yo_kv/set.rs
1//! A set, in whichever representation currently fits it.
2//!
3//! A set is one of an [`Intset`], a [`Listpack`], an [`Elements`] table or a
4//! [`Parts`] band, and which one is not a choice this file gets to make freely.
5//! `OBJECT ENCODING` has to answer `intset`, `listpack` or `hashtable` at exactly
6//! the sizes a real server answers them, because clients and test suites read it
7//! (`08` ยง1), so the promotion rules here are Redis's rules and they were read
8//! off `t_set.c` in the 8.10.1 tarball rather than reasoned out from what each
9//! structure is good at.
10//!
11//! ```text
12//! all integers small, any bytes everything else large
13//! +-------------+ +------------------+ +---------------+ +-----------+
14//! | intset |-->| listpack |-->| element table |-->| partition |
15//! | 2 B member | | ~2 B + payload | | one probe | | band |
16//! +-------------+ +------------------+ +---------------+ +-----------+
17//! to 512 members to 128 members past 262144
18//! <------- both say `hashtable` ---->
19//! ```
20//!
21//! The fourth step is the one Redis does not have, and it is invisible on
22//! purpose. A set past 262,144 members becomes several element tables rather than
23//! one, which is what makes the merges and the growth pauses bounded, and it
24//! still answers `hashtable` because a client that gets a fourth word back from
25//! `OBJECT ENCODING` is a client whose assertions break. What partitioning
26//! changes is how a large set is stored, not what it is. See [`crate::parts`].
27//!
28//! Promotion is one-way and upward, which is Y4. A set that has been a hash
29//! table does not go back to an intset when it shrinks, and neither does Redis's:
30//! a set that demoted on the way down would rewrite itself on every second
31//! operation for a workload that adds and removes across a threshold. The band
32//! follows the same rule for the same reason, so a set hovering at the partition
33//! threshold rehashes once rather than on every other `SREM`.
34//!
35//! # The rules, and the two that are not obvious
36//!
37//! Adding an integer to an intset keeps it an intset until it holds more than
38//! `set-max-intset-entries`, and then it becomes a **hash table** and not a
39//! listpack, because the intset ceiling is 512 and the listpack ceiling is 128
40//! and something over the first is well over the second.
41//!
42//! Adding a non integer to an intset is the asymmetric one. It becomes a
43//! listpack only if the intset is currently under the *listpack* ceiling of 128,
44//! so an intset of 200 integers that receives one string goes straight to a hash
45//! table and is never a listpack at all. Reading that off the source was worth
46//! more than reasoning about it, because the natural implementation converts to
47//! a listpack whenever the members would fit and gets a different encoding name
48//! than a real server for a shape a test suite actually builds.
49//!
50//! # Members
51//!
52//! A member comes back as a [`Member`], which is [`listpack::Entry`] under
53//! another name: either the bytes as they lie or an integer that has not been
54//! formatted yet. All three representations can produce one without copying, and
55//! the formatting happens once, into the reply buffer, at the moment the reply is
56//! built. That is Y18, and it is the same reason [`crate::value::Str`] has two
57//! arms.
58
59use crate::elem::Elements;
60use crate::frozen::{self, Broken};
61use crate::intset::Intset;
62use crate::listpack::{self, Listpack};
63use crate::parts::{PARTITION_AT, Parts, parts_for};
64use crate::scan::Cursor;
65use yo_common::num::{DIGITS_MAX, i64_digits, i64_len, parse_i64};
66
67/// A set member: bytes as they lie, or an integer not yet formatted.
68pub type Member<'a> = listpack::Entry<'a>;
69
70/// A single run intset, written as the Redis layout it already is.
71const FORM_INTSET: u8 = 1;
72/// A listpack, written as the layout it already is.
73const FORM_PACKED: u8 = 2;
74/// An intset that has split into runs, written as ascending gaps.
75const FORM_RUNS: u8 = 3;
76/// A table or a band, written as its members.
77const FORM_MEMBERS: u8 = 4;
78/// The top bit of the form byte, which carries [`Set::ints_past_limit`].
79///
80/// In the form byte rather than a byte of its own because it only ever means
81/// anything for the two integer forms, and a set that has to be told what to call
82/// itself has already been told what it is.
83const PAST_LIMIT: u8 = 0x80;
84
85/// A member on its way to being asked about, in every form the three
86/// representations want it in.
87///
88/// Set algebra walks one set and asks every other set the same question about
89/// each member, and the three representations do not want the question in the
90/// same shape. An intset wants a number, a listpack wants bytes and a number
91/// because it holds both kinds, and an element table wants bytes and their
92/// hash. Asking through [`Set::contains`] would redo all of that per question:
93/// a parse for the intset, another parse inside the listpack, and a hash per
94/// table. This does each once per member and then asks `k - 1` times.
95///
96/// The hash is computed whether or not any operand is a table, which is waste
97/// when none is. It is waste worth taking, because the sets where it is wasted
98/// are an intset or a listpack, which are capped at a few hundred members, and
99/// an operation over sets that small is finished before the saving could have
100/// been measured. The sets where the hash pays are the large ones, and those
101/// are tables by definition.
102#[derive(Debug, Clone, Copy)]
103pub struct Needle<'a> {
104 /// The member as bytes, which for an integer member is a caller's buffer.
105 bytes: &'a [u8],
106 /// The number it is, if it is one, under the same rule that decides whether
107 /// a set stores it as one.
108 int: Option<i64>,
109 /// What an element table would key it under.
110 hash: u64,
111}
112
113impl<'a> Needle<'a> {
114 /// A needle from bytes, which is what a command line argument is.
115 #[must_use]
116 pub fn new(bytes: &'a [u8]) -> Needle<'a> {
117 Needle {
118 bytes,
119 int: parse_i64(bytes),
120 hash: Elements::<()>::hash_of(bytes),
121 }
122 }
123
124 /// A needle from a member walked out of a set.
125 ///
126 /// `digits` is where an integer member's text goes, because an intset holds
127 /// the number and the digits do not exist anywhere until somebody writes
128 /// them. It is the caller's buffer rather than a field so that the needle
129 /// stays a borrow and the buffer is written once per member rather than
130 /// allocated once per member.
131 ///
132 /// A member that came out as bytes is still parsed, because the set being
133 /// asked may be an intset and `SINTER ints strings` has to find the members
134 /// they share. A member that came out as a number is not, which is the
135 /// whole saving.
136 #[must_use]
137 pub fn of(member: Member<'a>, digits: &'a mut [u8; DIGITS_MAX]) -> Needle<'a> {
138 match member {
139 Member::Str(s) => Needle::new(s),
140 Member::Int(n) => {
141 let bytes = i64_digits(digits, n);
142 Needle {
143 bytes,
144 int: Some(n),
145 hash: Elements::<()>::hash_of(bytes),
146 }
147 }
148 }
149 }
150
151 /// The member as bytes, which is what a caller collecting an answer wants.
152 #[must_use]
153 pub const fn bytes(&self) -> &'a [u8] {
154 self.bytes
155 }
156}
157
158/// Where the encodings change over.
159///
160/// These are `set-max-intset-entries`, `set-max-listpack-entries` and
161/// `set-max-listpack-value`, and they are runtime configuration in Redis, so
162/// they are a value passed in here rather than three constants. The defaults are
163/// Redis's defaults and a client that never touches `CONFIG SET` sees exactly
164/// the encodings a real server would give it.
165#[derive(Debug, Clone, Copy, PartialEq, Eq)]
166pub struct Limits {
167 /// Past this many members an all integer set stops being an intset.
168 pub max_intset_entries: usize,
169 /// At this many members a set stops being a listpack.
170 pub max_listpack_entries: usize,
171 /// A member longer than this cannot go in a listpack.
172 pub max_listpack_value: usize,
173}
174
175impl Limits {
176 /// Redis's defaults: 512, 128 and 64.
177 pub const DEFAULT: Limits = Limits {
178 max_intset_entries: 512,
179 max_listpack_entries: 128,
180 max_listpack_value: 64,
181 };
182}
183
184impl Default for Limits {
185 fn default() -> Limits {
186 Limits::DEFAULT
187 }
188}
189
190/// Which of the three a set is in, which is what `OBJECT ENCODING` reports.
191#[derive(Debug, Clone, Copy, PartialEq, Eq)]
192pub enum Encoding {
193 /// All members are integers and there are few enough of them.
194 Intset,
195 /// One packed blob, walked linearly.
196 Listpack,
197 /// The element table.
198 Hashtable,
199}
200
201impl Encoding {
202 /// The word `OBJECT ENCODING` replies with.
203 #[inline]
204 pub const fn name(self) -> &'static str {
205 match self {
206 Encoding::Intset => "intset",
207 Encoding::Listpack => "listpack",
208 Encoding::Hashtable => "hashtable",
209 }
210 }
211}
212
213/// The four representations, of which `OBJECT ENCODING` can see three.
214///
215/// [`Body::Split`] is the partitioned band and it is deliberately invisible from
216/// outside. Redis has three set encodings and a client that gets a fourth word
217/// back from `OBJECT ENCODING` is a client whose assertions break, so a split set
218/// answers `hashtable` like the table it was. What partitioning changes is how a
219/// large set is stored and merged, not what it is.
220#[derive(Debug, Clone)]
221enum Body {
222 Ints(Intset),
223 Packed(Listpack),
224 Table(Elements<()>),
225 Split(Parts<()>),
226}
227
228/// A set of members.
229#[derive(Debug, Clone)]
230pub struct Set {
231 body: Body,
232 /// Whether an all integer set has passed `set-max-intset-entries`.
233 ///
234 /// This is where Redis rehashes the members into a dictionary and starts
235 /// answering `hashtable` to `OBJECT ENCODING`. Nothing is rehashed here,
236 /// because [`Intset`] holds a large set in runs and stays at two to eight
237 /// bytes a member where a table would cost thirty, so the only thing the
238 /// ceiling still decides is the word.
239 ///
240 /// It is a flag and not a comparison against the length because the ceiling
241 /// is configurable and [`Set::encoding`] is not handed the configuration. It
242 /// is also one way, like every promotion here: a set that has been called a
243 /// hashtable once does not go back to being called an intset when members
244 /// are removed, which is Y4 and is what Redis does.
245 ///
246 /// Only ever true while [`Body::Ints`] is the body. Every other body already
247 /// knows what to answer.
248 ints_past_limit: bool,
249}
250
251impl Set {
252 /// An empty set, which starts as an intset.
253 ///
254 /// This is what `SADD` on a missing key creates when it has no size hint,
255 /// and the first member decides nothing: an intset that receives a string
256 /// converts on the spot, and it costs a conversion of nothing.
257 #[must_use]
258 pub fn new() -> Set {
259 Set {
260 body: Body::Ints(Intset::new()),
261 ints_past_limit: false,
262 }
263 }
264
265 /// An empty set sized for what is about to go in it.
266 ///
267 /// Redis's `setTypeCreate`, which picks the representation from the first
268 /// member and the count the caller expects, so that `SADD k a b c ...` with
269 /// a thousand arguments builds a table once rather than converting twice on
270 /// the way there. `hint` is only a hint and being wrong about it costs a
271 /// conversion and no correctness.
272 ///
273 /// An integer first member sends it to the intset even when the count is
274 /// past `set-max-intset-entries`, where Redis would go straight to a
275 /// dictionary. A large set of integers is the case the runs exist for, and
276 /// building a table and never leaving it would give up the whole saving on
277 /// the one call that said in advance it was going to matter. The listpack
278 /// band in between is still honoured, because a set that small has nothing
279 /// to save and a server configured that way expects a listpack.
280 #[must_use]
281 pub fn with_hint(first: &[u8], hint: usize, limits: &Limits) -> Set {
282 let ints = parse_i64(first).is_some()
283 && (hint <= limits.max_intset_entries || hint > limits.max_listpack_entries);
284 if ints {
285 Set {
286 body: Body::Ints(Intset::with_capacity(hint)),
287 ints_past_limit: hint > limits.max_intset_entries,
288 }
289 } else if hint <= limits.max_listpack_entries {
290 Set {
291 body: Body::Packed(Listpack::new()),
292 ints_past_limit: false,
293 }
294 } else if hint > PARTITION_AT {
295 // A caller that says up front it is about to load a million members
296 // should not build one table, fill it past the threshold and then
297 // rehash the lot into partitions. The hint is only a hint, and being
298 // wrong about it here costs a set with more partitions than it needs
299 // rather than anything incorrect.
300 Set {
301 body: Body::Split(Parts::with_parts(parts_for(hint))),
302 ints_past_limit: false,
303 }
304 } else {
305 Set {
306 body: Body::Table(Elements::with_capacity(hint)),
307 ints_past_limit: false,
308 }
309 }
310 }
311
312 /// Which representation this is in.
313 ///
314 /// Four bodies and three words, and the intset accounts for two of the
315 /// missing ones. The partitioned body answers `hashtable` because it is one,
316 /// and an intset past `set-max-intset-entries` answers `hashtable` because
317 /// that is what a real server would have turned into by then, even though
318 /// nothing here was rehashed.
319 #[inline]
320 #[must_use]
321 pub const fn encoding(&self) -> Encoding {
322 match self.body {
323 Body::Ints(_) if self.ints_past_limit => Encoding::Hashtable,
324 Body::Ints(_) => Encoding::Intset,
325 Body::Packed(_) => Encoding::Listpack,
326 Body::Table(_) | Body::Split(_) => Encoding::Hashtable,
327 }
328 }
329
330 /// The bytes behind a set that is stored the way Redis stores it.
331 ///
332 /// `DUMP` writes these straight out instead of walking the members, so this
333 /// exists for [`crate::rdb`] and for nothing else. The word this answers
334 /// against is [`Set::encoding`] and not the body, so a set that calls itself
335 /// a hashtable is walked even on the rare occasion its members are still in
336 /// one intset run. Keeping those two in step is what makes the rule sayable:
337 /// whatever `OBJECT ENCODING` says, that is the type byte the payload gets.
338 ///
339 /// `None` for a set with no such shape, which is the table, the partitioned
340 /// body and an intset that has split into runs.
341 #[inline]
342 pub(crate) fn packed_bytes(&self) -> Option<&[u8]> {
343 match &self.body {
344 Body::Ints(s) if !self.ints_past_limit => s.as_bytes(),
345 Body::Packed(lp) => Some(lp.as_bytes()),
346 _ => None,
347 }
348 }
349
350 /// Write the set out as the bytes it becomes when it leaves memory.
351 ///
352 /// One form byte and then whatever that form needs. The two representations
353 /// that already have a flat layout, the single run intset and the listpack,
354 /// are written as themselves and cost one byte, because those bytes are
355 /// exactly what has to come back. The other two are written as their members.
356 ///
357 /// The form carries enough to land back in the same representation, which
358 /// [`crate::rdb`] deliberately does not: `ints_past_limit` rides in the top
359 /// bit of the form byte, and a set of integers that has split into runs is
360 /// its own form rather than a list of members, so that a million integers do
361 /// not come back as a hash table at thirty bytes a member. A value that
362 /// changed encoding because it was quiet long enough to be demoted would be a
363 /// value whose `OBJECT ENCODING` depends on memory pressure.
364 ///
365 /// See [`Set::thaw`], which is the other half and reads exactly this.
366 pub fn freeze(&self, out: &mut Vec<u8>) {
367 let past = if self.ints_past_limit { PAST_LIMIT } else { 0 };
368 match &self.body {
369 // `as_bytes` and not `packed_bytes`, because the ceiling only decides
370 // what the encoding is called and the flag above carries that. A set
371 // over the ceiling still has the Redis layout while it is one run,
372 // and writing the layout it has beats writing its members.
373 Body::Ints(s) => match s.as_bytes() {
374 Some(bytes) => {
375 out.push(FORM_INTSET | past);
376 out.extend_from_slice(bytes);
377 }
378 // Split into runs, so there is no one blob. Ascending, so the
379 // gaps are what gets written: a dense set of integers is a byte
380 // a member here where its members as digits would be several.
381 None => {
382 out.push(FORM_RUNS | past);
383 frozen::put_uint(out, s.len() as u64);
384 let mut prev = 0u64;
385 for i in 0..s.len() {
386 let v = s.at(i) as u64;
387 frozen::put_uint(out, v.wrapping_sub(prev));
388 prev = v;
389 }
390 }
391 },
392 Body::Packed(lp) => {
393 out.push(FORM_PACKED);
394 out.extend_from_slice(lp.as_bytes());
395 }
396 Body::Table(_) | Body::Split(_) => {
397 out.push(FORM_MEMBERS);
398 frozen::put_uint(out, self.len() as u64);
399 let mut digits = [0u8; DIGITS_MAX];
400 for m in self.iter() {
401 match m {
402 Member::Str(bytes) => frozen::put_bytes(out, bytes),
403 // Neither of these two bodies stores an integer as one,
404 // so this arm is for the type and not for a case that
405 // arrives.
406 Member::Int(n) => frozen::put_bytes(out, i64_digits(&mut digits, n)),
407 }
408 }
409 }
410 }
411 }
412
413 /// Read back a set written by [`Set::freeze`].
414 ///
415 /// The count in the member form is a capacity and not a promise. It picks the
416 /// body and sizes it, and then the members decide the length, so a count that
417 /// disagrees with what follows costs a rehash and never a wrong set.
418 pub fn thaw(bytes: &[u8]) -> Result<Set, Broken> {
419 let mut cut = frozen::Cut::new(bytes);
420 let tag = cut.byte()?;
421 let ints_past_limit = tag & PAST_LIMIT != 0;
422 match tag & !PAST_LIMIT {
423 FORM_INTSET => Ok(Set {
424 body: Body::Ints(Intset::from_bytes(cut.rest()).map_err(|_| Broken::Body)?),
425 ints_past_limit,
426 }),
427 FORM_PACKED => Ok(Set {
428 body: Body::Packed(Listpack::from_bytes(cut.rest()).map_err(|_| Broken::Body)?),
429 ints_past_limit: false,
430 }),
431 FORM_RUNS => {
432 let n = usize::try_from(cut.uint()?).map_err(|_| Broken::Short)?;
433 let mut s = Intset::with_capacity(n);
434 let mut prev = 0u64;
435 for _ in 0..n {
436 prev = prev.wrapping_add(cut.uint()?);
437 s.add(prev as i64);
438 }
439 Ok(Set {
440 body: Body::Ints(s),
441 ints_past_limit,
442 })
443 }
444 FORM_MEMBERS => {
445 let n = usize::try_from(cut.uint()?).map_err(|_| Broken::Short)?;
446 // A member is at least one byte, so a count larger than what is
447 // left cannot be honest and is not worth an allocation.
448 if n > cut.rest().len() {
449 return Err(Broken::Body);
450 }
451 // The same threshold the live set splits at, so a band that was
452 // demoted comes back a band. One that has shrunk under the
453 // threshold since comes back a table, which is the one place a
454 // round trip changes the body, and both of them answer
455 // `hashtable` so nothing outside can tell.
456 let mut set = Set {
457 body: if n > PARTITION_AT {
458 Body::Split(Parts::with_parts(parts_for(n)))
459 } else {
460 Body::Table(Elements::with_capacity(n))
461 },
462 ints_past_limit: false,
463 };
464 for _ in 0..n {
465 let member = cut.bytes()?;
466 match &mut set.body {
467 Body::Table(t) => t.insert(member, ()).map_err(|_| Broken::Body)?,
468 Body::Split(p) => p.insert(member, ()).map_err(|_| Broken::Body)?,
469 _ => unreachable!("the body was just built as one of those two"),
470 };
471 }
472 Ok(set)
473 }
474 _ => Err(Broken::Form),
475 }
476 }
477
478 /// How many members. This is `SCARD`.
479 #[inline]
480 pub fn len(&self) -> usize {
481 match &self.body {
482 Body::Ints(s) => s.len(),
483 Body::Packed(lp) => lp.len(),
484 Body::Table(t) => t.len(),
485 Body::Split(p) => p.len(),
486 }
487 }
488
489 /// Whether there are none.
490 ///
491 /// An empty set does not exist in Redis, so the caller deletes the key when
492 /// this turns true rather than storing an empty one.
493 #[inline]
494 pub fn is_empty(&self) -> bool {
495 self.len() == 0
496 }
497
498 /// Whether `member` is in the set. This is `SISMEMBER`.
499 #[must_use]
500 pub fn contains(&self, member: &[u8]) -> bool {
501 match &self.body {
502 // A member that is not an integer cannot be in a set of integers,
503 // and answering that costs a parse rather than a search.
504 Body::Ints(s) => parse_i64(member).is_some_and(|v| s.contains(v)),
505 Body::Packed(lp) => lp.find(member, 1).is_some(),
506 Body::Table(t) => t.contains(member),
507 Body::Split(p) => p.contains(member),
508 }
509 }
510
511 /// The same question asked with the work already done. See [`Needle`].
512 ///
513 /// This is what set algebra probes with. Every arm is the arm
514 /// [`Set::contains`] would have taken, with the parse and the hash lifted
515 /// out of it, so the two cannot disagree about what a member is.
516 #[must_use]
517 #[inline]
518 pub fn has(&self, needle: &Needle<'_>) -> bool {
519 match &self.body {
520 Body::Ints(s) => needle.int.is_some_and(|v| s.contains(v)),
521 Body::Packed(lp) => lp.find_parsed(needle.bytes, needle.int, 1).is_some(),
522 Body::Table(t) => t.contains_hashed(needle.hash, needle.bytes),
523 Body::Split(p) => p.contains_hashed(needle.hash, needle.bytes),
524 }
525 }
526
527 /// The member at `index`, in whatever order the representation holds them.
528 ///
529 /// Ascending for an intset, insertion order for the other two. Redis makes
530 /// no promise about set order and neither does this, but a uniform draw
531 /// needs positions and this is what gives it them (K9).
532 #[must_use]
533 pub fn at(&self, index: usize) -> Option<Member<'_>> {
534 match &self.body {
535 Body::Ints(s) => s.get(index).map(Member::Int),
536 Body::Packed(lp) => lp.get(index),
537 Body::Table(t) => t.at(index).map(|(name, _)| Member::Str(name)),
538 Body::Split(p) => p.at(index).map(|(name, _)| Member::Str(name)),
539 }
540 }
541
542 /// Every member.
543 pub fn iter(&self) -> impl Iterator<Item = Member<'_>> {
544 (0..self.len()).map(|i| self.at(i).expect("index is under the length"))
545 }
546
547 /// The members as a sorted array of integers, when that is what this is.
548 ///
549 /// The one place the representation is not an implementation detail, and it
550 /// is here for [`crate::setops`]: two sorted arrays intersect by stepping
551 /// through both of them with no hash anywhere, which is a different order of
552 /// cost from asking a table a question per member. That was worth nothing
553 /// while an all integer set turned into a table at five hundred and twelve
554 /// members, and it is worth a great deal now that it does not.
555 #[inline]
556 #[must_use]
557 pub const fn ints(&self) -> Option<&Intset> {
558 match &self.body {
559 Body::Ints(s) => Some(s),
560 _ => None,
561 }
562 }
563
564 /// Walk part of the set and say where to resume. This is `SSCAN`.
565 ///
566 /// Only the table and the partitioned band walk in windows. An intset or a
567 /// listpack hands back
568 /// every member in one call and a cursor of [`Cursor::END`], ignoring the
569 /// cursor it was given, which is what Redis does for the same two encodings
570 /// and for the same reason: a hundred and twenty eight members is smaller
571 /// than the reply header arithmetic to split them up, and a set that small
572 /// cannot block the loop long enough for the split to be worth anything.
573 ///
574 /// Ignoring the cursor is safe rather than merely convenient, because
575 /// promotion is one way. A set that gave a client a listpack cursor is not
576 /// going to be a listpack again, so the only way to arrive at those two arms
577 /// with a cursor from somewhere else is for the key to have been deleted and
578 /// remade underneath the scan, and returning everything to that client
579 /// returns a member twice at worst, which the guarantee allows.
580 ///
581 /// A table cursor arriving at the band is the one crossing that does happen,
582 /// because a set can split part way through a client's scan. That is handled
583 /// rather than ignored: a table cursor names one partition, and
584 /// [`Cursor::rebase`] reads the widening and restarts the walk at the top of
585 /// the new layout, so the client sees some members a second time and misses
586 /// none. Repeats are what the `SCAN` guarantee gives up in exchange for
587 /// surviving a resize, and a set only splits once.
588 pub fn scan<F>(&self, cursor: Cursor, count: usize, mut f: F) -> Cursor
589 where
590 F: FnMut(Member<'_>),
591 {
592 match &self.body {
593 Body::Table(t) => t.scan(cursor, count, |name, ()| f(Member::Str(name))),
594 Body::Split(p) => p.scan(cursor, count, |name, ()| f(Member::Str(name))),
595 // An intset past the ceiling is the one thing outside the table
596 // band that a single reply cannot hold. Redis has a dictionary by
597 // this point and walks it in windows, and a set of a million
598 // integers answering `SSCAN` with a million members in one go would
599 // be a several megabyte reply and a loop iteration nobody could
600 // measure. So it walks in windows too, and by index, which the runs
601 // answer in a walk down their tree rather than a walk along them.
602 //
603 // Downward, which is the direction the element table walks and for
604 // the same reason. Positions in a sorted array shift when a member
605 // below them goes, so an upward walk would miss a member for every
606 // one removed behind it, and `SSCAN` followed by `SREM` on what it
607 // found is the commonest thing anyone does with this command.
608 // Walking down means those removals are all above the cursor, where
609 // they cost nothing.
610 Body::Ints(s) if self.ints_past_limit && !s.is_empty() => {
611 let top = s.len() - 1;
612 let mut at = match cursor.rebase(1).idx() {
613 Some(idx) => (idx as usize).min(top),
614 None => top,
615 };
616 for _ in 0..count.max(1) {
617 f(Member::Int(s.at(at)));
618 if at == 0 {
619 return Cursor::END;
620 }
621 at -= 1;
622 }
623 Cursor::at(1, 0, at as u64)
624 }
625 _ => {
626 for m in self.iter() {
627 f(m);
628 }
629 Cursor::END
630 }
631 }
632 }
633
634 /// Bytes held by whichever representation this is.
635 #[must_use]
636 pub fn memory_bytes(&self) -> usize {
637 match &self.body {
638 Body::Ints(s) => s.memory_bytes(),
639 Body::Packed(lp) => lp.byte_len(),
640 Body::Table(t) => t.memory_bytes(),
641 Body::Split(p) => p.memory_bytes(),
642 }
643 }
644
645 /// What the slot array costs on its own, or nothing if there is not one.
646 ///
647 /// The three of these split [`Set::memory_bytes`] into the arrays it is made
648 /// of, which is what turns an argument about where the memory went into a
649 /// number. An intset and a listpack are one allocation with no index over
650 /// it, so they answer all of it under names and nothing under the other two.
651 #[must_use]
652 pub fn slot_bytes(&self) -> usize {
653 match &self.body {
654 Body::Ints(_) | Body::Packed(_) => 0,
655 Body::Table(t) => t.slot_bytes(),
656 Body::Split(p) => p.slot_bytes(),
657 }
658 }
659
660 /// What the row array costs on its own, capacity and not length.
661 #[must_use]
662 pub fn row_bytes(&self) -> usize {
663 match &self.body {
664 Body::Ints(_) | Body::Packed(_) => 0,
665 Body::Table(t) => t.row_bytes(),
666 Body::Split(p) => p.row_bytes(),
667 }
668 }
669
670 /// What the member bytes cost, live ones and dead ones together.
671 #[must_use]
672 pub fn name_bytes(&self) -> usize {
673 match &self.body {
674 Body::Ints(s) => s.memory_bytes(),
675 Body::Packed(lp) => lp.byte_len(),
676 Body::Table(t) => t.name_bytes(),
677 Body::Split(p) => p.name_bytes(),
678 }
679 }
680
681 /// Add `member`, promoting if it no longer fits. Answers whether it was new.
682 ///
683 /// This is `setTypeAdd`, arm for arm.
684 pub fn add(&mut self, member: &[u8], limits: &Limits) -> bool {
685 match &mut self.body {
686 Body::Table(t) => {
687 let new = t.insert(member, ()).is_ok_and(|old| old.is_none());
688 // Checked after the insert rather than before, so the set that
689 // splits is the one that has actually outgrown a table and not
690 // the one that is about to.
691 if t.len() > PARTITION_AT {
692 self.become_split();
693 }
694 return new;
695 }
696 Body::Split(p) => {
697 let new = p.insert(member, ()).is_ok_and(|old| old.is_none());
698 // Asked rather than decided, because growing is a rehash of the
699 // whole set and the band leaves the timing to whoever knows
700 // whether this is one write or the middle of a bulk load.
701 if let Some(want) = p.wants_parts() {
702 p.grow_to(want);
703 }
704 return new;
705 }
706 Body::Packed(lp) => {
707 if lp.find(member, 1).is_some() {
708 return false;
709 }
710 if lp.len() < limits.max_listpack_entries
711 && member.len() <= limits.max_listpack_value
712 {
713 lp.push(member);
714 return true;
715 }
716 // Too many members, or one too long. It falls out to a table.
717 }
718 Body::Ints(s) => {
719 if let Some(v) = parse_i64(member) {
720 if !s.add(v) {
721 return false;
722 }
723 // Strictly greater, so the 512th member is still an intset
724 // and the 513th is what a real server would call a
725 // hashtable. Nothing is rewritten, only the word changes.
726 //
727 // Unless the ceilings have been configured the wrong way
728 // round, where a set past the intset ceiling is still under
729 // the listpack one and a real server puts it in a listpack.
730 // That set is a handful of members and there is no memory
731 // argument for keeping it here, so it goes where it would
732 // have gone.
733 if s.len() > limits.max_intset_entries {
734 if self.ints_fit_a_listpack_alone(limits) {
735 self.become_listpack();
736 } else {
737 self.ints_past_limit = true;
738 }
739 }
740 return true;
741 }
742 // Not an integer, so it is certainly not in a set of integers
743 // already. If the set is still small enough it becomes a
744 // listpack, and otherwise it falls out to a table.
745 if self.ints_fit_a_listpack(member, limits) {
746 self.become_listpack();
747 self.push_new(member);
748 return true;
749 }
750 }
751 }
752 self.become_table(1);
753 self.push_new(member);
754 true
755 }
756
757 /// Put in a member already known to be new and known to fit where it is.
758 ///
759 /// Only ever called on the far side of a promotion, where both of those are
760 /// facts the promotion established and not things worth establishing twice.
761 fn push_new(&mut self, member: &[u8]) {
762 match &mut self.body {
763 Body::Packed(lp) => lp.push(member),
764 Body::Table(t) => {
765 t.insert(member, ())
766 .expect("the table was sized for this one");
767 }
768 Body::Split(p) => {
769 p.insert(member, ())
770 .expect("the band was sized for this one");
771 }
772 Body::Ints(_) => unreachable!("no promotion ever lands on an intset"),
773 }
774 }
775
776 /// Remove `member`. Answers whether it was there.
777 ///
778 /// Never demotes, which is Y4's one-way rule and Redis's behaviour.
779 pub fn remove(&mut self, member: &[u8]) -> bool {
780 match &mut self.body {
781 Body::Ints(s) => parse_i64(member).is_some_and(|v| s.remove(v)),
782 Body::Packed(lp) => match lp.find(member, 1) {
783 Some(at) => lp.delete(at, 1),
784 None => false,
785 },
786 Body::Table(t) => t.remove(member).is_some(),
787 Body::Split(p) => p.remove(member).is_some(),
788 }
789 }
790
791 /// Take out the member at `index` and hand it back.
792 ///
793 /// This is what `SPOP` runs on top of. The table moves its last row into the
794 /// hole rather than shifting, so the position of every other member is
795 /// stable except for one; the other two shift. Neither is a promise a caller
796 /// can lean on, and `SPOP` does not need one because it draws again from the
797 /// new length each time.
798 pub fn remove_at(&mut self, index: usize) -> Option<Vec<u8>> {
799 match &mut self.body {
800 Body::Ints(s) => {
801 let v = s.get(index)?;
802 s.remove(v);
803 let mut out = Vec::with_capacity(i64_len(v));
804 Member::Int(v).write_to(&mut out);
805 Some(out)
806 }
807 Body::Packed(lp) => {
808 let out = lp.get(index)?.to_vec();
809 lp.delete(index, 1);
810 Some(out)
811 }
812 Body::Table(t) => t.take_at(index).map(|(name, ())| name),
813 Body::Split(p) => p.take_at(index).map(|(name, ())| name),
814 }
815 }
816
817 /// Take out the member at `index` without building it into a `Vec` first.
818 ///
819 /// The same removal as [`Set::remove_at`] for a caller that has already read
820 /// the member and does not need it handed back. That caller is `SPOP` on the
821 /// wire, which reads with [`Set::at`], writes the bytes straight into the
822 /// reply buffer, and only then calls this. It is an allocation a member
823 /// saved on the one command in the set whose whole cost is the allocating.
824 ///
825 /// [`Set::remove_at`] stays for the embedded API, where the caller wants the
826 /// bytes and has nowhere to put them.
827 pub fn drop_at(&mut self, index: usize) -> bool {
828 match &mut self.body {
829 Body::Ints(s) => match s.get(index) {
830 Some(v) => {
831 s.remove(v);
832 true
833 }
834 None => false,
835 },
836 Body::Packed(lp) => {
837 if index >= lp.len() {
838 return false;
839 }
840 lp.delete(index, 1);
841 true
842 }
843 Body::Table(t) => t.remove_at(index).is_some(),
844 Body::Split(p) => p.remove_at(index).is_some(),
845 }
846 }
847
848 /// Whether an intset plus one non integer member would still be a listpack.
849 ///
850 /// Three tests, and the first is the asymmetric one: the count is compared
851 /// against the *listpack* ceiling and not the intset one, so an intset with
852 /// two hundred members is already too big to become a listpack even though
853 /// it is a perfectly legal intset. The other two are the new member's length
854 /// and the longest existing member's length once it is written as digits,
855 /// which only bites when `set-max-listpack-value` has been turned down,
856 /// because no integer is more than twenty characters.
857 fn ints_fit_a_listpack(&self, member: &[u8], limits: &Limits) -> bool {
858 let Body::Ints(s) = &self.body else {
859 return false;
860 };
861 s.len() < limits.max_listpack_entries
862 && member.len() <= limits.max_listpack_value
863 && self.ints_are_short_enough(limits)
864 }
865
866 /// Whether this intset on its own would fit a listpack.
867 ///
868 /// The same question with no new member in it, which is what an intset that
869 /// has just passed `set-max-intset-entries` asks. It only ever answers yes
870 /// when the two ceilings have been configured the wrong way round, because
871 /// 512 is not under 128, and a server run that way expects a listpack there.
872 fn ints_fit_a_listpack_alone(&self, limits: &Limits) -> bool {
873 let Body::Ints(s) = &self.body else {
874 return false;
875 };
876 s.len() <= limits.max_listpack_entries && self.ints_are_short_enough(limits)
877 }
878
879 /// Whether every member, written as digits, is under the listpack ceiling.
880 fn ints_are_short_enough(&self, limits: &Limits) -> bool {
881 let Body::Ints(s) = &self.body else {
882 return false;
883 };
884 // The two ends bound the digits of everything between them, so there is
885 // nothing to walk.
886 let widest = s
887 .min()
888 .map(i64_len)
889 .unwrap_or(0)
890 .max(s.max().map(i64_len).unwrap_or(0));
891 widest <= limits.max_listpack_value
892 }
893
894 /// Rewrite as a listpack, which only an intset ever does.
895 fn become_listpack(&mut self) {
896 let Body::Ints(s) = &self.body else {
897 return;
898 };
899 let mut lp = Listpack::new();
900 let mut buf = Vec::with_capacity(20);
901 for v in s.iter() {
902 buf.clear();
903 Member::Int(v).write_to(&mut buf);
904 lp.push(&buf);
905 }
906 self.body = Body::Packed(lp);
907 }
908
909 /// Rewrite as an element table, with room for `extra` more members.
910 fn become_table(&mut self, extra: usize) {
911 let mut t = Elements::with_capacity(self.len() + extra);
912 let mut buf = Vec::with_capacity(20);
913 for m in self.iter() {
914 match m {
915 Member::Str(b) => {
916 t.insert(b, ()).expect("room, and every member was unique");
917 }
918 Member::Int(v) => {
919 buf.clear();
920 Member::Int(v).write_to(&mut buf);
921 t.insert(&buf, ())
922 .expect("room, and every member was unique");
923 }
924 }
925 }
926 self.body = Body::Table(t);
927 }
928
929 /// Spread an element table over partitions.
930 ///
931 /// One way, like every other promotion here. A set that drops back under the
932 /// threshold keeps its partitions, which is Y4's rule and Redis's behaviour
933 /// for the encodings it does expose: the cost of a representation is paid
934 /// when it is entered, and paying it again on the way back out turns one
935 /// `SREM` at the boundary into a rehash of the whole set.
936 fn become_split(&mut self) {
937 if let Body::Table(t) = &self.body {
938 let p = Parts::from_table(t, parts_for(t.len()));
939 self.body = Body::Split(p);
940 }
941 }
942}
943
944impl Default for Set {
945 fn default() -> Set {
946 Set::new()
947 }
948}
949
950#[cfg(test)]
951mod tests {
952 use super::*;
953
954 fn of(members: &[&str]) -> Set {
955 let mut s = Set::new();
956 for m in members {
957 assert!(s.add(m.as_bytes(), &Limits::DEFAULT), "{m} was new");
958 }
959 s
960 }
961
962 /// Freeze a set, thaw it, and check that what came back is the same set in
963 /// the same representation.
964 ///
965 /// Members are compared as bytes and as a sorted list, because two of the
966 /// four bodies hold them in insertion order and a round trip through the
967 /// member form does not promise to preserve it. What it does promise is the
968 /// membership and the word `OBJECT ENCODING` answers, and those are what this
969 /// checks.
970 fn round_trip(set: &Set) -> Set {
971 let mut out = Vec::new();
972 set.freeze(&mut out);
973 let back = Set::thaw(&out).expect("what freeze wrote, thaw reads");
974 assert_eq!(back.len(), set.len(), "member count");
975 assert_eq!(back.encoding(), set.encoding(), "encoding");
976 let mut was: Vec<Vec<u8>> = set.iter().map(bytes_of).collect();
977 let mut now: Vec<Vec<u8>> = back.iter().map(bytes_of).collect();
978 was.sort_unstable();
979 now.sort_unstable();
980 assert_eq!(was, now, "members");
981 back
982 }
983
984 fn bytes_of(m: Member<'_>) -> Vec<u8> {
985 match m {
986 Member::Str(b) => b.to_vec(),
987 Member::Int(n) => n.to_string().into_bytes(),
988 }
989 }
990
991 #[test]
992 fn a_frozen_set_comes_back_in_the_body_it_left() {
993 // An intset, which is one run and goes out as the Redis layout.
994 let mut ints = Set::new();
995 for i in 0..100i64 {
996 ints.add(i.to_string().as_bytes(), &Limits::DEFAULT);
997 }
998 assert_eq!(ints.encoding(), Encoding::Intset);
999 round_trip(&ints);
1000
1001 // A listpack, same story with the other flat layout.
1002 let packed = of(&["one", "two", "three"]);
1003 assert_eq!(packed.encoding(), Encoding::Listpack);
1004 round_trip(&packed);
1005
1006 // A table, which goes out as its members.
1007 let mut table = Set::new();
1008 for i in 0..500 {
1009 table.add(format!("member:{i}").as_bytes(), &Limits::DEFAULT);
1010 }
1011 assert_eq!(table.encoding(), Encoding::Hashtable);
1012 round_trip(&table);
1013
1014 // An empty set, which the keyspace never stores but which is one byte
1015 // and should not need a special case anywhere.
1016 round_trip(&Set::new());
1017 }
1018
1019 #[test]
1020 fn an_intset_past_its_ceiling_still_calls_itself_a_hashtable_after_a_round_trip() {
1021 let mut s = Set::new();
1022 for i in 0..2_000i64 {
1023 s.add(i.to_string().as_bytes(), &Limits::DEFAULT);
1024 }
1025 // Past `set-max-intset-entries`, so it answers `hashtable` while its
1026 // body is still an intset. That flag is the whole reason the form byte
1027 // has a spare bit, and losing it would change what a client sees.
1028 assert_eq!(s.encoding(), Encoding::Hashtable);
1029 let back = round_trip(&s);
1030 assert!(back.ints().is_some(), "still an intset underneath");
1031 }
1032
1033 #[test]
1034 fn a_large_set_of_integers_does_not_come_back_as_a_table() {
1035 // Past the point where the runs split, which is the case the gap form
1036 // exists for. Through a table it would be thirty bytes a member.
1037 let mut s = Set::new();
1038 for i in 0..200_000i64 {
1039 s.add((i * 3).to_string().as_bytes(), &Limits::DEFAULT);
1040 }
1041 assert!(s.ints().is_some_and(|i| i.as_bytes().is_none()), "split");
1042 let mut out = Vec::new();
1043 s.freeze(&mut out);
1044 let back = round_trip(&s);
1045 assert!(back.ints().is_some(), "came back as integers");
1046 // Gaps of three, so two bytes a member is the ceiling here and the
1047 // digits would have been six. The check is against the loose bound
1048 // because the point is the order of magnitude and not the byte.
1049 assert!(out.len() < s.len() * 3, "{} bytes", out.len());
1050 }
1051
1052 #[test]
1053 fn a_partitioned_set_comes_back_partitioned() {
1054 let mut s = Set::new();
1055 for i in 0..=PARTITION_AT {
1056 s.add(format!("member:{i}").as_bytes(), &Limits::DEFAULT);
1057 }
1058 assert!(matches!(s.body, Body::Split(_)), "split");
1059 let back = round_trip(&s);
1060 assert!(matches!(back.body, Body::Split(_)), "still split");
1061 }
1062
1063 #[test]
1064 fn a_frozen_set_that_arrives_damaged_is_an_error_and_not_a_panic() {
1065 let s = of(&["one", "two", "three"]);
1066 let mut out = Vec::new();
1067 s.freeze(&mut out);
1068 for cut in 1..out.len() {
1069 // Every prefix, because a short read off the store can end anywhere
1070 // and none of them may panic.
1071 let _ = Set::thaw(&out[..cut]);
1072 }
1073 assert_eq!(Set::thaw(&[]).err(), Some(Broken::Short));
1074 assert_eq!(Set::thaw(&[9]).err(), Some(Broken::Form));
1075 assert_eq!(Set::thaw(&[FORM_INTSET, 1, 2]).err(), Some(Broken::Body));
1076 }
1077
1078 /// What a set actually costs per member, which is half of M3's memory gate
1079 /// row and was an argument rather than a number until this was written.
1080 ///
1081 /// Run it with `cargo test -p yo-kv --release measure_bytes_per_member --
1082 /// --ignored --nocapture`. Ignored because a million members is not
1083 /// something every `cargo test` should pay for, and it prints rather than
1084 /// asserts because the number it prints is the thing being reported.
1085 ///
1086 /// Three shapes, because the gate names one of them and the other two are
1087 /// what most sets actually hold. Integers first, at every band an all
1088 /// integer set passes through, then the same counts of integers scattered
1089 /// over a wide range, then strings, which never see the intset at all.
1090 ///
1091 /// The scattered row is there because the dense one is the most favourable
1092 /// input this structure will ever be given and a gate read off it would be
1093 /// a gate read off the easy case. `0..n` is sorted, contiguous and arrives
1094 /// in order, so every run fills to `RUN_MAX` and never splits. Real integer
1095 /// sets are ids with holes in them, arriving in whatever order the writer
1096 /// had, and the number that matters is the one they produce.
1097 #[test]
1098 #[ignore = "a measurement, run it by name"]
1099 fn measure_bytes_per_member() {
1100 let limits = Limits::DEFAULT;
1101 for n in [512usize, 1_000, 100_000, 1_000_000] {
1102 let mut s = Set::new();
1103 for i in 0..n {
1104 s.add(i.to_string().as_bytes(), &limits);
1105 }
1106 println!(
1107 "int n={n:<9} band={:<10} total={:<10} per_member={:.2}",
1108 band(&s),
1109 s.memory_bytes(),
1110 s.memory_bytes() as f64 / n as f64
1111 );
1112 }
1113 // The same counts, scattered over a range sixteen times as wide and
1114 // arriving out of order. A cheap multiplicative shuffle rather than a
1115 // real generator, because what this needs is holes and disorder and not
1116 // statistical quality.
1117 for n in [512usize, 1_000, 100_000, 1_000_000] {
1118 let mut s = Set::new();
1119 for i in 0..n {
1120 let v = (i as u64).wrapping_mul(0x9e37_79b9_7f4a_7c15) % (n as u64 * 16);
1121 s.add(v.to_string().as_bytes(), &limits);
1122 }
1123 println!(
1124 "sparse n={n:<9} band={:<10} total={:<10} members={:<9} per_member={:.2}",
1125 band(&s),
1126 s.memory_bytes(),
1127 s.len(),
1128 s.memory_bytes() as f64 / s.len() as f64
1129 );
1130 }
1131 // Sixteen byte members, so the payload is a round number and the
1132 // overhead is whatever is above it.
1133 for n in [128usize, 1_000, 100_000, 1_000_000] {
1134 let mut s = Set::new();
1135 let mut payload = 0usize;
1136 for i in 0..n {
1137 let m = format!("member:{i:09}");
1138 payload += m.len();
1139 s.add(m.as_bytes(), &limits);
1140 }
1141 let total = s.memory_bytes();
1142 let per = |b: usize| b as f64 / n as f64;
1143 println!(
1144 "bytes n={n:<9} band={:<10} total={total:<10} payload={payload:<9} per_member={:.2} over_per_member={:.2} slots={:.2} rows={:.2} names={:.2} name_slack={:.2}",
1145 band(&s),
1146 per(total),
1147 per(total - payload),
1148 per(s.slot_bytes()),
1149 per(s.row_bytes()),
1150 per(s.name_bytes()),
1151 per(s.name_bytes() - payload)
1152 );
1153 }
1154 }
1155
1156 /// Which of the four a set is in, spelled out rather than through
1157 /// [`Set::encoding`], which folds the two table bands into one word because
1158 /// that is what `OBJECT ENCODING` has to say.
1159 fn band(s: &Set) -> &'static str {
1160 match &s.body {
1161 Body::Ints(_) => "intset",
1162 Body::Packed(_) => "listpack",
1163 Body::Table(_) => "table",
1164 Body::Split(_) => "split",
1165 }
1166 }
1167
1168 /// Everything about the partitioned band that needs a real set past the real
1169 /// threshold, in one test, because building 262,145 members is the expensive
1170 /// part and there is no reason to pay for it four times.
1171 #[test]
1172 fn a_set_past_the_threshold_splits_without_the_client_being_able_to_tell() {
1173 let limits = Limits::DEFAULT;
1174 let mut s = Set::new();
1175 // One short of the threshold. Still one table: the check is strictly
1176 // greater, so the set that splits is the one that has outgrown a table
1177 // and not the one that is about to.
1178 for i in 0..PARTITION_AT {
1179 assert!(s.add(format!("m{i}").as_bytes(), &limits));
1180 }
1181 assert!(matches!(s.body, Body::Table(_)));
1182 assert_eq!(s.len(), PARTITION_AT);
1183
1184 // The member that tips it over.
1185 assert!(s.add(b"tipping", &limits));
1186 assert!(matches!(s.body, Body::Split(_)), "it should have split");
1187 assert_eq!(s.len(), PARTITION_AT + 1);
1188
1189 // And the client cannot tell. This is the whole point: Redis has three
1190 // set encodings and a fourth word here breaks every suite that reads it.
1191 assert_eq!(s.encoding(), Encoding::Hashtable);
1192 assert_eq!(s.encoding().name(), "hashtable");
1193
1194 // Every member survived the rehash, asked both ways.
1195 assert!(s.contains(b"tipping"));
1196 assert!(s.contains(b"m0"));
1197 assert!(s.contains(b"m262143"));
1198 assert!(!s.contains(b"m262144"));
1199 assert!(s.has(&Needle::new(b"m1000")));
1200 assert!(!s.has(&Needle::new(b"nothing")));
1201
1202 // A rewrite is still not an add.
1203 assert!(!s.add(b"m0", &limits));
1204 assert_eq!(s.len(), PARTITION_AT + 1);
1205
1206 // Removing goes back through the same partition it went into, and the
1207 // band never demotes however far the set shrinks.
1208 assert!(s.remove(b"tipping"));
1209 assert!(!s.remove(b"tipping"));
1210 assert_eq!(s.len(), PARTITION_AT);
1211 assert!(matches!(s.body, Body::Split(_)), "promotion is one way");
1212 assert_eq!(s.encoding(), Encoding::Hashtable);
1213
1214 // The draw `SPOP` and `SRANDMEMBER` run on reaches the last position,
1215 // which is the one that has to land in the highest non empty partition.
1216 let last = s.at(s.len() - 1).expect("inside the set").to_vec();
1217 assert!(s.contains(&last));
1218 assert!(s.at(s.len()).is_none());
1219 assert_eq!(s.remove_at(s.len() - 1), Some(last.clone()));
1220 assert!(!s.contains(&last));
1221 assert!(s.drop_at(0));
1222 assert_eq!(s.len(), PARTITION_AT - 2);
1223
1224 // And a full scan sees every member exactly once.
1225 let mut seen = 0usize;
1226 let mut cursor = Cursor::START;
1227 let mut rounds = 0;
1228 loop {
1229 cursor = s.scan(cursor, 1_000, |_| seen += 1);
1230 rounds += 1;
1231 assert!(rounds < 100_000, "the scan is not finishing");
1232 if cursor.is_end() {
1233 break;
1234 }
1235 }
1236 assert_eq!(seen, s.len());
1237 }
1238
1239 /// The crossing that actually happens in production: a client holding a
1240 /// cursor from before the split. It has to see every member that stayed, and
1241 /// repeats are what the `SCAN` guarantee gives up in exchange.
1242 #[test]
1243 fn a_scan_survives_the_set_splitting_underneath_it() {
1244 let limits = Limits::DEFAULT;
1245 let mut s = Set::new();
1246 for i in 0..PARTITION_AT {
1247 s.add(format!("m{i}").as_bytes(), &limits);
1248 }
1249 assert!(matches!(s.body, Body::Table(_)));
1250
1251 let mut seen = Vec::new();
1252 let cursor = s.scan(Cursor::START, 5_000, |m| seen.push(m.to_vec()));
1253 assert!(!cursor.is_end(), "the scan should have stopped part way");
1254
1255 s.add(b"tipping", &limits);
1256 assert!(matches!(s.body, Body::Split(_)));
1257
1258 let mut cursor = cursor;
1259 let mut rounds = 0;
1260 loop {
1261 cursor = s.scan(cursor, 5_000, |m| seen.push(m.to_vec()));
1262 rounds += 1;
1263 assert!(rounds < 100_000, "the scan is not finishing");
1264 if cursor.is_end() {
1265 break;
1266 }
1267 }
1268 seen.sort_unstable();
1269 seen.dedup();
1270 assert_eq!(
1271 seen.len(),
1272 PARTITION_AT + 1,
1273 "the split lost a member the client was entitled to"
1274 );
1275 }
1276
1277 #[test]
1278 fn a_hint_past_the_threshold_builds_the_band_up_front() {
1279 let limits = Limits::DEFAULT;
1280 // A caller loading a million members should not fill one table, cross the
1281 // threshold and then rehash the lot.
1282 let s = Set::with_hint(b"first", 1_000_000, &limits);
1283 assert!(matches!(s.body, Body::Split(_)));
1284 assert_eq!(s.encoding(), Encoding::Hashtable);
1285 assert!(s.is_empty());
1286
1287 // A hint at the threshold is still one table, matching the add path.
1288 let s = Set::with_hint(b"first", PARTITION_AT, &limits);
1289 assert!(matches!(s.body, Body::Table(_)));
1290
1291 // And a hint is only a hint: the band takes members like anything else.
1292 let mut s = Set::with_hint(b"a", 1_000_000, &limits);
1293 assert!(s.add(b"a", &limits));
1294 assert!(!s.add(b"a", &limits));
1295 assert!(s.contains(b"a"));
1296 assert_eq!(s.len(), 1);
1297 assert_eq!(s.at(0).map(|m| m.to_vec()), Some(b"a".to_vec()));
1298 }
1299
1300 fn members(s: &Set) -> Vec<String> {
1301 let mut v: Vec<String> = s
1302 .iter()
1303 .map(|m| String::from_utf8(m.to_vec()).expect("utf8 in these tests"))
1304 .collect();
1305 // Order is not part of the contract and the three representations do not
1306 // agree on it, so every assertion here is against a sorted list.
1307 v.sort();
1308 v
1309 }
1310
1311 #[test]
1312 fn a_new_set_is_an_empty_intset() {
1313 let s = Set::new();
1314 assert_eq!(s.encoding(), Encoding::Intset);
1315 assert_eq!(s.len(), 0);
1316 assert!(s.is_empty());
1317 assert!(!s.contains(b"1"));
1318 assert_eq!(s.at(0), None);
1319 }
1320
1321 #[test]
1322 fn integers_stay_an_intset_and_come_back_as_members() {
1323 let s = of(&["1", "2", "3"]);
1324 assert_eq!(s.encoding(), Encoding::Intset);
1325 assert_eq!(members(&s), ["1", "2", "3"]);
1326 assert!(s.contains(b"2"));
1327 assert!(!s.contains(b"4"));
1328 assert!(!s.contains(b"two"));
1329 }
1330
1331 #[test]
1332 fn adding_the_same_member_twice_says_so_in_all_three() {
1333 let mut ints = of(&["1", "2"]);
1334 assert!(!ints.add(b"2", &Limits::DEFAULT));
1335 assert_eq!(ints.len(), 2);
1336
1337 let mut packed = of(&["a", "b"]);
1338 assert_eq!(packed.encoding(), Encoding::Listpack);
1339 assert!(!packed.add(b"b", &Limits::DEFAULT));
1340 assert_eq!(packed.len(), 2);
1341
1342 let mut table = of(&["a", "b"]);
1343 table.become_table(0);
1344 assert!(!table.add(b"b", &Limits::DEFAULT));
1345 assert_eq!(table.len(), 2);
1346 }
1347
1348 #[test]
1349 fn a_string_turns_a_small_intset_into_a_listpack() {
1350 let mut s = of(&["1", "2", "3"]);
1351 assert!(s.add(b"hello", &Limits::DEFAULT));
1352 assert_eq!(s.encoding(), Encoding::Listpack);
1353 assert_eq!(members(&s), ["1", "2", "3", "hello"]);
1354 assert!(s.contains(b"1"), "the integers survived the rewrite");
1355 assert!(s.contains(b"hello"));
1356 }
1357
1358 #[test]
1359 fn a_string_turns_a_big_intset_straight_into_a_table() {
1360 // The asymmetric rule, and the one a natural implementation gets wrong.
1361 // Two hundred integers is a legal intset and is already past the
1362 // listpack ceiling, so this never passes through the listpack at all.
1363 let mut s = Set::new();
1364 for i in 0..200 {
1365 s.add(i.to_string().as_bytes(), &Limits::DEFAULT);
1366 }
1367 assert_eq!(s.encoding(), Encoding::Intset);
1368 assert_eq!(s.len(), 200);
1369
1370 assert!(s.add(b"hello", &Limits::DEFAULT));
1371 assert_eq!(s.encoding(), Encoding::Hashtable);
1372 assert_eq!(s.len(), 201);
1373 assert!(s.contains(b"199"));
1374 assert!(s.contains(b"hello"));
1375 }
1376
1377 #[test]
1378 fn an_intset_holds_five_hundred_and_twelve_and_converts_at_the_next_one() {
1379 let mut s = Set::new();
1380 for i in 0..512 {
1381 s.add(i.to_string().as_bytes(), &Limits::DEFAULT);
1382 }
1383 assert_eq!(s.encoding(), Encoding::Intset, "512 is still an intset");
1384 assert_eq!(s.len(), 512);
1385
1386 s.add(b"512", &Limits::DEFAULT);
1387 assert_eq!(s.encoding(), Encoding::Hashtable, "513 is not");
1388 assert_eq!(s.len(), 513);
1389 // And not a listpack on the way, because 513 is well past 128.
1390 for i in 0..513 {
1391 assert!(s.contains(i.to_string().as_bytes()), "{i} survived");
1392 }
1393 }
1394
1395 #[test]
1396 fn a_listpack_converts_at_a_hundred_and_twenty_eight_members() {
1397 let mut s = of(&["x"]);
1398 assert_eq!(s.encoding(), Encoding::Listpack);
1399 for i in 0..127 {
1400 s.add(format!("m{i}").as_bytes(), &Limits::DEFAULT);
1401 }
1402 assert_eq!(s.len(), 128);
1403 assert_eq!(s.encoding(), Encoding::Listpack, "128 is still a listpack");
1404
1405 s.add(b"one more", &Limits::DEFAULT);
1406 assert_eq!(s.len(), 129);
1407 assert_eq!(s.encoding(), Encoding::Hashtable);
1408 assert!(s.contains(b"x"));
1409 assert!(s.contains(b"m126"));
1410 assert!(s.contains(b"one more"));
1411 }
1412
1413 #[test]
1414 fn a_long_member_converts_a_listpack_whatever_the_count() {
1415 let mut s = of(&["a"]);
1416 let long = vec![b'z'; 65];
1417 assert!(s.add(&long, &Limits::DEFAULT));
1418 assert_eq!(s.encoding(), Encoding::Hashtable, "65 is past 64");
1419 assert_eq!(s.len(), 2);
1420 assert!(s.contains(&long));
1421
1422 // And exactly at the boundary it does not.
1423 let mut ok = of(&["a"]);
1424 ok.add(&[b'z'; 64], &Limits::DEFAULT);
1425 assert_eq!(ok.encoding(), Encoding::Listpack, "64 fits");
1426 }
1427
1428 #[test]
1429 fn a_long_member_sends_an_intset_to_a_table_and_not_a_listpack() {
1430 let mut s = of(&["1", "2"]);
1431 assert!(s.add(&[b'z'; 65], &Limits::DEFAULT));
1432 assert_eq!(s.encoding(), Encoding::Hashtable);
1433 assert_eq!(s.len(), 3);
1434 }
1435
1436 #[test]
1437 fn the_limits_are_configuration_and_moving_them_moves_the_encodings() {
1438 let tight = Limits {
1439 max_intset_entries: 2,
1440 max_listpack_entries: 2,
1441 max_listpack_value: 3,
1442 };
1443 let mut s = Set::new();
1444 s.add(b"1", &tight);
1445 s.add(b"2", &tight);
1446 assert_eq!(s.encoding(), Encoding::Intset);
1447 s.add(b"3", &tight);
1448 assert_eq!(s.encoding(), Encoding::Hashtable, "three is past two");
1449
1450 // And a member longer than three characters cannot go in a listpack.
1451 let mut t = Set::new();
1452 t.add(b"abc", &tight);
1453 assert_eq!(t.encoding(), Encoding::Listpack, "three characters fit");
1454 t.add(b"defg", &tight);
1455 assert_eq!(t.encoding(), Encoding::Hashtable, "four do not");
1456 assert!(t.contains(b"abc"));
1457 assert!(t.contains(b"defg"));
1458
1459 // Including as the first member, which is a table from the off rather
1460 // than a listpack that converts on the next thing to arrive. Redis gets
1461 // to the same place by the other road: it creates the listpack, tries
1462 // the add, and converts before the reply.
1463 let mut u = Set::new();
1464 u.add(b"abcd", &tight);
1465 assert_eq!(u.encoding(), Encoding::Hashtable);
1466 assert!(u.contains(b"abcd"));
1467 }
1468
1469 #[test]
1470 fn with_hint_picks_the_representation_up_front() {
1471 let d = &Limits::DEFAULT;
1472 assert_eq!(
1473 Set::with_hint(b"1", 10, d).encoding(),
1474 Encoding::Intset,
1475 "an integer and few enough of them"
1476 );
1477 assert_eq!(
1478 Set::with_hint(b"1", 1000, d).encoding(),
1479 Encoding::Hashtable,
1480 "an integer and too many"
1481 );
1482 assert_eq!(
1483 Set::with_hint(b"x", 10, d).encoding(),
1484 Encoding::Listpack,
1485 "not an integer and few enough"
1486 );
1487 assert_eq!(
1488 Set::with_hint(b"x", 1000, d).encoding(),
1489 Encoding::Hashtable,
1490 "not an integer and too many"
1491 );
1492 }
1493
1494 #[test]
1495 fn removing_works_in_all_three_and_never_demotes() {
1496 let mut ints = of(&["1", "2", "3"]);
1497 assert!(ints.remove(b"2"));
1498 assert!(!ints.remove(b"2"));
1499 assert!(!ints.remove(b"nope"), "not an integer, so not a member");
1500 assert_eq!(members(&ints), ["1", "3"]);
1501 assert_eq!(ints.encoding(), Encoding::Intset);
1502
1503 let mut packed = of(&["a", "b", "c"]);
1504 assert!(packed.remove(b"b"));
1505 assert!(!packed.remove(b"b"));
1506 assert_eq!(members(&packed), ["a", "c"]);
1507 assert_eq!(packed.encoding(), Encoding::Listpack);
1508
1509 let mut table = of(&["a", "b", "c"]);
1510 table.become_table(0);
1511 assert!(table.remove(b"b"));
1512 assert!(!table.remove(b"b"));
1513 assert_eq!(members(&table), ["a", "c"]);
1514 assert_eq!(
1515 table.encoding(),
1516 Encoding::Hashtable,
1517 "down to two members and still a table"
1518 );
1519 }
1520
1521 #[test]
1522 fn a_set_can_be_emptied_a_member_at_a_time() {
1523 for mut s in [of(&["1", "2", "3"]), of(&["a", "b", "c"])] {
1524 let all: Vec<Vec<u8>> = s.iter().map(|m| m.to_vec()).collect();
1525 for m in &all {
1526 assert!(s.remove(m));
1527 }
1528 assert!(s.is_empty());
1529 assert_eq!(s.at(0), None);
1530 }
1531 }
1532
1533 #[test]
1534 fn removing_by_position_hands_the_member_back() {
1535 // What `SPOP` runs on. Drawing position zero every time has to empty the
1536 // set rather than run off the end or repeat a member, in all three.
1537 let mut table = of(&["a", "b", "c", "d"]);
1538 table.become_table(0);
1539 for mut s in [
1540 of(&["10", "20", "30", "40"]),
1541 of(&["a", "b", "c", "d"]),
1542 table,
1543 ] {
1544 let mut got = Vec::new();
1545 while !s.is_empty() {
1546 got.push(String::from_utf8(s.remove_at(0).expect("not empty")).expect("utf8"));
1547 }
1548 got.sort();
1549 assert_eq!(got.len(), 4, "four members and no repeats");
1550 assert_eq!(s.len(), 0);
1551 assert_eq!(s.remove_at(0), None);
1552 }
1553 }
1554
1555 #[test]
1556 fn an_integer_member_is_the_same_member_however_it_is_written() {
1557 // An intset holds 42 as a number, so `SADD s 42` twice is one member.
1558 // `042` does not parse as an integer, so it is a different member and it
1559 // converts the set, which is what a real server does too.
1560 let mut s = of(&["42"]);
1561 assert!(!s.add(b"42", &Limits::DEFAULT));
1562 assert_eq!(s.len(), 1);
1563 assert!(s.add(b"042", &Limits::DEFAULT));
1564 assert_eq!(s.encoding(), Encoding::Listpack);
1565 assert_eq!(members(&s), ["042", "42"]);
1566 assert!(s.contains(b"42"));
1567 assert!(s.contains(b"042"));
1568 }
1569
1570 #[test]
1571 fn the_small_bands_answer_a_scan_in_one_go() {
1572 for s in [of(&["1", "2", "3"]), of(&["a", "b", "c"])] {
1573 let mut seen = Vec::new();
1574 // A count of one, which the table band would honour and these two
1575 // do not, and a cursor from nowhere, which these two ignore.
1576 let next = s.scan(Cursor::at(1, 0, 99), 1, |m| seen.push(m.to_vec()));
1577 assert!(next.is_end(), "{:?} split a scan up", s.encoding());
1578 assert_eq!(seen.len(), 3);
1579 }
1580 }
1581
1582 #[test]
1583 fn the_table_band_walks_a_scan_in_windows_and_misses_nothing() {
1584 let mut s = Set::new();
1585 for i in 0..300 {
1586 s.add(format!("m{i}").as_bytes(), &Limits::DEFAULT);
1587 }
1588 assert_eq!(s.encoding(), Encoding::Hashtable);
1589
1590 let mut seen = Vec::new();
1591 let mut c = Cursor::START;
1592 let mut turns = 0;
1593 loop {
1594 c = s.scan(c, 7, |m| seen.push(m.to_vec()));
1595 turns += 1;
1596 assert!(turns < 100, "the scan did not finish");
1597 if c.is_end() {
1598 break;
1599 }
1600 }
1601 assert!(
1602 turns > 1,
1603 "a window of seven over three hundred took one turn"
1604 );
1605 seen.sort();
1606 seen.dedup();
1607 assert_eq!(seen.len(), 300, "every member came back at least once");
1608 }
1609
1610 #[test]
1611 fn a_conversion_loses_no_member_at_any_of_the_three_boundaries() {
1612 // One walk over each path out of an intset and out of a listpack, each
1613 // checking every member is still findable after the rewrite rather than
1614 // only checking the count.
1615 let mut wide = Set::new();
1616 for v in [i64::MIN, -1, 0, 1, i64::MAX] {
1617 wide.add(v.to_string().as_bytes(), &Limits::DEFAULT);
1618 }
1619 wide.add(b"str", &Limits::DEFAULT);
1620 assert_eq!(wide.encoding(), Encoding::Listpack);
1621 for v in [i64::MIN, -1, 0, 1, i64::MAX] {
1622 assert!(wide.contains(v.to_string().as_bytes()), "{v} survived");
1623 }
1624
1625 wide.add(&[b'q'; 100], &Limits::DEFAULT);
1626 assert_eq!(wide.encoding(), Encoding::Hashtable);
1627 for v in [i64::MIN, -1, 0, 1, i64::MAX] {
1628 assert!(
1629 wide.contains(v.to_string().as_bytes()),
1630 "{v} survived twice"
1631 );
1632 }
1633 assert!(wide.contains(b"str"));
1634 assert_eq!(wide.len(), 7);
1635 }
1636}