1use yo_common::Result;
31
32use crate::array::Array;
33use crate::hash::Hash;
34use crate::keyspace::Keyspace;
35use crate::list::List;
36use crate::set::Set;
37use crate::value::{self, Kind};
38use crate::zset::Zset;
39
40#[derive(Debug, Clone)]
46pub struct Record {
47 body: Body,
48 expire_at: Option<u64>,
51}
52
53impl Record {
54 #[must_use]
56 pub const fn kind(&self) -> Kind {
57 match self.body {
58 Body::String(_) => Kind::String,
59 Body::Set(_) => Kind::Set,
60 Body::Hash(_) => Kind::Hash,
61 Body::List(_) => Kind::List,
62 Body::Zset(_) => Kind::Zset,
63 Body::Array(_) => Kind::Array,
64 }
65 }
66
67 #[must_use]
69 pub const fn expire_at(&self) -> Option<u64> {
70 self.expire_at
71 }
72}
73
74#[derive(Debug, Clone)]
82enum Body {
83 String(Vec<u8>),
84 Set(Set),
85 Hash(Hash),
86 List(List),
87 Zset(Zset),
88 Array(Array),
89}
90
91#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub enum Moved {
94 Missing,
96 Taken,
98 Ok,
100}
101
102impl Keyspace {
103 pub fn export(&mut self, key: &[u8]) -> Option<Record> {
112 let addr = self.live_rec(key)?;
113 let rec = self.map.value_at(addr);
114 let expire_at = value::expire_at(rec);
115 let body = match value::kind(rec) {
119 Kind::String => Body::String(value::read(rec).to_vec()),
120 Kind::Set => Body::Set(
121 self.sets
122 .get(value::slot(rec))
123 .expect("the record points at its body")
124 .clone(),
125 ),
126 Kind::Hash => Body::Hash(
127 self.hashes
128 .get(value::slot(rec))
129 .expect("the record points at its body")
130 .clone(),
131 ),
132 Kind::List => Body::List(
133 self.lists
134 .get(value::slot(rec))
135 .expect("the record points at its body")
136 .clone(),
137 ),
138 Kind::Zset => Body::Zset(
139 self.zsets
140 .get(value::slot(rec))
141 .expect("the record points at its body")
142 .clone(),
143 ),
144 Kind::Array => Body::Array(
145 self.arrays
146 .get(value::slot(rec))
147 .expect("the record points at its body")
148 .clone(),
149 ),
150 Kind::Stream => unreachable!("nothing can store a stream yet"),
153 };
154 Some(Record { body, expire_at })
155 }
156
157 pub fn import(&mut self, key: &[u8], rec: Record) {
163 let at = rec.expire_at;
164 match rec.body {
165 Body::String(bytes) => self.store(key, &bytes, at),
168 Body::Set(set) => {
169 self.free_body(key);
170 let slot = self.sets.insert(set);
171 self.bodies += 1;
172 self.write_slot(key, Kind::Set, slot, at);
173 }
174 Body::Hash(hash) => {
175 self.free_body(key);
176 let slot = self.hashes.insert(hash);
177 self.bodies += 1;
178 self.write_slot(key, Kind::Hash, slot, at);
179 }
180 Body::List(list) => {
181 self.free_body(key);
182 let slot = self.lists.insert(list);
183 self.bodies += 1;
184 self.write_slot(key, Kind::List, slot, at);
185 }
186 Body::Zset(zset) => {
187 self.free_body(key);
188 let slot = self.zsets.insert(zset);
189 self.bodies += 1;
190 self.write_slot(key, Kind::Zset, slot, at);
191 }
192 Body::Array(array) => {
193 self.free_body(key);
194 let slot = self.arrays.insert(array);
195 self.bodies += 1;
196 self.write_slot(key, Kind::Array, slot, at);
197 }
198 }
199 }
200
201 pub fn rename(&mut self, src: &[u8], dst: &[u8], only_if_new: bool) -> Moved {
218 if self.live_rec(src).is_none() {
219 return Moved::Missing;
220 }
221 let same = src == dst;
222 if only_if_new && (same || self.live_rec(dst).is_some()) {
223 return Moved::Taken;
224 }
225 if same {
226 return Moved::Ok;
227 }
228 let addr = self.map.find(src).expect("it was live a line ago");
237 let mut bytes = std::mem::take(&mut self.scratch);
238 bytes.clear();
239 bytes.extend_from_slice(self.map.value_at(addr));
240 self.free_body(dst);
241 self.write_rec(dst, bytes.len(), |out| {
242 out.copy_from_slice(&bytes);
243 });
244 self.scratch = bytes;
245 self.map.del(src);
249 Moved::Ok
250 }
251
252 pub fn copy(&mut self, src: &[u8], dst: &[u8], replace: bool) -> Moved {
270 if self.live_rec(src).is_none() {
271 return Moved::Missing;
272 }
273 let same = src == dst;
274 if !replace && (same || self.live_rec(dst).is_some()) {
275 return Moved::Taken;
276 }
277 if same {
278 return Moved::Ok;
279 }
280 let addr = self.map.find(src).expect("it was live a line ago");
288 if value::kind(self.map.value_at(addr)) == Kind::String {
289 let mut bytes = std::mem::take(&mut self.scratch);
296 bytes.clear();
297 bytes.extend_from_slice(self.map.value_at(addr));
298 self.free_body(dst);
299 self.write_rec(dst, bytes.len(), |out| {
300 out.copy_from_slice(&bytes);
301 });
302 self.scratch = bytes;
303 return Moved::Ok;
304 }
305 let rec = self.export(src).expect("it was live a line ago");
308 self.import(dst, rec);
309 Moved::Ok
310 }
311
312 pub fn touch<'k>(&mut self, keys: impl Iterator<Item = &'k [u8]>) -> usize {
319 keys.filter(|key| self.exists(key)).count()
320 }
321
322 fn write_slot(&mut self, key: &[u8], kind: Kind, slot: u32, at: Option<u64>) {
326 let len = value::slot_record_len(at.is_some());
327 self.write_rec(key, len, |out| {
328 value::write_slot_record(out, kind, slot, at);
329 });
330 }
331}
332
333#[must_use]
339pub fn no_such_key() -> yo_common::Error {
340 yo_common::Error::new(yo_common::Code::Invalid, "no such key")
341}
342
343impl Moved {
348 pub fn found(self) -> Result<Moved> {
355 match self {
356 Moved::Missing => Err(no_such_key()),
357 other => Ok(other),
358 }
359 }
360}
361
362#[cfg(test)]
363mod tests {
364 use super::*;
365 use crate::Clock;
366 use crate::End;
367 use crate::zsets::ZAdd;
368
369 fn db() -> Keyspace {
370 Keyspace::with_clock(Clock::fixed(1_000_000))
371 }
372
373 fn members(d: &mut Keyspace, key: &[u8]) -> Vec<String> {
374 let mut out: Vec<String> = d
375 .smembers(key)
376 .expect("a set")
377 .expect("a key")
378 .map(|m| String::from_utf8(m.to_vec()).expect("utf8 in these tests"))
379 .collect();
380 out.sort();
381 out
382 }
383
384 fn put(d: &mut Keyspace, key: &[u8], val: &[u8]) {
385 d.set_plain(key, val).expect("room for a record");
386 }
387
388 fn read(d: &mut Keyspace, key: &[u8]) -> Vec<u8> {
389 d.get(key).expect("a string").expect("there").to_vec()
390 }
391
392 #[test]
393 fn a_rename_moves_the_value_and_leaves_nothing_behind() {
394 let mut d = db();
395 put(&mut d, b"a", b"v1");
396
397 assert_eq!(d.rename(b"a", b"b", false), Moved::Ok);
398 assert!(!d.exists(b"a"));
399 assert_eq!(read(&mut d, b"b"), b"v1");
400 }
401
402 #[test]
406 fn a_rename_does_not_allocate_to_carry_the_record_across() {
407 let mut d = db();
408 put(&mut d, b"a", b"v1");
409 for _ in 0..4 {
412 assert_eq!(d.rename(b"a", b"b", false), Moved::Ok);
413 assert_eq!(d.rename(b"b", b"a", false), Moved::Ok);
414 }
415 let (_, allocs) = crate::tally::counted(|| {
416 for _ in 0..50 {
417 assert_eq!(d.rename(b"a", b"b", false), Moved::Ok);
418 assert_eq!(d.rename(b"b", b"a", false), Moved::Ok);
419 }
420 });
421 assert_eq!(allocs, 0, "rename allocated {allocs} times in a hundred");
422 assert_eq!(read(&mut d, b"a"), b"v1");
423 }
424
425 #[test]
426 fn a_rename_with_no_source_is_the_one_error_in_this_file() {
427 let mut d = db();
428 assert_eq!(d.rename(b"a", b"b", false), Moved::Missing);
429 assert_eq!(d.rename(b"a", b"b", true), Moved::Missing);
430 assert_eq!(
431 d.copy(b"a", b"b", false),
432 Moved::Missing,
433 "copy just says 0"
434 );
435 }
436
437 #[test]
438 fn a_rename_carries_the_source_deadline_and_drops_the_destination_one() {
439 let mut d = db();
440 put(&mut d, b"a", b"v1");
441 d.set_expiry(b"a", Some(2_000_000));
442 put(&mut d, b"b", b"v2");
443 d.set_expiry(b"b", Some(1_500_000));
444
445 assert_eq!(d.rename(b"a", b"b", false), Moved::Ok);
446 assert_eq!(d.deadline_of(b"b"), crate::Ask::At(2_000_000));
447 }
448
449 #[test]
450 fn renaming_a_key_onto_itself_keeps_it_and_renamenx_refuses() {
451 let mut d = db();
452 put(&mut d, b"a", b"v1");
453 d.set_expiry(b"a", Some(2_000_000));
454
455 assert_eq!(d.rename(b"a", b"a", false), Moved::Ok);
456 assert_eq!(read(&mut d, b"a"), b"v1");
457 assert_eq!(d.deadline_of(b"a"), crate::Ask::At(2_000_000));
458 assert_eq!(d.rename(b"a", b"a", true), Moved::Taken);
459 }
460
461 #[test]
462 fn renamenx_writes_over_nothing() {
463 let mut d = db();
464 put(&mut d, b"a", b"v1");
465 put(&mut d, b"b", b"v2");
466
467 assert_eq!(d.rename(b"a", b"b", true), Moved::Taken);
468 assert_eq!(read(&mut d, b"a"), b"v1");
469 assert_eq!(read(&mut d, b"b"), b"v2");
470 assert_eq!(d.rename(b"a", b"c", true), Moved::Ok);
471 assert!(!d.exists(b"a"));
472 }
473
474 #[test]
475 fn renaming_a_set_moves_the_slot_and_not_the_members() {
476 let mut d = db();
477 d.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
478 .expect("a set");
479 let before = d.memory_bytes();
480
481 assert_eq!(d.rename(b"s", b"t", false), Moved::Ok);
482 assert_eq!(members(&mut d, b"t"), ["m1", "m2"]);
483 assert_eq!(d.kind_of(b"t"), Some(Kind::Set));
484 assert!(!d.exists(b"s"));
485 assert!(
488 d.memory_bytes().abs_diff(before) < 64,
489 "the members were not copied"
490 );
491 }
492
493 #[test]
494 fn renaming_over_a_set_frees_the_set_that_was_there() {
495 let mut d = db();
496 d.sadd(b"s", [b"m1".as_ref()].into_iter()).expect("a set");
497 d.sadd(b"t", [b"m2".as_ref()].into_iter()).expect("a set");
498 assert_eq!(d.sets.len(), 2);
499
500 assert_eq!(d.rename(b"s", b"t", false), Moved::Ok);
501 assert_eq!(d.sets.len(), 1, "the destination's body went with it");
502 assert_eq!(members(&mut d, b"t"), ["m1"]);
503 }
504
505 #[test]
506 fn a_copy_is_a_second_value_and_not_a_second_name() {
507 let mut d = db();
508 d.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
509 .expect("a set");
510
511 assert_eq!(d.copy(b"s", b"t", false), Moved::Ok);
512 d.sadd(b"t", [b"m3".as_ref()].into_iter()).expect("a set");
513 assert_eq!(
514 members(&mut d, b"s"),
515 ["m1", "m2"],
516 "the original is intact"
517 );
518 assert_eq!(members(&mut d, b"t"), ["m1", "m2", "m3"]);
519 }
520
521 #[test]
522 fn a_copy_refuses_a_destination_it_was_not_told_it_could_have() {
523 let mut d = db();
524 put(&mut d, b"a", b"v1");
525 put(&mut d, b"b", b"v2");
526
527 assert_eq!(d.copy(b"a", b"b", false), Moved::Taken);
528 assert_eq!(read(&mut d, b"b"), b"v2");
529 assert_eq!(d.copy(b"a", b"b", true), Moved::Ok);
530 assert_eq!(read(&mut d, b"b"), b"v1");
531 }
532
533 #[test]
536 fn a_copy_of_a_string_does_not_allocate() {
537 let mut d = db();
538 put(&mut d, b"a", b"a-value-of-some-length");
539 for _ in 0..4 {
542 assert_eq!(d.copy(b"a", b"b", true), Moved::Ok);
543 }
544 let (_, allocs) = crate::tally::counted(|| {
545 for _ in 0..50 {
546 assert_eq!(d.copy(b"a", b"b", true), Moved::Ok);
547 }
548 });
549 assert_eq!(allocs, 0, "copy allocated {allocs} times in fifty");
550 assert_eq!(read(&mut d, b"b"), b"a-value-of-some-length");
551 }
552
553 #[test]
558 fn a_copy_onto_itself_leaves_the_key_alone() {
559 let mut d = db();
560 d.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
561 .expect("a set");
562
563 assert_eq!(d.copy(b"s", b"s", false), Moved::Taken);
564 assert_eq!(d.copy(b"s", b"s", true), Moved::Ok);
565 assert_eq!(members(&mut d, b"s"), ["m1", "m2"]);
566 assert_eq!(d.sets.len(), 1, "no second body was made or lost");
567 }
568
569 #[test]
570 fn a_copy_carries_the_deadline() {
571 let mut d = db();
572 put(&mut d, b"a", b"v1");
573 d.set_expiry(b"a", Some(2_000_000));
574
575 assert_eq!(d.copy(b"a", b"b", false), Moved::Ok);
576 assert_eq!(d.deadline_of(b"b"), crate::Ask::At(2_000_000));
577 assert_eq!(d.deadline_of(b"a"), crate::Ask::At(2_000_000));
578 }
579
580 #[test]
581 fn a_destination_that_has_already_gone_counts_as_free() {
582 let mut d = db();
583 put(&mut d, b"a", b"v1");
584 put(&mut d, b"b", b"v2");
585 d.set_expiry(b"b", Some(999_999));
586
587 assert_eq!(d.copy(b"a", b"b", false), Moved::Ok, "b was already gone");
588 assert_eq!(read(&mut d, b"b"), b"v1");
589 }
590
591 #[test]
592 fn a_source_that_has_already_gone_is_not_a_source() {
593 let mut d = db();
594 put(&mut d, b"a", b"v1");
595 d.set_expiry(b"a", Some(999_999));
596
597 assert_eq!(d.rename(b"a", b"b", false), Moved::Missing);
598 assert_eq!(d.copy(b"a", b"b", false), Moved::Missing);
599 }
600
601 #[test]
602 fn a_record_taken_out_of_a_database_outlives_it() {
603 let mut from = db();
604 from.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
605 .expect("a set");
606 let rec = from.export(b"s").expect("a record");
607 assert_eq!(rec.kind(), Kind::Set);
608 from.clear();
609
610 let mut into = db();
611 into.import(b"s", rec);
612 assert_eq!(members(&mut into, b"s"), ["m1", "m2"]);
613 }
614
615 #[test]
616 fn importing_over_a_body_does_not_leave_it_in_the_slab() {
617 let mut d = db();
618 d.sadd(b"s", [b"m1".as_ref()].into_iter()).expect("a set");
619 d.sadd(b"t", [b"m2".as_ref()].into_iter()).expect("a set");
620 let rec = d.export(b"s").expect("a record");
621
622 d.import(b"t", rec);
623 assert_eq!(d.sets.len(), 2, "s and t, and not the one t used to hold");
624 assert_eq!(members(&mut d, b"t"), ["m1"]);
625 }
626
627 #[test]
628 fn importing_a_string_over_a_set_frees_the_set() {
629 let mut d = db();
630 put(&mut d, b"a", b"v1");
631 d.sadd(b"s", [b"m1".as_ref()].into_iter()).expect("a set");
632 assert_eq!(d.sets.len(), 1);
633
634 assert_eq!(d.copy(b"a", b"s", true), Moved::Ok);
635 assert_eq!(d.sets.len(), 0, "the set went when the string arrived");
636 assert_eq!(d.kind_of(b"s"), Some(Kind::String));
637 }
638
639 #[test]
651 fn a_list_can_be_copied_and_the_copy_is_its_own() {
652 let mut d = db();
653 d.push(b"l", End::Left, [b"a".as_ref(), b"b".as_ref()].into_iter())
654 .expect("a list");
655
656 assert_eq!(d.copy(b"l", b"m", false), Moved::Ok);
657 assert_eq!(d.kind_of(b"m"), Some(Kind::List));
658 assert_eq!(d.llen(b"m").expect("a list"), 2);
659
660 d.push(b"m", End::Left, [b"c".as_ref()].into_iter())
661 .expect("a list");
662 assert_eq!(d.llen(b"l").expect("a list"), 2, "the source did not grow");
663 assert_eq!(d.llen(b"m").expect("a list"), 3);
664 }
665
666 #[test]
668 fn a_zset_can_be_copied_and_the_copy_is_its_own() {
669 let mut d = db();
670 d.zadd(b"z", [(1.0, b"m1".as_ref())].into_iter(), ZAdd::default())
671 .expect("a zset");
672
673 assert_eq!(d.copy(b"z", b"y", false), Moved::Ok);
674 assert_eq!(d.kind_of(b"y"), Some(Kind::Zset));
675 assert_eq!(d.zscore(b"y", b"m1").expect("a zset"), Some(1.0));
676
677 d.zadd(b"y", [(2.0, b"m2".as_ref())].into_iter(), ZAdd::default())
678 .expect("a zset");
679 assert_eq!(d.zcard(b"z").expect("a zset"), 1, "the source did not grow");
680 assert_eq!(d.zcard(b"y").expect("a zset"), 2);
681 }
682
683 #[test]
690 fn copying_over_a_list_frees_the_list() {
691 let mut d = db();
692 put(&mut d, b"a", b"v1");
693 d.push(b"l", End::Left, [b"x".as_ref()].into_iter())
694 .expect("a list");
695
696 assert_eq!(d.copy(b"a", b"l", true), Moved::Ok);
697 assert_eq!(d.kind_of(b"l"), Some(Kind::String));
698 assert_eq!(read(&mut d, b"l"), b"v1");
699 }
700
701 #[test]
702 fn touch_counts_the_way_exists_counts() {
703 let mut d = db();
704 put(&mut d, b"a", b"v1");
705 put(&mut d, b"b", b"v2");
706
707 assert_eq!(d.touch([b"a".as_ref()].into_iter()), 1);
708 assert_eq!(d.touch([b"a".as_ref(), b"b".as_ref()].into_iter()), 2);
709 assert_eq!(d.touch([b"a".as_ref(), b"a".as_ref()].into_iter()), 2);
710 assert_eq!(d.touch([b"a".as_ref(), b"z".as_ref()].into_iter()), 1);
711 assert_eq!(d.touch([b"z".as_ref()].into_iter()), 0);
712 }
713}