1use std::num::NonZeroU32;
14use std::sync::atomic::{AtomicU32, Ordering};
15
16use serde::{Deserialize, Serialize};
17
18macro_rules! define_id_type {
19 ($name:ident) => {
20 #[derive(
21 Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize,
22 )]
23 #[serde(transparent)]
24 pub struct $name(NonZeroU32);
25
26 impl $name {
27 pub fn new(val: u32) -> Option<Self> {
30 NonZeroU32::new(val).map(Self)
31 }
32
33 pub fn to_raw(self) -> u32 {
35 self.0.get()
36 }
37 }
38
39 impl From<NonZeroU32> for $name {
40 fn from(val: NonZeroU32) -> Self {
41 Self(val)
42 }
43 }
44 };
45}
46
47define_id_type!(FileId);
48define_id_type!(SymbolId);
49define_id_type!(SnapshotId);
50define_id_type!(DataNodeId);
51
52#[derive(Debug)]
62pub struct IdGenerator<T> {
63 counter: AtomicU32,
64 _marker: std::marker::PhantomData<T>,
65}
66
67impl<T> IdGenerator<T> {
68 pub fn new() -> Self {
69 Self {
70 counter: AtomicU32::new(1),
71 _marker: std::marker::PhantomData,
72 }
73 }
74
75 pub fn with_start(start: u32) -> Self {
76 Self {
77 counter: AtomicU32::new(start.max(1)),
78 _marker: std::marker::PhantomData,
79 }
80 }
81
82 pub fn next(&self) -> T
83 where
84 T: From<NonZeroU32>,
85 {
86 let val = self.counter.fetch_add(1, Ordering::Relaxed);
87 let nz = NonZeroU32::new(val)
88 .expect("IdGenerator exhausted u32 space (allocated > u32::MAX ids)");
89 T::from(nz)
90 }
91
92 pub fn reserve(&self, count: u32) -> u32 {
98 if count == 0 {
99 return self.counter.load(Ordering::Relaxed);
100 }
101 let start = self.counter.fetch_add(count, Ordering::Relaxed);
102 assert!(
103 start != 0 && start.checked_add(count - 1).is_some(),
104 "IdGenerator exhausted u32 space (reserved > u32::MAX ids)"
105 );
106 start
107 }
108}
109
110impl<T> Default for IdGenerator<T> {
111 fn default() -> Self {
112 Self::new()
113 }
114}
115
116#[cfg(test)]
117mod tests {
118 use super::*;
119 use std::collections::HashSet;
120 use std::sync::Arc;
121
122 #[test]
123 fn file_id_sequential() {
124 let idgen = IdGenerator::<FileId>::new();
125 assert_eq!(idgen.next(), FileId::new(1).unwrap());
126 assert_eq!(idgen.next(), FileId::new(2).unwrap());
127 assert_eq!(idgen.next(), FileId::new(3).unwrap());
128 }
129
130 #[test]
131 fn symbol_id_sequential() {
132 let idgen = IdGenerator::<SymbolId>::new();
133 assert_eq!(idgen.next(), SymbolId::new(1).unwrap());
134 assert_eq!(idgen.next(), SymbolId::new(2).unwrap());
135 assert_eq!(idgen.next(), SymbolId::new(3).unwrap());
136 }
137
138 #[test]
139 fn snapshot_id_sequential() {
140 let idgen = IdGenerator::<SnapshotId>::new();
141 assert_eq!(idgen.next(), SnapshotId::new(1).unwrap());
142 assert_eq!(idgen.next(), SnapshotId::new(2).unwrap());
143 assert_eq!(idgen.next(), SnapshotId::new(3).unwrap());
144 }
145
146 #[test]
147 fn id_generator_thread_safe() {
148 let idgen = Arc::new(IdGenerator::<SymbolId>::new());
149 let mut handles = Vec::new();
150
151 for _ in 0..4 {
152 let idgen = Arc::clone(&idgen);
153 handles.push(std::thread::spawn(move || {
154 let mut ids = Vec::with_capacity(250);
155 for _ in 0..250 {
156 ids.push(idgen.next());
157 }
158 ids
159 }));
160 }
161
162 let all_ids: Vec<SymbolId> = handles
163 .into_iter()
164 .flat_map(|h| h.join().unwrap_or_default())
165 .collect();
166
167 let unique: HashSet<SymbolId> = all_ids.iter().copied().collect();
168 assert_eq!(unique.len(), 1000, "expected 1000 unique IDs");
169 assert!(
170 unique
171 .iter()
172 .all(|id| id.to_raw() >= 1 && id.to_raw() <= 1000),
173 "all IDs must be in range 1..=1000"
174 );
175 }
176
177 #[test]
178 fn file_id_serde_roundtrip() {
179 let original = FileId::new(42).unwrap();
180 let json = serde_json::to_string(&original).unwrap();
181 let roundtrip: FileId = serde_json::from_str(&json).unwrap();
182 assert_eq!(original, roundtrip);
183 }
184
185 #[test]
186 fn symbol_id_serde_roundtrip() {
187 let original = SymbolId::new(99).unwrap();
188 let json = serde_json::to_string(&original).unwrap();
189 let roundtrip: SymbolId = serde_json::from_str(&json).unwrap();
190 assert_eq!(original, roundtrip);
191 }
192
193 #[test]
194 fn file_id_to_raw() {
195 assert_eq!(FileId::new(7).unwrap().to_raw(), 7);
196 }
197
198 #[test]
199 fn id_generator_default() {
200 let idgen = IdGenerator::<FileId>::default();
201 assert_eq!(idgen.next(), FileId::new(1).unwrap());
202 }
203
204 #[test]
205 fn id_generator_with_start_zero_safeguard() {
206 let idgen = IdGenerator::<SymbolId>::with_start(0);
207 let first = idgen.next();
208 assert_eq!(
209 first.to_raw(),
210 1,
211 "with_start(0) must sanitize to start ID 1"
212 );
213 }
214
215 #[test]
216 fn zero_is_invalid() {
217 assert!(FileId::new(0).is_none());
218 assert!(SymbolId::new(0).is_none());
219 assert!(SnapshotId::new(0).is_none());
220 assert!(DataNodeId::new(0).is_none());
221 }
222
223 #[test]
224 fn zero_rejected_on_deserialize() {
225 let res: Result<FileId, _> = serde_json::from_str("0");
226 assert!(res.is_err(), "deserializing 0 must fail");
227 }
228
229 #[test]
230 fn option_id_is_niche_optimized() {
231 assert_eq!(
233 std::mem::size_of::<Option<FileId>>(),
234 std::mem::size_of::<FileId>()
235 );
236 assert_eq!(std::mem::size_of::<Option<FileId>>(), 4);
237 assert_eq!(std::mem::size_of::<Option<SymbolId>>(), 4);
238 assert_eq!(std::mem::size_of::<Option<SnapshotId>>(), 4);
239 assert_eq!(std::mem::size_of::<Option<DataNodeId>>(), 4);
240 }
241 #[test]
242 fn counters_hand_out_unique_ids_under_contention() {
243 let idgen = Arc::new(IdGenerator::<SymbolId>::new());
244 let handles: Vec<_> = (0..8)
245 .map(|_| {
246 let idgen = Arc::clone(&idgen);
247 std::thread::spawn(move || {
248 let mut ids = Vec::with_capacity(10_000);
249 for _ in 0..10_000 {
250 ids.push(idgen.next());
251 }
252 ids
253 })
254 })
255 .collect();
256
257 let all: Vec<SymbolId> = handles
258 .into_iter()
259 .flat_map(|handle| handle.join().unwrap())
260 .collect();
261 let unique: HashSet<SymbolId> = all.iter().copied().collect();
262 assert_eq!(unique.len(), 80_000, "every identifier is handed out once");
263 let raw: Vec<u32> = unique.iter().map(|id| id.to_raw()).collect();
264 assert_eq!(raw.iter().min(), Some(&1));
265 assert_eq!(raw.iter().max(), Some(&80_000));
266 }
267}