├── .gitignore ├── Cargo.lock ├── Cargo.toml ├── README.md ├── compat.sh ├── rustfmt.toml └── src ├── cmd.rs ├── lru.rs ├── main.rs ├── parser.rs ├── server.rs └── store.rs /.gitignore: -------------------------------------------------------------------------------- 1 | target 2 | *.bk 3 | -------------------------------------------------------------------------------- /Cargo.lock: -------------------------------------------------------------------------------- 1 | [root] 2 | name = "rustcached" 3 | version = "0.1.0" 4 | dependencies = [ 5 | "getopts 0.2.14 (registry+https://github.com/rust-lang/crates.io-index)", 6 | "nom 1.2.4 (registry+https://github.com/rust-lang/crates.io-index)", 7 | "regex 0.1.73 (registry+https://github.com/rust-lang/crates.io-index)", 8 | "time 0.1.35 (registry+https://github.com/rust-lang/crates.io-index)", 9 | ] 10 | 11 | [[package]] 12 | name = "aho-corasick" 13 | version = "0.5.2" 14 | source = "registry+https://github.com/rust-lang/crates.io-index" 15 | dependencies = [ 16 | "memchr 0.1.11 (registry+https://github.com/rust-lang/crates.io-index)", 17 | ] 18 | 19 | [[package]] 20 | name = "getopts" 21 | version = "0.2.14" 22 | source = "registry+https://github.com/rust-lang/crates.io-index" 23 | 24 | [[package]] 25 | name = "kernel32-sys" 26 | version = "0.2.2" 27 | source = "registry+https://github.com/rust-lang/crates.io-index" 28 | dependencies = [ 29 | "winapi 0.2.8 (registry+https://github.com/rust-lang/crates.io-index)", 30 | "winapi-build 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)", 31 | ] 32 | 33 | [[package]] 34 | name = "libc" 35 | version = "0.2.15" 36 | source = "registry+https://github.com/rust-lang/crates.io-index" 37 | 38 | [[package]] 39 | name = "memchr" 40 | version = "0.1.11" 41 | source = "registry+https://github.com/rust-lang/crates.io-index" 42 | dependencies = [ 43 | "libc 0.2.15 (registry+https://github.com/rust-lang/crates.io-index)", 44 | ] 45 | 46 | [[package]] 47 | name = "nom" 48 | version = "1.2.4" 49 | source = "registry+https://github.com/rust-lang/crates.io-index" 50 | 51 | [[package]] 52 | name = "regex" 53 | version = "0.1.73" 54 | source = "registry+https://github.com/rust-lang/crates.io-index" 55 | dependencies = [ 56 | "aho-corasick 0.5.2 (registry+https://github.com/rust-lang/crates.io-index)", 57 | "memchr 0.1.11 (registry+https://github.com/rust-lang/crates.io-index)", 58 | "regex-syntax 0.3.4 (registry+https://github.com/rust-lang/crates.io-index)", 59 | "thread_local 0.2.6 (registry+https://github.com/rust-lang/crates.io-index)", 60 | "utf8-ranges 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)", 61 | ] 62 | 63 | [[package]] 64 | name = "regex-syntax" 65 | version = "0.3.4" 66 | source = "registry+https://github.com/rust-lang/crates.io-index" 67 | 68 | [[package]] 69 | name = "thread-id" 70 | version = "2.0.0" 71 | source = "registry+https://github.com/rust-lang/crates.io-index" 72 | dependencies = [ 73 | "kernel32-sys 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)", 74 | "libc 0.2.15 (registry+https://github.com/rust-lang/crates.io-index)", 75 | ] 76 | 77 | [[package]] 78 | name = "thread_local" 79 | version = "0.2.6" 80 | source = "registry+https://github.com/rust-lang/crates.io-index" 81 | dependencies = [ 82 | "thread-id 2.0.0 (registry+https://github.com/rust-lang/crates.io-index)", 83 | ] 84 | 85 | [[package]] 86 | name = "time" 87 | version = "0.1.35" 88 | source = "registry+https://github.com/rust-lang/crates.io-index" 89 | dependencies = [ 90 | "kernel32-sys 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)", 91 | "libc 0.2.15 (registry+https://github.com/rust-lang/crates.io-index)", 92 | "winapi 0.2.8 (registry+https://github.com/rust-lang/crates.io-index)", 93 | ] 94 | 95 | [[package]] 96 | name = "utf8-ranges" 97 | version = "0.1.3" 98 | source = "registry+https://github.com/rust-lang/crates.io-index" 99 | 100 | [[package]] 101 | name = "winapi" 102 | version = "0.2.8" 103 | source = "registry+https://github.com/rust-lang/crates.io-index" 104 | 105 | [[package]] 106 | name = "winapi-build" 107 | version = "0.1.1" 108 | source = "registry+https://github.com/rust-lang/crates.io-index" 109 | 110 | [metadata] 111 | "checksum aho-corasick 0.5.2 (registry+https://github.com/rust-lang/crates.io-index)" = "2b3fb52b09c1710b961acb35390d514be82e4ac96a9969a8e38565a29b878dc9" 112 | "checksum getopts 0.2.14 (registry+https://github.com/rust-lang/crates.io-index)" = "d9047cfbd08a437050b363d35ef160452c5fe8ea5187ae0a624708c91581d685" 113 | "checksum kernel32-sys 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)" = "7507624b29483431c0ba2d82aece8ca6cdba9382bff4ddd0f7490560c056098d" 114 | "checksum libc 0.2.15 (registry+https://github.com/rust-lang/crates.io-index)" = "23e3757828fa702a20072c37ff47938e9dd331b92fac6e223d26d4b7a55f7ee2" 115 | "checksum memchr 0.1.11 (registry+https://github.com/rust-lang/crates.io-index)" = "d8b629fb514376c675b98c1421e80b151d3817ac42d7c667717d282761418d20" 116 | "checksum nom 1.2.4 (registry+https://github.com/rust-lang/crates.io-index)" = "a5b8c256fd9471521bcb84c3cdba98921497f1a331cbc15b8030fc63b82050ce" 117 | "checksum regex 0.1.73 (registry+https://github.com/rust-lang/crates.io-index)" = "56b7ee9f764ecf412c6e2fff779bca4b22980517ae335a21aeaf4e32625a5df2" 118 | "checksum regex-syntax 0.3.4 (registry+https://github.com/rust-lang/crates.io-index)" = "31040aad7470ad9d8c46302dcffba337bb4289ca5da2e3cd6e37b64109a85199" 119 | "checksum thread-id 2.0.0 (registry+https://github.com/rust-lang/crates.io-index)" = "a9539db560102d1cef46b8b78ce737ff0bb64e7e18d35b2a5688f7d097d0ff03" 120 | "checksum thread_local 0.2.6 (registry+https://github.com/rust-lang/crates.io-index)" = "55dd963dbaeadc08aa7266bf7f91c3154a7805e32bb94b820b769d2ef3b4744d" 121 | "checksum time 0.1.35 (registry+https://github.com/rust-lang/crates.io-index)" = "3c7ec6d62a20df54e07ab3b78b9a3932972f4b7981de295563686849eb3989af" 122 | "checksum utf8-ranges 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)" = "a1ca13c08c41c9c3e04224ed9ff80461d97e121589ff27c753a16cb10830ae0f" 123 | "checksum winapi 0.2.8 (registry+https://github.com/rust-lang/crates.io-index)" = "167dc9d6949a9b857f3451275e911c3f44255842c1f7a76f33c55103a909087a" 124 | "checksum winapi-build 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)" = "2d315eee3b34aca4797b2da6b13ed88266e6d612562a0c46390af8299fc699bc" 125 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "rustcached" 3 | version = "0.1.0" 4 | authors = ["David King "] 5 | 6 | [dependencies] 7 | nom = "~1.2.0" 8 | time = "0.1" 9 | getopts = "0.2" 10 | regex = "0.1" 11 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | A memcached-compatible clone written in rust 2 | 3 | # Features: 4 | 5 | * Aside from what's listed in Missing below, we support all memcached commands and are fully compatible 6 | 7 | # Missing: 8 | 9 | * binary protocol 10 | * `stats` 11 | * UDP 12 | * unix sockets 13 | * `delete` with expires (memcached dropped this support in 1.4) 14 | * `flush_all` with expires 15 | * `verbosity` is recognised but ignored 16 | * SASL 17 | 18 | # Performance 19 | 20 | rustcache is about half as fast as memcached according to [memslap](http://docs.libmemcached.org/bin/memslap.html): 21 | 22 | rustcached: 23 | 24 | $ memslap --servers=localhost:11211 --test=set --concurrency=10 --debug --execute-number=10000 --flag --flush 25 | Threads connecting to servers 10 26 | Took 3.145 seconds to load data 27 | $ memslap --servers=localhost:11211 --test=get --concurrency=10 --debug --execute-number=10000 --flag --flush 28 | Threads connecting to servers 10 29 | Took 4.976 seconds to read data 30 | 31 | memcached: 32 | 33 | $ memslap --servers=localhost:11212 --test=set --concurrency=10 --debug --execute-number=10000 --flag --flush 34 | Threads connecting to servers 10 35 | Took 2.736 seconds to load data 36 | $ memslap --servers=localhost:11212 --test=get --concurrency=10 --debug --execute-number=10000 --flag --flush 37 | Threads connecting to servers 10 38 | Took 3.933 seconds to read data 39 | 40 | # Code organisation: 41 | 42 | * `cmd.rs`: control starts here, command line arguments parsed, and the server started 43 | * `store.rs`: houses the memcached application logic (e.g. what does "add" mean and how do I apply it?) 44 | * `lru.rs`: the LRU cache 45 | * `parser.rs`: protocol parsing 46 | * `server.rs`: socket handling and response writing 47 | 48 | # Todo: 49 | 50 | * We're using the default rust allocator (jemalloc) for everything, instead of a slab allocator like memcached 51 | * you can exhaust memory in the parser with an unlimited key/value length 52 | * you can't `set` values larger than the memcached length limits, but you can `append`/`prepend` past them 53 | * We copy a lot of stuff around right now that we don't have to, especially in the response builder 54 | * `lru.rs` uses `Arc` but there's probably a way to just use `Rc` when we go single-threaded 55 | 56 | # Future features: 57 | 58 | * we use LRU right now like memcached does, but I want to try ARC in a real environment 59 | * eager deletion of expired items during idle periods 60 | -------------------------------------------------------------------------------- /compat.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh -e 2 | 3 | sep() { 4 | echo ---------------------------- 5 | } 6 | 7 | echo killing 8 | 9 | pkill memcached||true 10 | pkill rustcache||true 11 | 12 | echo building 13 | cargo test 14 | cargo clean 15 | 16 | echo building for release 17 | cargo build --release 18 | sep 19 | 20 | echo starting servers 21 | memcached -p 11212 & 22 | cargo run --release -- -p 11211 & 23 | sep 24 | 25 | # I'd like to use these tests but they fail intermittently on real memcached 26 | # echo memcached memcapable 27 | # memcapable -t1 -a -p 11212 -v || true 28 | # sep 29 | # echo rustcache memcapable 30 | # memcapable -t1 -a -p 11211 -v || true 31 | # sep 32 | 33 | echo memcached memslap 34 | time memslap --servers=localhost:11212 --concurrency=10 --execute-number=100000 || true 35 | time memslap --servers=localhost:11212 --concurrency=10 --execute-number=100000 || true 36 | echo rustcache memslap 37 | time memslap --servers=localhost:11211 --concurrency=10 --execute-number=100000 || true 38 | time memslap --servers=localhost:11211 --concurrency=10 --execute-number=100000 || true 39 | 40 | kill %1 %2 41 | wait 42 | -------------------------------------------------------------------------------- /rustfmt.toml: -------------------------------------------------------------------------------- 1 | max_width = 80 2 | ideal_width = 80 3 | format_strings = false # to leave sql strings alone -------------------------------------------------------------------------------- /src/cmd.rs: -------------------------------------------------------------------------------- 1 | use std::env; 2 | use std::process; 3 | use std::str::FromStr; 4 | use std::io::Write; 5 | 6 | use getopts::Options; 7 | 8 | use server; 9 | use parser::parse_size; 10 | 11 | macro_rules! println_stderr( 12 | ($($arg:tt)*) => ( 13 | match writeln!(&mut ::std::io::stderr(), $($arg)* ) { 14 | Ok(_) => {}, 15 | Err(x) => panic!("Unable to write to stderr: {}", x), 16 | } 17 | ) 18 | ); 19 | 20 | pub fn main() { 21 | let mut port = 11211; 22 | let mut capacity = 64 * 1024 * 1024; 23 | let mut verbose = false; 24 | 25 | let args: Vec = env::args().collect(); 26 | let program = args[0].clone(); 27 | 28 | let mut opts = Options::new(); 29 | opts.optopt("p", "port", "port to listen on (default: 11211)", "PORT"); 30 | opts.optopt("m", "memory", "port to listen on (default: 64mb)", "MEMORY"); 31 | opts.optflag("v", "verbose", "be really verbose"); 32 | opts.optflag("h", "help", "print help and exit"); 33 | 34 | let print_usage_and_die = |exit_code: i32| -> ! { 35 | let brief = format!("Usage: {} [options]", program); 36 | println_stderr!("{}", opts.usage(&brief)); 37 | process::exit(exit_code); 38 | }; 39 | 40 | let matches = match opts.parse(&args[1..]) { 41 | Ok(m) => m, 42 | Err(f) => { 43 | println_stderr!("{}", f); 44 | return print_usage_and_die(1); 45 | } 46 | }; 47 | 48 | if matches.opt_present("h") || !matches.free.is_empty() { 49 | return print_usage_and_die(1); 50 | } 51 | 52 | if let Some(digits) = matches.opt_str("p") { 53 | if let Result::Ok(port_num) = FromStr::from_str(&digits) { 54 | port = port_num; 55 | } else { 56 | println_stderr!("couldn't parse port num {}", digits); 57 | return print_usage_and_die(1); 58 | } 59 | } 60 | 61 | if let Some(size_spec) = matches.opt_str("m") { 62 | if let Some(size) = parse_size(&size_spec) { 63 | capacity = size; 64 | } else { 65 | println_stderr!("couldn't parse size {}", size_spec); 66 | return print_usage_and_die(1); 67 | } 68 | } 69 | 70 | if matches.opt_present("v") { 71 | verbose = true; 72 | } 73 | 74 | server::start(port, capacity, verbose); 75 | } 76 | -------------------------------------------------------------------------------- /src/lru.rs: -------------------------------------------------------------------------------- 1 | /// The LRU storage engine 2 | 3 | use std::cmp::Ord; 4 | use std::collections::HashMap; 5 | use std::collections::BTreeSet; 6 | use std::mem; 7 | use std::hash::Hash; 8 | use std::sync::Arc; 9 | 10 | pub type Weight = usize; 11 | pub type Timestamp = u32; 12 | 13 | type LruEntryUsed = (Timestamp, Arc); 14 | type LruEntryExpires = (Timestamp, Arc); 15 | 16 | #[derive(Debug)] 17 | pub struct LruCache { 18 | map: HashMap, LruEntry>, 19 | lru: BTreeSet>, 20 | expires: BTreeSet>, 21 | capacity: Weight, 22 | weight: Weight, // TODO store this? 23 | } 24 | 25 | pub trait HasWeight { 26 | fn weight(&self) -> Weight; 27 | } 28 | 29 | #[derive(Debug)] 30 | pub struct LruEntry { 31 | pub data: V, 32 | pub key: Arc, 33 | used: Timestamp, 34 | pub expires: Option, 35 | weight: Weight, 36 | } 37 | 38 | impl LruCache { 39 | pub fn new(capacity: Weight) -> LruCache { 40 | LruCache { 41 | map: HashMap::new(), 42 | lru: BTreeSet::new(), 43 | expires: BTreeSet::new(), 44 | capacity: capacity, 45 | weight: 0, 46 | } 47 | } 48 | 49 | pub fn clear(&mut self) { 50 | self.map.clear(); 51 | self.lru.clear(); 52 | self.expires.clear(); 53 | self.weight = 0; 54 | } 55 | 56 | pub fn get_full_entry(&mut self, 57 | key: &K, 58 | now: Timestamp) 59 | -> Option<&LruEntry> { 60 | let entry = self._get_full_entry(key, now); 61 | entry.map(|e| &*e) 62 | } 63 | 64 | fn _get_full_entry(&mut self, 65 | key: &K, 66 | now: Timestamp) 67 | -> Option<&mut LruEntry> { 68 | match self.map.get_mut(key) { 69 | None => Option::None, 70 | 71 | Some(ref entry) if expired((*entry).expires, now) => { 72 | // we found it, but it's expired. we could theoretically 73 | // pre-emptively remove it on discovering this, but for the 74 | // moment we'll leave it there and clean it up during the normal 75 | // cleaup process (thereby keeping our reads fast and paying the 76 | // cost on writes instead) 77 | Option::None 78 | } 79 | 80 | Some(entry) => { 81 | // we found it and it hasn't expired. Since it's being used now 82 | // we need to update its position in the LRU 83 | 84 | // use the key we found inside of the structure instead of the 85 | // one that was passed in, so we can save memory by just 86 | // incrementing the reference count on that one instead of 87 | // copying it 88 | let inner_key = (*entry).key.clone(); 89 | 90 | if ((*entry).used) != now { 91 | // only update it if the value would change 92 | let old_lru_key = ((*entry).used, inner_key.clone()); 93 | self.lru.remove(&old_lru_key); 94 | 95 | entry.used = now; 96 | 97 | let new_lru_key = (now, inner_key.clone()); 98 | self.lru.insert(new_lru_key); 99 | } 100 | 101 | Some(entry) 102 | } 103 | } 104 | } 105 | 106 | pub fn get(&mut self, key: &K, now: Timestamp) -> Option<&V> { 107 | self.get_full_entry(key, now).map(|entry| &entry.data) 108 | } 109 | 110 | pub fn set(&mut self, 111 | key: K, 112 | value: V, 113 | expires: Option, 114 | now: Timestamp) 115 | -> bool { 116 | if expired(expires, now) { 117 | // if it's already expired there's no need to store it 118 | return false; 119 | } 120 | 121 | // if it's already in here, we need to get rid of it 122 | self.delete(&key); 123 | 124 | let weight = compute_weight(&key, &value); 125 | 126 | if weight > self.capacity { 127 | // we'll never be able to store this 128 | return false; 129 | } 130 | 131 | // free up any space that we need to in order to fit this 132 | let capacity = self.capacity; 133 | self.deweight(capacity - weight, now); 134 | 135 | let k2 = Arc::new(key.clone()); 136 | 137 | let entry = LruEntry { 138 | key: k2.clone(), 139 | data: value, 140 | expires: expires, 141 | weight: weight, 142 | used: now, 143 | }; 144 | 145 | self.map.insert(k2.clone(), entry); 146 | self.weight += weight; 147 | 148 | let lru_key = (now, k2.clone()); 149 | self.lru.insert(lru_key); 150 | 151 | if let Some(expires_ts) = expires { 152 | // if it expires, add it to the expiration queue 153 | let expires_key = (expires_ts, k2.clone()); 154 | self.expires.insert(expires_key); 155 | } 156 | 157 | // TODO Store always ignores this return value. Do we care? 158 | true 159 | } 160 | 161 | pub fn fast_get(&self, key: &K, now: Timestamp) -> Option<&V> { 162 | // fetch the value of a key without updating the LRU 163 | match self.map.get(key) { 164 | Some(entry) if expired((*entry).expires, now) => None, 165 | Some(entry) => Some(&(*entry).data), 166 | None => None, 167 | } 168 | } 169 | 170 | pub fn contains(&self, key: &K, now: Timestamp) -> bool { 171 | // checks for the presence of a key (without updating the LRU) 172 | self.fast_get(key, now).is_some() 173 | } 174 | 175 | pub fn touch(&mut self, 176 | key: &K, 177 | expires: Option, 178 | now: Timestamp) 179 | -> bool { 180 | // update the timestamp and last-used field of a row without copying the 181 | // whole contents 182 | let (old_key, old_expires, old_used) = match self._get_full_entry(key, now) { 183 | None => { 184 | // just bail, it was never in here anyway 185 | return false; 186 | } 187 | Some(full_entry) => { 188 | // change it in-place 189 | let old_key = (*full_entry).key.clone(); 190 | let old_expires = (*full_entry).expires; 191 | let old_used = (*full_entry).used; 192 | (*full_entry).expires = expires; 193 | (*full_entry).used = now; 194 | (old_key, old_expires, old_used) 195 | } 196 | }; 197 | 198 | // update our data structures 199 | self._touch(old_key, old_expires, expires, old_used, now); 200 | true 201 | } 202 | 203 | fn _touch(&mut self, 204 | key: Arc, 205 | old_expires: Option, 206 | new_expires: Option, 207 | old_used: Timestamp, 208 | now: Timestamp) { 209 | 210 | if old_expires != new_expires { 211 | if let Some(old_expires_ts) = old_expires { 212 | // if it expired before, we have to remove it 213 | let old_expires_key = (old_expires_ts, key.clone()); 214 | self.expires.remove(&old_expires_key); 215 | } 216 | 217 | if let Some(expires_ts) = new_expires { 218 | // if it expires now, we have to add it 219 | let expires_key = (expires_ts, key.clone()); 220 | self.expires.insert(expires_key); 221 | } 222 | } 223 | 224 | if old_used != now { 225 | let old_lru_key = (old_used, key.clone()); 226 | self.lru.remove(&old_lru_key); 227 | let new_lru_key = (now, key.clone()); 228 | self.lru.insert(new_lru_key); 229 | } 230 | } 231 | 232 | pub fn delete(&mut self, key: &K) -> bool { 233 | let found = { 234 | match self.map.get(key) { 235 | None => None, 236 | Some(entry) => { 237 | Some(((*entry).key.clone(), 238 | (*entry).expires, 239 | (*entry).used, 240 | (*entry).weight)) 241 | } 242 | } 243 | }; 244 | 245 | if let Some((old_key, expires, used, weight)) = found { 246 | self.map.remove(&old_key); 247 | let lru_key = (used, old_key.clone()); 248 | self.lru.remove(&lru_key); 249 | if let Some(expires_ts) = expires { 250 | let expires_key = (expires_ts, old_key.clone()); 251 | self.lru.remove(&expires_key); 252 | } 253 | self.weight -= weight; 254 | true 255 | } else { 256 | false 257 | } 258 | } 259 | 260 | fn deweight(&mut self, target_weight: Weight, now: Timestamp) { 261 | // we're trying to add more data, but there isn't room for it. We need 262 | // to delete at least `weight` worth of data to fit this new entry 263 | 264 | while self.weight > target_weight && !self.map.is_empty() { 265 | self.deweight_once(now); 266 | } 267 | 268 | assert!(self.capacity >= target_weight); 269 | } 270 | 271 | fn deweight_once(&mut self, now: Timestamp) { 272 | if self.map.is_empty() { 273 | // nothing we can delete if it's already empty 274 | return; 275 | } 276 | 277 | let expired_key = { 278 | // check the expiration queue for stuff that's already expired that 279 | // we can just delete 280 | let maybe_expirable = &self.expires; 281 | let mut maybe_expirable = maybe_expirable.into_iter(); 282 | let maybe_expirable = maybe_expirable.next(); 283 | match maybe_expirable { 284 | None => None, 285 | Some(found_tuple) => { 286 | let (ref expired_ts, ref expired_key) = *found_tuple; 287 | if _expired(*expired_ts, now) { 288 | Some(expired_key.clone()) 289 | } else { 290 | None 291 | } 292 | } 293 | } 294 | }; 295 | 296 | if let Some(key_ref) = expired_key { 297 | self.delete(&*key_ref); 298 | return; 299 | } 300 | 301 | // otherwise we have to use the LRU 302 | let lru_key = { 303 | let lru = &self.lru; 304 | let mut lru = lru.into_iter(); 305 | let lru = lru.next(); 306 | match lru { 307 | None => None, 308 | Some(found_tuple) => { 309 | let (_, ref lru_key) = *found_tuple; 310 | Some(lru_key.clone()) 311 | } 312 | } 313 | }; 314 | 315 | if let Some(key_ref) = lru_key { 316 | self.delete(&*key_ref); 317 | return; 318 | } 319 | 320 | unreachable!("there's nothing on the LRU?"); 321 | } 322 | 323 | #[cfg(test)] 324 | pub fn all_keys(&self, now: Timestamp) -> Vec { 325 | // very expensive operation that fetches a full list of all of the keys 326 | // that we know about that aren't expired 327 | let mut ret = Vec::new(); 328 | for (ref key, ref value) in &self.map { 329 | if !expired((*value).expires, now) { 330 | // more stars is better, right? 331 | let copied = (*(**key)).clone(); 332 | ret.push(copied); 333 | } 334 | } 335 | ret.sort(); 336 | ret 337 | } 338 | } 339 | 340 | fn expired(timestamp: Option, now: Timestamp) -> bool { 341 | match timestamp { 342 | Some(ts) if _expired(ts, now) => true, 343 | _ => false, 344 | } 345 | } 346 | 347 | fn _expired(timestamp: Timestamp, now: Timestamp) -> bool { 348 | timestamp < now 349 | } 350 | 351 | pub fn compute_weight(key: &K, 352 | value: &V) 353 | -> Weight { 354 | // this isn't perfect because it ignores some hashtable and btreeset 355 | // overhead, but it's a pretty good guess at the memory usage of an entry 356 | let mut sum = 0; 357 | sum += key.weight(); 358 | sum += 3 * mem::size_of::>(); 359 | sum += value.weight(); 360 | sum += mem::size_of::(); 361 | sum += 2 * mem::size_of::(); 362 | sum += 2 * mem::size_of::>(); 363 | sum 364 | } 365 | 366 | impl HasWeight for Vec { 367 | fn weight(&self) -> Weight { 368 | self.len() 369 | } 370 | } 371 | 372 | #[cfg(test)] 373 | mod tests { 374 | use super::*; 375 | 376 | const NOW: Timestamp = 100; 377 | const FUTURE: Timestamp = NOW + 1; 378 | const FUTURE2: Timestamp = NOW + 2; 379 | const PAST: Timestamp = NOW - 1; 380 | const CAPACITY: Weight = 300; 381 | 382 | #[test] 383 | fn basic_set() { 384 | let mut store = make_store(); 385 | 386 | store.set(b("foo"), b("data"), None, NOW); 387 | assert!(store.contains(&b("foo"), NOW)); 388 | assert_eq!(store.all_keys(NOW), vec![b("foo")]); 389 | } 390 | 391 | #[test] 392 | fn set_already_expired() { 393 | let mut store = make_store(); 394 | 395 | store.set(b("foo"), b("data"), Some(PAST), NOW); 396 | assert!(store.all_keys(NOW).is_empty()); 397 | } 398 | 399 | #[test] 400 | fn set_expires() { 401 | let mut store = make_store(); 402 | 403 | store.set(b("foo"), b("data"), Some(PAST), NOW); 404 | assert!(store.all_keys(FUTURE).is_empty()); 405 | } 406 | 407 | #[test] 408 | fn too_big() { 409 | let mut store = make_store(); 410 | 411 | store.set(b("foo1"), b("data"), None, PAST); 412 | 413 | let big = make_big(CAPACITY * 2); 414 | 415 | store.set(b("foo2"), big, None, PAST); 416 | 417 | // setting that big item should have been rejected, so the old data 418 | // should still be there 419 | assert!(store.contains(&b("foo1"), NOW)); 420 | assert_eq!(store.all_keys(NOW), vec![b("foo1")]); 421 | } 422 | 423 | #[test] 424 | fn outgrow() { 425 | let mut store = make_store(); 426 | 427 | store.set(b("foo1"), b("data"), None, PAST); 428 | 429 | // these should push that guy out 430 | store.set(b("foo2"), make_big(30), None, NOW); 431 | store.set(b("foo3"), make_big(30), None, FUTURE); 432 | 433 | assert!(!store.contains(&b("foo1"), FUTURE2)); 434 | assert!(store.contains(&b("foo2"), FUTURE2)); 435 | assert!(store.contains(&b("foo3"), FUTURE2)); 436 | } 437 | 438 | #[test] 439 | fn prefer_expired() { 440 | // make sure we prefer to remove expired members over live members 441 | let mut store = make_store(); 442 | 443 | store.set(b("foo1"), make_big(30), None, PAST); 444 | store.set(b("foo2"), make_big(30), Some(NOW), NOW); 445 | 446 | // this has to push one of them out 447 | store.set(b("foo3"), make_big(30), None, FUTURE); 448 | 449 | assert!(store.contains(&b("foo1"), FUTURE2)); 450 | assert!(!store.contains(&b("foo2"), FUTURE2)); 451 | assert!(store.contains(&b("foo3"), FUTURE2)); 452 | } 453 | 454 | #[test] 455 | fn clear() { 456 | let mut store = make_store(); 457 | 458 | store.set(b("foo1"), b("data"), None, PAST); 459 | store.clear(); 460 | 461 | assert!(!store.contains(&b("foo1"), NOW)); 462 | } 463 | 464 | 465 | fn make_store() -> LruCache, Vec> { 466 | let store = LruCache::new(CAPACITY); 467 | store 468 | } 469 | 470 | fn make_big(size: usize) -> Vec { 471 | let mut big = Vec::new(); 472 | for x in 0..size * 2 { 473 | big.push(x as u8); 474 | } 475 | big 476 | } 477 | 478 | fn b(inp: &'static str) -> Vec { 479 | // syntactic sugar for tests 480 | let mut s = String::new(); 481 | s.push_str(inp); 482 | s.into_bytes() 483 | } 484 | } 485 | -------------------------------------------------------------------------------- /src/main.rs: -------------------------------------------------------------------------------- 1 | #[macro_use] 2 | extern crate nom; 3 | extern crate time; 4 | extern crate getopts; 5 | extern crate regex; 6 | 7 | mod parser; 8 | mod store; 9 | mod lru; 10 | mod server; 11 | mod cmd; 12 | 13 | pub fn main() { 14 | cmd::main() 15 | } 16 | -------------------------------------------------------------------------------- /src/parser.rs: -------------------------------------------------------------------------------- 1 | /// The parser. Takes a stream of bytes and turns it into a series of parsed 2 | /// commands ready to be send to the store 3 | 4 | use std::str::from_utf8; 5 | use std::str::FromStr; 6 | 7 | use nom::{crlf, space, digit}; 8 | use regex::Regex; // used for the size parser 9 | 10 | pub use nom::{IResult, Needed}; 11 | 12 | use store::ServerCommand; 13 | use store::IncrementerType; 14 | use store::GetterType; 15 | use store::SetterType; 16 | 17 | #[derive(Debug,PartialEq,Eq)] 18 | pub struct CommandConfig<'a> { 19 | pub should_reply: bool, 20 | pub command: ServerCommand<'a>, 21 | } 22 | 23 | named!(key_parser<&[u8], &[u8]>, is_not!(" \t\r\n\0")); 24 | 25 | fn unwrap_noreply(tag: Option<&[u8]>) -> bool { 26 | match tag { 27 | Some(b"noreply") => false, 28 | Some(_) => panic!(format!("can't unwrap noreply tag {:?}", tag)), 29 | None => true, 30 | } 31 | } 32 | 33 | named!(u32_digit, 34 | map_res!( 35 | map_res!( 36 | digit, 37 | from_utf8 38 | ), 39 | FromStr::from_str 40 | ) 41 | ); 42 | 43 | named!(usize_digit, 44 | map_res!( 45 | map_res!( 46 | digit, 47 | from_utf8 48 | ), 49 | FromStr::from_str 50 | ) 51 | ); 52 | 53 | 54 | named!(u64_digit, 55 | map_res!( 56 | map_res!( 57 | digit, 58 | from_utf8 59 | ), 60 | FromStr::from_str 61 | ) 62 | ); 63 | 64 | fn map_setter_name(res: &[u8]) -> SetterType { 65 | match res { 66 | b"set" => SetterType::Set, 67 | b"add" => SetterType::Add, 68 | b"prepend" => SetterType::Prepend, 69 | b"replace" => SetterType::Replace, 70 | b"append" => SetterType::Append, 71 | _ => panic!(format!("unknown setter mapped? {:?}", res)), 72 | } 73 | } 74 | 75 | named!(parse_setter_name, 76 | alt!( 77 | tag!("set") | 78 | tag!("add") | 79 | tag!("prepend") | 80 | tag!("replace") | 81 | tag!("append") 82 | ) 83 | ); 84 | 85 | // cas [noreply]\r\n 86 | named!(cmd_cas<&[u8], CommandConfig>, 87 | chain!( 88 | tag!("cas") ~ 89 | space ~ 90 | key: key_parser ~ 91 | space ~ 92 | flags: u32_digit ~ 93 | space ~ 94 | ttl: u32_digit ~ 95 | space ~ 96 | bytes: usize_digit ~ 97 | space ~ 98 | cas_unique: u64_digit ~ 99 | noreply: chain!(space ~ x: tag!("noreply"), || {x})? ~ 100 | crlf ~ 101 | payload: take!(bytes) ~ 102 | crlf, 103 | || { 104 | CommandConfig { 105 | should_reply: unwrap_noreply(noreply), 106 | command: ServerCommand::Setter{ 107 | setter: SetterType::Cas(cas_unique), 108 | key: key, 109 | data: payload, 110 | ttl: ttl, 111 | flags: flags, 112 | } 113 | } 114 | } 115 | ) 116 | ); 117 | 118 | // setters: 119 | // [noreply]\r\npayload\r\n 120 | named!(cmd_set<&[u8], CommandConfig>, 121 | chain!( 122 | setter_name: parse_setter_name ~ 123 | space ~ 124 | key: key_parser ~ 125 | space ~ 126 | flags: u32_digit ~ 127 | space ~ 128 | ttl: u32_digit ~ 129 | space ~ 130 | bytes: u32_digit ~ 131 | noreply: chain!(space ~ x: tag!("noreply"), || {x})? ~ 132 | crlf ~ 133 | payload: take!(bytes) ~ // assuming this is where the payload is 134 | crlf, 135 | || { 136 | let setter = map_setter_name(setter_name); 137 | 138 | CommandConfig { 139 | should_reply: unwrap_noreply(noreply), 140 | command: ServerCommand::Setter { 141 | setter: setter, 142 | key: key, 143 | data: payload, 144 | ttl: ttl, 145 | flags: flags, 146 | } 147 | } 148 | } 149 | ) 150 | ); 151 | 152 | fn map_getter_name(res: &[u8]) -> GetterType { 153 | match res { 154 | b"get" => GetterType::Get, 155 | b"gets" => GetterType::Gets, 156 | _ => panic!(format!("unknown getter mapped? {:?}", res)), 157 | } 158 | } 159 | 160 | named!(parse_getter_name, 161 | alt!( 162 | tag!("gets") | 163 | tag!("get") 164 | ) 165 | ); 166 | 167 | 168 | // get *\r\n 169 | // gets *\r\n 170 | named!(cmd_get<&[u8], CommandConfig>, 171 | chain!( 172 | getter_name: parse_getter_name ~ 173 | space ~ 174 | keys: separated_nonempty_list!(space, key_parser) ~ 175 | crlf, 176 | || { 177 | CommandConfig { 178 | should_reply: true, 179 | command: ServerCommand::Getter { 180 | getter: map_getter_name(getter_name), 181 | keys: keys 182 | } 183 | } 184 | } 185 | ) 186 | ); 187 | 188 | // delete [noreply]\r\n 189 | // TODO there's a rumour that this can take a time? 190 | named!(cmd_delete<&[u8], CommandConfig>, 191 | chain!( 192 | tag!("delete") ~ 193 | space ~ 194 | key: key_parser ~ 195 | noreply: chain!(space ~ x: tag!("noreply"), || {x})? ~ 196 | crlf, 197 | || { 198 | CommandConfig { 199 | should_reply: unwrap_noreply(noreply), 200 | command: ServerCommand::Delete { 201 | key: key 202 | } 203 | } 204 | } 205 | ) 206 | ); 207 | 208 | // touch [noreply]\r\n 209 | named!(cmd_touch<&[u8], CommandConfig>, 210 | chain!( 211 | tag!("touch") ~ 212 | space ~ 213 | key: key_parser ~ 214 | space ~ 215 | ttl: u32_digit ~ 216 | noreply: chain!(space ~ x: tag!("noreply"), || {x})? ~ 217 | crlf, 218 | || { 219 | CommandConfig { 220 | should_reply: unwrap_noreply(noreply), 221 | command: ServerCommand::Touch { 222 | key: key, 223 | ttl: ttl, 224 | } 225 | } 226 | } 227 | ) 228 | ); 229 | 230 | fn map_incr_name(res: &[u8]) -> IncrementerType { 231 | match res { 232 | b"incr" => IncrementerType::Incr, 233 | b"decr" => IncrementerType::Decr, 234 | _ => panic!(format!("unknown getter mapped? {:?}", res)), 235 | } 236 | } 237 | 238 | named!(parse_incr_name, 239 | alt!( 240 | tag!("incr") | 241 | tag!("decr") 242 | ) 243 | ); 244 | 245 | // incr [noreply]\r\n 246 | // decr [noreply]\r\n 247 | named!(cmd_incr<&[u8], CommandConfig>, 248 | chain!( 249 | incr_name: parse_incr_name ~ 250 | space ~ 251 | key: key_parser ~ 252 | space ~ 253 | value: u64_digit ~ 254 | noreply: chain!(space ~ x: tag!("noreply"), || {x})? ~ 255 | crlf, 256 | || { 257 | CommandConfig { 258 | should_reply: unwrap_noreply(noreply), 259 | command: ServerCommand::Incrementer { 260 | incrementer: map_incr_name(incr_name), 261 | key: key, 262 | value: value, 263 | } 264 | } 265 | } 266 | ) 267 | ); 268 | 269 | // verbosity \r\n 270 | named!(cmd_verbosity<&[u8], CommandConfig>, 271 | chain!( 272 | tag!("verbosity") ~ 273 | space ~ 274 | u32_digit ~ 275 | noreply: chain!(space ~ x: tag!("noreply"), || {x})? ~ 276 | crlf, 277 | || { 278 | CommandConfig { 279 | should_reply: unwrap_noreply(noreply), 280 | command: ServerCommand::Verbosity 281 | } 282 | } 283 | ) 284 | ); 285 | 286 | // version\r\n 287 | named!(cmd_version<&[u8], CommandConfig>, 288 | chain!( 289 | tag!("version") ~ 290 | crlf, 291 | || { 292 | CommandConfig { 293 | should_reply: true, 294 | command: ServerCommand::Version 295 | } 296 | } 297 | ) 298 | ); 299 | 300 | // quit\r\n 301 | named!(cmd_quit<&[u8], CommandConfig>, 302 | chain!( 303 | tag!("quit") ~ 304 | crlf, 305 | || { 306 | CommandConfig { 307 | should_reply: true, 308 | command: ServerCommand::Quit 309 | } 310 | } 311 | ) 312 | ); 313 | 314 | // flush_all\r\n 315 | named!(cmd_flushall<&[u8], CommandConfig>, 316 | chain!( 317 | tag!("flush_all") ~ 318 | noreply: chain!(space ~ x: tag!("noreply"), || {x})? ~ 319 | crlf, 320 | || { 321 | CommandConfig { 322 | should_reply: unwrap_noreply(noreply), 323 | command: ServerCommand::FlushAll 324 | } 325 | } 326 | ) 327 | ); 328 | 329 | // anything else is a malformed command 330 | named!(cmd_bad<&[u8], CommandConfig>, 331 | chain!( 332 | bad_stuff: is_not!("\r\n")? ~ 333 | crlf, 334 | || { 335 | CommandConfig { 336 | should_reply: true, 337 | command: ServerCommand::Bad(bad_stuff.unwrap_or(b"")) 338 | } 339 | } 340 | ) 341 | ); 342 | 343 | named!(pub parse_command<&[u8], CommandConfig>, 344 | alt!( 345 | // these short ones need to go first to work around a bug in nom where 346 | // it thinks it needs more data than it does 347 | cmd_quit | cmd_version | cmd_flushall | cmd_verbosity 348 | | cmd_set | cmd_cas | cmd_get | cmd_delete | cmd_incr | cmd_touch 349 | | cmd_bad 350 | ) 351 | ); 352 | 353 | pub fn parse_size(size_str: &str) -> Option { 354 | let re = Regex::new(r"^(\d+)([kmgt]?)b?$").unwrap(); 355 | match re.captures(size_str) { 356 | None => None, 357 | Some(matches) => { 358 | let digits = matches.at(1).unwrap(); 359 | let number: usize = FromStr::from_str(digits).unwrap(); 360 | let suffix = matches.at(2); 361 | let mult = match suffix { 362 | None | Some("b") | Some("") => 1, 363 | Some("k") => 1024, 364 | Some("m") => 1024 * 1024, 365 | Some("g") => 1024 * 1024 * 1024, 366 | Some("t") => 1024 * 1024 * 1024 * 1024, 367 | bad_mult => { 368 | unreachable!(format!("weird suffix {:?}", bad_mult)) 369 | } 370 | }; 371 | Some(number * mult) 372 | } 373 | } 374 | } 375 | 376 | #[cfg(test)] 377 | mod tests { 378 | use super::*; 379 | use store::ServerCommand; 380 | use store::IncrementerType; 381 | use store::GetterType; 382 | use store::SetterType; 383 | 384 | #[test] 385 | pub fn commands() { 386 | let tests: Vec<(&str, IResult<&[u8], CommandConfig>)> = vec![ 387 | ("set foo 12 34 5\r\ndata!\r\n", 388 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Setter {setter: SetterType::Set, key: b"foo", data: b"data!", ttl: 34, flags: 12 } })), 389 | ("set foo 12 34 5 noreply\r\ndata!\r\n", 390 | IResult::Done(b"", CommandConfig { should_reply: false, command: ServerCommand::Setter { setter: SetterType::Set, key: b"foo", data: b"data!", ttl: 34, flags: 12 } })), 391 | ("add foo 12 34 5\r\ndata!\r\n", 392 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Setter {setter: SetterType::Add, key: b"foo", data: b"data!", ttl: 34, flags: 12 } })), 393 | ("add foo 12 34 5 noreply\r\ndata!\r\n", 394 | IResult::Done(b"", CommandConfig { should_reply: false, command: ServerCommand::Setter { setter: SetterType::Add, key: b"foo", data: b"data!", ttl: 34, flags: 12 } })), 395 | ("append foo 12 34 5\r\ndata!\r\n", 396 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Setter { setter: SetterType::Append, key: b"foo", data: b"data!", ttl: 34, flags: 12 } })), 397 | ("append foo 12 34 5 noreply\r\ndata!\r\n", 398 | IResult::Done(b"", CommandConfig { should_reply: false, command: ServerCommand::Setter { setter: SetterType::Append, key: b"foo", data: b"data!", ttl: 34, flags: 12 } })), 399 | ("prepend foo 12 34 5\r\ndata!\r\n", 400 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Setter { setter: SetterType::Prepend, key: b"foo", data: b"data!", ttl: 34, flags: 12 } })), 401 | ("prepend foo 12 34 5 noreply\r\ndata!\r\n", 402 | IResult::Done(b"", CommandConfig { should_reply: false, command: ServerCommand::Setter { setter: SetterType::Prepend, key: b"foo", data: b"data!", ttl: 34, flags: 12 } })), 403 | ("replace foo 12 34 5 noreply\r\ndata!\r\n", 404 | IResult::Done(b"", CommandConfig { should_reply: false, command: ServerCommand::Setter { setter: SetterType::Replace, key: b"foo", data: b"data!", ttl: 34, flags: 12 } })), 405 | ("replace foo 12 34 5 noreply\r\ndata!\r\n", 406 | IResult::Done(b"", CommandConfig { should_reply: false, command: ServerCommand::Setter { setter: SetterType::Replace, key: b"foo", data: b"data!", ttl: 34, flags: 12 } })), 407 | 408 | ("cas foo 12 34 5 89\r\ndata!\r\n", 409 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Setter { setter: SetterType::Cas(89), key: b"foo", data: b"data!", ttl: 34, flags: 12 } })), 410 | ("cas foo 12 34 5 89 noreply\r\ndata!\r\n", 411 | IResult::Done(b"", CommandConfig { should_reply: false, command: ServerCommand::Setter { setter: SetterType::Cas(89), key: b"foo", data: b"data!", ttl: 34, flags: 12 } })), 412 | 413 | ("get foo\r\n", 414 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Getter { getter: GetterType::Get, keys: vec![b"foo"] } })), 415 | ("get foo1 foo2\r\n", 416 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Getter { getter: GetterType::Get, keys: vec![b"foo1", b"foo2"] } })), 417 | ("gets foo\r\n", 418 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Getter { getter: GetterType::Gets, keys: vec![b"foo"] } })), 419 | ("gets foo1 foo2\r\n", 420 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Getter { getter: GetterType::Gets, keys: vec![b"foo1", b"foo2"] } })), 421 | 422 | ("delete foo\r\n", 423 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Delete { key: b"foo" } })), 424 | ("delete foo noreply\r\n", 425 | IResult::Done(b"", CommandConfig { should_reply: false, command: ServerCommand::Delete { key: b"foo" } })), 426 | 427 | ("incr foo 5\r\n", 428 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Incrementer { incrementer: IncrementerType::Incr, key: b"foo", value: 5 } })), 429 | ("incr foo 5 noreply\r\n", 430 | IResult::Done(b"", CommandConfig { should_reply: false, command: ServerCommand::Incrementer { incrementer: IncrementerType::Incr, key: b"foo", value: 5 } })), 431 | ("decr foo 5\r\n", 432 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Incrementer { incrementer: IncrementerType::Decr, key: b"foo", value: 5 } })), 433 | ("decr foo 5 noreply\r\n", 434 | IResult::Done(b"", CommandConfig { should_reply: false, command: ServerCommand::Incrementer { incrementer: IncrementerType::Decr, key: b"foo", value: 5 } })), 435 | 436 | ("touch foo 5\r\n", 437 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Touch { key: b"foo", ttl: 5 } })), 438 | ("touch foo 5 noreply\r\n", 439 | IResult::Done(b"", CommandConfig { should_reply: false, command: ServerCommand::Touch { key: b"foo", ttl: 5 } })), 440 | 441 | ("flush_all\r\n", 442 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::FlushAll })), 443 | ("flush_all noreply\r\n", 444 | IResult::Done(b"", CommandConfig { should_reply: false, command: ServerCommand::FlushAll })), 445 | ("version\r\n", 446 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Version })), 447 | ("quit\r\n", 448 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Quit })), 449 | ("verbosity 10\r\n", 450 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Verbosity })), 451 | ("verbosity 10 noreply\r\n", 452 | IResult::Done(b"", CommandConfig { should_reply: false, command: ServerCommand::Verbosity })), 453 | 454 | ("foo bar\r\n", 455 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Bad(b"foo bar") })), 456 | ("version foo bar\r\n", 457 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Bad(b"version foo bar") })), 458 | ("\r\n", 459 | IResult::Done(b"", CommandConfig { should_reply: true, command: ServerCommand::Bad(b"") } )), 460 | 461 | ]; 462 | 463 | for &(ref command, ref expected_result) in &tests { 464 | // let (command, expected_result) = item; 465 | println!("command: {:?}", *command); 466 | let parsed = parse_command(command.as_bytes()); 467 | println!("expect: {:?}", expected_result); 468 | println!("got: {:?}", parsed); 469 | if *expected_result == parsed { 470 | println!("good!"); 471 | } else { 472 | println!("bad :("); 473 | } 474 | assert_eq!(*expected_result, parsed); 475 | } 476 | } 477 | 478 | #[test] 479 | pub fn parse_sizes() { 480 | let tests = vec![ 481 | ("0", Some(0)), 482 | ("1", Some(1)), 483 | ("1b", Some(1)), 484 | ("10", Some(10)), 485 | ("100", Some(100)), 486 | ("1k", Some(1024)), 487 | ("2k", Some(2048)), 488 | ("1m", Some(1024*1024)), 489 | ("2m", Some(2*1024*1024)), 490 | ("2mb", Some(2*1024*1024)), 491 | ("garbage", None), 492 | ("1.5gb", None), // might be nice to support this some day 493 | ]; 494 | 495 | for &(ref text, ref expected_result) in &tests { 496 | println!("Parsing {:?}", text); 497 | println!("Expect {:?}", expected_result); 498 | let parsed = parse_size(text); 499 | println!("Got {:?}", parsed); 500 | assert_eq!(*expected_result, parsed); 501 | } 502 | } 503 | } 504 | -------------------------------------------------------------------------------- /src/server.rs: -------------------------------------------------------------------------------- 1 | use std::thread::spawn; 2 | use std::sync::{Arc, Mutex}; 3 | use std::net::{TcpListener, TcpStream}; 4 | 5 | use std::io; 6 | use std::io::{Read, Write}; 7 | 8 | use store::Store; 9 | use store::Response; 10 | use store::ServerCommand; 11 | use parser::CommandConfig; 12 | use parser; 13 | 14 | pub const NAME: &'static [u8] = b"rustcache"; 15 | pub const VERSION: &'static [u8] = b"0.1.0"; 16 | 17 | fn format_response(response: Response, socket: &mut Write) -> io::Result<()> { 18 | match response { 19 | Response::Data { responses } => { 20 | for response in &responses { 21 | try!(socket.write(b"VALUE ")); 22 | try!(socket.write(response.key)); 23 | try!(socket.write(b" ")); 24 | try!(socket.write(format!("{}", response.flags).as_bytes())); 25 | try!(socket.write(b" ")); 26 | try!(socket.write(format!("{}", response.data.len()).as_bytes())); 27 | try!(socket.write(b"\r\n")); 28 | try!(socket.write(&response.data)); 29 | try!(socket.write(b"\r\n")); 30 | } 31 | try!(socket.write(b"END\r\n")); 32 | } 33 | Response::Gets { responses } => { 34 | for response in &responses { 35 | try!(socket.write(b"VALUE ")); 36 | try!(socket.write(response.key)); 37 | try!(socket.write(b" ")); 38 | try!(socket.write(format!("{}", response.flags).as_bytes())); 39 | try!(socket.write(b" ")); 40 | try!(socket.write(format!("{}", response.data.len()).as_bytes())); 41 | try!(socket.write(b" ")); 42 | try!(socket.write(format!("{}", response.unique).as_bytes())); 43 | try!(socket.write(b" ")); 44 | try!(socket.write(b"\r\n")); 45 | try!(socket.write(&response.data)); 46 | try!(socket.write(b"\r\n")); 47 | } 48 | try!(socket.write(b"END\r\n")); 49 | } 50 | Response::Incr { value } => { 51 | try!(socket.write(format!("{}", value).as_bytes())); 52 | try!(socket.write(b"\r\n")); 53 | } 54 | Response::Deleted => { 55 | try!(socket.write(b"DELETED\r\n")); 56 | } 57 | Response::Touched => { 58 | try!(socket.write(b"TOUCHED\r\n")); 59 | } 60 | Response::Ok => { 61 | try!(socket.write(b"OK\r\n")); 62 | } 63 | Response::Stored => { 64 | try!(socket.write(b"STORED\r\n")); 65 | } 66 | Response::NotStored => { 67 | try!(socket.write(b"NOT_STORED\r\n")); 68 | } 69 | Response::Exists => { 70 | try!(socket.write(b"EXISTS\r\n")); 71 | } 72 | Response::NotFound => { 73 | try!(socket.write(b"NOT_FOUND\r\n")); 74 | } 75 | Response::Error => { 76 | try!(socket.write(b"ERROR\r\n")); 77 | } 78 | Response::ClientError { message } => { 79 | try!(socket.write(b"CLIENT_ERROR ")); 80 | try!(socket.write(message)); 81 | try!(socket.write(b"\r\n")); 82 | }, 83 | Response::ServerError{message} => { 84 | try!(socket.write(b"SERVER_ERROR ")); 85 | try!(socket.write(message)); 86 | try!(socket.write(b"\r\n")); 87 | } 88 | Response::TooBig => { 89 | try!(socket.write(b"SERVER_ERROR object too large for cache")); 90 | } 91 | Response::Version => { 92 | try!(socket.write(b"VERSION ")); 93 | try!(socket.write(NAME)); 94 | try!(socket.write(b" ")); 95 | try!(socket.write(VERSION)); 96 | try!(socket.write(b"\r\n")); 97 | } 98 | } 99 | 100 | try!(socket.flush()); 101 | 102 | Ok(()) 103 | } 104 | 105 | fn client(locked_store: Arc>, 106 | mut socket: TcpStream, 107 | verbose: bool) { 108 | if verbose { 109 | println!("client connect"); 110 | } 111 | 112 | // this buffer on our stack is the largest amount that we can read from the 113 | // wire in a single go. bigger means fewer copies but more memory used per 114 | // client connection 115 | let mut buff: [u8; 10240] = [0; 10240]; 116 | 117 | // the accumulated data that's been read but not parsed yet. TODO we can 118 | // avoid a lot of copies here by trying to use buff directly when possible 119 | // and only spilling onto the heap when necessary. TODO this can be become 120 | // infinite in size. We need provisions for booting clients that grow it too 121 | // big, and for shrinking it occasionally so every client doesn't have 122 | // megabytes of buffer just because they used that much once in the past 123 | let mut parse_state: Vec = Vec::with_capacity(buff.len()); 124 | 125 | loop { 126 | match socket.read(&mut buff) { 127 | Err(err) => { 128 | if verbose { 129 | println!("client err: {:?}", err) 130 | } 131 | return; 132 | } 133 | Ok(size) if size == 0 => { 134 | if verbose { 135 | println!("client disconnect"); 136 | } 137 | return; // eof 138 | } 139 | Ok(size) => { 140 | parse_state.extend_from_slice(&buff[0..size]); 141 | 142 | // TODO this is all sorts of slow. we hold the lock until the 143 | // client is done receiving all of our bits! 144 | 145 | match parser::parse_command(&parse_state.to_vec()) { // TODO copy 146 | parser::IResult::Done(remaining, command_config) => { 147 | let CommandConfig { should_reply, command } = 148 | command_config; 149 | 150 | let response = match command { 151 | ServerCommand::Quit => { 152 | // no response, just disconnect them and quit 153 | return; 154 | } 155 | ServerCommand::Bad(text) => { 156 | if verbose { 157 | println!("bad client command: {:?}", 158 | String::from_utf8_lossy(text)) 159 | } 160 | Response::Error 161 | } 162 | _ => { 163 | // all others must be sent to the store 164 | let mut unlocked_store = locked_store.lock() 165 | .unwrap(); 166 | unlocked_store.apply(command) 167 | } 168 | }; 169 | if should_reply { 170 | match format_response(response, &mut socket) { 171 | Result::Ok(_) => (), 172 | Result::Err(err) => { 173 | if verbose { 174 | println!("client write error {:?}", err); 175 | } 176 | // TODO right now we just disconnect them 177 | return; 178 | } 179 | } 180 | } 181 | // TODO this does all sorts of copying 182 | parse_state.clear(); 183 | parse_state.extend_from_slice(remaining); 184 | } 185 | parser::IResult::Error(err) => { 186 | if verbose { 187 | println!("parser error? {:?}", err); 188 | } 189 | // TODO can we recover from this? 190 | return; 191 | } 192 | parser::IResult::Incomplete(_needed) => { 193 | continue; 194 | } 195 | } 196 | } 197 | } 198 | } 199 | } 200 | 201 | fn start_client(locked_store: Arc>, 202 | socket: TcpStream, 203 | verbose: bool) { 204 | spawn(move || client(locked_store, socket, verbose)); 205 | } 206 | 207 | pub fn start(port: u16, capacity: usize, verbose: bool) { 208 | let store = Store::new(capacity); 209 | let locked_store = Arc::new(Mutex::new(store)); 210 | let uri = format!("0.0.0.0:{}", port); 211 | let uri: &str = &uri; 212 | 213 | if verbose { 214 | println!("starting server"); 215 | } 216 | 217 | for client_stream in TcpListener::bind(&uri).unwrap().incoming() { 218 | match client_stream { 219 | Ok(client_stream) => { 220 | start_client(locked_store.clone(), client_stream, verbose); 221 | } 222 | Err(err) => { 223 | println!("client accept error: {:?}", err); 224 | } 225 | } 226 | } 227 | } 228 | -------------------------------------------------------------------------------- /src/store.rs: -------------------------------------------------------------------------------- 1 | /// The storage server. This implements all of the memcached semantics, 2 | /// including mapping parsed commands into actual fetches and mutations on an 3 | /// LRU that it contains 4 | #[allow(unused_imports)] 5 | use time; 6 | 7 | use std::str; 8 | use std::mem; 9 | 10 | use lru; 11 | 12 | // Keys as we get them from the client 13 | pub type Key<'a> = &'a [u8]; 14 | // Keys as we store them 15 | pub type StoredKey = Vec; 16 | // Keys as we return them to the client (for now we're only returning them keys 17 | // that they gave us, so we're using the same type to return it to them so we 18 | // can just use that same pointer instead of copying it) 19 | pub type ReturnedKey<'a> = &'a [u8]; 20 | 21 | // Data as we get it from the client 22 | pub type Data<'a> = &'a [u8]; 23 | // Data as we store it (copied from the client connection's memory) 24 | pub type StoredData = Vec; 25 | // Data as we return it to a client (copied for now) 26 | pub type ReturnedData = Vec; 27 | 28 | pub type Ttl = u32; 29 | pub type Flags = u32; 30 | pub type CasUnique = u64; 31 | pub type IncrValue = u64; 32 | 33 | pub type Capacity = usize; 34 | 35 | #[derive(Debug,PartialEq)] 36 | struct DataContainer { 37 | data: StoredData, 38 | flags: Flags, 39 | 40 | // TODO we are updating these by hand in the individual Setter handlers that 41 | // can change it, but we'll want to change that 42 | unique: CasUnique, 43 | } 44 | 45 | #[derive(Debug,PartialEq,Eq)] 46 | pub enum SetterType { 47 | Set, 48 | Add, 49 | Replace, 50 | Append, 51 | Prepend, 52 | Cas(CasUnique), 53 | } 54 | 55 | #[derive(Debug,PartialEq,Eq)] 56 | pub enum GetterType { 57 | Get, 58 | Gets, 59 | } 60 | 61 | #[derive(Debug,PartialEq,Eq)] 62 | pub enum IncrementerType { 63 | Incr, 64 | Decr, 65 | } 66 | 67 | #[derive(Debug,PartialEq,Eq)] 68 | pub enum ServerCommand<'a> { 69 | Setter { 70 | setter: SetterType, 71 | key: Key<'a>, 72 | data: Data<'a>, 73 | ttl: Ttl, 74 | flags: Flags, 75 | }, 76 | Getter { 77 | getter: GetterType, 78 | keys: Vec>, 79 | }, 80 | Delete { 81 | key: Key<'a>, 82 | }, 83 | Touch { 84 | key: Key<'a>, 85 | ttl: Ttl, 86 | }, 87 | Incrementer { 88 | incrementer: IncrementerType, 89 | key: Key<'a>, 90 | value: IncrValue, 91 | }, 92 | FlushAll, 93 | Bad(&'a [u8]), 94 | Quit, 95 | Version, 96 | Verbosity, 97 | } 98 | 99 | #[derive(Debug,PartialEq,Eq)] 100 | pub struct SingleGetResponse<'a> { 101 | pub key: ReturnedKey<'a>, 102 | pub data: ReturnedData, 103 | pub flags: Flags, 104 | pub unique: CasUnique, 105 | } 106 | 107 | #[derive(Debug,PartialEq,Eq)] 108 | pub enum Response<'a> { 109 | // Data and Gets share the SingleGetResponse format for simplicity. If 110 | // copying a bunch of unneeded unique values turns out to be a problem we 111 | // can revisit but this makes the response builder much simpler 112 | Data { 113 | responses: Vec>, 114 | }, 115 | Gets { 116 | responses: Vec>, 117 | }, 118 | Incr { 119 | value: IncrValue, 120 | }, 121 | Deleted, 122 | Touched, 123 | Ok, 124 | Stored, 125 | NotStored, 126 | Exists, 127 | NotFound, 128 | Error, 129 | ClientError { 130 | message: &'a [u8], 131 | }, 132 | ServerError { 133 | message: &'a [u8], 134 | }, 135 | Version, 136 | TooBig, 137 | } 138 | 139 | fn forgetful_parse_int(current_data: &StoredData) -> Option { 140 | // try to interpret it as an int 141 | let as_string = str::from_utf8(current_data); 142 | if as_string.is_err() { 143 | return None; 144 | } 145 | let as_string = as_string.unwrap(); 146 | let as_int = as_string.parse::(); 147 | if as_int.is_err() { 148 | return None; 149 | } 150 | 151 | Option::Some(as_int.unwrap()) 152 | } 153 | 154 | enum _IncrSubResult { 155 | NotFound, 156 | BadInt, 157 | NewValue(IncrValue, Option, Flags), 158 | } 159 | 160 | // the number of seconds in a TTL after which we start recognising it as a 161 | // timestamp instead. This is a magic number used by memcached so we're cloning 162 | // its behaviour here. used by wrap_ttl 163 | const MAGIC_DATE: Ttl = 60 * 60 * 24 * 30; 164 | 165 | // right now these are enforced here but it would be nice if the parser could do 166 | // some of it too 167 | const MAX_KEY: usize = 255; 168 | const MAX_DATA: usize = 1024 * 1024; // 1MB 169 | 170 | pub fn wrap_ttl(ttl: Ttl, now: Ttl) -> Option { 171 | if ttl == 0 { 172 | None 173 | } else if ttl < MAGIC_DATE { 174 | Option::Some(now + ttl) 175 | } else { 176 | Option::Some(ttl) 177 | } 178 | } 179 | 180 | #[cfg(not(test))] 181 | fn epoch_time() -> Ttl { 182 | time::get_time().sec as Ttl 183 | } 184 | #[cfg(test)] 185 | pub fn epoch_time() -> Ttl { 186 | // so the tests can use a known value 187 | 1455082881 188 | } 189 | 190 | #[derive(Debug)] 191 | pub struct Store { 192 | store: lru::LruCache, 193 | last_cas_id: CasUnique, 194 | } 195 | 196 | impl Store { 197 | pub fn new(capacity: Capacity) -> Store { 198 | Store { 199 | store: lru::LruCache::new(capacity), 200 | last_cas_id: 0, 201 | } 202 | } 203 | 204 | fn make_cas_id(&mut self) -> CasUnique { 205 | self.last_cas_id += 1; 206 | self.last_cas_id 207 | } 208 | 209 | pub fn apply<'a>(&mut self, command: ServerCommand<'a>) -> Response<'a> { 210 | let now = epoch_time(); // TODO lazy? 211 | 212 | match command { 213 | ServerCommand::Setter { key: ckey, data: cdata, .. } 214 | if ckey.len() > MAX_KEY || 215 | cdata.len() > MAX_DATA => Response::TooBig, 216 | ServerCommand::Setter { setter, 217 | key: ckey, 218 | data: cdata, 219 | ttl: cttl, 220 | flags } => { 221 | let new_cas = self.make_cas_id(); // TODO too many IDs 222 | let ttl = wrap_ttl(cttl, now); 223 | let skey = ckey.to_vec(); 224 | 225 | let container = |data: &[u8], flags| { 226 | DataContainer { 227 | data: data.to_vec(), // does a copy 228 | flags: flags, 229 | unique: new_cas, 230 | } 231 | }; 232 | 233 | match setter { 234 | SetterType::Add if self.store.contains(&skey, now) => { 235 | Response::NotStored 236 | } 237 | SetterType::Add | SetterType::Set => { 238 | self.store.set(skey, container(cdata, flags), ttl, now); 239 | Response::Stored 240 | } 241 | SetterType::Replace if self.store.contains(&skey, now) => { 242 | self.store.set(skey, container(cdata, flags), ttl, now); 243 | Response::Stored 244 | } 245 | SetterType::Replace => Response::NotStored, 246 | SetterType::Append if self.store.contains(&skey, now) => { 247 | // this is pretty slow because we use immutable data in 248 | // the lru. it's possible to make this faster by using 249 | // mutable data structures instead that we can just 250 | // directly modify, but then we'd need to make sure to 251 | // keep the weights and stuff in sync and that's a pain 252 | let (new_vec, old_ttl, old_flags) = { 253 | let current_entry = 254 | self.store.get_full_entry(&skey, now).unwrap(); 255 | let current_container = ¤t_entry.data; 256 | let new_size = cdata.len() + 257 | current_container.data.len(); 258 | let mut new_vec = Vec::with_capacity(new_size); 259 | new_vec.extend_from_slice(¤t_container.data); 260 | new_vec.extend_from_slice(cdata); 261 | (new_vec, 262 | current_entry.expires, 263 | current_container.flags) 264 | }; 265 | self.store.set(skey, 266 | container(&new_vec, old_flags), 267 | old_ttl, 268 | now); 269 | Response::Stored 270 | } 271 | SetterType::Append => Response::NotStored, 272 | SetterType::Prepend if self.store.contains(&skey, now) => { 273 | let (new_vec, old_ttl, old_flags) = { 274 | let current_entry = 275 | self.store.get_full_entry(&skey, now).unwrap(); 276 | let current_container = ¤t_entry.data; 277 | let new_size = cdata.len() + 278 | current_container.data.len(); 279 | let mut new_vec = Vec::with_capacity(new_size); 280 | new_vec.extend_from_slice(cdata); 281 | new_vec.extend_from_slice(¤t_container.data); 282 | (new_vec, 283 | current_entry.expires, 284 | current_container.flags) 285 | }; 286 | self.store.set(skey, 287 | container(&new_vec, old_flags), 288 | old_ttl, 289 | now); 290 | Response::Stored 291 | } 292 | SetterType::Prepend => Response::NotStored, 293 | SetterType::Cas(_) if !self.store.contains(&skey, now) => { 294 | Response::NotFound 295 | } 296 | SetterType::Cas(unique) if (self.store 297 | .fast_get(&skey, now) 298 | .map(|cont| cont.unique) == 299 | Some(unique)) => { 300 | self.store.set(skey, container(cdata, flags), ttl, now); 301 | Response::Stored 302 | } 303 | SetterType::Cas(_) => { 304 | // n.b. failed cas updates don't update the lru 305 | Response::Exists 306 | } 307 | } 308 | } 309 | 310 | ServerCommand::Getter { getter, keys } => { 311 | let mut found = Vec::with_capacity(keys.len()); 312 | for ckey in keys { 313 | let skey = ckey.to_vec(); 314 | if let Some(item) = self.store.get(&skey, now) { 315 | found.push(SingleGetResponse { 316 | key: ckey, 317 | data: item.data.clone(), // does a copy 318 | flags: item.flags, 319 | unique: item.unique, 320 | }); 321 | } 322 | } 323 | // and turn that into the right result format for the request 324 | // (does this really have to be this repetetive?) 325 | match getter { 326 | GetterType::Get => Response::Data { responses: found }, 327 | GetterType::Gets => Response::Gets { responses: found }, 328 | } 329 | } 330 | ServerCommand::Delete { key: ckey } => { 331 | let skey = ckey.to_vec(); 332 | 333 | if self.store.delete(&skey) { 334 | Response::Deleted 335 | } else { 336 | Response::NotFound 337 | } 338 | } 339 | ServerCommand::Touch { key: ckey, ttl: cttl } => { 340 | let skey = ckey.to_vec(); 341 | let ttl = wrap_ttl(cttl, now); 342 | 343 | if self.store.contains(&skey, now) { 344 | self.store.touch(&skey, ttl, now); 345 | Response::Touched 346 | } else { 347 | Response::NotFound 348 | } 349 | } 350 | ServerCommand::Incrementer { incrementer, key: ckey, value } => { 351 | let new_cas = self.make_cas_id(); 352 | let skey = ckey.to_vec(); 353 | 354 | let isr = match self.store.get_full_entry(&skey, now) { 355 | None => _IncrSubResult::NotFound, 356 | Some(full_entry) => { 357 | let item = &(*full_entry).data; 358 | let current_data = &(*item).data; 359 | let as_int = forgetful_parse_int(current_data); 360 | match as_int { 361 | None => _IncrSubResult::BadInt, 362 | Some(current_int) => { 363 | let new_int = match incrementer { 364 | // memcached is saturating in the negative 365 | // direction 366 | IncrementerType::Decr => { 367 | current_int.saturating_sub(value) 368 | } 369 | // ...but wrapping in the positive direction 370 | IncrementerType::Incr => { 371 | current_int.wrapping_add(value) 372 | } 373 | }; 374 | _IncrSubResult::NewValue(new_int, 375 | full_entry.expires, 376 | item.flags) 377 | } 378 | } 379 | } 380 | }; 381 | match isr { 382 | _IncrSubResult::NotFound => Response::NotFound, 383 | _IncrSubResult::BadInt => { 384 | Response::ClientError { 385 | message: b"cannot increment or decrement non-numeric value", 386 | } 387 | } 388 | _IncrSubResult::NewValue(new_int, sttl, flags) => { 389 | let re_str = new_int.to_string(); 390 | let re_bytes = re_str.as_bytes(); 391 | let new_data = re_bytes.to_vec(); 392 | let new_container = DataContainer { 393 | data: new_data.to_vec(), 394 | flags: flags, 395 | unique: new_cas, 396 | }; 397 | self.store.set(skey, new_container, sttl, now); 398 | Response::Incr { value: new_int } 399 | } 400 | } 401 | } 402 | ServerCommand::FlushAll => { 403 | self.store.clear(); // weeeeee 404 | Response::Ok 405 | } 406 | ServerCommand::Bad(_) => Response::Error, 407 | ServerCommand::Version => Response::Version, 408 | // we ignore this, we just support it to make memcapable happy 409 | ServerCommand::Verbosity => Response::Ok, 410 | ServerCommand::Quit => { 411 | unreachable!("this should have been handled by the server dispatch loop") 412 | } 413 | 414 | } 415 | } 416 | 417 | #[cfg(test)] 418 | pub fn simple_get(&mut self, key: &str) -> Option { 419 | let as_bytes = key.as_bytes(); 420 | let as_vec = as_bytes.to_vec(); 421 | match self.store.fast_get(&as_vec, epoch_time()) { 422 | None => None, 423 | Some(container) => { 424 | let ref container_data = container.data; 425 | let container_as_string = 426 | String::from_utf8_lossy(&container_data); 427 | let mut new_string = String::new(); 428 | new_string.push_str(&container_as_string); 429 | Some(new_string) 430 | } 431 | } 432 | } 433 | 434 | #[cfg(test)] 435 | pub fn simple_get_flags(&mut self, key: &str) -> Option { 436 | let as_bytes = key.as_bytes(); 437 | let as_vec = as_bytes.to_vec(); 438 | self.store.get(&as_vec, epoch_time()).map(|c| c.flags) 439 | } 440 | 441 | #[cfg(test)] 442 | pub fn simple_get_ttl(&mut self, key: &str) -> Option { 443 | let as_bytes = key.as_bytes(); 444 | let as_vec = as_bytes.to_vec(); 445 | match self.store.get_full_entry(&as_vec, epoch_time()) { 446 | None => None, 447 | Some(entry) => (*entry).expires, 448 | } 449 | } 450 | 451 | #[cfg(test)] 452 | pub fn simple_set(&mut self, key: &str, data: &str) { 453 | self.simple_set_cas(key, data, 0); 454 | } 455 | 456 | #[cfg(test)] 457 | pub fn simple_set_cas(&mut self, 458 | key: &str, 459 | data: &str, 460 | unique: CasUnique) { 461 | let mut key_vec: Vec = Vec::new(); 462 | key_vec.extend_from_slice(key.as_bytes()); 463 | let mut data_vec: Vec = Vec::new(); 464 | data_vec.extend_from_slice(data.as_bytes()); 465 | 466 | self.store.set(key_vec, 467 | DataContainer { 468 | data: data_vec, 469 | flags: 0, 470 | unique: unique, 471 | }, 472 | Option::None, 473 | epoch_time()); 474 | } 475 | } 476 | 477 | impl lru::HasWeight for DataContainer { 478 | fn weight(&self) -> lru::Weight { 479 | (self.data.capacity() + mem::size_of::() + 480 | mem::size_of::()) 481 | } 482 | } 483 | 484 | #[cfg(test)] 485 | mod tests { 486 | use super::*; 487 | 488 | #[test] 489 | pub fn set() { 490 | let mut store = Store::new(200); 491 | let res = store.apply(ServerCommand::Setter { 492 | setter: SetterType::Set, 493 | key: b"foo", 494 | data: b"bar", 495 | flags: 0, 496 | ttl: 0, 497 | }); 498 | assert_eq!(Response::Stored, res); 499 | assert_eq!(Some("bar".to_string()), store.simple_get("foo")); 500 | } 501 | 502 | #[test] 503 | pub fn add_not_present() { 504 | let mut store = Store::new(200); 505 | let res = store.apply(ServerCommand::Setter { 506 | setter: SetterType::Add, 507 | key: b"foo", 508 | data: b"bar", 509 | flags: 0, 510 | ttl: 0, 511 | }); 512 | assert_eq!(Response::Stored, res); 513 | assert_eq!(Some("bar".to_string()), store.simple_get("foo")); 514 | } 515 | 516 | #[test] 517 | pub fn add_present() { 518 | let mut store = Store::new(200); 519 | store.simple_set("foo", "bar"); 520 | let res = store.apply(ServerCommand::Setter { 521 | setter: SetterType::Add, 522 | key: b"foo", 523 | data: b"baz", 524 | flags: 0, 525 | ttl: 0, 526 | }); 527 | assert_eq!(Response::NotStored, res); 528 | assert_eq!(Some("bar".to_string()), store.simple_get("foo")); 529 | } 530 | 531 | #[test] 532 | pub fn replace_not_present() { 533 | let mut store = Store::new(200); 534 | let res = store.apply(ServerCommand::Setter { 535 | setter: SetterType::Replace, 536 | key: b"foo", 537 | data: b"baz", 538 | flags: 0, 539 | ttl: 0, 540 | }); 541 | assert_eq!(Response::NotStored, res); 542 | assert_eq!(None, store.simple_get("foo")); 543 | } 544 | 545 | #[test] 546 | pub fn replace_present() { 547 | let mut store = Store::new(200); 548 | store.simple_set("foo", "bar"); 549 | let res = store.apply(ServerCommand::Setter { 550 | setter: SetterType::Replace, 551 | key: b"foo", 552 | data: b"baz", 553 | flags: 0, 554 | ttl: 0, 555 | }); 556 | assert_eq!(Response::Stored, res); 557 | assert_eq!(Some("baz".to_string()), store.simple_get("foo")); 558 | } 559 | 560 | #[test] 561 | pub fn append_not_present() { 562 | let mut store = Store::new(200); 563 | let res = store.apply(ServerCommand::Setter { 564 | setter: SetterType::Append, 565 | key: b"foo", 566 | data: b"baz", 567 | flags: 0, 568 | ttl: 0, 569 | }); 570 | assert_eq!(Response::NotStored, res); 571 | assert_eq!(None, store.simple_get("foo")); 572 | } 573 | 574 | #[test] 575 | pub fn append_present() { 576 | let mut store = Store::new(200); 577 | store.simple_set("foo", "bar"); 578 | let res = store.apply(ServerCommand::Setter { 579 | setter: SetterType::Append, 580 | key: b"foo", 581 | data: b"baz", 582 | flags: 12, 583 | ttl: 34, 584 | }); 585 | assert_eq!(Response::Stored, res); 586 | assert_eq!(Some("barbaz".to_string()), store.simple_get("foo")); 587 | // make sure we didn't update the flags or ttl 588 | assert_eq!(Some(0), store.simple_get_flags("foo")); 589 | assert_eq!(None, store.simple_get_ttl("bar")); 590 | } 591 | 592 | #[test] 593 | pub fn prepend_not_present() { 594 | let mut store = Store::new(200); 595 | let res = store.apply(ServerCommand::Setter { 596 | setter: SetterType::Prepend, 597 | key: b"foo", 598 | data: b"baz", 599 | flags: 0, 600 | ttl: 0, 601 | }); 602 | assert_eq!(Response::NotStored, res); 603 | assert_eq!(None, store.simple_get("foo")); 604 | } 605 | 606 | #[test] 607 | pub fn prepend_present() { 608 | let mut store = Store::new(200); 609 | store.simple_set("foo", "bar"); 610 | let res = store.apply(ServerCommand::Setter { 611 | setter: SetterType::Prepend, 612 | key: b"foo", 613 | data: b"baz", 614 | flags: 0, 615 | ttl: 0, 616 | }); 617 | assert_eq!(Response::Stored, res); 618 | assert_eq!(Some("bazbar".to_string()), store.simple_get("foo")); 619 | // make sure we didn't update the flags or ttl 620 | assert_eq!(Some(0), store.simple_get_flags("foo")); 621 | assert_eq!(None, store.simple_get_ttl("bar")); 622 | } 623 | 624 | #[test] 625 | pub fn cas_not_present() { 626 | let mut store = Store::new(200); 627 | let res = store.apply(ServerCommand::Setter { 628 | setter: SetterType::Cas(50), 629 | key: b"foo", 630 | data: b"baz", 631 | flags: 0, 632 | ttl: 0, 633 | }); 634 | assert_eq!(Response::NotFound, res); 635 | assert_eq!(None, store.simple_get("foo")); 636 | } 637 | 638 | #[test] 639 | pub fn cas_wrong() { 640 | let mut store = Store::new(200); 641 | store.simple_set_cas("foo", "bar", 100); 642 | let res = store.apply(ServerCommand::Setter { 643 | setter: SetterType::Cas(200), 644 | key: b"foo", 645 | data: b"baz", 646 | flags: 0, 647 | ttl: 0, 648 | }); 649 | assert_eq!(Response::Exists, res); 650 | assert_eq!(Some("bar".to_string()), store.simple_get("foo")); 651 | } 652 | 653 | #[test] 654 | pub fn cas_right() { 655 | let mut store = Store::new(200); 656 | store.simple_set_cas("foo", "bar", 100); 657 | let res = store.apply(ServerCommand::Setter { 658 | setter: SetterType::Cas(100), 659 | key: b"foo", 660 | data: b"baz", 661 | flags: 0, 662 | ttl: 0, 663 | }); 664 | assert_eq!(Response::Stored, res); 665 | assert_eq!(Some("baz".to_string()), store.simple_get("foo")); 666 | } 667 | 668 | #[test] 669 | pub fn cas_refreshes() { 670 | let mut store = Store::new(200); 671 | store.simple_set_cas("foo", "bar", 100); 672 | store.simple_set("foo", "quux"); 673 | let res = store.apply(ServerCommand::Setter { 674 | setter: SetterType::Cas(100), 675 | key: b"foo", 676 | data: b"baz", 677 | flags: 0, 678 | ttl: 0, 679 | }); 680 | assert_eq!(Response::Exists, res); 681 | assert_eq!(Some("quux".to_string()), store.simple_get("foo")); 682 | } 683 | 684 | #[test] 685 | pub fn get() { 686 | let mut store = Store::new(200); 687 | store.simple_set("foo", "bar"); 688 | let res = store.apply(ServerCommand::Getter { 689 | getter: GetterType::Get, 690 | keys: vec!["foo".as_bytes()], 691 | }); 692 | assert_eq!(res, 693 | Response::Data { 694 | responses: vec![SingleGetResponse { 695 | key: "foo".as_bytes(), 696 | data: b("bar"), 697 | flags: 0, 698 | unique: 0, 699 | }], 700 | }); 701 | } 702 | 703 | #[test] 704 | pub fn get_multi() { 705 | let mut store = Store::new(200); 706 | store.simple_set("foo1", "bar1"); 707 | store.simple_set("foo2", "bar2"); 708 | let res = store.apply(ServerCommand::Getter { 709 | getter: GetterType::Get, 710 | keys: vec!["foo1".as_bytes(), "foo2".as_bytes(), "foo3".as_bytes()], 711 | }); 712 | assert_eq!(res, 713 | Response::Data { 714 | responses: vec![SingleGetResponse { 715 | key: "foo1".as_bytes(), 716 | data: b("bar1"), 717 | flags: 0, 718 | unique: 0, 719 | }, 720 | SingleGetResponse { 721 | key: "foo2".as_bytes(), 722 | data: b("bar2"), 723 | flags: 0, 724 | unique: 0, 725 | }], 726 | }); 727 | } 728 | 729 | #[test] 730 | pub fn gets() { 731 | let mut store = Store::new(200); 732 | store.simple_set("foo", "bar"); 733 | let res = store.apply(ServerCommand::Getter { 734 | getter: GetterType::Get, 735 | keys: vec!["foo".as_bytes()], 736 | }); 737 | assert_eq!(res, 738 | Response::Data { 739 | responses: vec![SingleGetResponse { 740 | key: "foo".as_bytes(), 741 | data: b("bar"), 742 | flags: 0, 743 | unique: 0, 744 | }], 745 | }); 746 | } 747 | 748 | #[test] 749 | pub fn gets_multi() { 750 | let mut store = Store::new(200); 751 | store.simple_set_cas("foo1", "bar1", 100); 752 | store.simple_set_cas("foo2", "bar2", 100); 753 | let res = store.apply(ServerCommand::Getter { 754 | getter: GetterType::Gets, 755 | keys: vec!["foo1".as_bytes(), "foo2".as_bytes(), "foo3".as_bytes()], 756 | }); 757 | assert_eq!(res, 758 | Response::Gets { 759 | responses: vec![SingleGetResponse { 760 | key: "foo1".as_bytes(), 761 | data: b("bar1"), 762 | flags: 0, 763 | unique: 100, 764 | }, 765 | SingleGetResponse { 766 | key: "foo2".as_bytes(), 767 | data: b("bar2"), 768 | flags: 0, 769 | unique: 100, 770 | }], 771 | }); 772 | } 773 | 774 | #[test] 775 | pub fn incr_present_and_good() { 776 | let mut store = Store::new(200); 777 | store.simple_set("foo", "1"); 778 | let res = store.apply(ServerCommand::Incrementer { 779 | incrementer: IncrementerType::Incr, 780 | key: b"foo", 781 | value: 5, 782 | }); 783 | assert_eq!(res, Response::Incr { value: 6 }); 784 | } 785 | 786 | #[test] 787 | pub fn incr_present_and_bad() { 788 | let mut store = Store::new(200); 789 | store.simple_set("foo", "bar"); 790 | let res = store.apply(ServerCommand::Incrementer { 791 | incrementer: IncrementerType::Incr, 792 | key: b"foo", 793 | value: 5, 794 | }); 795 | assert_eq!(res, 796 | Response::ClientError { 797 | message: b"cannot increment or decrement non-numeric value", 798 | }); 799 | } 800 | 801 | #[test] 802 | pub fn incr_not_present() { 803 | let mut store = Store::new(200); 804 | let res = store.apply(ServerCommand::Incrementer { 805 | incrementer: IncrementerType::Incr, 806 | key: b"foo", 807 | value: 5, 808 | }); 809 | assert_eq!(res, Response::NotFound); 810 | } 811 | 812 | #[test] 813 | pub fn incr_refreshes_cas() { 814 | let mut store = Store::new(200); 815 | store.simple_set_cas("foo", "20", 100); 816 | let res = store.apply(ServerCommand::Incrementer { 817 | incrementer: IncrementerType::Incr, 818 | key: b"foo", 819 | value: 5, 820 | }); 821 | assert_eq!(Response::Incr { value: 25 }, res); 822 | let res = store.apply(ServerCommand::Setter { 823 | setter: SetterType::Cas(100), 824 | key: b"foo", 825 | data: b"30", 826 | flags: 0, 827 | ttl: 0, 828 | }); 829 | assert_eq!(Response::Exists, res); 830 | assert_eq!(Some("25".to_string()), store.simple_get("foo")); 831 | } 832 | 833 | 834 | #[test] 835 | pub fn decr() { 836 | let mut store = Store::new(200); 837 | store.simple_set("foo", "20"); 838 | let res = store.apply(ServerCommand::Incrementer { 839 | incrementer: IncrementerType::Decr, 840 | key: b"foo", 841 | value: 5, 842 | }); 843 | assert_eq!(res, Response::Incr { value: 15 }); 844 | } 845 | 846 | #[test] 847 | pub fn decr_saturates() { 848 | let mut store = Store::new(200); 849 | store.simple_set("foo", "20"); 850 | let res = store.apply(ServerCommand::Incrementer { 851 | incrementer: IncrementerType::Decr, 852 | key: b"foo", 853 | value: 100, 854 | }); 855 | assert_eq!(res, Response::Incr { value: 0 }); 856 | } 857 | 858 | #[test] 859 | pub fn incr_wraps() { 860 | let mut store = Store::new(200); 861 | store.simple_set("foo", "18446744073709551615"); 862 | let res = store.apply(ServerCommand::Incrementer { 863 | incrementer: IncrementerType::Incr, 864 | key: b"foo", 865 | value: 2, 866 | }); 867 | assert_eq!(res, Response::Incr { value: 1 }); 868 | } 869 | 870 | #[test] 871 | pub fn delete_present() { 872 | let mut store = Store::new(200); 873 | store.simple_set("foo", "bar"); 874 | let res = store.apply(ServerCommand::Delete { key: b"foo" }); 875 | assert_eq!(res, Response::Deleted); 876 | } 877 | 878 | #[test] 879 | pub fn delete_not_present() { 880 | let mut store = Store::new(200); 881 | let res = store.apply(ServerCommand::Delete { key: b"foo" }); 882 | assert_eq!(res, Response::NotFound); 883 | } 884 | 885 | #[test] 886 | pub fn touch_not_present() { 887 | let mut store = Store::new(200); 888 | 889 | let res = store.apply(ServerCommand::Touch { 890 | key: b"foo", 891 | ttl: 0, 892 | }); 893 | assert_eq!(Response::NotFound, res); 894 | assert_eq!(None, store.simple_get("foo")); 895 | } 896 | 897 | #[test] 898 | pub fn touch() { 899 | let mut store = Store::new(1000); 900 | 901 | // will get the version marked cfg(test)! 902 | let now: Ttl = epoch_time(); 903 | 904 | store.simple_set("foo", "bar"); 905 | assert_eq!(store.simple_get_ttl("foo"), None); 906 | 907 | let res = store.apply(ServerCommand::Touch { 908 | key: b"foo", 909 | ttl: 0, 910 | }); 911 | assert_eq!(Response::Touched, res); 912 | assert_eq!(store.simple_get_ttl("foo"), None); 913 | assert_eq!(Some("bar".to_string()), store.simple_get("foo")); 914 | 915 | let res = store.apply(ServerCommand::Touch { 916 | key: b"foo", 917 | ttl: 100, 918 | }); 919 | assert_eq!(Response::Touched, res); 920 | assert_eq!(Some("bar".to_string()), store.simple_get("foo")); 921 | assert_eq!(store.simple_get_ttl("foo"), Some(now + 100)); 922 | 923 | // make sure we cna set it back to 0 924 | let res = store.apply(ServerCommand::Touch { 925 | key: b"foo", 926 | ttl: 0, 927 | }); 928 | assert_eq!(Response::Touched, res); 929 | assert_eq!(store.simple_get_ttl("foo"), None); 930 | } 931 | 932 | #[test] 933 | pub fn wrapping_ttl() { 934 | // memcached accepts timestamps in seconds-in-the-future or in absolute 935 | // epoch seconds. It uses a heuristic magic number (MAGIC_DATE above) to 936 | // guess which one to use, so we have to make sure we properly support 937 | // this 938 | let mut store = Store::new(200); 939 | 940 | // will get the version marked cfg(test)! 941 | let now: Ttl = epoch_time(); 942 | 943 | assert_eq!(wrap_ttl(0, now), None); 944 | assert_eq!(wrap_ttl(1, now), Some(now + 1)); 945 | assert_eq!(wrap_ttl(2, now), Some(now + 2)); 946 | assert_eq!(wrap_ttl(now, now), Some(now)); 947 | assert_eq!(wrap_ttl(now + 1, now), Some(now + 1)); 948 | 949 | store.simple_set("foo", "bar"); 950 | 951 | let res = store.apply(ServerCommand::Touch { 952 | key: b"foo", 953 | ttl: 100, 954 | }); 955 | assert_eq!(Response::Touched, res); 956 | assert_eq!(Some("bar".to_string()), store.simple_get("foo")); 957 | assert_eq!(store.simple_get_ttl("foo"), Some(now + 100)); 958 | 959 | let res = store.apply(ServerCommand::Touch { 960 | key: b"foo", 961 | ttl: now + 200, 962 | }); 963 | assert_eq!(Response::Touched, res); 964 | assert_eq!(Some("bar".to_string()), store.simple_get("foo")); 965 | assert_eq!(store.simple_get_ttl("foo"), Some(now + 200)); 966 | 967 | let res = store.apply(ServerCommand::Setter { 968 | setter: SetterType::Set, 969 | key: b"foo", 970 | data: b"bar", 971 | ttl: now + 300, 972 | flags: 0, 973 | }); 974 | assert_eq!(Response::Stored, res); 975 | assert_eq!(Some("bar".to_string()), store.simple_get("foo")); 976 | assert_eq!(store.simple_get_ttl("foo"), Some(now + 300)); 977 | } 978 | 979 | #[test] 980 | pub fn flushall() { 981 | let mut store = Store::new(200); 982 | store.simple_set("foo", "bar"); 983 | assert_eq!(Some("bar".to_string()), store.simple_get("foo")); 984 | 985 | let res = store.apply(ServerCommand::FlushAll); 986 | assert_eq!(res, Response::Ok); 987 | 988 | assert_eq!(None, store.simple_get("foo")); 989 | } 990 | 991 | fn b(inp: &'static str) -> Vec { 992 | // syntactic sugar for tests 993 | let mut s = String::new(); 994 | s.push_str(inp); 995 | s.into_bytes() 996 | } 997 | } 998 | --------------------------------------------------------------------------------