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.del_rec(src);
252 Moved::Ok
253 }
254
255 pub fn copy(&mut self, src: &[u8], dst: &[u8], replace: bool) -> Moved {
273 if self.live_rec(src).is_none() {
274 return Moved::Missing;
275 }
276 let same = src == dst;
277 if !replace && (same || self.live_rec(dst).is_some()) {
278 return Moved::Taken;
279 }
280 if same {
281 return Moved::Ok;
282 }
283 let addr = self.map.find(src).expect("it was live a line ago");
291 if value::kind(self.map.value_at(addr)) == Kind::String {
292 let mut bytes = std::mem::take(&mut self.scratch);
299 bytes.clear();
300 bytes.extend_from_slice(self.map.value_at(addr));
301 self.free_body(dst);
302 self.write_rec(dst, bytes.len(), |out| {
303 out.copy_from_slice(&bytes);
304 });
305 self.scratch = bytes;
306 return Moved::Ok;
307 }
308 let rec = self.export(src).expect("it was live a line ago");
311 self.import(dst, rec);
312 Moved::Ok
313 }
314
315 pub fn touch<'k>(&mut self, keys: impl Iterator<Item = &'k [u8]>) -> usize {
322 keys.filter(|key| self.exists(key)).count()
323 }
324
325 fn write_slot(&mut self, key: &[u8], kind: Kind, slot: u32, at: Option<u64>) {
329 let len = value::slot_record_len(at.is_some());
330 self.write_rec(key, len, |out| {
331 value::write_slot_record(out, kind, slot, at);
332 });
333 }
334}
335
336#[must_use]
342pub fn no_such_key() -> yo_common::Error {
343 yo_common::Error::new(yo_common::Code::Invalid, "no such key")
344}
345
346impl Moved {
351 pub fn found(self) -> Result<Moved> {
358 match self {
359 Moved::Missing => Err(no_such_key()),
360 other => Ok(other),
361 }
362 }
363}
364
365#[cfg(test)]
366mod tests {
367 use super::*;
368 use crate::Clock;
369 use crate::End;
370 use crate::zsets::ZAdd;
371
372 fn db() -> Keyspace {
373 Keyspace::with_clock(Clock::fixed(1_000_000))
374 }
375
376 fn members(d: &mut Keyspace, key: &[u8]) -> Vec<String> {
377 let mut out: Vec<String> = d
378 .smembers(key)
379 .expect("a set")
380 .expect("a key")
381 .map(|m| String::from_utf8(m.to_vec()).expect("utf8 in these tests"))
382 .collect();
383 out.sort();
384 out
385 }
386
387 fn put(d: &mut Keyspace, key: &[u8], val: &[u8]) {
388 d.set_plain(key, val).expect("room for a record");
389 }
390
391 fn read(d: &mut Keyspace, key: &[u8]) -> Vec<u8> {
392 d.get(key).expect("a string").expect("there").to_vec()
393 }
394
395 #[test]
396 fn a_rename_moves_the_value_and_leaves_nothing_behind() {
397 let mut d = db();
398 put(&mut d, b"a", b"v1");
399
400 assert_eq!(d.rename(b"a", b"b", false), Moved::Ok);
401 assert!(!d.exists(b"a"));
402 assert_eq!(read(&mut d, b"b"), b"v1");
403 }
404
405 #[test]
409 fn a_rename_does_not_allocate_to_carry_the_record_across() {
410 let mut d = db();
411 put(&mut d, b"a", b"v1");
412 for _ in 0..4 {
415 assert_eq!(d.rename(b"a", b"b", false), Moved::Ok);
416 assert_eq!(d.rename(b"b", b"a", false), Moved::Ok);
417 }
418 let (_, allocs) = crate::tally::counted(|| {
419 for _ in 0..50 {
420 assert_eq!(d.rename(b"a", b"b", false), Moved::Ok);
421 assert_eq!(d.rename(b"b", b"a", false), Moved::Ok);
422 }
423 });
424 assert_eq!(allocs, 0, "rename allocated {allocs} times in a hundred");
425 assert_eq!(read(&mut d, b"a"), b"v1");
426 }
427
428 #[test]
429 fn a_rename_with_no_source_is_the_one_error_in_this_file() {
430 let mut d = db();
431 assert_eq!(d.rename(b"a", b"b", false), Moved::Missing);
432 assert_eq!(d.rename(b"a", b"b", true), Moved::Missing);
433 assert_eq!(
434 d.copy(b"a", b"b", false),
435 Moved::Missing,
436 "copy just says 0"
437 );
438 }
439
440 #[test]
441 fn a_rename_carries_the_source_deadline_and_drops_the_destination_one() {
442 let mut d = db();
443 put(&mut d, b"a", b"v1");
444 d.set_expiry(b"a", Some(2_000_000));
445 put(&mut d, b"b", b"v2");
446 d.set_expiry(b"b", Some(1_500_000));
447
448 assert_eq!(d.rename(b"a", b"b", false), Moved::Ok);
449 assert_eq!(d.deadline_of(b"b"), crate::Ask::At(2_000_000));
450 }
451
452 #[test]
453 fn renaming_a_key_onto_itself_keeps_it_and_renamenx_refuses() {
454 let mut d = db();
455 put(&mut d, b"a", b"v1");
456 d.set_expiry(b"a", Some(2_000_000));
457
458 assert_eq!(d.rename(b"a", b"a", false), Moved::Ok);
459 assert_eq!(read(&mut d, b"a"), b"v1");
460 assert_eq!(d.deadline_of(b"a"), crate::Ask::At(2_000_000));
461 assert_eq!(d.rename(b"a", b"a", true), Moved::Taken);
462 }
463
464 #[test]
465 fn renamenx_writes_over_nothing() {
466 let mut d = db();
467 put(&mut d, b"a", b"v1");
468 put(&mut d, b"b", b"v2");
469
470 assert_eq!(d.rename(b"a", b"b", true), Moved::Taken);
471 assert_eq!(read(&mut d, b"a"), b"v1");
472 assert_eq!(read(&mut d, b"b"), b"v2");
473 assert_eq!(d.rename(b"a", b"c", true), Moved::Ok);
474 assert!(!d.exists(b"a"));
475 }
476
477 #[test]
478 fn renaming_a_set_moves_the_slot_and_not_the_members() {
479 let mut d = db();
480 d.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
481 .expect("a set");
482 let before = d.memory_bytes();
483
484 assert_eq!(d.rename(b"s", b"t", false), Moved::Ok);
485 assert_eq!(members(&mut d, b"t"), ["m1", "m2"]);
486 assert_eq!(d.kind_of(b"t"), Some(Kind::Set));
487 assert!(!d.exists(b"s"));
488 assert!(
491 d.memory_bytes().abs_diff(before) < 64,
492 "the members were not copied"
493 );
494 }
495
496 #[test]
497 fn renaming_over_a_set_frees_the_set_that_was_there() {
498 let mut d = db();
499 d.sadd(b"s", [b"m1".as_ref()].into_iter()).expect("a set");
500 d.sadd(b"t", [b"m2".as_ref()].into_iter()).expect("a set");
501 assert_eq!(d.sets.len(), 2);
502
503 assert_eq!(d.rename(b"s", b"t", false), Moved::Ok);
504 assert_eq!(d.sets.len(), 1, "the destination's body went with it");
505 assert_eq!(members(&mut d, b"t"), ["m1"]);
506 }
507
508 #[test]
509 fn a_copy_is_a_second_value_and_not_a_second_name() {
510 let mut d = db();
511 d.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
512 .expect("a set");
513
514 assert_eq!(d.copy(b"s", b"t", false), Moved::Ok);
515 d.sadd(b"t", [b"m3".as_ref()].into_iter()).expect("a set");
516 assert_eq!(
517 members(&mut d, b"s"),
518 ["m1", "m2"],
519 "the original is intact"
520 );
521 assert_eq!(members(&mut d, b"t"), ["m1", "m2", "m3"]);
522 }
523
524 #[test]
525 fn a_copy_refuses_a_destination_it_was_not_told_it_could_have() {
526 let mut d = db();
527 put(&mut d, b"a", b"v1");
528 put(&mut d, b"b", b"v2");
529
530 assert_eq!(d.copy(b"a", b"b", false), Moved::Taken);
531 assert_eq!(read(&mut d, b"b"), b"v2");
532 assert_eq!(d.copy(b"a", b"b", true), Moved::Ok);
533 assert_eq!(read(&mut d, b"b"), b"v1");
534 }
535
536 #[test]
539 fn a_copy_of_a_string_does_not_allocate() {
540 let mut d = db();
541 put(&mut d, b"a", b"a-value-of-some-length");
542 for _ in 0..4 {
545 assert_eq!(d.copy(b"a", b"b", true), Moved::Ok);
546 }
547 let (_, allocs) = crate::tally::counted(|| {
548 for _ in 0..50 {
549 assert_eq!(d.copy(b"a", b"b", true), Moved::Ok);
550 }
551 });
552 assert_eq!(allocs, 0, "copy allocated {allocs} times in fifty");
553 assert_eq!(read(&mut d, b"b"), b"a-value-of-some-length");
554 }
555
556 #[test]
561 fn a_copy_onto_itself_leaves_the_key_alone() {
562 let mut d = db();
563 d.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
564 .expect("a set");
565
566 assert_eq!(d.copy(b"s", b"s", false), Moved::Taken);
567 assert_eq!(d.copy(b"s", b"s", true), Moved::Ok);
568 assert_eq!(members(&mut d, b"s"), ["m1", "m2"]);
569 assert_eq!(d.sets.len(), 1, "no second body was made or lost");
570 }
571
572 #[test]
573 fn a_copy_carries_the_deadline() {
574 let mut d = db();
575 put(&mut d, b"a", b"v1");
576 d.set_expiry(b"a", Some(2_000_000));
577
578 assert_eq!(d.copy(b"a", b"b", false), Moved::Ok);
579 assert_eq!(d.deadline_of(b"b"), crate::Ask::At(2_000_000));
580 assert_eq!(d.deadline_of(b"a"), crate::Ask::At(2_000_000));
581 }
582
583 #[test]
584 fn a_destination_that_has_already_gone_counts_as_free() {
585 let mut d = db();
586 put(&mut d, b"a", b"v1");
587 put(&mut d, b"b", b"v2");
588 d.set_expiry(b"b", Some(999_999));
589
590 assert_eq!(d.copy(b"a", b"b", false), Moved::Ok, "b was already gone");
591 assert_eq!(read(&mut d, b"b"), b"v1");
592 }
593
594 #[test]
595 fn a_source_that_has_already_gone_is_not_a_source() {
596 let mut d = db();
597 put(&mut d, b"a", b"v1");
598 d.set_expiry(b"a", Some(999_999));
599
600 assert_eq!(d.rename(b"a", b"b", false), Moved::Missing);
601 assert_eq!(d.copy(b"a", b"b", false), Moved::Missing);
602 }
603
604 #[test]
605 fn a_record_taken_out_of_a_database_outlives_it() {
606 let mut from = db();
607 from.sadd(b"s", [b"m1".as_ref(), b"m2".as_ref()].into_iter())
608 .expect("a set");
609 let rec = from.export(b"s").expect("a record");
610 assert_eq!(rec.kind(), Kind::Set);
611 from.clear();
612
613 let mut into = db();
614 into.import(b"s", rec);
615 assert_eq!(members(&mut into, b"s"), ["m1", "m2"]);
616 }
617
618 #[test]
619 fn importing_over_a_body_does_not_leave_it_in_the_slab() {
620 let mut d = db();
621 d.sadd(b"s", [b"m1".as_ref()].into_iter()).expect("a set");
622 d.sadd(b"t", [b"m2".as_ref()].into_iter()).expect("a set");
623 let rec = d.export(b"s").expect("a record");
624
625 d.import(b"t", rec);
626 assert_eq!(d.sets.len(), 2, "s and t, and not the one t used to hold");
627 assert_eq!(members(&mut d, b"t"), ["m1"]);
628 }
629
630 #[test]
631 fn importing_a_string_over_a_set_frees_the_set() {
632 let mut d = db();
633 put(&mut d, b"a", b"v1");
634 d.sadd(b"s", [b"m1".as_ref()].into_iter()).expect("a set");
635 assert_eq!(d.sets.len(), 1);
636
637 assert_eq!(d.copy(b"a", b"s", true), Moved::Ok);
638 assert_eq!(d.sets.len(), 0, "the set went when the string arrived");
639 assert_eq!(d.kind_of(b"s"), Some(Kind::String));
640 }
641
642 #[test]
654 fn a_list_can_be_copied_and_the_copy_is_its_own() {
655 let mut d = db();
656 d.push(b"l", End::Left, [b"a".as_ref(), b"b".as_ref()].into_iter())
657 .expect("a list");
658
659 assert_eq!(d.copy(b"l", b"m", false), Moved::Ok);
660 assert_eq!(d.kind_of(b"m"), Some(Kind::List));
661 assert_eq!(d.llen(b"m").expect("a list"), 2);
662
663 d.push(b"m", End::Left, [b"c".as_ref()].into_iter())
664 .expect("a list");
665 assert_eq!(d.llen(b"l").expect("a list"), 2, "the source did not grow");
666 assert_eq!(d.llen(b"m").expect("a list"), 3);
667 }
668
669 #[test]
671 fn a_zset_can_be_copied_and_the_copy_is_its_own() {
672 let mut d = db();
673 d.zadd(b"z", [(1.0, b"m1".as_ref())].into_iter(), ZAdd::default())
674 .expect("a zset");
675
676 assert_eq!(d.copy(b"z", b"y", false), Moved::Ok);
677 assert_eq!(d.kind_of(b"y"), Some(Kind::Zset));
678 assert_eq!(d.zscore(b"y", b"m1").expect("a zset"), Some(1.0));
679
680 d.zadd(b"y", [(2.0, b"m2".as_ref())].into_iter(), ZAdd::default())
681 .expect("a zset");
682 assert_eq!(d.zcard(b"z").expect("a zset"), 1, "the source did not grow");
683 assert_eq!(d.zcard(b"y").expect("a zset"), 2);
684 }
685
686 #[test]
693 fn copying_over_a_list_frees_the_list() {
694 let mut d = db();
695 put(&mut d, b"a", b"v1");
696 d.push(b"l", End::Left, [b"x".as_ref()].into_iter())
697 .expect("a list");
698
699 assert_eq!(d.copy(b"a", b"l", true), Moved::Ok);
700 assert_eq!(d.kind_of(b"l"), Some(Kind::String));
701 assert_eq!(read(&mut d, b"l"), b"v1");
702 }
703
704 #[test]
705 fn touch_counts_the_way_exists_counts() {
706 let mut d = db();
707 put(&mut d, b"a", b"v1");
708 put(&mut d, b"b", b"v2");
709
710 assert_eq!(d.touch([b"a".as_ref()].into_iter()), 1);
711 assert_eq!(d.touch([b"a".as_ref(), b"b".as_ref()].into_iter()), 2);
712 assert_eq!(d.touch([b"a".as_ref(), b"a".as_ref()].into_iter()), 2);
713 assert_eq!(d.touch([b"a".as_ref(), b"z".as_ref()].into_iter()), 1);
714 assert_eq!(d.touch([b"z".as_ref()].into_iter()), 0);
715 }
716}