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}