tinyufo/
buckets.rs

1// Copyright 2025 Cloudflare, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Concurrent storage backend
16
17use super::{Bucket, Key};
18use ahash::RandomState;
19use crossbeam_skiplist::{map::Entry, SkipMap};
20use flurry::HashMap;
21
22/// N-shard skip list. Memory efficient, constant time lookup on average, but a bit slower
23/// than hash map
24pub struct Compact<T>(Box<[SkipMap<Key, Bucket<T>>]>);
25
26impl<T: Send + 'static> Compact<T> {
27    /// Create a new [Compact]
28    pub fn new(total_items: usize, items_per_shard: usize) -> Self {
29        assert!(items_per_shard > 0);
30
31        let shards = std::cmp::max(total_items / items_per_shard, 1);
32        let mut shard_array = vec![];
33        for _ in 0..shards {
34            shard_array.push(SkipMap::new());
35        }
36        Self(shard_array.into_boxed_slice())
37    }
38
39    pub fn get(&self, key: &Key) -> Option<Entry<Key, Bucket<T>>> {
40        let shard = *key as usize % self.0.len();
41        self.0[shard].get(key)
42    }
43
44    pub fn get_map<V, F: FnOnce(Entry<Key, Bucket<T>>) -> V>(&self, key: &Key, f: F) -> Option<V> {
45        let v = self.get(key);
46        v.map(f)
47    }
48
49    fn insert(&self, key: Key, value: Bucket<T>) -> Option<()> {
50        let shard = key as usize % self.0.len();
51        let removed = self.0[shard].remove(&key);
52        self.0[shard].insert(key, value);
53        removed.map(|_| ())
54    }
55
56    fn remove(&self, key: &Key) {
57        let shard = *key as usize % self.0.len();
58        (&self.0)[shard].remove(key);
59    }
60}
61
62// Concurrent hash map, fast but use more memory
63pub struct Fast<T>(HashMap<Key, Bucket<T>, RandomState>);
64
65impl<T: Send + Sync> Fast<T> {
66    pub fn new(total_items: usize) -> Self {
67        Self(HashMap::with_capacity_and_hasher(
68            total_items,
69            RandomState::new(),
70        ))
71    }
72
73    pub fn get_map<V, F: FnOnce(&Bucket<T>) -> V>(&self, key: &Key, f: F) -> Option<V> {
74        let pinned = self.0.pin();
75        let v = pinned.get(key);
76        v.map(f)
77    }
78
79    fn insert(&self, key: Key, value: Bucket<T>) -> Option<()> {
80        let pinned = self.0.pin();
81        pinned.insert(key, value).map(|_| ())
82    }
83
84    fn remove(&self, key: &Key) {
85        let pinned = self.0.pin();
86        pinned.remove(key);
87    }
88}
89
90pub enum Buckets<T> {
91    Fast(Box<Fast<T>>),
92    Compact(Compact<T>),
93}
94
95impl<T: Send + Sync + 'static> Buckets<T> {
96    pub fn new_fast(items: usize) -> Self {
97        Self::Fast(Box::new(Fast::new(items)))
98    }
99
100    pub fn new_compact(items: usize, items_per_shard: usize) -> Self {
101        Self::Compact(Compact::new(items, items_per_shard))
102    }
103
104    pub fn insert(&self, key: Key, value: Bucket<T>) -> Option<()> {
105        match self {
106            Self::Compact(c) => c.insert(key, value),
107            Self::Fast(f) => f.insert(key, value),
108        }
109    }
110
111    pub fn remove(&self, key: &Key) {
112        match self {
113            Self::Compact(c) => c.remove(key),
114            Self::Fast(f) => f.remove(key),
115        }
116    }
117
118    pub fn get_map<V, F: FnOnce(&Bucket<T>) -> V>(&self, key: &Key, f: F) -> Option<V> {
119        match self {
120            Self::Compact(c) => c.get_map(key, |v| f(v.value())),
121            Self::Fast(c) => c.get_map(key, f),
122        }
123    }
124
125    #[cfg(test)]
126    pub fn get_queue(&self, key: &Key) -> Option<bool> {
127        self.get_map(key, |v| v.queue.is_main())
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134
135    #[test]
136    fn test_fast() {
137        let fast = Buckets::new_fast(10);
138
139        assert!(fast.get_map(&1, |_| ()).is_none());
140
141        let bucket = Bucket {
142            queue: crate::Location::new_small(),
143            weight: 1,
144            uses: Default::default(),
145            data: 1,
146        };
147        fast.insert(1, bucket);
148
149        assert_eq!(fast.get_map(&1, |v| v.data), Some(1));
150
151        fast.remove(&1);
152        assert!(fast.get_map(&1, |_| ()).is_none());
153    }
154
155    #[test]
156    fn test_compact() {
157        let compact = Buckets::new_compact(10, 2);
158
159        assert!(compact.get_map(&1, |_| ()).is_none());
160
161        let bucket = Bucket {
162            queue: crate::Location::new_small(),
163            weight: 1,
164            uses: Default::default(),
165            data: 1,
166        };
167        compact.insert(1, bucket);
168
169        assert_eq!(compact.get_map(&1, |v| v.data), Some(1));
170
171        compact.remove(&1);
172        assert!(compact.get_map(&1, |_| ()).is_none());
173    }
174}
tinyufo/buckets.rs

tinyufo/
buckets.rs
