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