Skip to main content

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}