From e88f6f6ee978ba6ce5df3a999e7efe2717e8532e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Tue, 4 Mar 2025 15:09:46 +0100
Subject: [PATCH] Add property-based testing for `RadixAllocator` (#3068)

---
 Cargo.lock               |   1 +
 backends/v3/Cargo.toml   |   1 +
 backends/v3/src/radix.rs | 138 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 140 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index 4603f77d2..16f8c80de 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4847,6 +4847,7 @@ dependencies = [
  "rand 0.8.5",
  "regex",
  "reqwest 0.11.27",
+ "rustc-hash 2.1.1",
  "serde",
  "serde_json",
  "slotmap",
diff --git a/backends/v3/Cargo.toml b/backends/v3/Cargo.toml
index 996290ed3..588a2716f 100644
--- a/backends/v3/Cargo.toml
+++ b/backends/v3/Cargo.toml
@@ -71,6 +71,7 @@ prost-build = "0.12.1"
 [dev-dependencies]
 criterion = "0.3"
 itertools = "0.13"
+rustc-hash = "2"
 
 [features]
 default = ["ngrok"]
diff --git a/backends/v3/src/radix.rs b/backends/v3/src/radix.rs
index 81ce61d1d..aea69693f 100644
--- a/backends/v3/src/radix.rs
+++ b/backends/v3/src/radix.rs
@@ -635,6 +635,12 @@ fn shared_prefix(left: &[u32], right: &[u32], block_size: usize) -> usize {
 mod tests {
     use std::sync::Arc;
 
+    use rand::{
+        distributions::Uniform, prelude::Distribution, rngs::SmallRng, seq::SliceRandom,
+        SeedableRng,
+    };
+    use rustc_hash::FxHashSet;
+
     use super::*;
 
     #[test]
@@ -900,4 +906,136 @@ mod tests {
         assert_eq!(blocks, vec![0, 1]);
         assert_eq!(node_id, trie.find(&[0, 1], &mut blocks))
     }
+
+    struct AllocationWithInfo {
+        allocation: BlockAllocation,
+        // We are doing a lot of set operations and `FxBuildHasher` is
+        // muc faster for a set of integers.
+        blockset: FxHashSet<u32>,
+        non_prefix_blocks: FxHashSet<u32>,
+    }
+
+    #[test]
+    fn invariants_hold_on_many_operations_remove_all() {
+        invariants_hold_on_many_insertions(true);
+    }
+
+    #[test]
+    fn invariants_hold_on_many_operations_remove_subset() {
+        invariants_hold_on_many_insertions(false);
+    }
+
+    fn invariants_hold_on_many_insertions(remove_all: bool) {
+        // Small vocabulary sizes lead to violations more quickly due to
+        // prefix sharing, etc.
+        const VOCAB_SIZE: u32 = 2;
+        const DATA_LEN: usize = 1_000;
+
+        const MAX_PREFILL_LEN: usize = 8;
+        const MAX_DECODE_LEN: usize = 8;
+
+        let vocab_range = Uniform::new(0, VOCAB_SIZE);
+        let data_range = Uniform::new(0, DATA_LEN);
+        let prefill_len_range = Uniform::new(0, MAX_PREFILL_LEN);
+        let decode_len_range = Uniform::new(0, MAX_DECODE_LEN);
+
+        let mut rng = SmallRng::seed_from_u64(64);
+        let data = (0..DATA_LEN)
+            .map(|_| vocab_range.sample(&mut rng))
+            .collect::<Vec<_>>();
+        let mut allocator = RadixAllocator::new(1, 100, None);
+
+        let mut allocations = Vec::new();
+
+        for i in 0..100_000 {
+            // Allocate until all blocks are used.
+            'allocation: loop {
+                // Use offset 0 half of the times for prefix sharing.
+                let prefill_offset = data_range.sample(&mut rng);
+                let prefill_len = prefill_len_range.sample(&mut rng);
+                let decode_len = decode_len_range.sample(&mut rng);
+
+                let prefill =
+                    data[prefill_offset..data.len().min(prefill_offset + prefill_len)].to_vec();
+
+                let allocation = match allocator
+                    .allocate((prefill.len() + decode_len) as u32, Some(Arc::new(prefill)))
+                {
+                    Some(allocation) => allocation,
+                    None => break 'allocation,
+                };
+                let non_prefix_blocks = allocation.blocks[allocation.prefix_len as usize..]
+                    .iter()
+                    .copied()
+                    .collect::<FxHashSet<_>>();
+                let blockset = allocation.blocks.iter().copied().collect::<FxHashSet<_>>();
+
+                // No duplicate blocks in an allocation.
+                assert_eq!(
+                    allocation.blocks.len(),
+                    blockset.len(),
+                    "Duplicate blocks in allocation"
+                );
+
+                allocations.push(AllocationWithInfo {
+                    allocation,
+                    blockset,
+                    non_prefix_blocks,
+                });
+            }
+
+            // Check invariants. Skip first iteration, since there is no prefix sharing yet.
+            if i > 1 {
+                check_allocation_invariants(&allocations);
+            }
+
+            // Remove 20% of the allocations, randomly.
+            if remove_all {
+                allocations.into_iter().for_each(|allocation| {
+                    allocator.free(
+                        allocation.allocation.blocks.clone(),
+                        allocation.allocation.allocation_id,
+                    )
+                });
+                allocations = Vec::new();
+            } else {
+                allocations.shuffle(&mut rng);
+                let remove_index = (allocations.len() as f64 * 0.8) as usize;
+                for allocation in allocations.drain(remove_index..) {
+                    allocator.free(
+                        allocation.allocation.blocks.clone(),
+                        allocation.allocation.allocation_id,
+                    );
+                }
+            }
+        }
+    }
+
+    fn check_allocation_invariants(allocations: &[AllocationWithInfo]) {
+        for i in 0..allocations.len() {
+            let allocation = &allocations[i];
+
+            // 0 is used for health checks, must not be used.
+            assert!(
+                !allocation.blockset.contains(&0),
+                "Block 0 must not be allocated"
+            );
+
+            // No duplicate blocks in an allocation.
+            assert_eq!(
+                allocation.allocation.blocks.len(),
+                allocation.blockset.len(),
+                "Duplicate blocks in allocation"
+            );
+
+            for other_allocation in &allocations[i + 1..] {
+                assert!(
+                    other_allocation
+                        .non_prefix_blocks
+                        .is_disjoint(&allocation.non_prefix_blocks),
+                    "Allocations share non-prefix blocks"
+                )
+            }
+        }
+    }
 }