Skip to main content

yo_index/
map.rs

1//! The raw map: an index and an arena wired together.
2//!
3//! This is the smallest thing that is actually a key value store, and it is the
4//! thing M0's exit gate measures against aki's `f1raw` numbers. There is no
5//! record header yet beyond two lengths, no TTL, no type byte, no version. All
6//! of that arrives in M1 and replaces [`Record`] without the index noticing,
7//! which is the point of keeping the two crates apart.
8//!
9//! Layout of one record in the arena:
10//!
11//! ```text
12//! +--------+--------+-----------+-------------+
13//! | klen   | vlen   | key bytes | value bytes |
14//! | u32 LE | u32 LE | klen      | vlen        |
15//! +--------+--------+-----------+-------------+
16//! ```
17//!
18//! Key and value live in one allocation so that a hit is one cache miss for the
19//! bucket and one for the record, not three.
20
21use crate::index::{Index, Keys};
22use crate::scan::Cursor;
23use crate::tagged::Tagged;
24use yo_arena::Arena;
25use yo_common::{Addr, Space, bytes_eq, wyhash};
26
27/// Bytes of length prefix in front of a record.
28const HDR: usize = 8;
29
30/// The least a single [`RawMap::compact_step`] walks.
31///
32/// A segment is two megabytes and evacuating one in a single call was a stop
33/// the world pause in the middle of a batch. At 64 byte values that is around
34/// twenty six thousand records, each one an index probe, a copy and an index
35/// write, and the replies behind it wait for all of them. It is why the write
36/// rows had a p99 of 3.9 milliseconds against Redis at 0.8 while the p50 was
37/// in line: the median command paid nothing and one command in a few thousand
38/// paid for the whole segment.
39///
40/// Sixty four kilobytes is a thirty second of a segment, which puts the worst
41/// call at a few hundred records. Smaller would be smoother and would spend
42/// more of the total on the fixed cost of picking up where the last call left
43/// off; this is the smallest size at which that overhead is still noise.
44///
45/// The budget is spent on how far the cursor moves and not on how many records
46/// move, because a segment can be entirely dead. Charging only for records
47/// that move would let one call walk two megabytes of headers for free, which
48/// is the pause this exists to prevent, just without the copying.
49const EVAC_FLOOR: usize = 64 * 1024;
50
51/// The most, which is a whole segment.
52///
53/// The cap is here so that the scaling below has an end, not because a segment
54/// is a good amount of work to do at once. Reaching it means the collector is
55/// sixteen times past the line it starts at, at which point the pause is the
56/// smaller problem.
57const EVAC_CEILING: usize = yo_arena::SEGMENT_SIZE;
58
59/// A segment that is partway through being evacuated, and how far it got.
60#[derive(Clone, Copy)]
61struct Evac {
62    seg: usize,
63    off: usize,
64}
65
66struct Record;
67
68impl Record {
69    #[inline]
70    fn lens(bytes: &[u8]) -> (usize, usize) {
71        let k = u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) as usize;
72        let v = u32::from_le_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]) as usize;
73        (k, v)
74    }
75}
76
77/// Arena backed record access, which is what the index probes through.
78struct Records<'a> {
79    arena: &'a Arena,
80}
81
82impl Keys for Records<'_> {
83    #[inline]
84    fn hash_at(&self, addr: Addr) -> u64 {
85        let (klen, _) = Record::lens(self.arena.get(addr, HDR));
86        let bytes = self.arena.get(addr, HDR + klen);
87        wyhash(&bytes[HDR..], 0)
88    }
89
90    #[inline]
91    fn eq_at(&self, addr: Addr, key: &[u8]) -> bool {
92        let bytes = self.arena.get(addr, HDR);
93        let (klen, _) = Record::lens(bytes);
94        if klen != key.len() {
95            return false;
96        }
97        let bytes = self.arena.get(addr, HDR + klen);
98        bytes_eq(&bytes[HDR..], key)
99    }
100}
101
102/// A single shard's key value map: bytes in, bytes out, nothing else.
103///
104/// Not `Sync`, and deliberately so. One of these belongs to one shard thread
105/// and is reached through `ShardLocal`, which is `05` section 1's whole
106/// argument: single ownership means no atomics on the hot path.
107///
108/// ```
109/// let mut m = yo_index::RawMap::new();
110/// assert_eq!(m.set(b"k", b"v"), None);
111/// assert_eq!(m.get(b"k"), Some(&b"v"[..]));
112/// assert_eq!(m.set(b"k", b"w").is_some(), true);
113/// assert_eq!(m.get(b"k"), Some(&b"w"[..]));
114/// assert_eq!(m.del(b"k"), true);
115/// assert_eq!(m.get(b"k"), None);
116/// ```
117pub struct RawMap {
118    index: Index,
119    arena: Arena,
120    /// Where the last `compact_step` stopped, if it stopped partway.
121    evac: Option<Evac>,
122    /// How many times anything in here has been written to.
123    ///
124    /// A caller that resolved a key once and wants to skip resolving it again
125    /// needs to know whether anything could have moved in between, and the
126    /// honest answer is any write at all. Every method that takes `&mut self`
127    /// bumps this, including the in place ones, so the question a caller asks is
128    /// "has this map been written since" and not "has this map been written in a
129    /// way I thought would matter".
130    ///
131    /// It lives here rather than in the caller because there are eleven places
132    /// in `yo-kv` that write to a map and one place here that could be missed,
133    /// and a missed invalidation is a stale answer rather than a slow one.
134    ///
135    /// [`RawMap::value_at_mut`] is the one exception and it is argued for where
136    /// it is written. Everything else, including the in place ones, bumps this.
137    writes: u64,
138    /// The records the caller marked when it wrote them.
139    ///
140    /// A second index of a subset of the keys, which exists so that a caller
141    /// looking for one of them does not have to walk past the ones it is not
142    /// looking for. The only thing that uses it is expiry: a key with a deadline
143    /// is rare in most databases, and both the active expire cycle and the
144    /// `volatile-*` eviction policies were sampling the whole map to find one.
145    ///
146    /// It is here and not in `yo-kv` because this is the only thing that knows
147    /// where a record is. An overwrite can move one, a delete takes one away,
148    /// and compaction moves them between segments, and all three are in this
149    /// file. A set of addresses kept anywhere else would go stale on the third.
150    ///
151    /// What "marked" means is entirely the caller's business. This holds
152    /// addresses and has never heard of a deadline.
153    tagged: Tagged,
154}
155
156impl RawMap {
157    /// An empty map.
158    pub fn new() -> RawMap {
159        RawMap {
160            index: Index::new(),
161            arena: Arena::new(),
162            evac: None,
163            writes: 0,
164            tagged: Tagged::new(),
165        }
166    }
167
168    /// How many times this map has been written to.
169    ///
170    /// Two reads of this with the same value either side of some work mean
171    /// nothing in the map moved, so an address or a slot resolved before the
172    /// first read is still the right one after the second. It never goes
173    /// backwards, including across [`RawMap::clear`].
174    #[inline]
175    #[must_use]
176    pub const fn writes(&self) -> u64 {
177        self.writes
178    }
179
180    /// How many keys are stored.
181    #[inline]
182    pub fn len(&self) -> usize {
183        self.index.len()
184    }
185
186    /// Whether the map is empty.
187    #[inline]
188    pub fn is_empty(&self) -> bool {
189        self.index.is_empty()
190    }
191
192    /// Throw everything away and give the memory back.
193    ///
194    /// A fresh index and a fresh arena rather than a walk that deletes each key
195    /// in turn. Deleting one at a time would leave an arena the size of the
196    /// data that used to be in it and an index still grown to fit it, and the
197    /// one thing a client that has just said `FLUSHALL` is entitled to expect is
198    /// the memory back.
199    pub fn clear(&mut self) {
200        // Carried across the reset and bumped, because a counter that went back
201        // to zero here could land on a value a memo was already holding and
202        // read as "nothing moved" on the one call where everything did.
203        let writes = self.writes;
204        *self = RawMap::new();
205        self.writes = writes + 1;
206    }
207
208    /// The hash this map files `key` under.
209    ///
210    /// Public because the batch walk in `04` section 3 hashes on the first walk
211    /// and looks up on the second, and the alternative is hashing every key
212    /// twice to keep the seed a private detail.
213    #[inline]
214    #[must_use]
215    pub fn hash_of(key: &[u8]) -> u64 {
216        wyhash(key, 0)
217    }
218
219    /// Ask the cache for the bucket `hash` will be looked up in.
220    #[inline]
221    pub fn prefetch(&self, hash: u64) {
222        self.index.prefetch(hash);
223    }
224
225    /// The value stored under `key`.
226    #[inline]
227    pub fn get(&self, key: &[u8]) -> Option<&[u8]> {
228        self.get_hashed(Self::hash_of(key), key)
229    }
230
231    /// The value stored under `key`, whose hash the caller already has.
232    ///
233    /// The second walk's entry point. `hash` has to be [`RawMap::hash_of`] of
234    /// this key: a hash from somewhere else is not unsafe, it just misses.
235    #[inline]
236    pub fn get_hashed(&self, hash: u64, key: &[u8]) -> Option<&[u8]> {
237        let addr = self.index.get(hash, key, &Records { arena: &self.arena })?;
238        Some(self.value_at(addr))
239    }
240
241    /// Where `key`'s record is, for a caller that has to look at it twice.
242    ///
243    /// A `GET` has to know whether the key is past its deadline before it can
244    /// answer, and then has to read the value it just decided about. Asking
245    /// [`RawMap::get`] twice is two hashes and two probes for one record, and a
246    /// probe is the expensive half of a command. This hands back the address
247    /// instead, and [`RawMap::value_at`] reads it with no probe at all.
248    ///
249    /// The address is good until the next write to this map. Anything that
250    /// inserts, deletes or compacts can move a record, and an address held
251    /// across one of those reads whatever is at that spot now. Hold it for the
252    /// length of one command and no longer.
253    #[inline]
254    pub fn find(&self, key: &[u8]) -> Option<Addr> {
255        self.find_hashed(Self::hash_of(key), key)
256    }
257
258    /// [`RawMap::find`] for a caller that already hashed the key.
259    #[inline]
260    pub fn find_hashed(&self, hash: u64, key: &[u8]) -> Option<Addr> {
261        self.index.get(hash, key, &Records { arena: &self.arena })
262    }
263
264    /// The value at an address this map handed out, with no probe.
265    ///
266    /// See [`RawMap::find`] for how long an address is worth holding.
267    #[inline]
268    #[must_use]
269    pub fn value_at(&self, addr: Addr) -> &[u8] {
270        let (klen, vlen) = Record::lens(self.arena.get(addr, HDR));
271        &self.arena.get(addr, HDR + klen + vlen)[HDR + klen..]
272    }
273
274    /// The value at an address, to be overwritten in place, without counting as
275    /// a write.
276    ///
277    /// This is the one method taking a mutable borrow that leaves
278    /// [`RawMap::writes`] where it was, and that is a deliberate exception to
279    /// the rule stated on the counter rather than an oversight in it.
280    ///
281    /// It is sound because nothing moves. The record already exists, the caller
282    /// already holds its address, there is no allocation and no index write, so
283    /// every address and every number read out of a record before the call is
284    /// still right afterwards. That is a stronger guarantee than the counter is
285    /// asking about, and it is one this method can actually make.
286    ///
287    /// It exists because the conservative answer costs more here than it
288    /// protects. The eviction clock is written back on nearly every read, under
289    /// eight of the ten policies including the default, so counting it as a write
290    /// would invalidate the caller's memo on every single command rather than on
291    /// every write. That is a measured nineteen nanoseconds a command on single
292    /// key `SADD`, given up to avoid thinking once about three bytes written
293    /// inside a record that is not going anywhere.
294    ///
295    /// The length cannot change, for the same reason it cannot in
296    /// [`RawMap::value_mut`], and an address is only good until the next real
297    /// write, for the same reason it is in [`RawMap::find`].
298    #[inline]
299    pub fn value_at_mut(&mut self, addr: Addr) -> &mut [u8] {
300        let (klen, vlen) = Record::lens(self.arena.get(addr, HDR));
301        &mut self.arena.get_mut(addr, HDR + klen + vlen)[HDR + klen..]
302    }
303
304    /// The value stored under `key`, to be overwritten where it lies.
305    ///
306    /// The length cannot change, which is the whole reason this is safe to
307    /// offer. `INCR` on an integer encoded string is a probe, an add and a
308    /// store, and the store is eight bytes back into the record it came from
309    /// (`08` section 2). Going through [`RawMap::set`] instead would write a
310    /// fresh record and free the old one on every increment, which is an arena
311    /// append and a dead byte per operation for a value whose size never moves.
312    ///
313    /// There is no reader to tear. A map belongs to one shard thread and is not
314    /// `Sync`, so the only code that can observe a half written value is the
315    /// code doing the writing. When a replica stream or a snapshot reader starts
316    /// walking the arena from another thread, this becomes an epoch question and
317    /// the write becomes an install rather than an overwrite.
318    #[inline]
319    pub fn value_mut(&mut self, key: &[u8]) -> Option<&mut [u8]> {
320        self.value_mut_hashed(Self::hash_of(key), key)
321    }
322
323    /// [`RawMap::value_mut`] for a caller that already hashed the key.
324    #[inline]
325    pub fn value_mut_hashed(&mut self, hash: u64, key: &[u8]) -> Option<&mut [u8]> {
326        self.writes += 1;
327        let addr = self.index.get(hash, key, &Records { arena: &self.arena })?;
328        let (klen, vlen) = Record::lens(self.arena.get(addr, HDR));
329        Some(&mut self.arena.get_mut(addr, HDR + klen + vlen)[HDR + klen..])
330    }
331
332    /// Store `val` under `key`, returning the length of the value it replaced.
333    pub fn set(&mut self, key: &[u8], val: &[u8]) -> Option<usize> {
334        self.set_with(
335            key,
336            val.len(),
337            |_| {},
338            |buf| {
339                buf.copy_from_slice(val);
340                false
341            },
342        )
343    }
344
345    /// The largest record this map can store, key and value and header together.
346    ///
347    /// A value past this belongs in the log region rather than the arena, which
348    /// is `06` section 2's business and not this crate's.
349    #[inline]
350    #[must_use]
351    pub const fn max_record() -> usize {
352        yo_arena::MAX_ALLOC
353    }
354
355    /// Bytes of record header in front of the key.
356    #[inline]
357    #[must_use]
358    pub const fn header_len() -> usize {
359        HDR
360    }
361
362    /// Store a `vlen` byte value under `key`, written by `fill`.
363    ///
364    /// The same thing [`RawMap::set`] does, except that the caller writes
365    /// straight into the record instead of building the value somewhere else
366    /// first and having it copied in. A string with a one byte encoding tag in
367    /// front of it would otherwise be assembled in a scratch buffer and then
368    /// memcpy'd again, and two copies for one `SET` is one too many on a path
369    /// that is trying to be ten times faster than Redis.
370    ///
371    /// `fill` is handed exactly `vlen` bytes of uninitialised-looking storage.
372    /// It is arena memory that has been handed out before and freed, so its
373    /// contents are arbitrary and every byte of it must be written. What it
374    /// answers is whether this record should be marked, which is what
375    /// [`RawMap::sample_tagged`] later draws from. A caller with no use for that
376    /// answers `false` and pays a branch.
377    ///
378    /// `peek` is handed the value that was already under `key`, if there was
379    /// one, before anything is written over it. It exists because the caller
380    /// keeps counts that depend on what the old value was, and this is the only
381    /// place those bytes can be read for free: both paths through here have
382    /// already loaded the old record's header to find out how long it is, so the
383    /// value is in cache and would otherwise cost a second lookup to see. A
384    /// caller with nothing to ask passes an empty closure and pays nothing.
385    ///
386    /// # Panics
387    ///
388    /// If the whole record would exceed [`RawMap::max_record`].
389    pub fn set_with<P, F>(&mut self, key: &[u8], vlen: usize, peek: P, fill: F) -> Option<usize>
390    where
391        P: FnOnce(&[u8]),
392        F: FnOnce(&mut [u8]) -> bool,
393    {
394        self.writes += 1;
395        assert!(key.len() <= u32::MAX as usize, "key too long");
396        assert!(vlen <= u32::MAX as usize, "value too long");
397        let total = HDR + key.len() + vlen;
398        let h = wyhash(key, 0);
399
400        // A key that is already here, in a record exactly the size the new value
401        // needs, is written over where it lies. No allocation, no dead bytes, no
402        // index write, and nothing for compaction to collect later.
403        //
404        // This used to say the in place path had to wait for epochs, because a
405        // reader that had already resolved the address would see a torn value.
406        // That was never a rule this map kept: `value_mut` is the same write and
407        // `INCR` has been doing it since the day it was written, for the same
408        // reason given there. A map belongs to one shard thread and is not
409        // `Sync`, so the only code that can see a half written value is the code
410        // writing it. When a replica stream or a snapshot reader starts walking
411        // the arena from another thread, both of these become an install rather
412        // than an overwrite, together.
413        //
414        // Exactly the size and not merely small enough. A shorter value in a
415        // longer record would leave the header disagreeing with the space the
416        // record occupies, and compaction walks a segment by stepping over each
417        // record by the length in its header, so the walk would land in the
418        // middle of the next one.
419        //
420        // Overwriting a key with a value the same size as the last one is what
421        // half of the world's caches do, and it is what every SET benchmark
422        // does. On gamingpc it was 25 percent of SET throughput at pipeline 16
423        // and 37 percent of MSET, all of it spent making garbage and then
424        // collecting it.
425        if let Some(addr) = self.index.get(h, key, &Records { arena: &self.arena }) {
426            let (klen, old_vlen) = Record::lens(self.arena.get(addr, HDR));
427            debug_assert_eq!(klen, key.len(), "the index matched a different key");
428            // Before `fill`, because the in place path writes over exactly the
429            // bytes `peek` is being handed. Once, and here rather than next to
430            // the free below, because this is the branch that knows the key was
431            // there and both paths out of it go past this line.
432            peek(&self.arena.get(addr, HDR + klen + old_vlen)[HDR + klen..]);
433            if old_vlen == vlen {
434                let rec = self.arena.get_mut(addr, total);
435                let tag = fill(&mut rec[HDR + klen..]);
436                // The record did not move, so this is the only thing that can
437                // have changed about where it stands: `PERSIST` on a key whose
438                // value is the same length is exactly this branch.
439                self.retag(addr, tag);
440                return Some(vlen);
441            }
442        }
443
444        let (addr, buf) = self
445            .arena
446            .alloc(total)
447            .expect("record is larger than a segment");
448        buf[0..4].copy_from_slice(&(key.len() as u32).to_le_bytes());
449        buf[4..8].copy_from_slice(&(vlen as u32).to_le_bytes());
450        // The arena hands back a run padded up to its alignment, so index to
451        // `total` rather than to the end of the slice.
452        buf[HDR..HDR + key.len()].copy_from_slice(key);
453        let tag = fill(&mut buf[HDR + key.len()..total]);
454
455        let old = {
456            let recs = Records { arena: &self.arena };
457            self.index.insert(h, key, addr, &recs)
458        };
459        // After the insert and not before, because the address the old record
460        // was at is only known once the index has handed it back, and tagging
461        // the new one first would put both in the set for the width of the call
462        // if they happened to be the same address, which they cannot be, but the
463        // order that does not depend on that is the one to write.
464        if let Some(prev) = old {
465            self.tagged.remove(prev);
466        }
467        if tag {
468            self.tagged.insert(addr);
469        } else {
470            // Nothing to take out. `addr` is a run the arena has just handed
471            // back, and nothing is ever freed while it is still marked: a delete
472            // unmarks before it frees, an overwrite unmarks the record it
473            // replaces on the line above, and compaction moves the mark before
474            // it frees the copy it moved from. So a fresh address is never in
475            // the set, and this is the common path, which is every `SET` on a
476            // database that has any deadline in it at all.
477            debug_assert!(
478                !self.tagged.contains(addr),
479                "the arena handed out an address that is still marked"
480            );
481        }
482        match old {
483            Some(prev) => {
484                let (pk, pv) = Record::lens(self.arena.get(prev, HDR));
485                self.arena.free(prev, HDR + pk + pv);
486                Some(pv)
487            }
488            None => None,
489        }
490    }
491
492    /// Put `addr` in the marked set, or take it out, to match `tag`.
493    ///
494    /// For the in place path, which is the one where the record was already
495    /// there and could already have been marked. It cannot tell whether the mark
496    /// changed without asking, because a deadline is eight bytes in the record
497    /// and a value eight bytes shorter with a deadline is the same length as a
498    /// value without one, so a write that lands in place is not proof that the
499    /// mark stayed put.
500    ///
501    /// On a database where nothing is marked the ask is one comparison against a
502    /// zero length, which is what the overwhelming majority of servers pay.
503    #[inline]
504    fn retag(&mut self, addr: Addr, tag: bool) {
505        if tag {
506            self.tagged.insert(addr);
507        } else {
508            self.tagged.remove(addr);
509        }
510    }
511
512    /// Remove `key`, returning whether it was there.
513    #[inline]
514    pub fn del(&mut self, key: &[u8]) -> bool {
515        self.del_with(key, |_| {})
516    }
517
518    /// Remove `key`, showing its value to `peek` first, and return whether it
519    /// was there.
520    ///
521    /// The sibling of [`RawMap::set_with`], and it exists for the same reason.
522    /// This already reads the record's header to find out how long it is before
523    /// handing the bytes back to the arena, so the value is in cache and a
524    /// caller who keeps a count that depends on what was removed can read it
525    /// here for the price of a closure call. Asking with a [`RawMap::get`] first
526    /// would be a second lookup for a question this one already knows the answer
527    /// to. `peek` is not called when the key was not there.
528    pub fn del_with<P: FnOnce(&[u8])>(&mut self, key: &[u8], peek: P) -> bool {
529        self.writes += 1;
530        let h = wyhash(key, 0);
531        let addr = {
532            let recs = Records { arena: &self.arena };
533            self.index.remove(h, key, &recs)
534        };
535        match addr {
536            Some(a) => {
537                let (k, v) = Record::lens(self.arena.get(a, HDR));
538                peek(&self.arena.get(a, HDR + k + v)[HDR + k..]);
539                self.tagged.remove(a);
540                self.arena.free(a, HDR + k + v);
541                true
542            }
543            None => false,
544        }
545    }
546
547    /// Whether `key` is present.
548    #[inline]
549    pub fn contains(&self, key: &[u8]) -> bool {
550        let h = wyhash(key, 0);
551        self.index.contains(h, key, &Records { arena: &self.arena })
552    }
553
554    /// The key and the value at an address this map handed out.
555    ///
556    /// The pair rather than either one alone, because they are one contiguous
557    /// read: the header says how long the key is and the value starts where the
558    /// key ends, so asking for both costs what asking for one costs.
559    #[inline]
560    #[must_use]
561    pub fn entry_at(&self, addr: Addr) -> (&[u8], &[u8]) {
562        let (klen, vlen) = Record::lens(self.arena.get(addr, HDR));
563        let bytes = self.arena.get(addr, HDR + klen + vlen);
564        (&bytes[HDR..HDR + klen], &bytes[HDR + klen..])
565    }
566
567    /// Walk a batch of the map, and say where the next batch starts.
568    ///
569    /// This is `SCAN`. `budget` is how many entries the caller would like, and
570    /// it is a floor and not a ceiling: the walk stops at the first bucket
571    /// boundary past it, so a batch of ten can come back with fifteen. Redis's
572    /// `COUNT` behaves the same way and for the same reason, which is that a
573    /// bucket is the smallest unit a cursor can name.
574    ///
575    /// A budget of zero still does one bucket, so a caller that keeps passing
576    /// the cursor back always finishes rather than spinning on the same number.
577    ///
578    /// The guarantee, in full: a key that is present for the whole walk is
579    /// handed to `out` at least once. A key added or removed partway through may
580    /// or may not appear, and a key may appear twice. The reasoning is in
581    /// [`Cursor`], and the part worth knowing here is that none of it depends on
582    /// the map holding still between calls.
583    pub fn scan(&self, from: Cursor, budget: usize, mut out: impl FnMut(&[u8], &[u8])) -> Cursor {
584        // The index and the arena are separate fields, so the walk can hold one
585        // and the closure the other. That is what keeps this allocation free:
586        // there is no list of addresses in between.
587        let arena = &self.arena;
588        let mut at = from;
589        let mut seen = 0usize;
590        loop {
591            at = self.index.scan(at, |addr| {
592                let (klen, vlen) = Record::lens(arena.get(addr, HDR));
593                let bytes = arena.get(addr, HDR + klen + vlen);
594                out(&bytes[HDR..HDR + klen], &bytes[HDR + klen..]);
595                seen += 1;
596            });
597            if at.is_end() || seen >= budget {
598                return at;
599            }
600        }
601    }
602
603    /// Entries picked at random, for eviction sampling, until `out` says stop.
604    ///
605    /// The key, the value and the address of each, because a caller choosing a
606    /// victim needs all three: the value to score it, the key to delete it, and
607    /// the address to delete it by without a second probe. `out` answers whether
608    /// to keep going. [`Index::sample`] is where the argument for all of it lives,
609    /// including why the budget is the caller's and why this can hand back
610    /// nothing at all.
611    pub fn sample(&self, r: u64, mut out: impl FnMut(&[u8], &[u8], Addr) -> bool) {
612        let arena = &self.arena;
613        self.index.sample(r, |addr| {
614            let (klen, vlen) = Record::lens(arena.get(addr, HDR));
615            let bytes = arena.get(addr, HDR + klen + vlen);
616            out(&bytes[HDR..HDR + klen], &bytes[HDR + klen..], addr)
617        });
618    }
619
620    /// The index, for stats and for compaction.
621    pub fn index(&self) -> &Index {
622        &self.index
623    }
624
625    /// The arena, for stats and for compaction.
626    pub fn arena(&self) -> &Arena {
627        &self.arena
628    }
629
630    /// Bytes held by index structure plus arena segments.
631    pub fn memory_bytes(&self) -> usize {
632        self.index.memory_bytes()
633            + self.arena.reserved_bytes() as usize
634            + self.tagged.memory_bytes()
635    }
636
637    /// How many records are marked.
638    ///
639    /// Exact, and kept exact by every write path, so a caller can branch on a
640    /// zero here rather than starting a sweep that was never going to find
641    /// anything.
642    #[inline]
643    #[must_use]
644    pub fn tagged_len(&self) -> usize {
645        self.tagged.len()
646    }
647
648    /// Whether the record at `addr` is marked.
649    ///
650    /// For a test and for a debug assertion. Nothing on a hot path asks this:
651    /// the mark is written from the record's own bytes, so anything holding the
652    /// record already knows.
653    #[must_use]
654    pub fn is_tagged(&self, addr: Addr) -> bool {
655        self.tagged.contains(addr)
656    }
657
658    /// Walk marked records from wherever `r` lands, until `out` says stop.
659    ///
660    /// [`RawMap::sample`] for the marked subset, and the reason the subset
661    /// exists. A database of ten million keys where a thousand carry a deadline
662    /// gives the expire cycle a thousand candidates to draw from instead of ten
663    /// million, and the cycle stops costing anything at all in the case that
664    /// matters most, which is the one where the answer is that there is nothing
665    /// to do.
666    pub fn sample_tagged(&self, r: u64, mut out: impl FnMut(&[u8], &[u8], Addr) -> bool) {
667        let arena = &self.arena;
668        self.tagged.sample(r, |addr| {
669            let (klen, vlen) = Record::lens(arena.get(addr, HDR));
670            let bytes = arena.get(addr, HDR + klen + vlen);
671            out(&bytes[HDR..HDR + klen], &bytes[HDR + klen..], addr)
672        });
673    }
674
675    /// Move every live record out of `seg` and into the current segment, then
676    /// put the segment back on the arena's free list.
677    ///
678    /// Copy, rewrite the index entry, done. No forwarding pointers and no read
679    /// barrier, which is the F2 shape from `05` section 3.2 and is what an
680    /// allocation having exactly one referent buys.
681    ///
682    /// The walk is over the segment and not over the index. Both find the same
683    /// records, and the index walk is the one written in the spec, but it reads
684    /// the whole index to compact two megabytes: fine when this only ran in a
685    /// test, wrong once the event loop calls it, because the pause would then
686    /// grow with the size of the database rather than with the size of a
687    /// segment. Walking the segment costs one index probe per record in it and
688    /// does not care how many keys exist elsewhere.
689    ///
690    /// Records sit back to back from the header to the segment's bump, each one
691    /// rounded up to the arena's alignment, and every arena allocation is a
692    /// record, so the next one is always a known distance away. A record is
693    /// live when the index still points at this copy of it, and dead when it
694    /// points somewhere else or at nothing, which is exactly what an overwrite
695    /// and a delete leave behind.
696    ///
697    /// The reclaim at the end is the part that makes the space usable again.
698    /// Moving the records out only makes a segment empty, and an empty segment
699    /// that nothing ever bumps through again is still two megabytes the process
700    /// is holding.
701    pub fn compact_segment(&mut self, seg: usize) -> usize {
702        self.writes += 1;
703        if seg == self.arena.current_segment() {
704            // Its bump is a cursor, not a checkpoint, and reclaiming it would
705            // take the ground out from under the next allocation.
706            return 0;
707        }
708        let (moved, _) = self.evacuate(seg, yo_arena::HEADER_SIZE, usize::MAX);
709        self.arena.reclaim(seg);
710        moved
711    }
712
713    /// Walk `seg` from `from`, moving live records out, and stop once the walk
714    /// has covered `budget` bytes of it. Says how many records moved and where
715    /// to start again.
716    ///
717    /// The record that straddles the budget is finished rather than cut in
718    /// half, so the walk can go a little past what was asked for. The overrun
719    /// is one record and the budget is thousands of bytes.
720    ///
721    /// Nothing here reclaims. A segment is only empty once the walk reaches the
722    /// bump, and the caller is the one that knows whether it did.
723    fn evacuate(&mut self, seg: usize, from: usize, budget: usize) -> (usize, usize) {
724        let base = (seg as u64) << yo_arena::SEGMENT_SHIFT;
725        let bump = self.arena.recorded_bump(seg) as usize;
726        let stop = from.saturating_add(budget).min(bump);
727
728        let mut moved = 0;
729        let mut off = from;
730        while off < stop {
731            let old = Addr::new(Space::Arena, base + off as u64);
732            let (klen, vlen) = Record::lens(self.arena.get(old, HDR));
733            let total = HDR + klen + vlen;
734            off += total.next_multiple_of(yo_arena::ALIGN);
735
736            let hash = {
737                let bytes = self.arena.get(old, HDR + klen);
738                wyhash(&bytes[HDR..], 0)
739            };
740            let live = {
741                let bytes = self.arena.get(old, HDR + klen);
742                let key = &bytes[HDR..];
743                let recs = Records { arena: &self.arena };
744                self.index.get(hash, key, &recs) == Some(old)
745            };
746            if !live {
747                continue;
748            }
749
750            let new = self.arena.copy_within(old, total);
751            let bytes = self.arena.get(new, HDR + klen);
752            let key = &bytes[HDR..];
753            let recs = Records { arena: &self.arena };
754            let ok = self.index.relocate(hash, key, new, &recs);
755            debug_assert!(ok, "compaction lost an entry the index just handed us");
756            // The one place a record moves without anybody writing to it, and
757            // therefore the one place the tagged set would go stale if this line
758            // were not here.
759            if self.tagged.remove(old) {
760                self.tagged.insert(new);
761            }
762            self.arena.free(old, total);
763            moved += 1;
764        }
765        (moved, off)
766    }
767
768    /// How much to walk on this call, given how far behind the collector is.
769    ///
770    /// A fixed budget has to be either a good pause or a good collection rate
771    /// and it cannot be both. At 64 kilobytes a segment takes thirty two calls,
772    /// and a pipelined flood of writes makes garbage faster than one call per
773    /// batch gets it back: measured with variable sized values at pipeline 16,
774    /// the tail came down from 2.6 milliseconds to 1.6 and the process held 18
775    /// MB more, because segments queued up waiting their turn to be walked.
776    ///
777    /// So the floor is what a command can be asked to wait for, and the depth
778    /// of that queue is what says how much more than the floor is needed to
779    /// keep up. One candidate is a store that is keeping up and pays the floor.
780    /// Nine is a store nine segments behind, and it walks nine slices.
781    ///
782    /// The queue and not the dead byte total. Dead bytes were tried first,
783    /// measured against the point compaction starts at, and that ratio cannot
784    /// see a backlog at all: the threshold is a fraction of what the arena
785    /// holds, so a collector that falls behind grows the arena, which raises
786    /// the threshold, which puts the ratio back where it was. It sat at the
787    /// floor through the whole flood and the 18 MB stayed exactly where it was.
788    /// A count of segments has no such denominator.
789    ///
790    /// Linear in the depth and not squared. This is a controller in a loop with
791    /// its own input, and a term that grows faster than the error is how one of
792    /// those starts to oscillate.
793    fn budget(&self) -> usize {
794        let behind = self.arena.candidate_count().max(1);
795        EVAC_FLOOR.saturating_mul(behind).min(EVAC_CEILING)
796    }
797
798    /// Do one bounded slice of compaction, and say how many records moved.
799    ///
800    /// `None` means there was no candidate and there is nothing in flight. It
801    /// is not the same as `Some(0)`, which is a slice that walked only records
802    /// that had already been overwritten: that one made progress and cost
803    /// something, and a caller deciding whether to go round again needs to be
804    /// told so.
805    ///
806    /// This is the whole maintenance contract: a bounded amount of work per
807    /// call, so a caller that runs it once per batch never pays for a full pass
808    /// over the arena and never pays for a whole segment either. Finding out
809    /// there is nothing to do is one comparison against the running dead byte
810    /// total.
811    ///
812    /// A segment takes as many calls as it takes. Each one picks up where the
813    /// last stopped and only the call that reaches the end gives the two
814    /// megabytes back, so the space comes back in one lump at the end while the
815    /// cost of getting it back is spread over the batches in between. That is
816    /// the trade: a segment stays around a little longer than it used to, and
817    /// no single command waits for the whole of it.
818    ///
819    /// The segment in flight is finished before another is chosen, rather than
820    /// asking which segment is worst on every call. Otherwise a segment that is
821    /// three quarters evacuated could be put down in favour of a worse one and
822    /// never picked up, and the arena would fill with segments that are nearly
823    /// empty and never reclaimed.
824    pub fn compact_step(&mut self) -> Option<usize> {
825        self.compact(false)
826    }
827
828    /// One slice of compaction for a store that has run out of room.
829    ///
830    /// The same work, choosing between segments the way
831    /// [`Arena::any_candidate`](yo_arena::Arena::any_candidate) chooses rather
832    /// than the way [`Arena::worst_candidate`](yo_arena::Arena::worst_candidate)
833    /// does, so a segment holding a single dead record is still worth
834    /// evacuating. The reason is written on `any_candidate`: at a memory limit
835    /// the copying is cheaper than the alternative, which is telling a client no.
836    ///
837    /// A segment already in flight is finished first either way, so switching
838    /// between this and [`RawMap::compact_step`] cannot leave a segment half
839    /// evacuated forever.
840    pub fn compact_hard(&mut self) -> Option<usize> {
841        self.compact(true)
842    }
843
844    fn compact(&mut self, hard: bool) -> Option<usize> {
845        self.writes += 1;
846        let (seg, from) = match self.evac {
847            Some(e) => (e.seg, e.off),
848            None => {
849                let pick = if hard {
850                    self.arena.any_candidate()?
851                } else {
852                    self.arena.worst_candidate()?
853                };
854                (pick, yo_arena::HEADER_SIZE)
855            }
856        };
857        if seg == self.arena.current_segment() {
858            self.evac = None;
859            return Some(0);
860        }
861
862        // After the choice and not before it. The count is a walk over the
863        // segment headers, and a store with nothing to collect should not pay
864        // for one on every batch to be told there is nothing to collect.
865        let budget = self.budget();
866        let (moved, off) = self.evacuate(seg, from, budget);
867        if off >= self.arena.recorded_bump(seg) as usize {
868            self.arena.reclaim(seg);
869            self.evac = None;
870        } else {
871            self.evac = Some(Evac { seg, off });
872        }
873        Some(moved)
874    }
875}
876
877impl Default for RawMap {
878    fn default() -> RawMap {
879        RawMap::new()
880    }
881}
882
883#[cfg(test)]
884mod tests {
885    use super::*;
886    use std::collections::{HashMap, HashSet};
887
888    /// `key:` and the index zero padded to twelve digits.
889    ///
890    /// Written out by hand rather than with `format!`, which produces the same
891    /// bytes. Formatting is a lot of machinery for twelve digits, and Miri pays
892    /// per operation rather than per instruction, so under the interpreter one
893    /// `format!` costs a couple of milliseconds. `grows_through_many_splits`
894    /// calls this once per set, get, delete and contains, which is ten thousand
895    /// calls on its own, and that is twenty seconds of the ninety five this
896    /// crate's Miri shard used to take.
897    fn key(i: usize) -> Vec<u8> {
898        let mut k = *b"key:000000000000";
899        let mut n = i;
900        let mut p = k.len() - 1;
901        while n > 0 {
902            k[p] = b'0' + (n % 10) as u8;
903            n /= 10;
904            p -= 1;
905        }
906        k.to_vec()
907    }
908
909    /// `v` and the index, unpadded, which is what `format!("v{i}")` gives.
910    fn val(i: usize) -> Vec<u8> {
911        let mut v = vec![b'v'];
912        if i == 0 {
913            v.push(b'0');
914            return v;
915        }
916        let start = v.len();
917        let mut n = i;
918        while n > 0 {
919            v.push(b'0' + (n % 10) as u8);
920            n /= 10;
921        }
922        v[start..].reverse();
923        v
924    }
925
926    // Miri is a few hundred times slower than the machine, so the counts below
927    // shrink under it. They stay large enough to force directory doublings,
928    // segment splits and overflow chains, which is what these tests are for.
929    // Only the scale goes away, not the coverage.
930    // Three thousand and not fewer. `splits() > 4` is the assertion and the
931    // splits go 1, 1, 2, 3, 3, 5 at 800, 1200, 1500, 2000, 2500 and 3000 keys,
932    // so this is already the smallest count that grows the directory the number
933    // of times the test asks about.
934    #[cfg(miri)]
935    const GROW_N: usize = 3_000;
936    #[cfg(not(miri))]
937    const GROW_N: usize = 200_000;
938
939    #[cfg(miri)]
940    const ADVERSARIAL_N: u64 = 1_000;
941    #[cfg(not(miri))]
942    const ADVERSARIAL_N: u64 = 50_000;
943
944    // Big values so that a handful of records fills a 2 MiB segment and
945    // compaction has something to do without a hundred thousand writes.
946    #[cfg(miri)]
947    const COMPACT_VAL: usize = 65_536;
948    #[cfg(miri)]
949    const COMPACT_N: usize = 200;
950    #[cfg(not(miri))]
951    const COMPACT_VAL: usize = 1024;
952    #[cfg(not(miri))]
953    const COMPACT_N: usize = 8_000;
954
955    #[test]
956    fn set_get_del() {
957        let mut m = RawMap::new();
958        assert!(m.is_empty());
959        assert_eq!(m.set(b"a", b"1"), None);
960        assert_eq!(m.get(b"a"), Some(&b"1"[..]));
961        assert_eq!(m.len(), 1);
962        assert_eq!(m.set(b"a", b"22"), Some(1));
963        assert_eq!(m.get(b"a"), Some(&b"22"[..]));
964        assert_eq!(m.len(), 1);
965        assert!(m.del(b"a"));
966        assert!(!m.del(b"a"));
967        assert_eq!(m.get(b"a"), None);
968        assert!(m.is_empty());
969    }
970
971    #[test]
972    fn a_value_can_be_overwritten_where_it_lies() {
973        let mut m = RawMap::new();
974        m.set(b"n", &7u64.to_le_bytes());
975        m.set(b"other", b"untouched");
976        let before = m.arena().live_bytes();
977
978        let v = m.value_mut(b"n").expect("the key is there");
979        v.copy_from_slice(&8u64.to_le_bytes());
980
981        assert_eq!(m.get(b"n"), Some(&8u64.to_le_bytes()[..]));
982        assert_eq!(m.get(b"other"), Some(&b"untouched"[..]));
983        // The point of the whole method: no second record and nothing dead.
984        assert_eq!(m.arena().live_bytes(), before);
985        assert_eq!(m.len(), 2);
986
987        assert!(m.value_mut(b"missing").is_none());
988    }
989
990    /// A key overwritten with a value the same size stays in the record it is
991    /// already in, and one overwritten with a different size does not.
992    ///
993    /// The first is the shape every SET benchmark and half the world's caches
994    /// have: the same keys, the same value size, over and over. Writing a fresh
995    /// record for each of those makes a dead one to go with it, and compaction
996    /// then spends a quarter of the server's write throughput copying live
997    /// records out from between them.
998    #[test]
999    fn an_overwrite_of_the_same_size_makes_no_garbage() {
1000        let mut m = RawMap::new();
1001        m.set(b"k", b"12345678");
1002        m.set(b"other", b"untouched");
1003        let live = m.arena().live_bytes();
1004        let dead = m.arena().dead_bytes_total();
1005
1006        for i in 0..1000u32 {
1007            let v = format!("{i:08}");
1008            assert_eq!(m.set(b"k", v.as_bytes()), Some(8));
1009        }
1010
1011        assert_eq!(m.get(b"k"), Some(&b"00000999"[..]));
1012        assert_eq!(m.get(b"other"), Some(&b"untouched"[..]));
1013        assert_eq!(m.len(), 2);
1014        assert_eq!(m.arena().live_bytes(), live, "a thousand writes, no growth");
1015        assert_eq!(m.arena().dead_bytes_total(), dead, "and nothing dead");
1016
1017        // A different length cannot go in the same hole, because the record has
1018        // to be as long as its header says it is.
1019        assert_eq!(m.set(b"k", b"123456789"), Some(8));
1020        assert_eq!(m.get(b"k"), Some(&b"123456789"[..]));
1021        assert!(
1022            m.arena().dead_bytes_total() > dead,
1023            "the old record is dead"
1024        );
1025    }
1026
1027    /// An expiring value and a plain one are different record lengths, so the
1028    /// one does not get written over the other.
1029    ///
1030    /// This is the case the in place path has to refuse rather than the case it
1031    /// is for, and it is the one that would corrupt a record if it took it: the
1032    /// value here is a keyspace record, whose deadline is inside the value, so
1033    /// two values of the same visible length are two different record lengths.
1034    #[test]
1035    fn a_longer_value_moves_and_the_index_follows_it() {
1036        let mut m = RawMap::new();
1037        m.set(b"k", b"aaaa");
1038        let first = m
1039            .index()
1040            .get(RawMap::hash_of(b"k"), b"k", &Records { arena: m.arena() });
1041
1042        m.set(b"k", b"aaaaaaaa");
1043        let second = m
1044            .index()
1045            .get(RawMap::hash_of(b"k"), b"k", &Records { arena: m.arena() });
1046
1047        assert_ne!(first, second, "a longer value needs a new record");
1048        assert_eq!(m.get(b"k"), Some(&b"aaaaaaaa"[..]));
1049    }
1050
1051    #[test]
1052    fn empty_key_and_empty_value() {
1053        let mut m = RawMap::new();
1054        m.set(b"", b"");
1055        assert_eq!(m.get(b""), Some(&b""[..]));
1056        m.set(b"x", b"");
1057        assert_eq!(m.get(b"x"), Some(&b""[..]));
1058        assert_eq!(m.len(), 2);
1059    }
1060
1061    #[test]
1062    fn grows_through_many_splits() {
1063        let mut m = RawMap::new();
1064        const N: usize = GROW_N;
1065        for i in 0..N {
1066            m.set(&key(i), &val(i));
1067        }
1068        assert_eq!(m.len(), N);
1069        assert!(
1070            m.index().splits() > 4,
1071            "expected real growth, saw {} splits",
1072            m.index().splits()
1073        );
1074        for i in 0..N {
1075            assert_eq!(
1076                m.get(&key(i)),
1077                Some(val(i).as_slice()),
1078                "lost key {i} after {} splits",
1079                m.index().splits()
1080            );
1081        }
1082        for i in (0..N).step_by(3) {
1083            assert!(m.del(&key(i)), "delete missed key {i}");
1084        }
1085        for i in 0..N {
1086            assert_eq!(
1087                m.contains(&key(i)),
1088                i % 3 != 0,
1089                "wrong presence for key {i}"
1090            );
1091        }
1092    }
1093
1094    #[test]
1095    fn compaction_preserves_everything() {
1096        let mut m = RawMap::new();
1097        // Enough to fill several arena segments with 1 KiB values.
1098        let val = vec![b'z'; COMPACT_VAL];
1099        const N: usize = COMPACT_N;
1100        for i in 0..N {
1101            m.set(&key(i), &val);
1102        }
1103        // Kill half, which pushes the early segments over the dead ratio.
1104        for i in (0..N).step_by(2) {
1105            m.del(&key(i));
1106        }
1107        let candidates = m.arena().compaction_candidates();
1108        assert!(
1109            !candidates.is_empty(),
1110            "expected at least one segment past the dead ratio"
1111        );
1112        for seg in candidates {
1113            m.compact_segment(seg);
1114        }
1115        for i in 0..N {
1116            let want = if i % 2 == 0 { None } else { Some(val.clone()) };
1117            assert_eq!(m.get(&key(i)).map(|v| v.to_vec()), want, "key {i}");
1118        }
1119    }
1120
1121    /// A mark follows its record wherever the record goes.
1122    ///
1123    /// The whole reason the marked set lives in this file. Compaction moves a
1124    /// record to a new address without anybody writing to it, so a set of
1125    /// addresses kept by a caller would be pointing at freed space afterwards,
1126    /// and the sample would read whatever the arena handed out next.
1127    #[test]
1128    fn compaction_carries_the_marks_with_it() {
1129        let mut m = RawMap::new();
1130        let val = vec![b'z'; COMPACT_VAL];
1131        const N: usize = COMPACT_N;
1132        for i in 0..N {
1133            m.set_with(
1134                &key(i),
1135                val.len(),
1136                |_| {},
1137                |b| {
1138                    b.copy_from_slice(&val);
1139                    i % 3 == 0
1140                },
1141            );
1142        }
1143        let want = (0..N).filter(|i| i % 3 == 0).count();
1144        assert_eq!(m.tagged_len(), want);
1145
1146        for i in (0..N).step_by(2) {
1147            m.del(&key(i));
1148        }
1149        let want = (0..N).filter(|i| i % 3 == 0 && i % 2 == 1).count();
1150        assert_eq!(m.tagged_len(), want, "a delete takes the mark with it");
1151
1152        for seg in m.arena().compaction_candidates() {
1153            m.compact_segment(seg);
1154        }
1155        assert_eq!(
1156            m.tagged_len(),
1157            want,
1158            "and compaction moves it rather than losing it"
1159        );
1160
1161        // Every mark points at a record that is still there and is one of the
1162        // ones that was marked, which is what a stale address would fail.
1163        let mut seen = 0;
1164        m.sample_tagged(0, |k, _, addr| {
1165            assert!(m.get(k).is_some(), "a mark on a key that is gone");
1166            let i: usize = std::str::from_utf8(&k[4..]).unwrap().parse().unwrap();
1167            assert!(
1168                i.is_multiple_of(3) && !i.is_multiple_of(2),
1169                "key {i} was never marked"
1170            );
1171            assert!(m.is_tagged(addr));
1172            seen += 1;
1173            true
1174        });
1175        assert_eq!(seen, want);
1176    }
1177
1178    /// A mark goes on and comes off with the record's own bytes, which is how
1179    /// PERSIST works: the value is the same length, so the record does not move
1180    /// and only the mark changes.
1181    #[test]
1182    fn a_mark_goes_on_and_comes_off_in_place() {
1183        let mut m = RawMap::new();
1184        let mark = |m: &mut RawMap, on: bool| {
1185            m.set_with(
1186                b"k",
1187                1,
1188                |_| {},
1189                |b| {
1190                    b[0] = b'v';
1191                    on
1192                },
1193            )
1194        };
1195        mark(&mut m, true);
1196        assert_eq!(m.tagged_len(), 1);
1197        mark(&mut m, true);
1198        assert_eq!(m.tagged_len(), 1, "marking twice is marking once");
1199        mark(&mut m, false);
1200        assert_eq!(m.tagged_len(), 0);
1201        mark(&mut m, true);
1202        assert_eq!(m.tagged_len(), 1);
1203        assert!(m.del(b"k"));
1204        assert_eq!(m.tagged_len(), 0);
1205    }
1206
1207    /// The bug this exists for: overwriting a key writes a new record and only
1208    /// counts the old one dead, so without compaction a server that rewrites
1209    /// the same keys holds every version of every one of them forever. Measured
1210    /// on a real server before this, 400000 sets over 100000 keys came to 742
1211    /// bytes a key for 64 byte values.
1212    #[test]
1213    fn rewriting_the_same_keys_stops_growing() {
1214        let mut m = RawMap::new();
1215        let val = vec![b'z'; COMPACT_VAL];
1216        const N: usize = COMPACT_N;
1217
1218        for i in 0..N {
1219            m.set(&key(i), &val);
1220            m.compact_step();
1221        }
1222        let after_first_pass = m.arena().reserved_bytes();
1223
1224        // Nine more passes over the same keys, writing the same amount of data
1225        // nine more times and keeping exactly as much of it.
1226        for _ in 0..9 {
1227            for i in 0..N {
1228                m.set(&key(i), &val);
1229                m.compact_step();
1230            }
1231        }
1232        let after_ten = m.arena().reserved_bytes();
1233
1234        assert!(
1235            after_ten <= after_first_pass * 2,
1236            "held {after_ten} after ten passes against {after_first_pass} after one, \
1237             which is the grow forever shape"
1238        );
1239        assert!(
1240            after_ten < m.arena().live_bytes() * 2,
1241            "held {after_ten} for {} live, which is more than the ratio allows",
1242            m.arena().live_bytes()
1243        );
1244        for i in 0..N {
1245            assert_eq!(
1246                m.get(&key(i)).map(<[u8]>::to_vec),
1247                Some(val.clone()),
1248                "key {i}"
1249            );
1250        }
1251    }
1252
1253    /// A segment is evacuated over several calls, and it comes back only on the
1254    /// call whose walk reaches the end of it.
1255    ///
1256    /// This is what the budget is for. One call used to copy every live record
1257    /// in two megabytes, around twenty six thousand of them at 64 byte values,
1258    /// and the whole batch of replies queued behind it waited for all of them.
1259    /// That is where a p99 of 3.9 milliseconds on the write rows came from
1260    /// while the p50 was in line with Redis: the median command paid nothing
1261    /// and one command in a few thousand paid for a segment.
1262    ///
1263    /// The loop is also what catches a walk that restarts instead of resuming.
1264    /// A restart would move records and look like progress, and it would spend
1265    /// every call re-walking the dead space it made on the last one, so the
1266    /// cursor would never reach the bump and the segment would never come back.
1267    #[test]
1268    fn a_segment_comes_back_over_several_calls() {
1269        let mut m = RawMap::new();
1270        let val = vec![b'z'; COMPACT_VAL];
1271        const N: usize = COMPACT_N;
1272        for i in 0..N {
1273            m.set(&key(i), &val);
1274        }
1275        // Every other key, so the early segments are well past the dead ratio
1276        // and there is still a live half to copy out.
1277        for i in (0..N).step_by(2) {
1278            m.del(&key(i));
1279        }
1280
1281        let rec = (HDR + key(0).len() + COMPACT_VAL).next_multiple_of(yo_arena::ALIGN);
1282        let per_call = m.budget() / rec + 1;
1283        let free = m.arena().free_segments();
1284
1285        let moved = m.compact_step().expect("half of it is dead");
1286        assert!(
1287            moved <= per_call,
1288            "one call moved {moved} records and the budget is {per_call}"
1289        );
1290        assert_eq!(
1291            m.arena().free_segments(),
1292            free,
1293            "a segment came back before the walk reached the end of it"
1294        );
1295
1296        let mut calls = 1;
1297        while m.arena().free_segments() == free {
1298            m.compact_step()
1299                .expect("the segment in flight is not finished");
1300            calls += 1;
1301            assert!(calls < 1000, "the walk is not getting any further along");
1302        }
1303        assert!(calls > 2, "the whole segment came back in {calls} calls");
1304
1305        for i in 0..N {
1306            let want = if i % 2 == 0 { None } else { Some(val.clone()) };
1307            assert_eq!(m.get(&key(i)).map(<[u8]>::to_vec), want, "key {i}");
1308        }
1309    }
1310
1311    /// A store barely holding any garbage collects nothing until it is asked to.
1312    ///
1313    /// The two ratios are the whole reason [`RawMap::compact_hard`] exists. A
1314    /// server under a memory limit needs the pages back whatever the ratios
1315    /// think of the trade, and a server that is not under one should not pay for
1316    /// copying that buys it a few kilobytes.
1317    #[test]
1318    fn a_store_with_little_dead_in_it_only_collects_when_pushed() {
1319        let mut m = RawMap::new();
1320        let val = vec![b'z'; COMPACT_VAL];
1321        const N: usize = COMPACT_N;
1322        for i in 0..N {
1323            m.set(&key(i), &val);
1324        }
1325        // One key in fifty, which is well under the eighth of everything held
1326        // that compaction normally waits for.
1327        for i in (0..N).step_by(50) {
1328            m.del(&key(i));
1329        }
1330
1331        assert_eq!(m.compact_step(), None, "not worth collecting");
1332        let free = m.arena().free_segments();
1333        let mut calls = 0;
1334        while m.arena().free_segments() == free {
1335            assert!(
1336                m.compact_hard().is_some(),
1337                "there is a segment holding something dead"
1338            );
1339            calls += 1;
1340            assert!(calls < 1000, "the walk is not getting any further along");
1341        }
1342        // Everything still reads back, which is the thing that matters: the
1343        // records that were live in the segment that came back were moved and
1344        // their index entries were moved with them.
1345        for i in 0..N {
1346            let want = if i % 50 == 0 { None } else { Some(val.clone()) };
1347            assert_eq!(m.get(&key(i)).map(<[u8]>::to_vec), want, "key {i}");
1348        }
1349    }
1350
1351    /// The budget grows with how far behind the collector is.
1352    ///
1353    /// A store with one segment waiting pays the floor, which is the pause a
1354    /// command can be asked to wait for. One with a queue of them walks a slice
1355    /// per segment in the queue, which is what keeps a pipelined write flood
1356    /// from outrunning one call per batch and leaving the process holding the
1357    /// segments that never got their turn.
1358    #[test]
1359    fn the_budget_scales_with_the_backlog() {
1360        let mut m = RawMap::new();
1361        let val = vec![b'z'; COMPACT_VAL];
1362        const N: usize = COMPACT_N;
1363        for i in 0..N {
1364            m.set(&key(i), &val);
1365        }
1366        assert_eq!(m.budget(), EVAC_FLOOR, "nothing is waiting yet");
1367
1368        for i in 0..N {
1369            m.del(&key(i));
1370        }
1371        let flooded = m.budget();
1372        assert!(
1373            flooded >= EVAC_FLOOR * m.arena().candidate_count(),
1374            "{} segments are waiting and the budget is {flooded}",
1375            m.arena().candidate_count()
1376        );
1377        assert!(
1378            flooded > EVAC_FLOOR,
1379            "every segment is dead and the budget is still the floor"
1380        );
1381        assert!(flooded <= EVAC_CEILING, "walked past a whole segment");
1382    }
1383
1384    /// A segment that is partway through being evacuated is finished before a
1385    /// worse one is started.
1386    ///
1387    /// Writes keep coming while a segment is being walked and they make dead
1388    /// space elsewhere, so the answer to "which segment is worst" moves around
1389    /// underneath a walk that takes thirty calls. Asking it again on every call
1390    /// would let a segment be put down at nine tenths done in favour of one
1391    /// that is slightly worse, and the arena would fill up with segments that
1392    /// are nearly empty and never reclaimed.
1393    ///
1394    /// Here the first quarter of the keyspace is deleted so that the segment at
1395    /// the front is the only candidate, one call starts on it, and then the
1396    /// back half goes too so that another segment ties with it mid walk. The
1397    /// tie goes to the later segment, so a walk that asked again would move to
1398    /// it and leave the first one part done.
1399    #[test]
1400    fn the_segment_in_flight_is_finished_first() {
1401        let mut m = RawMap::new();
1402        let val = vec![b'z'; COMPACT_VAL];
1403        const N: usize = COMPACT_N;
1404        for i in 0..N {
1405            m.set(&key(i), &val);
1406        }
1407        for i in 0..N / 4 {
1408            m.del(&key(i));
1409        }
1410
1411        let free = m.arena().free_segments();
1412        let first = m.arena().worst_candidate().expect("the front is all dead");
1413        m.compact_step().expect("there is a candidate");
1414
1415        for i in N / 2..N {
1416            m.del(&key(i));
1417        }
1418        let worse = m.arena().worst_candidate().expect("the back is all dead");
1419        assert_ne!(worse, first, "the test needs the answer to have moved");
1420
1421        while m.arena().free_segments() == free {
1422            m.compact_step()
1423                .expect("the segment in flight is not finished");
1424        }
1425        assert!(
1426            m.arena().is_free(first),
1427            "the segment that was in flight is not the one that came back"
1428        );
1429        assert!(
1430            !m.arena().is_free(worse),
1431            "the walk moved to the segment that tied with it partway through"
1432        );
1433    }
1434
1435    /// A segment that compaction emptied is bumped through again rather than
1436    /// sitting there holding two megabytes.
1437    #[test]
1438    fn an_emptied_segment_is_used_again() {
1439        let mut m = RawMap::new();
1440        let val = vec![b'z'; COMPACT_VAL];
1441        const N: usize = COMPACT_N;
1442        for i in 0..N {
1443            m.set(&key(i), &val);
1444        }
1445        for i in (0..N).step_by(2) {
1446            m.del(&key(i));
1447        }
1448
1449        let before = m.arena().segment_count();
1450        let seg = m.arena().worst_candidate().expect("half of it is dead");
1451        m.compact_segment(seg);
1452        assert_eq!(
1453            m.arena().free_segments(),
1454            1,
1455            "the segment did not come back"
1456        );
1457
1458        // Write until the free segment has to be taken, and the count is where
1459        // it was rather than one higher.
1460        for i in N..N * 2 {
1461            m.set(&key(i), &val);
1462            if m.arena().free_segments() == 0 {
1463                break;
1464            }
1465        }
1466        assert_eq!(
1467            m.arena().segment_count(),
1468            before,
1469            "asked the system for memory while holding an empty segment"
1470        );
1471    }
1472
1473    #[test]
1474    fn adversarial_keys_that_share_low_bits() {
1475        // Keys chosen so that many land in the same bucket index. The point is
1476        // that overflow chaining and splitting both still work when the hash is
1477        // not being kind.
1478        let mut m = RawMap::new();
1479        let mut inserted = Vec::new();
1480        for i in 0..ADVERSARIAL_N {
1481            let k = i.to_le_bytes().to_vec();
1482            m.set(&k, b"v");
1483            inserted.push(k);
1484        }
1485        for k in &inserted {
1486            assert_eq!(m.get(k), Some(&b"v"[..]));
1487        }
1488        assert_eq!(m.len(), inserted.len());
1489    }
1490
1491    /// Whatever memoizes against this counter is only correct if every way of
1492    /// moving something in the map moves it too. A method that mutates and does
1493    /// not is not a slow memo, it is a wrong answer, so this asserts on the whole
1494    /// `&mut self` surface rather than on the ones that look like they matter.
1495    ///
1496    /// The single exception is pinned by the test below this one, so a method
1497    /// added without a decision about which side it falls on fails here.
1498    #[test]
1499    fn every_way_of_writing_moves_the_counter() {
1500        let mut m = RawMap::new();
1501        let mut last = m.writes();
1502        let mut moved = |m: &RawMap, what: &str| {
1503            assert!(m.writes() > last, "{what} did not move the counter");
1504            last = m.writes();
1505        };
1506
1507        m.set(b"k", b"v");
1508        moved(&m, "set");
1509        m.set_with(
1510            b"k",
1511            1,
1512            |_| {},
1513            |b| {
1514                b[0] = b'w';
1515                false
1516            },
1517        );
1518        moved(&m, "set_with");
1519        m.value_mut(b"k");
1520        moved(&m, "value_mut");
1521        m.value_mut_hashed(RawMap::hash_of(b"k"), b"k");
1522        moved(&m, "value_mut_hashed");
1523        m.compact_step();
1524        moved(&m, "compact_step");
1525        m.compact_segment(0);
1526        moved(&m, "compact_segment");
1527        m.del(b"k");
1528        moved(&m, "del");
1529    }
1530
1531    /// The exception, pinned so that it stays a decision rather than becoming a
1532    /// habit. An in place stamp leaves the counter alone, and everything the
1533    /// caller resolved before it is still right after it.
1534    #[test]
1535    fn sampling_hands_back_real_entries_and_stops_when_told() {
1536        let mut m = RawMap::new();
1537        for i in 0..2000u32 {
1538            m.set(format!("k{i}").as_bytes(), format!("v{i}").as_bytes());
1539        }
1540
1541        // Whatever it hands over is really in the map, key and value together,
1542        // and the address it gives is the address that key resolves to.
1543        let mut count = 0usize;
1544        m.sample(0x1234_5678_9abc_def0, |key, val, addr| {
1545            assert_eq!(m.get(key), Some(val));
1546            assert_eq!(m.find(key), Some(addr));
1547            count += 1;
1548            count < 5
1549        });
1550        assert_eq!(count, 5, "it did not stop when it was told to");
1551
1552        // A caller that never says stop still terminates, because the segment is
1553        // the bound and not the caller.
1554        let mut all = 0usize;
1555        m.sample(0, |_, _, _| {
1556            all += 1;
1557            true
1558        });
1559        assert!(all > 0, "it found nothing in a map of two thousand keys");
1560        assert!(
1561            all < m.len(),
1562            "one segment and not the whole map, got {all} of {}",
1563            m.len()
1564        );
1565    }
1566
1567    #[test]
1568    fn sampling_a_sparse_map_still_finds_something() {
1569        // The case a sampler that looked in one bucket would get wrong. Two keys
1570        // in a map sized for two thousand is sixty two empty buckets for every
1571        // two that are worth looking in.
1572        let mut m = RawMap::new();
1573        for i in 0..2000u32 {
1574            m.set(format!("k{i}").as_bytes(), b"v");
1575        }
1576        for i in 0..1998u32 {
1577            m.del(format!("k{i}").as_bytes());
1578        }
1579        assert_eq!(m.len(), 2);
1580
1581        // Not every draw lands in the segment those two are in, so this is about
1582        // whether it ever finds them rather than whether it always does.
1583        let mut found = 0usize;
1584        for r in 0..200u64 {
1585            m.sample(r.wrapping_mul(0x9e37_79b9_7f4a_7c15), |_, _, _| {
1586                found += 1;
1587                true
1588            });
1589        }
1590        assert!(found > 0, "two hundred draws and it never found either key");
1591    }
1592
1593    #[test]
1594    fn stamping_a_value_in_place_is_not_a_write() {
1595        let mut m = RawMap::new();
1596        m.set(b"k", b"hello");
1597        let addr = m.find(b"k").expect("just stored");
1598        let before = m.writes();
1599
1600        m.value_at_mut(addr)[0] = b'j';
1601
1602        assert_eq!(m.writes(), before, "a stamp counted as a write");
1603        assert_eq!(m.get(b"k"), Some(&b"jello"[..]));
1604        // And the address the caller was holding still means what it meant, which
1605        // is the guarantee the counter would otherwise be asked about.
1606        assert_eq!(m.find(b"k"), Some(addr));
1607        assert_eq!(m.value_at(addr), b"jello");
1608    }
1609
1610    /// `clear` replaces the map with a fresh one, and a fresh one starts at
1611    /// zero. A memo taken at write 3 against a map that went back to 0 and
1612    /// climbed to 3 again would read as still valid on the one call where every
1613    /// key in the map had been thrown away.
1614    #[test]
1615    fn clearing_does_not_send_the_counter_backwards() {
1616        let mut m = RawMap::new();
1617        for i in 0..10u32 {
1618            m.set(&i.to_le_bytes(), b"v");
1619        }
1620        let before = m.writes();
1621        m.clear();
1622        assert!(m.writes() > before, "clear went backwards or stood still");
1623    }
1624
1625    /// Enough keys to have split several times, so a walk crosses segments of
1626    /// different local depths rather than staying inside one.
1627    #[cfg(miri)]
1628    const SCAN_N: usize = 400;
1629    #[cfg(not(miri))]
1630    const SCAN_N: usize = 20_000;
1631
1632    #[test]
1633    fn a_walk_of_an_empty_map_ends_on_the_first_call() {
1634        let m = RawMap::new();
1635        let mut seen = 0;
1636        let at = m.scan(Cursor::START, 1000, |_, _| seen += 1);
1637        assert_eq!(seen, 0);
1638        assert!(
1639            at.is_end(),
1640            "an empty map took more than one call to finish"
1641        );
1642    }
1643
1644    /// The plain case, and the one every other guarantee is stated against: no
1645    /// writes during the walk, so every key comes back once and no key comes
1646    /// back twice.
1647    #[test]
1648    fn a_quiet_walk_returns_every_key_exactly_once() {
1649        let mut m = RawMap::new();
1650        for i in 0..SCAN_N {
1651            m.set(&key(i), &val(i));
1652        }
1653
1654        let mut counts: HashMap<Vec<u8>, usize> = HashMap::new();
1655        let mut at = Cursor::START;
1656        let mut calls = 0;
1657        loop {
1658            at = m.scan(at, 1, |k, v| {
1659                // Both borrows are shared, so the walk can look the key up
1660                // while it is handing it over. The pair arriving together is
1661                // the point: a bucket walk that read the header of one record
1662                // and the body of the next would still pass a key only check.
1663                assert_eq!(m.get(k), Some(v), "the value came back on the wrong key");
1664                *counts.entry(k.to_vec()).or_default() += 1;
1665            });
1666            calls += 1;
1667            assert!(calls < 1_000_000, "the cursor is not advancing");
1668            if at.is_end() {
1669                break;
1670            }
1671        }
1672
1673        assert_eq!(
1674            counts.len(),
1675            SCAN_N,
1676            "the walk missed keys or invented them"
1677        );
1678        for i in 0..SCAN_N {
1679            assert_eq!(counts.get(&key(i)).copied(), Some(1), "key {i}");
1680        }
1681    }
1682
1683    /// A budget is a floor and not a ceiling, and asking for everything at once
1684    /// is one call.
1685    #[test]
1686    fn a_budget_big_enough_finishes_in_one_call() {
1687        let mut m = RawMap::new();
1688        for i in 0..SCAN_N {
1689            m.set(&key(i), &val(i));
1690        }
1691
1692        let mut seen = 0;
1693        let at = m.scan(Cursor::START, usize::MAX, |_, _| seen += 1);
1694        assert_eq!(seen, SCAN_N);
1695        assert!(at.is_end());
1696    }
1697
1698    /// The guarantee that matters: the map grows underneath the walk, the
1699    /// directory doubles and segments split, and a key that was there the whole
1700    /// time still comes back.
1701    ///
1702    /// Written the way a client uses it, which is a cursor held across calls
1703    /// with other work happening in between, because the failure this is looking
1704    /// for is a cursor that means one thing before a split and another after.
1705    #[test]
1706    fn a_walk_survives_the_map_growing_underneath_it() {
1707        let mut m = RawMap::new();
1708        // The keys that are there throughout. Named apart from the ones added
1709        // during the walk so the two are easy to tell apart in the assertion.
1710        for i in 0..SCAN_N {
1711            m.set(&key(i), &val(i));
1712        }
1713        let depth_before = m.index().global_depth();
1714
1715        let mut seen: HashSet<Vec<u8>> = HashSet::new();
1716        let mut at = Cursor::START;
1717        let mut added = SCAN_N;
1718        loop {
1719            at = m.scan(at, 8, |k, _| {
1720                seen.insert(k.to_vec());
1721            });
1722            if at.is_end() {
1723                break;
1724            }
1725            // Between one call and the next, which is where a client would be.
1726            for _ in 0..64 {
1727                m.set(&key(added), &val(added));
1728                added += 1;
1729            }
1730        }
1731
1732        assert!(
1733            m.index().global_depth() > depth_before,
1734            "the directory never doubled, so this test proved nothing"
1735        );
1736        for i in 0..SCAN_N {
1737            assert!(
1738                seen.contains(&key(i)),
1739                "key {i} was there throughout and never came back"
1740            );
1741        }
1742    }
1743
1744    /// Deletes during a walk are the other half of the same guarantee. A key
1745    /// that survives to the end still comes back, whatever happened to its
1746    /// neighbours.
1747    #[test]
1748    fn a_walk_survives_keys_being_deleted_underneath_it() {
1749        let mut m = RawMap::new();
1750        for i in 0..SCAN_N {
1751            m.set(&key(i), &val(i));
1752        }
1753
1754        let mut seen: HashSet<Vec<u8>> = HashSet::new();
1755        let mut at = Cursor::START;
1756        let mut next_gone = 1;
1757        loop {
1758            at = m.scan(at, 8, |k, _| {
1759                seen.insert(k.to_vec());
1760            });
1761            if at.is_end() {
1762                break;
1763            }
1764            // Every odd key goes, a few at a time. The even ones are what the
1765            // assertion is about.
1766            for _ in 0..16 {
1767                if next_gone < SCAN_N {
1768                    m.del(&key(next_gone));
1769                    next_gone += 2;
1770                }
1771            }
1772        }
1773
1774        for i in (0..SCAN_N).step_by(2) {
1775            assert!(
1776                seen.contains(&key(i)),
1777                "key {i} was never deleted and never came back"
1778            );
1779        }
1780    }
1781
1782    /// A cursor names a place in the keyspace and not a place in memory, so a
1783    /// walk started partway through returns everything from there on.
1784    ///
1785    /// The prefix is what says where that is. Starting at prefix `p` resumes in
1786    /// the segment holding `p`, which begins at or before it, so every key whose
1787    /// own prefix is `p` or higher is still ahead of the walk.
1788    #[test]
1789    fn a_walk_that_starts_partway_returns_everything_from_there_on() {
1790        let mut m = RawMap::new();
1791        for i in 0..SCAN_N {
1792            m.set(&key(i), &val(i));
1793        }
1794
1795        let half = 1u64 << (crate::scan::PREFIX_BITS - 1);
1796        let mut seen: HashSet<Vec<u8>> = HashSet::new();
1797        let at = m.scan(Cursor::at(half, 0), usize::MAX, |k, _| {
1798            seen.insert(k.to_vec());
1799        });
1800        assert!(at.is_end());
1801
1802        let mut expected = 0;
1803        for i in 0..SCAN_N {
1804            let k = key(i);
1805            if Cursor::prefix_of(RawMap::hash_of(&k)) >= half {
1806                expected += 1;
1807                assert!(
1808                    seen.contains(&k),
1809                    "key {i} is past the cursor and did not come back"
1810                );
1811            }
1812        }
1813        // Both halves of the keyspace have keys in them, or the assertion above
1814        // is checking nothing.
1815        assert!(
1816            expected > 0 && expected < SCAN_N,
1817            "the split point was degenerate"
1818        );
1819    }
1820}