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