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_mut().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}