yo_kv/keys.rs
1//! Moving a key, copying one, and touching one.
2//!
3//! Three of the four commands here move whole values around, and all three of
4//! them are careful about the same thing: a value lives in two places at once.
5//! A string lives entirely in its record, and a set or a hash lives in a slab
6//! with the record holding nothing but a slot number. So there is no one way to
7//! move a value, and a command that forgets which case it is in either drops
8//! members on the floor or leaves a body in the slab that nothing points at.
9//!
10//! [`Keyspace::rename`] moves the record's bytes and leaves the body exactly
11//! where it is, because a slot number that moves to a different key is still
12//! the same slot. Renaming a set of a million members writes thirteen bytes.
13//!
14//! [`Keyspace::copy`] cannot do that, since two records pointing at one slot
15//! would be one set that answers to two names and `SADD` to either would show
16//! up in both. So the body is cloned, which is the one thing here that costs
17//! what the value is worth. That is Redis's cost too and there is no version of
18//! `COPY` that avoids it.
19//!
20//! # Why export and import are separate and public
21//!
22//! `COPY key dst DB n` puts a value in a database this one cannot reach. The
23//! wire layer holds every database and this one holds none of them, so the two
24//! halves are separate calls and the caller is what joins them up.
25//!
26//! It also makes the pair the answer for `MOVE`, `DUMP` and `RESTORE`, which
27//! want exactly this: a value lifted out of a database, standing on its own with
28//! its deadline attached.
29//!
30//! There are two ways to lift one out. [`Keyspace::export`] clones the body and
31//! leaves the key where it is, which is what `COPY` needs, and
32//! [`Keyspace::take`] pulls the body out of the slab and deletes the key, which
33//! is what `MOVE` needs. `MOVE` through `export` would clone a set of a million
34//! members and then throw the original away a line later, so the two are
35//! separate calls rather than one call with a flag.
36//!
37//! # And the same pair again, with bytes in the middle
38//!
39//! `DUMP` and `RESTORE` are the same shape one step further out. A record is a
40//! value standing on its own inside this process, and a payload is a value
41//! standing on its own outside it, so [`Keyspace::dump`] is an export followed
42//! by [`crate::rdb`] and [`Keyspace::restore`] is `rdb` followed by an import.
43//! The deadline is the one thing that does not make the trip, because `DUMP`
44//! drops it and `RESTORE` is given a fresh one.
45
46use yo_common::Result;
47
48use crate::array::Array;
49use crate::foreign::Foreign;
50use crate::hash::Hash;
51use crate::keyspace::Keyspace;
52use crate::list::List;
53use crate::rdb;
54use crate::set::Set;
55use crate::stream::Stream;
56use crate::value::{self, Kind};
57use crate::zset::Zset;
58
59/// Everything under one key, lifted out so it can be put somewhere else.
60///
61/// It owns what it holds. A record taken out of a database survives that
62/// database being written to, flushed or dropped, which is what makes it safe
63/// to carry between two of them.
64#[derive(Debug, Clone)]
65pub struct Record {
66 body: Body,
67 /// The deadline, which travels with the value. `COPY` and `RENAME` both
68 /// keep it, and a copy of a key with ten seconds left has ten seconds left.
69 expire_at: Option<u64>,
70}
71
72impl Record {
73 /// A record built from parts, for a caller that has both.
74 ///
75 /// [`crate::rdb`] is that caller and there is no other. A record normally
76 /// comes out of a database and this is the one way to make one that never
77 /// was in a database, which is what a payload arriving from a client is.
78 pub(crate) const fn new(body: Body, expire_at: Option<u64>) -> Record {
79 Record { body, expire_at }
80 }
81
82 /// What it holds, for the code that has to write it down.
83 pub(crate) const fn body(&self) -> &Body {
84 &self.body
85 }
86
87 /// What type this is, which the caller usually knows and sometimes does not.
88 #[must_use]
89 pub const fn kind(&self) -> Kind {
90 match self.body {
91 Body::String(_) => Kind::String,
92 Body::Set(_) => Kind::Set,
93 Body::Hash(_) => Kind::Hash,
94 Body::List(_) => Kind::List,
95 Body::Zset(_) => Kind::Zset,
96 Body::Array(_) => Kind::Array,
97 Body::Stream(_) => Kind::Stream,
98 Body::Foreign(_) => Kind::Foreign,
99 }
100 }
101
102 /// When it goes away, if anything says.
103 #[must_use]
104 pub const fn expire_at(&self) -> Option<u64> {
105 self.expire_at
106 }
107}
108
109/// The eight things a record can be, owned rather than borrowed.
110///
111/// One variant per type that a key can hold, and that is the point: the day an
112/// eighth type lands, the compiler names this file. It did not before, because
113/// the match in [`Keyspace::export`] had a catch all arm at the bottom, and a
114/// catch all in front of an enum the rest of the crate keeps growing is a hole
115/// that reports itself as a panic on a live server rather than as a build error.
116#[derive(Debug)]
117pub(crate) enum Body {
118 String(Vec<u8>),
119 Set(Set),
120 Hash(Hash),
121 List(List),
122 Zset(Zset),
123 Array(Array),
124 Stream(Stream),
125 Foreign(Box<dyn Foreign>),
126}
127
128/// Every body but the foreign one can be copied.
129///
130/// Written out rather than derived so that the one variant which cannot is a
131/// named arm here instead of a `Clone` bound the escape could never satisfy.
132/// Nothing reaches it: [`Keyspace::export`] is the only thing that clones a
133/// body and it answers `None` for a foreign one before it gets this far, so
134/// this is the assertion of that rather than a case to handle.
135impl Clone for Body {
136 fn clone(&self) -> Body {
137 match self {
138 Body::String(v) => Body::String(v.clone()),
139 Body::Set(v) => Body::Set(v.clone()),
140 Body::Hash(v) => Body::Hash(v.clone()),
141 Body::List(v) => Body::List(v.clone()),
142 Body::Zset(v) => Body::Zset(v.clone()),
143 Body::Array(v) => Body::Array(v.clone()),
144 Body::Stream(v) => Body::Stream(v.clone()),
145 Body::Foreign(_) => unreachable!("a foreign body never reaches a clone"),
146 }
147 }
148}
149
150/// What a rename or a copy did.
151#[derive(Debug, Clone, Copy, PartialEq, Eq)]
152pub enum Moved {
153 /// There was no source key, so there was nothing to move.
154 Missing,
155 /// The destination was there and the caller said not to write over it.
156 Taken,
157 /// The source holds something there is no way to copy.
158 ///
159 /// A foreign body is owned by the engine above this crate and there is no
160 /// generic way to ask one for a duplicate of itself. A graph could grow a
161 /// deep copy and a vector index probably should not have one at all, so the
162 /// decision belongs to whichever of them is under the key rather than here.
163 /// Answered rather than panicked so the wire can say so in a sentence.
164 Unsupported,
165 /// It happened.
166 Ok,
167}
168
169impl Keyspace {
170 /// Take a copy of everything under `key`, deadline included.
171 ///
172 /// `None` for a key that is not there, and for one whose deadline has gone,
173 /// which is reaped on the way through the same as every other read.
174 ///
175 /// This clones the body, so exporting a set of a million members costs a set
176 /// of a million members. [`Keyspace::rename`] exists so that the one case
177 /// which does not need a copy does not pay for one.
178 pub fn export(&mut self, key: &[u8]) -> Option<Record> {
179 let mut addr = self.live_rec(key)?;
180 // A value on the file is read back but not put back. `DUMP` over a
181 // whole database is the scan the doorkeeper is there for: a backup
182 // should not be able to pull everything into memory on its way past. A
183 // chain that will not read back answers as a key that is not there,
184 // which is the only answer this signature has room for.
185 if value::cold(self.map.value_at(addr)).is_some() {
186 if self.warm(key).is_err() {
187 return None;
188 }
189 addr = self.map.find(key)?;
190 }
191 let rec = self.map.value_at(addr);
192 let expire_at = value::expire_at(rec);
193 // The slot is read inside the arms and not before them. A string record
194 // holds the string and not a slot, so reading four bytes where the slot
195 // would be reads off the end of a short one.
196 let body = match value::kind(rec) {
197 Kind::String => Body::String(self.value_of(key, rec).to_vec()),
198 Kind::Set => Body::Set(
199 self.sets
200 .get(value::slot(rec))
201 .expect("the record points at its body")
202 .clone(),
203 ),
204 Kind::Hash => Body::Hash(
205 self.hashes
206 .get(value::slot(rec))
207 .expect("the record points at its body")
208 .clone(),
209 ),
210 Kind::List => Body::List(
211 self.lists
212 .get(value::slot(rec))
213 .expect("the record points at its body")
214 .clone(),
215 ),
216 Kind::Zset => Body::Zset(
217 self.zsets
218 .get(value::slot(rec))
219 .expect("the record points at its body")
220 .clone(),
221 ),
222 Kind::Array => Body::Array(
223 self.arrays
224 .get(value::slot(rec))
225 .expect("the record points at its body")
226 .clone(),
227 ),
228 Kind::Stream => Body::Stream(
229 self.streams
230 .get(value::slot(rec))
231 .expect("the record points at its body")
232 .clone(),
233 ),
234 // A copy is the one thing a foreign body cannot be asked for. See
235 // [`Moved::Unsupported`]. `None` here reads the same as a missing
236 // key to a caller that only wanted the record, which is why `COPY`
237 // and `DUMP` both check the kind themselves before they get here
238 // rather than reporting a graph as absent.
239 Kind::Foreign => return None,
240 };
241 Some(Record { body, expire_at })
242 }
243
244 /// Lift everything under `key` out and leave the key gone.
245 ///
246 /// The same answer [`Keyspace::export`] gives, without the clone. A body in
247 /// the slab is already a value standing on its own, so a caller that is
248 /// about to delete the source can have that body itself rather than a copy
249 /// of it, and taking a set of a million members costs a slot number.
250 ///
251 /// This is what `MOVE` wants and what `COPY` cannot have. The difference is
252 /// that a move leaves nothing behind, so there is never a moment where two
253 /// records point at one slot.
254 ///
255 /// The record is removed here rather than by the caller, because the body is
256 /// out of the slab by then and a record still pointing at a slot that has
257 /// been freed is the one state this file exists to prevent. A `del` on top
258 /// of this would free the body a second time and underflow the count of keys
259 /// that hold one.
260 pub fn take(&mut self, key: &[u8]) -> Option<Record> {
261 let mut addr = self.live_rec(key)?;
262 // The same read as [`Keyspace::export`] does, for the same reason,
263 // except that here the record is about to go anyway. What the caller
264 // does with the bytes decides where they end up, and on a `RENAME` that
265 // is a resident record under the new name with the old chunks left for
266 // the log's compaction to collect.
267 if value::cold(self.map.value_at(addr)).is_some() {
268 if self.warm(key).is_err() {
269 return None;
270 }
271 addr = self.map.find(key)?;
272 }
273 let rec = self.map.value_at(addr);
274 let expire_at = value::expire_at(rec);
275 let kind = value::kind(rec);
276 // A string record is the value, so there is nothing in the slab to take
277 // and the bytes have to be copied out before the record goes. It leaves
278 // early because the slot below is not there to read on this one.
279 if kind == Kind::String {
280 let bytes = self.value_of(key, rec).to_vec();
281 self.del_rec(key);
282 return Some(Record {
283 body: Body::String(bytes),
284 expire_at,
285 });
286 }
287 let slot = value::slot(rec);
288 let gone = "the record points at its body";
289 let body = match kind {
290 Kind::Set => Body::Set(self.sets.remove(slot).expect(gone)),
291 Kind::Hash => Body::Hash(self.hashes.remove(slot).expect(gone)),
292 Kind::List => Body::List(self.lists.remove(slot).expect(gone)),
293 Kind::Zset => Body::Zset(self.zsets.remove(slot).expect(gone)),
294 Kind::Array => Body::Array(self.arrays.remove(slot).expect(gone)),
295 Kind::Stream => Body::Stream(self.streams.remove(slot).expect(gone)),
296 // A move is the one of the two that a foreign body can do, because
297 // it hands the box over rather than asking for a second one.
298 Kind::Foreign => Body::Foreign(self.foreign.remove(slot).expect(gone)),
299 // Handled above, and named rather than caught, as in `export`.
300 Kind::String => unreachable!("handled above"),
301 };
302 self.bodies -= 1;
303 self.del_rec(key);
304 Some(Record { body, expire_at })
305 }
306
307 /// Put `rec` under `key`, over whatever was there.
308 ///
309 /// The caller has already decided that writing over the destination is
310 /// allowed, which is why this answers nothing. Whatever was under `key` is
311 /// taken away first, record and body both, so this cannot leak a slab slot.
312 ///
313 /// The record goes rather than being written over because this is a key
314 /// arriving and not a value changing. `RESTORE`, `COPY` and `MOVE` all land
315 /// here, and all three of them put a key somewhere it was not, even when
316 /// the name was taken and they were told to take it. The store forms are
317 /// the other case and they go through `Keyspace::put_set` and its
318 /// neighbours, which keep the record where it stands. A client watching for
319 /// keys that were not there before can tell the two apart, so they have to
320 /// be told apart here.
321 pub fn import(&mut self, key: &[u8], rec: Record) {
322 let at = rec.expire_at;
323 self.drop_key(key);
324 match rec.body {
325 Body::String(bytes) => self.store(key, &bytes, at),
326 Body::Set(set) => {
327 let slot = self.sets.insert(set);
328 self.bodies += 1;
329 self.write_slot(key, Kind::Set, slot, at);
330 }
331 Body::Hash(hash) => {
332 let slot = self.hashes.insert(hash);
333 self.bodies += 1;
334 self.write_slot(key, Kind::Hash, slot, at);
335 }
336 Body::List(list) => {
337 let slot = self.lists.insert(list);
338 self.bodies += 1;
339 self.write_slot(key, Kind::List, slot, at);
340 }
341 Body::Zset(zset) => {
342 let slot = self.zsets.insert(zset);
343 self.bodies += 1;
344 self.write_slot(key, Kind::Zset, slot, at);
345 }
346 Body::Array(array) => {
347 let slot = self.arrays.insert(array);
348 self.bodies += 1;
349 self.write_slot(key, Kind::Array, slot, at);
350 }
351 Body::Stream(stream) => {
352 let slot = self.streams.insert(stream);
353 self.bodies += 1;
354 self.write_slot(key, Kind::Stream, slot, at);
355 }
356 Body::Foreign(body) => {
357 let slot = self.foreign.insert(body);
358 self.bodies += 1;
359 self.write_slot(key, Kind::Foreign, slot, at);
360 }
361 }
362 }
363
364 /// `DUMP key`, which is a value on its own with a checksum on the end.
365 ///
366 /// `None` for a key that is not there, and for a key holding something with
367 /// no RDB shape, which today is only the sparse array and which no command
368 /// on the wire can create. Both answer the null bulk that `DUMP` gives for a
369 /// missing key, so a client cannot tell them apart and there is nothing here
370 /// for it to tell apart yet.
371 ///
372 /// The deadline is deliberately left behind. Redis's `DUMP` does the same
373 /// and the reason is that a payload has no idea how long it will be in
374 /// flight, so carrying an absolute deadline would arrive already expired and
375 /// carrying a relative one would quietly extend it. `RESTORE` takes the ttl
376 /// as an argument instead, which puts the decision on whoever knows.
377 pub fn dump(&mut self, key: &[u8]) -> Option<Vec<u8>> {
378 let rec = self.export(key)?;
379 rdb::dump(&rec)
380 }
381
382 /// `RESTORE key ttl payload`, with `replace` for the `REPLACE` option.
383 ///
384 /// [`Moved::Taken`] for a key that is already there without `REPLACE`, which
385 /// is checked before the payload is looked at because that is the order
386 /// Redis checks in and a busy key should not depend on whether the bytes
387 /// behind it happened to be good.
388 ///
389 /// The clone in `export` is not paid here. The payload is parsed straight
390 /// into a body and that body goes into the slab, so restoring a set of a
391 /// million members builds one set.
392 ///
393 /// # Errors
394 ///
395 /// [`rdb::Bad::Footer`] when the version is from the future or the checksum
396 /// does not match, and [`rdb::Bad::Format`] when the bytes were intact and
397 /// still did not describe anything this server can hold. The wire layer has
398 /// a different message for each and clients depend on the difference.
399 pub fn restore(
400 &mut self,
401 key: &[u8],
402 payload: &[u8],
403 expire_at: Option<u64>,
404 replace: bool,
405 ) -> std::result::Result<Moved, rdb::Bad> {
406 if !replace && self.exists(key) {
407 return Ok(Moved::Taken);
408 }
409 let limits = rdb::Limits {
410 set: &self.limits,
411 hash: &self.hash_limits,
412 list: &self.list_limits,
413 zset: &self.zset_limits,
414 };
415 let now = self.clock.now_ms();
416 let body = rdb::load(payload, limits, now)?;
417 // A deadline that has already gone means there is nothing to create, and
418 // the payload is still parsed first rather than skipped. A client that
419 // sent bad bytes and a stale deadline should be told about the bytes,
420 // and finding out only when the deadline is fixed is a bad afternoon.
421 if expire_at.is_some_and(|at| at <= now) {
422 // A no op unless `REPLACE` was given, since a key that was there
423 // without it has already been refused above.
424 self.del(key);
425 return Ok(Moved::Ok);
426 }
427 self.import(key, Record::new(body, expire_at));
428 Ok(Moved::Ok)
429 }
430
431 /// `RENAME src dst`, and `RENAMENX` when `only_if_new`.
432 ///
433 /// The body never moves. A set or a hash is a slot number in a record, and a
434 /// slot number under a different key is the same set, so this writes the
435 /// source's record bytes under the destination and deletes the source
436 /// record without freeing anything. That is why renaming a large collection
437 /// is the same call as renaming a short string.
438 ///
439 /// The deadline travels with the source and the destination's own deadline
440 /// goes with the value it belonged to, which falls out of moving the whole
441 /// record rather than being a rule applied on top of it.
442 ///
443 /// Renaming a key onto itself is allowed and does nothing, which is Redis's
444 /// answer. `RENAMENX` on the same key answers [`Moved::Taken`] instead,
445 /// because the destination does exist, and a key is not new because it is
446 /// the one you already had.
447 pub fn rename(&mut self, src: &[u8], dst: &[u8], only_if_new: bool) -> Moved {
448 if self.live_rec(src).is_none() {
449 return Moved::Missing;
450 }
451 let same = src == dst;
452 if only_if_new && (same || self.live_rec(dst).is_some()) {
453 return Moved::Taken;
454 }
455 if same {
456 return Moved::Ok;
457 }
458 // The record and not the value: a tag, a deadline and then either the
459 // string itself or four bytes saying which slot the body is in. Copying
460 // it out ends the borrow of the map so the write below can begin.
461 //
462 // Into the database's scratch buffer rather than a fresh `Vec`, because
463 // a record under a collection key is nine bytes and `RENAME` is not
464 // rare enough to pay a malloc and a free for nine bytes. Taken out and
465 // put back, so the map is free to be borrowed in between.
466 let addr = self.map.find(src).expect("it was live a line ago");
467 let mut bytes = std::mem::take(&mut self.scratch);
468 bytes.clear();
469 bytes.extend_from_slice(self.map.value_at(addr));
470 // The whole key and not just its body, for the reason
471 // [`Keyspace::import`] gives: what lands on the destination is a key
472 // arriving, whether or not the name was taken.
473 self.drop_key(dst);
474 self.write_rec(dst, bytes.len(), |out| {
475 out.copy_from_slice(&bytes);
476 });
477 self.scratch = bytes;
478 // `del_rec` and not `drop_key`, which is the whole point. The body under
479 // the source belongs to the destination now and freeing it here would
480 // take it away from the key that just gained it. It still goes through
481 // `del_rec` rather than straight at the map, because the record is going
482 // away either way and the count of keys with deadlines has to hear about
483 // it.
484 self.del_rec(src);
485 Moved::Ok
486 }
487
488 /// `COPY src dst`, within one database.
489 ///
490 /// Across two databases the caller runs [`Keyspace::export`] on one and
491 /// [`Keyspace::import`] on the other, because a database cannot see its
492 /// neighbours from in here.
493 ///
494 /// A destination whose deadline has gone counts as free, so this answers
495 /// [`Moved::Ok`] without `replace` on a key that has technically expired and
496 /// not yet been collected. That is Redis's behaviour and it is the only one
497 /// that is consistent with `EXISTS` saying zero for the same key.
498 /// A key copied onto itself answers [`Moved::Ok`] and does nothing, and
499 /// without `replace` it answers [`Moved::Taken`], which is the same pair of
500 /// answers [`Keyspace::rename`] gives. The wire never asks: Redis refuses
501 /// `COPY k k` with an error and so does the dispatch. This is for the
502 /// embedded caller, who can ask, and for whom freeing the body and then
503 /// writing a record that points at it would be the worst of the answers
504 /// available.
505 pub fn copy(&mut self, src: &[u8], dst: &[u8], replace: bool) -> Moved {
506 if self.live_rec(src).is_none() {
507 return Moved::Missing;
508 }
509 let same = src == dst;
510 if !replace && (same || self.live_rec(dst).is_some()) {
511 return Moved::Taken;
512 }
513 if same {
514 return Moved::Ok;
515 }
516 // Asked before anything is written, so a refused copy leaves both keys
517 // exactly as they were rather than freeing the destination first. A
518 // rename does not need the same guard, because it moves the record and
519 // the body under it travels with the record. Only a copy needs a second
520 // body, and a foreign one cannot be asked for one.
521 if self.kind_of(src) == Some(Kind::Foreign) {
522 return Moved::Unsupported;
523 }
524 // The destination is settled before anything is copied, which is the
525 // difference between a refused copy of a million member set costing
526 // nothing and costing the set.
527 //
528 // Both keys have been reaped by now, so the address below stays good
529 // for as long as it is held. It is read after the reaping and not
530 // before, because a reap can move records around.
531 let addr = self.map.find(src).expect("it was live a line ago");
532 if value::kind(self.map.value_at(addr)) == Kind::String {
533 // A string record is the value, deadline and all, so copying the
534 // record is copying the key. That is [`Keyspace::rename`]'s trick,
535 // except the source stays where it is, and it goes through the
536 // database's scratch buffer for the same reason: the borrow of the
537 // map has to end before the write can begin, and a short string is
538 // not worth a malloc and a free.
539 let mut bytes = std::mem::take(&mut self.scratch);
540 bytes.clear();
541 bytes.extend_from_slice(self.map.value_at(addr));
542 self.drop_key(dst);
543 self.write_rec(dst, bytes.len(), |out| {
544 out.copy_from_slice(&bytes);
545 });
546 self.scratch = bytes;
547 return Moved::Ok;
548 }
549 // A collection is a clone and there is no way around that: the
550 // destination has to end up owning a set of its own.
551 let rec = self.export(src).expect("it was live a line ago");
552 self.import(dst, rec);
553 Moved::Ok
554 }
555
556 /// `TOUCH key [key ...]`. Answers how many of them are there.
557 ///
558 /// The same answer `EXISTS` gives, including a key named twice counting
559 /// twice. On a real server the difference is that this moves the key up the
560 /// eviction order, and there is no eviction here yet, so for now the two are
561 /// the same walk and the day eviction lands this is where the bump goes.
562 pub fn touch<'k>(&mut self, keys: impl Iterator<Item = &'k [u8]>) -> usize {
563 keys.filter(|key| self.exists(key)).count()
564 }
565
566 /// The record a set or a hash gets: a tag, a slot number and maybe a
567 /// deadline. Both arms of [`Keyspace::import`] want it and neither wants to
568 /// spell it out.
569 fn write_slot(&mut self, key: &[u8], kind: Kind, slot: u32, at: Option<u64>) {
570 let len = value::slot_record_len(at.is_some());
571 self.write_rec(key, len, |out| {
572 value::write_slot_record(out, kind, slot, at);
573 });
574 }
575}
576
577/// The error `RENAME` and `RENAMENX` answer for a source that is not there.
578///
579/// It is the same sentence for both and it is an error and not a zero, which is
580/// unusual enough among the keyspace commands to be worth its own name: every
581/// other command here treats a missing key as an ordinary answer.
582#[must_use]
583pub fn no_such_key() -> yo_common::Error {
584 yo_common::Error::new(yo_common::Code::Invalid, "no such key")
585}
586
587/// So that a caller can write `?` on a rename without unpacking the enum.
588///
589/// [`Moved::Taken`] is not an error here, because for `RENAMENX` it is the whole
590/// answer and for `RENAME` it cannot happen.
591impl Moved {
592 /// The source was there, or the error `RENAME` gives when it was not.
593 ///
594 /// # Errors
595 ///
596 /// [`yo_common::Code::Invalid`] with Redis's `no such key` for
597 /// [`Moved::Missing`].
598 pub fn found(self) -> Result<Moved> {
599 match self {
600 Moved::Missing => Err(no_such_key()),
601 other => Ok(other),
602 }
603 }
604}
605
606#[cfg(test)]
607mod tests {
608 use super::*;
609 use crate::Clock;
610 use crate::End;
611 use crate::zsets::ZAdd;
612 use crate::{Applied, Cond};
613
614 fn db() -> Keyspace {
615 Keyspace::with_clock(Clock::fixed(1_000_000))
616 }
617
618 fn members(d: &mut Keyspace, key: &[u8]) -> Vec<String> {
619 let mut out: Vec<String> = d
620 .smembers(key)
621 .expect("a set")
622 .expect("a key")
623 .map(|m| String::from_utf8(m.to_vec()).expect("utf8 in these tests"))
624 .collect();
625 out.sort();
626 out
627 }
628
629 fn put(d: &mut Keyspace, key: &[u8], val: &[u8]) {
630 d.set_plain(key, val).expect("room for a record");
631 }
632
633 fn read(d: &mut Keyspace, key: &[u8]) -> Vec<u8> {
634 d.get(key).expect("a string").expect("there").to_vec()
635 }
636
637 #[test]
638 fn a_rename_moves_the_value_and_leaves_nothing_behind() {
639 let mut d = db();
640 put(&mut d, b"a", b"v1");
641
642 assert_eq!(d.rename(b"a", b"b", false), Moved::Ok);
643 assert!(!d.exists(b"a"));
644 assert_eq!(read(&mut d, b"b"), b"v1");
645 }
646
647 /// `RENAME` used to copy the source record into a fresh `Vec` so it could
648 /// let go of the map before writing, and that record is nine bytes when the
649 /// key holds a collection.
650 #[test]
651 fn a_rename_does_not_allocate_to_carry_the_record_across() {
652 let mut d = db();
653 put(&mut d, b"a", b"v1");
654 // Both names get used before the count starts, so the map has already
655 // made room for them and the loop below is renames and nothing else.
656 for _ in 0..4 {
657 assert_eq!(d.rename(b"a", b"b", false), Moved::Ok);
658 assert_eq!(d.rename(b"b", b"a", false), Moved::Ok);
659 }
660 let (_, allocs) = crate::tally::counted(|| {
661 for _ in 0..50 {
662 assert_eq!(d.rename(b"a", b"b", false), Moved::Ok);
663 assert_eq!(d.rename(b"b", b"a", false), Moved::Ok);
664 }
665 });
666 assert_eq!(allocs, 0, "rename allocated {allocs} times in a hundred");
667 assert_eq!(read(&mut d, b"a"), b"v1");
668 }
669
670 #[test]
671 fn a_rename_with_no_source_is_the_one_error_in_this_file() {
672 let mut d = db();
673 assert_eq!(d.rename(b"a", b"b", false), Moved::Missing);
674 assert_eq!(d.rename(b"a", b"b", true), Moved::Missing);
675 assert_eq!(
676 d.copy(b"a", b"b", false),
677 Moved::Missing,
678 "copy just says 0"
679 );
680 }
681
682 #[test]
683 fn a_rename_carries_the_source_deadline_and_drops_the_destination_one() {
684 let mut d = db();
685 put(&mut d, b"a", b"v1");
686 d.set_expiry(b"a", Some(2_000_000));
687 put(&mut d, b"b", b"v2");
688 d.set_expiry(b"b", Some(1_500_000));
689
690 assert_eq!(d.rename(b"a", b"b", false), Moved::Ok);
691 assert_eq!(d.deadline_of(b"b"), crate::Ask::At(2_000_000));
692 }
693
694 #[test]
695 fn renaming_a_key_onto_itself_keeps_it_and_renamenx_refuses() {
696 let mut d = db();
697 put(&mut d, b"a", b"v1");
698 d.set_expiry(b"a", Some(2_000_000));
699
700 assert_eq!(d.rename(b"a", b"a", false), Moved::Ok);
701 assert_eq!(read(&mut d, b"a"), b"v1");
702 assert_eq!(d.deadline_of(b"a"), crate::Ask::At(2_000_000));
703 assert_eq!(d.rename(b"a", b"a", true), Moved::Taken);
704 }
705
706 #[test]
707 fn renamenx_writes_over_nothing() {
708 let mut d = db();
709 put(&mut d, b"a", b"v1");
710 put(&mut d, b"b", b"v2");
711
712 assert_eq!(d.rename(b"a", b"b", true), Moved::Taken);
713 assert_eq!(read(&mut d, b"a"), b"v1");
714 assert_eq!(read(&mut d, b"b"), b"v2");
715 assert_eq!(d.rename(b"a", b"c", true), Moved::Ok);
716 assert!(!d.exists(b"a"));
717 }
718
719 #[test]
720 fn renaming_a_set_moves_the_slot_and_not_the_members() {
721 let mut d = db();
722 d.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
723 .expect("a set");
724 let before = d.memory_bytes();
725
726 assert_eq!(d.rename(b"s", b"t", false), Moved::Ok);
727 assert_eq!(members(&mut d, b"t"), ["m1", "m2"]);
728 assert_eq!(d.kind_of(b"t"), Some(Kind::Set));
729 assert!(!d.exists(b"s"));
730 // The record moved and the body did not, so the only thing that can
731 // have changed size is the record itself.
732 assert!(
733 d.memory_bytes().abs_diff(before) < 64,
734 "the members were not copied"
735 );
736 }
737
738 #[test]
739 fn renaming_over_a_set_frees_the_set_that_was_there() {
740 let mut d = db();
741 d.sadd(b"s", [b"m1".as_ref()].into_iter()).expect("a set");
742 d.sadd(b"t", [b"m2".as_ref()].into_iter()).expect("a set");
743 assert_eq!(d.sets.len(), 2);
744
745 assert_eq!(d.rename(b"s", b"t", false), Moved::Ok);
746 assert_eq!(d.sets.len(), 1, "the destination's body went with it");
747 assert_eq!(members(&mut d, b"t"), ["m1"]);
748 }
749
750 #[test]
751 fn a_copy_is_a_second_value_and_not_a_second_name() {
752 let mut d = db();
753 d.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
754 .expect("a set");
755
756 assert_eq!(d.copy(b"s", b"t", false), Moved::Ok);
757 d.sadd(b"t", [b"m3".as_ref()].into_iter()).expect("a set");
758 assert_eq!(
759 members(&mut d, b"s"),
760 ["m1", "m2"],
761 "the original is intact"
762 );
763 assert_eq!(members(&mut d, b"t"), ["m1", "m2", "m3"]);
764 }
765
766 #[test]
767 fn a_copy_refuses_a_destination_it_was_not_told_it_could_have() {
768 let mut d = db();
769 put(&mut d, b"a", b"v1");
770 put(&mut d, b"b", b"v2");
771
772 assert_eq!(d.copy(b"a", b"b", false), Moved::Taken);
773 assert_eq!(read(&mut d, b"b"), b"v2");
774 assert_eq!(d.copy(b"a", b"b", true), Moved::Ok);
775 assert_eq!(read(&mut d, b"b"), b"v1");
776 }
777
778 /// `COPY` of a string used to go through `export`, which builds a `Vec` of
779 /// the value so that `import` can copy it into the map and drop it.
780 #[test]
781 fn a_copy_of_a_string_does_not_allocate() {
782 let mut d = db();
783 put(&mut d, b"a", b"a-value-of-some-length");
784 // Warmed up, so the map has already made room for both names and the
785 // loop below is copies and nothing else.
786 for _ in 0..4 {
787 assert_eq!(d.copy(b"a", b"b", true), Moved::Ok);
788 }
789 let (_, allocs) = crate::tally::counted(|| {
790 for _ in 0..50 {
791 assert_eq!(d.copy(b"a", b"b", true), Moved::Ok);
792 }
793 });
794 assert_eq!(allocs, 0, "copy allocated {allocs} times in fifty");
795 assert_eq!(read(&mut d, b"b"), b"a-value-of-some-length");
796 }
797
798 /// The embedded caller can ask for this and the wire cannot, because the
799 /// dispatch turns it into an error before it gets here. Freeing the body
800 /// and then writing a record that still points at it would be the way to
801 /// get this wrong.
802 #[test]
803 fn a_copy_onto_itself_leaves_the_key_alone() {
804 let mut d = db();
805 d.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
806 .expect("a set");
807
808 assert_eq!(d.copy(b"s", b"s", false), Moved::Taken);
809 assert_eq!(d.copy(b"s", b"s", true), Moved::Ok);
810 assert_eq!(members(&mut d, b"s"), ["m1", "m2"]);
811 assert_eq!(d.sets.len(), 1, "no second body was made or lost");
812 }
813
814 #[test]
815 fn a_copy_carries_the_deadline() {
816 let mut d = db();
817 put(&mut d, b"a", b"v1");
818 d.set_expiry(b"a", Some(2_000_000));
819
820 assert_eq!(d.copy(b"a", b"b", false), Moved::Ok);
821 assert_eq!(d.deadline_of(b"b"), crate::Ask::At(2_000_000));
822 assert_eq!(d.deadline_of(b"a"), crate::Ask::At(2_000_000));
823 }
824
825 #[test]
826 fn a_destination_that_has_already_gone_counts_as_free() {
827 let mut d = db();
828 put(&mut d, b"a", b"v1");
829 put(&mut d, b"b", b"v2");
830 d.set_expiry(b"b", Some(999_999));
831
832 assert_eq!(d.copy(b"a", b"b", false), Moved::Ok, "b was already gone");
833 assert_eq!(read(&mut d, b"b"), b"v1");
834 }
835
836 #[test]
837 fn a_source_that_has_already_gone_is_not_a_source() {
838 let mut d = db();
839 put(&mut d, b"a", b"v1");
840 d.set_expiry(b"a", Some(999_999));
841
842 assert_eq!(d.rename(b"a", b"b", false), Moved::Missing);
843 assert_eq!(d.copy(b"a", b"b", false), Moved::Missing);
844 }
845
846 #[test]
847 fn a_record_taken_out_of_a_database_outlives_it() {
848 let mut from = db();
849 from.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
850 .expect("a set");
851 let rec = from.export(b"s").expect("a record");
852 assert_eq!(rec.kind(), Kind::Set);
853 from.clear();
854
855 let mut into = db();
856 into.import(b"s", rec);
857 assert_eq!(members(&mut into, b"s"), ["m1", "m2"]);
858 }
859
860 #[test]
861 fn importing_over_a_body_does_not_leave_it_in_the_slab() {
862 let mut d = db();
863 d.sadd(b"s", [b"m1".as_ref()].into_iter()).expect("a set");
864 d.sadd(b"t", [b"m2".as_ref()].into_iter()).expect("a set");
865 let rec = d.export(b"s").expect("a record");
866
867 d.import(b"t", rec);
868 assert_eq!(d.sets.len(), 2, "s and t, and not the one t used to hold");
869 assert_eq!(members(&mut d, b"t"), ["m1"]);
870 }
871
872 #[test]
873 fn importing_a_string_over_a_set_frees_the_set() {
874 let mut d = db();
875 put(&mut d, b"a", b"v1");
876 d.sadd(b"s", [b"m1".as_ref()].into_iter()).expect("a set");
877 assert_eq!(d.sets.len(), 1);
878
879 assert_eq!(d.copy(b"a", b"s", true), Moved::Ok);
880 assert_eq!(d.sets.len(), 0, "the set went when the string arrived");
881 assert_eq!(d.kind_of(b"s"), Some(Kind::String));
882 }
883
884 /// `COPY` of a list, which used to take the server down with it.
885 ///
886 /// The catch all arm at the bottom of `export` was written when a set and a
887 /// hash were the only bodies there were, and the list and the sorted set
888 /// arrived past it without anybody coming back here. So `COPY mylist other`
889 /// reached `unreachable!` and panicked the shard, from a command any client
890 /// can send, against a type the server otherwise supports completely.
891 ///
892 /// The copy has to be a copy and not a second name for the same body, which
893 /// is the other half of what this checks: pushing to the destination must
894 /// not show up in the source.
895 #[test]
896 fn a_list_can_be_copied_and_the_copy_is_its_own() {
897 let mut d = db();
898 d.push(b"l", End::Left, [b"a".as_ref(), b"b".as_ref()].into_iter())
899 .expect("a list");
900
901 assert_eq!(d.copy(b"l", b"m", false), Moved::Ok);
902 assert_eq!(d.kind_of(b"m"), Some(Kind::List));
903 assert_eq!(d.llen(b"m").expect("a list"), 2);
904
905 d.push(b"m", End::Left, [b"c".as_ref()].into_iter())
906 .expect("a list");
907 assert_eq!(d.llen(b"l").expect("a list"), 2, "the source did not grow");
908 assert_eq!(d.llen(b"m").expect("a list"), 3);
909 }
910
911 /// The same for a sorted set, which had the same hole for the same reason.
912 #[test]
913 fn a_zset_can_be_copied_and_the_copy_is_its_own() {
914 let mut d = db();
915 d.zadd(b"z", [(1.0, b"m1".as_ref())].into_iter(), ZAdd::default())
916 .expect("a zset");
917
918 assert_eq!(d.copy(b"z", b"y", false), Moved::Ok);
919 assert_eq!(d.kind_of(b"y"), Some(Kind::Zset));
920 assert_eq!(d.zscore(b"y", b"m1").expect("a zset"), Some(1.0));
921
922 d.zadd(b"y", [(2.0, b"m2".as_ref())].into_iter(), ZAdd::default())
923 .expect("a zset");
924 assert_eq!(d.zcard(b"z").expect("a zset"), 1, "the source did not grow");
925 assert_eq!(d.zcard(b"y").expect("a zset"), 2);
926 }
927
928 /// A copy over a key that held a list gives the list back.
929 ///
930 /// The leak this guards against is the same one the set version guards
931 /// against: a record written over a body that nothing freed leaves a slab
932 /// slot reachable and never reused, and nothing about the server looks wrong
933 /// afterwards.
934 #[test]
935 fn copying_over_a_list_frees_the_list() {
936 let mut d = db();
937 put(&mut d, b"a", b"v1");
938 d.push(b"l", End::Left, [b"x".as_ref()].into_iter())
939 .expect("a list");
940
941 assert_eq!(d.copy(b"a", b"l", true), Moved::Ok);
942 assert_eq!(d.kind_of(b"l"), Some(Kind::String));
943 assert_eq!(read(&mut d, b"l"), b"v1");
944 }
945
946 /// The whole reason `take` exists: the body arrives without being cloned and
947 /// the slab it came out of is empty afterwards.
948 #[test]
949 fn taking_a_set_empties_the_slab_and_the_key() {
950 let mut d = db();
951 d.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
952 .expect("a set");
953 assert_eq!(d.sets.len(), 1);
954
955 let rec = d.take(b"s").expect("a record");
956 assert_eq!(rec.kind(), Kind::Set);
957 assert_eq!(d.sets.len(), 0, "the body left with the record");
958 assert!(!d.exists(b"s"), "and so did the key");
959
960 let mut into = db();
961 into.import(b"s", rec);
962 assert_eq!(members(&mut into, b"s"), ["m1", "m2"]);
963 }
964
965 /// A string has no slab slot, so the bytes are copied and the count is left
966 /// alone. Taking one and then taking it again answers nothing the second
967 /// time, which is the check that the record went too.
968 #[test]
969 fn taking_a_string_takes_the_record_with_it() {
970 let mut d = db();
971 put(&mut d, b"a", b"v1");
972
973 let rec = d.take(b"a").expect("a record");
974 assert_eq!(rec.kind(), Kind::String);
975 assert!(d.take(b"a").is_none());
976 assert_eq!(d.len(), 0);
977 }
978
979 /// The deadline travels, the same as it does through `export`.
980 #[test]
981 fn a_taken_key_keeps_the_time_it_had_left() {
982 let mut d = db();
983 put(&mut d, b"a", b"v1");
984 assert_eq!(d.expire(b"a", 2_000_000, Cond::Always), Applied::Ok);
985
986 let rec = d.take(b"a").expect("a record");
987 assert_eq!(rec.expire_at(), Some(2_000_000));
988 }
989
990 /// A key past its deadline is not there to take, which is the reaping every
991 /// other read does and not a special case here.
992 #[test]
993 fn a_dead_key_cannot_be_taken() {
994 let mut d = db();
995 d.sadd(b"s", [b"m1".as_ref()].into_iter()).expect("a set");
996 assert_eq!(d.expire(b"s", 1_000_001, Cond::Always), Applied::Ok);
997 d.clock().advance(10);
998
999 assert!(d.take(b"s").is_none());
1000 assert_eq!(d.sets.len(), 0, "and the body did not stay behind");
1001 }
1002
1003 #[test]
1004 fn touch_counts_the_way_exists_counts() {
1005 let mut d = db();
1006 put(&mut d, b"a", b"v1");
1007 put(&mut d, b"b", b"v2");
1008
1009 assert_eq!(d.touch([b"a".as_ref()].into_iter()), 1);
1010 assert_eq!(d.touch([b"a".as_ref(), b"b".as_ref()].into_iter()), 2);
1011 assert_eq!(d.touch([b"a".as_ref(), b"a".as_ref()].into_iter()), 2);
1012 assert_eq!(d.touch([b"a".as_ref(), b"z".as_ref()].into_iter()), 1);
1013 assert_eq!(d.touch([b"z".as_ref()].into_iter()), 0);
1014 }
1015}