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    /// freed first, body and all, so this cannot leak a slab slot.
312    pub fn import(&mut self, key: &[u8], rec: Record) {
313        let at = rec.expire_at;
314        match rec.body {
315            // The string path frees the old body itself, because every string
316            // write has to and this is not the place to make it special.
317            Body::String(bytes) => self.store(key, &bytes, at),
318            Body::Set(set) => {
319                self.free_body(key);
320                let slot = self.sets.insert(set);
321                self.bodies += 1;
322                self.write_slot(key, Kind::Set, slot, at);
323            }
324            Body::Hash(hash) => {
325                self.free_body(key);
326                let slot = self.hashes.insert(hash);
327                self.bodies += 1;
328                self.write_slot(key, Kind::Hash, slot, at);
329            }
330            Body::List(list) => {
331                self.free_body(key);
332                let slot = self.lists.insert(list);
333                self.bodies += 1;
334                self.write_slot(key, Kind::List, slot, at);
335            }
336            Body::Zset(zset) => {
337                self.free_body(key);
338                let slot = self.zsets.insert(zset);
339                self.bodies += 1;
340                self.write_slot(key, Kind::Zset, slot, at);
341            }
342            Body::Array(array) => {
343                self.free_body(key);
344                let slot = self.arrays.insert(array);
345                self.bodies += 1;
346                self.write_slot(key, Kind::Array, slot, at);
347            }
348            Body::Stream(stream) => {
349                self.free_body(key);
350                let slot = self.streams.insert(stream);
351                self.bodies += 1;
352                self.write_slot(key, Kind::Stream, slot, at);
353            }
354            Body::Foreign(body) => {
355                self.free_body(key);
356                let slot = self.foreign.insert(body);
357                self.bodies += 1;
358                self.write_slot(key, Kind::Foreign, slot, at);
359            }
360        }
361    }
362
363    /// `DUMP key`, which is a value on its own with a checksum on the end.
364    ///
365    /// `None` for a key that is not there, and for a key holding something with
366    /// no RDB shape, which today is only the sparse array and which no command
367    /// on the wire can create. Both answer the null bulk that `DUMP` gives for a
368    /// missing key, so a client cannot tell them apart and there is nothing here
369    /// for it to tell apart yet.
370    ///
371    /// The deadline is deliberately left behind. Redis's `DUMP` does the same
372    /// and the reason is that a payload has no idea how long it will be in
373    /// flight, so carrying an absolute deadline would arrive already expired and
374    /// carrying a relative one would quietly extend it. `RESTORE` takes the ttl
375    /// as an argument instead, which puts the decision on whoever knows.
376    pub fn dump(&mut self, key: &[u8]) -> Option<Vec<u8>> {
377        let rec = self.export(key)?;
378        rdb::dump(&rec)
379    }
380
381    /// `RESTORE key ttl payload`, with `replace` for the `REPLACE` option.
382    ///
383    /// [`Moved::Taken`] for a key that is already there without `REPLACE`, which
384    /// is checked before the payload is looked at because that is the order
385    /// Redis checks in and a busy key should not depend on whether the bytes
386    /// behind it happened to be good.
387    ///
388    /// The clone in `export` is not paid here. The payload is parsed straight
389    /// into a body and that body goes into the slab, so restoring a set of a
390    /// million members builds one set.
391    ///
392    /// # Errors
393    ///
394    /// [`rdb::Bad::Footer`] when the version is from the future or the checksum
395    /// does not match, and [`rdb::Bad::Format`] when the bytes were intact and
396    /// still did not describe anything this server can hold. The wire layer has
397    /// a different message for each and clients depend on the difference.
398    pub fn restore(
399        &mut self,
400        key: &[u8],
401        payload: &[u8],
402        expire_at: Option<u64>,
403        replace: bool,
404    ) -> std::result::Result<Moved, rdb::Bad> {
405        if !replace && self.exists(key) {
406            return Ok(Moved::Taken);
407        }
408        let limits = rdb::Limits {
409            set: &self.limits,
410            hash: &self.hash_limits,
411            list: &self.list_limits,
412            zset: &self.zset_limits,
413        };
414        let now = self.clock.now_ms();
415        let body = rdb::load(payload, limits, now)?;
416        // A deadline that has already gone means there is nothing to create, and
417        // the payload is still parsed first rather than skipped. A client that
418        // sent bad bytes and a stale deadline should be told about the bytes,
419        // and finding out only when the deadline is fixed is a bad afternoon.
420        if expire_at.is_some_and(|at| at <= now) {
421            // A no op unless `REPLACE` was given, since a key that was there
422            // without it has already been refused above.
423            self.del(key);
424            return Ok(Moved::Ok);
425        }
426        self.import(key, Record::new(body, expire_at));
427        Ok(Moved::Ok)
428    }
429
430    /// `RENAME src dst`, and `RENAMENX` when `only_if_new`.
431    ///
432    /// The body never moves. A set or a hash is a slot number in a record, and a
433    /// slot number under a different key is the same set, so this writes the
434    /// source's record bytes under the destination and deletes the source
435    /// record without freeing anything. That is why renaming a large collection
436    /// is the same call as renaming a short string.
437    ///
438    /// The deadline travels with the source and the destination's own deadline
439    /// goes with the value it belonged to, which falls out of moving the whole
440    /// record rather than being a rule applied on top of it.
441    ///
442    /// Renaming a key onto itself is allowed and does nothing, which is Redis's
443    /// answer. `RENAMENX` on the same key answers [`Moved::Taken`] instead,
444    /// because the destination does exist, and a key is not new because it is
445    /// the one you already had.
446    pub fn rename(&mut self, src: &[u8], dst: &[u8], only_if_new: bool) -> Moved {
447        if self.live_rec(src).is_none() {
448            return Moved::Missing;
449        }
450        let same = src == dst;
451        if only_if_new && (same || self.live_rec(dst).is_some()) {
452            return Moved::Taken;
453        }
454        if same {
455            return Moved::Ok;
456        }
457        // The record and not the value: a tag, a deadline and then either the
458        // string itself or four bytes saying which slot the body is in. Copying
459        // it out ends the borrow of the map so the write below can begin.
460        //
461        // Into the database's scratch buffer rather than a fresh `Vec`, because
462        // a record under a collection key is nine bytes and `RENAME` is not
463        // rare enough to pay a malloc and a free for nine bytes. Taken out and
464        // put back, so the map is free to be borrowed in between.
465        let addr = self.map.find(src).expect("it was live a line ago");
466        let mut bytes = std::mem::take(&mut self.scratch);
467        bytes.clear();
468        bytes.extend_from_slice(self.map.value_at(addr));
469        self.free_body(dst);
470        self.write_rec(dst, bytes.len(), |out| {
471            out.copy_from_slice(&bytes);
472        });
473        self.scratch = bytes;
474        // `del_rec` and not `drop_key`, which is the whole point. The body under
475        // the source belongs to the destination now and freeing it here would
476        // take it away from the key that just gained it. It still goes through
477        // `del_rec` rather than straight at the map, because the record is going
478        // away either way and the count of keys with deadlines has to hear about
479        // it.
480        self.del_rec(src);
481        Moved::Ok
482    }
483
484    /// `COPY src dst`, within one database.
485    ///
486    /// Across two databases the caller runs [`Keyspace::export`] on one and
487    /// [`Keyspace::import`] on the other, because a database cannot see its
488    /// neighbours from in here.
489    ///
490    /// A destination whose deadline has gone counts as free, so this answers
491    /// [`Moved::Ok`] without `replace` on a key that has technically expired and
492    /// not yet been collected. That is Redis's behaviour and it is the only one
493    /// that is consistent with `EXISTS` saying zero for the same key.
494    /// A key copied onto itself answers [`Moved::Ok`] and does nothing, and
495    /// without `replace` it answers [`Moved::Taken`], which is the same pair of
496    /// answers [`Keyspace::rename`] gives. The wire never asks: Redis refuses
497    /// `COPY k k` with an error and so does the dispatch. This is for the
498    /// embedded caller, who can ask, and for whom freeing the body and then
499    /// writing a record that points at it would be the worst of the answers
500    /// available.
501    pub fn copy(&mut self, src: &[u8], dst: &[u8], replace: bool) -> Moved {
502        if self.live_rec(src).is_none() {
503            return Moved::Missing;
504        }
505        let same = src == dst;
506        if !replace && (same || self.live_rec(dst).is_some()) {
507            return Moved::Taken;
508        }
509        if same {
510            return Moved::Ok;
511        }
512        // Asked before anything is written, so a refused copy leaves both keys
513        // exactly as they were rather than freeing the destination first. A
514        // rename does not need the same guard, because it moves the record and
515        // the body under it travels with the record. Only a copy needs a second
516        // body, and a foreign one cannot be asked for one.
517        if self.kind_of(src) == Some(Kind::Foreign) {
518            return Moved::Unsupported;
519        }
520        // The destination is settled before anything is copied, which is the
521        // difference between a refused copy of a million member set costing
522        // nothing and costing the set.
523        //
524        // Both keys have been reaped by now, so the address below stays good
525        // for as long as it is held. It is read after the reaping and not
526        // before, because a reap can move records around.
527        let addr = self.map.find(src).expect("it was live a line ago");
528        if value::kind(self.map.value_at(addr)) == Kind::String {
529            // A string record is the value, deadline and all, so copying the
530            // record is copying the key. That is [`Keyspace::rename`]'s trick,
531            // except the source stays where it is, and it goes through the
532            // database's scratch buffer for the same reason: the borrow of the
533            // map has to end before the write can begin, and a short string is
534            // not worth a malloc and a free.
535            let mut bytes = std::mem::take(&mut self.scratch);
536            bytes.clear();
537            bytes.extend_from_slice(self.map.value_at(addr));
538            self.free_body(dst);
539            self.write_rec(dst, bytes.len(), |out| {
540                out.copy_from_slice(&bytes);
541            });
542            self.scratch = bytes;
543            return Moved::Ok;
544        }
545        // A collection is a clone and there is no way around that: the
546        // destination has to end up owning a set of its own.
547        let rec = self.export(src).expect("it was live a line ago");
548        self.import(dst, rec);
549        Moved::Ok
550    }
551
552    /// `TOUCH key [key ...]`. Answers how many of them are there.
553    ///
554    /// The same answer `EXISTS` gives, including a key named twice counting
555    /// twice. On a real server the difference is that this moves the key up the
556    /// eviction order, and there is no eviction here yet, so for now the two are
557    /// the same walk and the day eviction lands this is where the bump goes.
558    pub fn touch<'k>(&mut self, keys: impl Iterator<Item = &'k [u8]>) -> usize {
559        keys.filter(|key| self.exists(key)).count()
560    }
561
562    /// The record a set or a hash gets: a tag, a slot number and maybe a
563    /// deadline. Both arms of [`Keyspace::import`] want it and neither wants to
564    /// spell it out.
565    fn write_slot(&mut self, key: &[u8], kind: Kind, slot: u32, at: Option<u64>) {
566        let len = value::slot_record_len(at.is_some());
567        self.write_rec(key, len, |out| {
568            value::write_slot_record(out, kind, slot, at);
569        });
570    }
571}
572
573/// The error `RENAME` and `RENAMENX` answer for a source that is not there.
574///
575/// It is the same sentence for both and it is an error and not a zero, which is
576/// unusual enough among the keyspace commands to be worth its own name: every
577/// other command here treats a missing key as an ordinary answer.
578#[must_use]
579pub fn no_such_key() -> yo_common::Error {
580    yo_common::Error::new(yo_common::Code::Invalid, "no such key")
581}
582
583/// So that a caller can write `?` on a rename without unpacking the enum.
584///
585/// [`Moved::Taken`] is not an error here, because for `RENAMENX` it is the whole
586/// answer and for `RENAME` it cannot happen.
587impl Moved {
588    /// The source was there, or the error `RENAME` gives when it was not.
589    ///
590    /// # Errors
591    ///
592    /// [`yo_common::Code::Invalid`] with Redis's `no such key` for
593    /// [`Moved::Missing`].
594    pub fn found(self) -> Result<Moved> {
595        match self {
596            Moved::Missing => Err(no_such_key()),
597            other => Ok(other),
598        }
599    }
600}
601
602#[cfg(test)]
603mod tests {
604    use super::*;
605    use crate::Clock;
606    use crate::End;
607    use crate::zsets::ZAdd;
608    use crate::{Applied, Cond};
609
610    fn db() -> Keyspace {
611        Keyspace::with_clock(Clock::fixed(1_000_000))
612    }
613
614    fn members(d: &mut Keyspace, key: &[u8]) -> Vec<String> {
615        let mut out: Vec<String> = d
616            .smembers(key)
617            .expect("a set")
618            .expect("a key")
619            .map(|m| String::from_utf8(m.to_vec()).expect("utf8 in these tests"))
620            .collect();
621        out.sort();
622        out
623    }
624
625    fn put(d: &mut Keyspace, key: &[u8], val: &[u8]) {
626        d.set_plain(key, val).expect("room for a record");
627    }
628
629    fn read(d: &mut Keyspace, key: &[u8]) -> Vec<u8> {
630        d.get(key).expect("a string").expect("there").to_vec()
631    }
632
633    #[test]
634    fn a_rename_moves_the_value_and_leaves_nothing_behind() {
635        let mut d = db();
636        put(&mut d, b"a", b"v1");
637
638        assert_eq!(d.rename(b"a", b"b", false), Moved::Ok);
639        assert!(!d.exists(b"a"));
640        assert_eq!(read(&mut d, b"b"), b"v1");
641    }
642
643    /// `RENAME` used to copy the source record into a fresh `Vec` so it could
644    /// let go of the map before writing, and that record is nine bytes when the
645    /// key holds a collection.
646    #[test]
647    fn a_rename_does_not_allocate_to_carry_the_record_across() {
648        let mut d = db();
649        put(&mut d, b"a", b"v1");
650        // Both names get used before the count starts, so the map has already
651        // made room for them and the loop below is renames and nothing else.
652        for _ in 0..4 {
653            assert_eq!(d.rename(b"a", b"b", false), Moved::Ok);
654            assert_eq!(d.rename(b"b", b"a", false), Moved::Ok);
655        }
656        let (_, allocs) = crate::tally::counted(|| {
657            for _ in 0..50 {
658                assert_eq!(d.rename(b"a", b"b", false), Moved::Ok);
659                assert_eq!(d.rename(b"b", b"a", false), Moved::Ok);
660            }
661        });
662        assert_eq!(allocs, 0, "rename allocated {allocs} times in a hundred");
663        assert_eq!(read(&mut d, b"a"), b"v1");
664    }
665
666    #[test]
667    fn a_rename_with_no_source_is_the_one_error_in_this_file() {
668        let mut d = db();
669        assert_eq!(d.rename(b"a", b"b", false), Moved::Missing);
670        assert_eq!(d.rename(b"a", b"b", true), Moved::Missing);
671        assert_eq!(
672            d.copy(b"a", b"b", false),
673            Moved::Missing,
674            "copy just says 0"
675        );
676    }
677
678    #[test]
679    fn a_rename_carries_the_source_deadline_and_drops_the_destination_one() {
680        let mut d = db();
681        put(&mut d, b"a", b"v1");
682        d.set_expiry(b"a", Some(2_000_000));
683        put(&mut d, b"b", b"v2");
684        d.set_expiry(b"b", Some(1_500_000));
685
686        assert_eq!(d.rename(b"a", b"b", false), Moved::Ok);
687        assert_eq!(d.deadline_of(b"b"), crate::Ask::At(2_000_000));
688    }
689
690    #[test]
691    fn renaming_a_key_onto_itself_keeps_it_and_renamenx_refuses() {
692        let mut d = db();
693        put(&mut d, b"a", b"v1");
694        d.set_expiry(b"a", Some(2_000_000));
695
696        assert_eq!(d.rename(b"a", b"a", false), Moved::Ok);
697        assert_eq!(read(&mut d, b"a"), b"v1");
698        assert_eq!(d.deadline_of(b"a"), crate::Ask::At(2_000_000));
699        assert_eq!(d.rename(b"a", b"a", true), Moved::Taken);
700    }
701
702    #[test]
703    fn renamenx_writes_over_nothing() {
704        let mut d = db();
705        put(&mut d, b"a", b"v1");
706        put(&mut d, b"b", b"v2");
707
708        assert_eq!(d.rename(b"a", b"b", true), Moved::Taken);
709        assert_eq!(read(&mut d, b"a"), b"v1");
710        assert_eq!(read(&mut d, b"b"), b"v2");
711        assert_eq!(d.rename(b"a", b"c", true), Moved::Ok);
712        assert!(!d.exists(b"a"));
713    }
714
715    #[test]
716    fn renaming_a_set_moves_the_slot_and_not_the_members() {
717        let mut d = db();
718        d.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
719            .expect("a set");
720        let before = d.memory_bytes();
721
722        assert_eq!(d.rename(b"s", b"t", false), Moved::Ok);
723        assert_eq!(members(&mut d, b"t"), ["m1", "m2"]);
724        assert_eq!(d.kind_of(b"t"), Some(Kind::Set));
725        assert!(!d.exists(b"s"));
726        // The record moved and the body did not, so the only thing that can
727        // have changed size is the record itself.
728        assert!(
729            d.memory_bytes().abs_diff(before) < 64,
730            "the members were not copied"
731        );
732    }
733
734    #[test]
735    fn renaming_over_a_set_frees_the_set_that_was_there() {
736        let mut d = db();
737        d.sadd(b"s", [b"m1".as_ref()].into_iter()).expect("a set");
738        d.sadd(b"t", [b"m2".as_ref()].into_iter()).expect("a set");
739        assert_eq!(d.sets.len(), 2);
740
741        assert_eq!(d.rename(b"s", b"t", false), Moved::Ok);
742        assert_eq!(d.sets.len(), 1, "the destination's body went with it");
743        assert_eq!(members(&mut d, b"t"), ["m1"]);
744    }
745
746    #[test]
747    fn a_copy_is_a_second_value_and_not_a_second_name() {
748        let mut d = db();
749        d.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
750            .expect("a set");
751
752        assert_eq!(d.copy(b"s", b"t", false), Moved::Ok);
753        d.sadd(b"t", [b"m3".as_ref()].into_iter()).expect("a set");
754        assert_eq!(
755            members(&mut d, b"s"),
756            ["m1", "m2"],
757            "the original is intact"
758        );
759        assert_eq!(members(&mut d, b"t"), ["m1", "m2", "m3"]);
760    }
761
762    #[test]
763    fn a_copy_refuses_a_destination_it_was_not_told_it_could_have() {
764        let mut d = db();
765        put(&mut d, b"a", b"v1");
766        put(&mut d, b"b", b"v2");
767
768        assert_eq!(d.copy(b"a", b"b", false), Moved::Taken);
769        assert_eq!(read(&mut d, b"b"), b"v2");
770        assert_eq!(d.copy(b"a", b"b", true), Moved::Ok);
771        assert_eq!(read(&mut d, b"b"), b"v1");
772    }
773
774    /// `COPY` of a string used to go through `export`, which builds a `Vec` of
775    /// the value so that `import` can copy it into the map and drop it.
776    #[test]
777    fn a_copy_of_a_string_does_not_allocate() {
778        let mut d = db();
779        put(&mut d, b"a", b"a-value-of-some-length");
780        // Warmed up, so the map has already made room for both names and the
781        // loop below is copies and nothing else.
782        for _ in 0..4 {
783            assert_eq!(d.copy(b"a", b"b", true), Moved::Ok);
784        }
785        let (_, allocs) = crate::tally::counted(|| {
786            for _ in 0..50 {
787                assert_eq!(d.copy(b"a", b"b", true), Moved::Ok);
788            }
789        });
790        assert_eq!(allocs, 0, "copy allocated {allocs} times in fifty");
791        assert_eq!(read(&mut d, b"b"), b"a-value-of-some-length");
792    }
793
794    /// The embedded caller can ask for this and the wire cannot, because the
795    /// dispatch turns it into an error before it gets here. Freeing the body
796    /// and then writing a record that still points at it would be the way to
797    /// get this wrong.
798    #[test]
799    fn a_copy_onto_itself_leaves_the_key_alone() {
800        let mut d = db();
801        d.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
802            .expect("a set");
803
804        assert_eq!(d.copy(b"s", b"s", false), Moved::Taken);
805        assert_eq!(d.copy(b"s", b"s", true), Moved::Ok);
806        assert_eq!(members(&mut d, b"s"), ["m1", "m2"]);
807        assert_eq!(d.sets.len(), 1, "no second body was made or lost");
808    }
809
810    #[test]
811    fn a_copy_carries_the_deadline() {
812        let mut d = db();
813        put(&mut d, b"a", b"v1");
814        d.set_expiry(b"a", Some(2_000_000));
815
816        assert_eq!(d.copy(b"a", b"b", false), Moved::Ok);
817        assert_eq!(d.deadline_of(b"b"), crate::Ask::At(2_000_000));
818        assert_eq!(d.deadline_of(b"a"), crate::Ask::At(2_000_000));
819    }
820
821    #[test]
822    fn a_destination_that_has_already_gone_counts_as_free() {
823        let mut d = db();
824        put(&mut d, b"a", b"v1");
825        put(&mut d, b"b", b"v2");
826        d.set_expiry(b"b", Some(999_999));
827
828        assert_eq!(d.copy(b"a", b"b", false), Moved::Ok, "b was already gone");
829        assert_eq!(read(&mut d, b"b"), b"v1");
830    }
831
832    #[test]
833    fn a_source_that_has_already_gone_is_not_a_source() {
834        let mut d = db();
835        put(&mut d, b"a", b"v1");
836        d.set_expiry(b"a", Some(999_999));
837
838        assert_eq!(d.rename(b"a", b"b", false), Moved::Missing);
839        assert_eq!(d.copy(b"a", b"b", false), Moved::Missing);
840    }
841
842    #[test]
843    fn a_record_taken_out_of_a_database_outlives_it() {
844        let mut from = db();
845        from.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
846            .expect("a set");
847        let rec = from.export(b"s").expect("a record");
848        assert_eq!(rec.kind(), Kind::Set);
849        from.clear();
850
851        let mut into = db();
852        into.import(b"s", rec);
853        assert_eq!(members(&mut into, b"s"), ["m1", "m2"]);
854    }
855
856    #[test]
857    fn importing_over_a_body_does_not_leave_it_in_the_slab() {
858        let mut d = db();
859        d.sadd(b"s", [b"m1".as_ref()].into_iter()).expect("a set");
860        d.sadd(b"t", [b"m2".as_ref()].into_iter()).expect("a set");
861        let rec = d.export(b"s").expect("a record");
862
863        d.import(b"t", rec);
864        assert_eq!(d.sets.len(), 2, "s and t, and not the one t used to hold");
865        assert_eq!(members(&mut d, b"t"), ["m1"]);
866    }
867
868    #[test]
869    fn importing_a_string_over_a_set_frees_the_set() {
870        let mut d = db();
871        put(&mut d, b"a", b"v1");
872        d.sadd(b"s", [b"m1".as_ref()].into_iter()).expect("a set");
873        assert_eq!(d.sets.len(), 1);
874
875        assert_eq!(d.copy(b"a", b"s", true), Moved::Ok);
876        assert_eq!(d.sets.len(), 0, "the set went when the string arrived");
877        assert_eq!(d.kind_of(b"s"), Some(Kind::String));
878    }
879
880    /// `COPY` of a list, which used to take the server down with it.
881    ///
882    /// The catch all arm at the bottom of `export` was written when a set and a
883    /// hash were the only bodies there were, and the list and the sorted set
884    /// arrived past it without anybody coming back here. So `COPY mylist other`
885    /// reached `unreachable!` and panicked the shard, from a command any client
886    /// can send, against a type the server otherwise supports completely.
887    ///
888    /// The copy has to be a copy and not a second name for the same body, which
889    /// is the other half of what this checks: pushing to the destination must
890    /// not show up in the source.
891    #[test]
892    fn a_list_can_be_copied_and_the_copy_is_its_own() {
893        let mut d = db();
894        d.push(b"l", End::Left, [b"a".as_ref(), b"b".as_ref()].into_iter())
895            .expect("a list");
896
897        assert_eq!(d.copy(b"l", b"m", false), Moved::Ok);
898        assert_eq!(d.kind_of(b"m"), Some(Kind::List));
899        assert_eq!(d.llen(b"m").expect("a list"), 2);
900
901        d.push(b"m", End::Left, [b"c".as_ref()].into_iter())
902            .expect("a list");
903        assert_eq!(d.llen(b"l").expect("a list"), 2, "the source did not grow");
904        assert_eq!(d.llen(b"m").expect("a list"), 3);
905    }
906
907    /// The same for a sorted set, which had the same hole for the same reason.
908    #[test]
909    fn a_zset_can_be_copied_and_the_copy_is_its_own() {
910        let mut d = db();
911        d.zadd(b"z", [(1.0, b"m1".as_ref())].into_iter(), ZAdd::default())
912            .expect("a zset");
913
914        assert_eq!(d.copy(b"z", b"y", false), Moved::Ok);
915        assert_eq!(d.kind_of(b"y"), Some(Kind::Zset));
916        assert_eq!(d.zscore(b"y", b"m1").expect("a zset"), Some(1.0));
917
918        d.zadd(b"y", [(2.0, b"m2".as_ref())].into_iter(), ZAdd::default())
919            .expect("a zset");
920        assert_eq!(d.zcard(b"z").expect("a zset"), 1, "the source did not grow");
921        assert_eq!(d.zcard(b"y").expect("a zset"), 2);
922    }
923
924    /// A copy over a key that held a list gives the list back.
925    ///
926    /// The leak this guards against is the same one the set version guards
927    /// against: a record written over a body that nothing freed leaves a slab
928    /// slot reachable and never reused, and nothing about the server looks wrong
929    /// afterwards.
930    #[test]
931    fn copying_over_a_list_frees_the_list() {
932        let mut d = db();
933        put(&mut d, b"a", b"v1");
934        d.push(b"l", End::Left, [b"x".as_ref()].into_iter())
935            .expect("a list");
936
937        assert_eq!(d.copy(b"a", b"l", true), Moved::Ok);
938        assert_eq!(d.kind_of(b"l"), Some(Kind::String));
939        assert_eq!(read(&mut d, b"l"), b"v1");
940    }
941
942    /// The whole reason `take` exists: the body arrives without being cloned and
943    /// the slab it came out of is empty afterwards.
944    #[test]
945    fn taking_a_set_empties_the_slab_and_the_key() {
946        let mut d = db();
947        d.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
948            .expect("a set");
949        assert_eq!(d.sets.len(), 1);
950
951        let rec = d.take(b"s").expect("a record");
952        assert_eq!(rec.kind(), Kind::Set);
953        assert_eq!(d.sets.len(), 0, "the body left with the record");
954        assert!(!d.exists(b"s"), "and so did the key");
955
956        let mut into = db();
957        into.import(b"s", rec);
958        assert_eq!(members(&mut into, b"s"), ["m1", "m2"]);
959    }
960
961    /// A string has no slab slot, so the bytes are copied and the count is left
962    /// alone. Taking one and then taking it again answers nothing the second
963    /// time, which is the check that the record went too.
964    #[test]
965    fn taking_a_string_takes_the_record_with_it() {
966        let mut d = db();
967        put(&mut d, b"a", b"v1");
968
969        let rec = d.take(b"a").expect("a record");
970        assert_eq!(rec.kind(), Kind::String);
971        assert!(d.take(b"a").is_none());
972        assert_eq!(d.len(), 0);
973    }
974
975    /// The deadline travels, the same as it does through `export`.
976    #[test]
977    fn a_taken_key_keeps_the_time_it_had_left() {
978        let mut d = db();
979        put(&mut d, b"a", b"v1");
980        assert_eq!(d.expire(b"a", 2_000_000, Cond::Always), Applied::Ok);
981
982        let rec = d.take(b"a").expect("a record");
983        assert_eq!(rec.expire_at(), Some(2_000_000));
984    }
985
986    /// A key past its deadline is not there to take, which is the reaping every
987    /// other read does and not a special case here.
988    #[test]
989    fn a_dead_key_cannot_be_taken() {
990        let mut d = db();
991        d.sadd(b"s", [b"m1".as_ref()].into_iter()).expect("a set");
992        assert_eq!(d.expire(b"s", 1_000_001, Cond::Always), Applied::Ok);
993        d.clock().advance(10);
994
995        assert!(d.take(b"s").is_none());
996        assert_eq!(d.sets.len(), 0, "and the body did not stay behind");
997    }
998
999    #[test]
1000    fn touch_counts_the_way_exists_counts() {
1001        let mut d = db();
1002        put(&mut d, b"a", b"v1");
1003        put(&mut d, b"b", b"v2");
1004
1005        assert_eq!(d.touch([b"a".as_ref()].into_iter()), 1);
1006        assert_eq!(d.touch([b"a".as_ref(), b"b".as_ref()].into_iter()), 2);
1007        assert_eq!(d.touch([b"a".as_ref(), b"a".as_ref()].into_iter()), 2);
1008        assert_eq!(d.touch([b"a".as_ref(), b"z".as_ref()].into_iter()), 1);
1009        assert_eq!(d.touch([b"z".as_ref()].into_iter()), 0);
1010    }
1011}