mz_ore/
id_gen.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
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 in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! ID generation utilities.
17
18use hibitset::BitSet;
19use serde::{Serialize, Serializer};
20use std::borrow::Borrow;
21use std::fmt;
22use std::hash::Hash;
23use std::marker::PhantomData;
24use std::ops::{AddAssign, Sub};
25use std::sync::atomic::{AtomicU64, Ordering};
26use std::sync::{Arc, Mutex};
27use uuid::Uuid;
28
29use rand::RngExt;
30use rand::rngs::StdRng;
31
32use crate::cast::CastFrom;
33
34/// Manages the allocation of unique IDs.
35#[derive(Debug, Clone)]
36pub struct Gen<Id> {
37    id: u64,
38    phantom: PhantomData<Id>,
39}
40
41impl<Id> Default for Gen<Id> {
42    fn default() -> Self {
43        Self {
44            id: 0,
45            phantom: PhantomData,
46        }
47    }
48}
49
50impl<Id: From<u64>> Gen<Id> {
51    /// Allocates a new identifier of type `Id` and advances the generator.
52    pub fn allocate_id(&mut self) -> Id {
53        let id = self.id;
54        self.id += 1;
55        id.into()
56    }
57}
58
59/// A generator of u64-bit IDs.
60pub type IdGen = Gen<u64>;
61
62/// Manages the allocation of unique IDs.
63///
64/// Atomic version of `Gen`, for sharing between threads.
65#[derive(Debug)]
66pub struct AtomicGen<Id> {
67    id: AtomicU64,
68    phantom: PhantomData<Id>,
69}
70
71impl<Id> Default for AtomicGen<Id> {
72    fn default() -> Self {
73        Self {
74            id: AtomicU64::new(0),
75            phantom: PhantomData,
76        }
77    }
78}
79
80impl<Id: From<u64> + Default> AtomicGen<Id> {
81    /// Allocates a new identifier of type `Id` and advances the generator.
82    pub fn allocate_id(&self) -> Id {
83        // The only purpose of the atomic here is to ensure every caller receives a distinct ID,
84        // there are no requirements on the order in which IDs are produced and there is no other
85        // state protected by this atomic. `Relaxed` ordering is therefore sufficient.
86        let id = self.id.fetch_add(1, Ordering::Relaxed);
87        id.into()
88    }
89}
90
91/// A generator of u64-bit IDs.
92///
93/// Atomic version of `IdGen`, for sharing between threads.
94pub type AtomicIdGen = AtomicGen<u64>;
95
96/// IdAllocator common traits.
97pub trait IdGenerator:
98    From<u8> + AddAssign + Sub + PartialOrd + Copy + Eq + Hash + Ord + Serialize + fmt::Display
99{
100}
101
102impl<T> IdGenerator for T where
103    T: From<u8> + AddAssign + Sub + PartialOrd + Copy + Eq + Hash + Ord + Serialize + fmt::Display
104{
105}
106
107/// Manages allocation of numeric IDs.
108#[derive(Debug)]
109pub struct IdAllocator<A: IdAllocatorInner>(pub Arc<Mutex<A>>);
110
111// `derive(Clone)` would require `A: Clone`, but clones share the inner allocator.
112impl<A: IdAllocatorInner> Clone for IdAllocator<A> {
113    fn clone(&self) -> Self {
114        IdAllocator(Arc::clone(&self.0))
115    }
116}
117
118/// Common trait for id allocators.
119pub trait IdAllocatorInner: std::fmt::Debug + Send {
120    /// Name of the allocator.
121    const NAME: &'static str;
122    /// Construct an allocator with the given range. Returned ids will be OR'd with `mask`. `mask`
123    /// must not have any bits that could be set by a number <= `max`.
124    fn new(min: u32, max: u32, mask: u32) -> Self;
125    /// Allocate a new id.
126    fn alloc(&mut self) -> Option<u32>;
127    /// Deallocate a used id, making it available for reuse.
128    fn remove(&mut self, id: u32);
129}
130
131/// IdAllocator using a HiBitSet.
132#[derive(Debug)]
133pub struct IdAllocatorInnerBitSet {
134    next: StdRng,
135    min: u32,
136    max: u32,
137    mask: u32,
138    used: BitSet,
139}
140
141impl IdAllocatorInner for IdAllocatorInnerBitSet {
142    const NAME: &'static str = "hibitset";
143
144    fn new(min: u32, max: u32, mask: u32) -> Self {
145        let total = usize::cast_from(max - min);
146        assert!(total < BitSet::BITS_PER_USIZE.pow(4));
147        IdAllocatorInnerBitSet {
148            next: rand::make_rng(),
149            min,
150            max,
151            mask,
152            used: BitSet::new(),
153        }
154    }
155
156    fn alloc(&mut self) -> Option<u32> {
157        let range = self.min..=self.max;
158        let init = self.next.random_range(range);
159        let mut next = init;
160        loop {
161            // Because hibitset has a hard maximum of 64**4 (~16 million), subtract the min in case
162            // max is above that. This is safe because we already asserted above that `max - min <
163            // 64**4`.
164            let stored = next - self.min;
165            if !self.used.add(stored) {
166                assert!(
167                    next & self.mask == 0,
168                    "chosen ID must not intersect with mask:\n{:#034b}\n{:#034b}",
169                    next,
170                    self.mask
171                );
172                return Some(next | self.mask);
173            }
174            // Existing value, increment and try again. Wrap once we hit max back to min.
175            next = if next == self.max { self.min } else { next + 1 };
176            // We fully wrapped around. BitSet doesn't have a rank or count method, so we can't
177            // compute this early.
178            if next == init {
179                return None;
180            }
181        }
182    }
183
184    fn remove(&mut self, id: u32) {
185        let id = (!self.mask) & id;
186        let stored = id - self.min;
187        self.used.remove(stored);
188    }
189}
190
191impl<A: IdAllocatorInner> IdAllocator<A> {
192    /// Creates a new `IdAllocator` that will assign IDs between `min` and
193    /// `max`, both inclusive.
194    pub fn new(min: u32, max: u32, mask: u32) -> IdAllocator<A> {
195        assert!(min <= max);
196        if mask != 0 && max > 0 {
197            // mask_check is all 1s in any bit set by all numbers >= max. Assert that the mask
198            // doesn't share any bits with those.
199            let mask_check = (1 << (max.ilog2() + 1)) - 1;
200            assert_eq!(mask & mask_check, 0, "max and mask share bits");
201        }
202        let inner = A::new(min, max, mask);
203        IdAllocator(Arc::new(Mutex::new(inner)))
204    }
205
206    /// Allocates a new ID randomly distributed between min and max.
207    ///
208    /// Returns `None` if the allocator is exhausted.
209    ///
210    /// The ID associated with the [`IdHandle`] will be freed when all of the
211    /// outstanding [`IdHandle`]s have been dropped.
212    pub fn alloc(&self) -> Option<IdHandle<u32, A>> {
213        let inner = Arc::new(internal::IdHandleInner::new(self)?);
214        Some(IdHandle::Dynamic(inner))
215    }
216
217    // Attempt to allocate a new ID. We want the ID to be randomly distributed in the range. To do
218    // this, choose a random candidate. Check if it's already in the used set, and increment in a
219    // loop until an unused slot is found. Ideally we could ask used set what its next used or
220    // unused id is after some given X, but that's not part of the API. This means that when the
221    // range is large and the used set is near full, we will spend a lot of cycles looking for an
222    // open slot. However, we limit the number of connections to the thousands, and connection IDs
223    // have 20 bits (~1 million) of space, so it is not currently possible to enter that state.
224    fn alloc_internal(&self) -> Option<u32> {
225        let mut inner = self.0.lock().expect("lock poisoned");
226        inner.alloc()
227    }
228
229    fn free_internal(&self, id: u32) {
230        let mut inner = self.0.lock().expect("lock poisoned");
231        inner.remove(id);
232    }
233}
234
235/// A clone-able owned reference to an ID.
236///
237/// Once all of the [`IdHandle`]s referencing an ID have been dropped, we will then free the ID
238/// for later re-use.
239#[derive(Debug)]
240pub enum IdHandle<T, A: IdAllocatorInner> {
241    /// An ID "allocated" at compile time.
242    ///
243    /// Note: It is *entirely* up to the caller to make sure the provided ID is
244    /// not used by a dynamic ID allocator.
245    Static(T),
246    /// An ID allocated at runtime, gets freed once all handles have been dropped.
247    Dynamic(Arc<internal::IdHandleInner<T, A>>),
248}
249
250impl<T: Clone, A: IdAllocatorInner> Clone for IdHandle<T, A> {
251    fn clone(&self) -> Self {
252        match self {
253            IdHandle::Static(t) => IdHandle::Static(t.clone()),
254            IdHandle::Dynamic(handle) => IdHandle::Dynamic(Arc::clone(handle)),
255        }
256    }
257}
258
259impl<T: IdGenerator, A: IdAllocatorInner> IdHandle<T, A> {
260    /// Returns the raw ID inside of this handle.
261    ///
262    /// Use with caution! It is easy for a raw ID to outlive the handle from
263    /// which it came. You are responsible for ensuring that your use of the raw
264    /// ID does not lead to ID reuse bugs.
265    pub fn unhandled(&self) -> T {
266        *self.borrow()
267    }
268}
269
270impl<T: IdGenerator, A: IdAllocatorInner> PartialEq for IdHandle<T, A> {
271    fn eq(&self, other: &Self) -> bool {
272        self.unhandled() == other.unhandled()
273    }
274}
275impl<T: IdGenerator, A: IdAllocatorInner> Eq for IdHandle<T, A> {}
276
277impl<T: IdGenerator, A: IdAllocatorInner> PartialOrd for IdHandle<T, A> {
278    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
279        Some(self.cmp(other))
280    }
281}
282
283impl<T: IdGenerator, A: IdAllocatorInner> Ord for IdHandle<T, A> {
284    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
285        self.unhandled().cmp(&other.unhandled())
286    }
287}
288
289impl<T, A: IdAllocatorInner> Borrow<T> for IdHandle<T, A> {
290    fn borrow(&self) -> &T {
291        match self {
292            IdHandle::Static(id) => id,
293            IdHandle::Dynamic(inner) => &inner.id,
294        }
295    }
296}
297
298impl<T: IdGenerator, A: IdAllocatorInner> fmt::Display for IdHandle<T, A> {
299    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
300        self.unhandled().fmt(f)
301    }
302}
303
304impl<T: IdGenerator, A: IdAllocatorInner> Serialize for IdHandle<T, A> {
305    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
306    where
307        S: Serializer,
308    {
309        self.unhandled().serialize(serializer)
310    }
311}
312
313mod internal {
314    use std::fmt::Debug;
315    use std::sync::Arc;
316
317    use crate::cast::CastFrom;
318    use crate::id_gen::{IdAllocator, IdAllocatorInner};
319
320    pub struct IdHandleInner<T, A: IdAllocatorInner> {
321        /// A handle to the [`IdAllocator`] used to allocated the provided id.
322        pub(super) allocator: IdAllocator<A>,
323        /// The actual ID that was allocated.
324        pub(super) id: T,
325        stored: u32,
326    }
327
328    impl<T: Debug, A: IdAllocatorInner> Debug for IdHandleInner<T, A> {
329        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
330            f.debug_struct("IdHandleInner")
331                .field("id", &self.id)
332                .field("stored", &self.stored)
333                .finish_non_exhaustive()
334        }
335    }
336
337    impl<T, A: IdAllocatorInner> IdHandleInner<T, A>
338    where
339        T: CastFrom<u32>,
340    {
341        pub fn new(allocator: &IdAllocator<A>) -> Option<Self> {
342            let stored = allocator.alloc_internal()?;
343            Some(IdHandleInner {
344                allocator: IdAllocator(Arc::clone(&allocator.0)),
345                id: T::cast_from(stored),
346                stored,
347            })
348        }
349    }
350
351    impl<T, A: IdAllocatorInner> Drop for IdHandleInner<T, A> {
352        fn drop(&mut self) {
353            // Release our ID for later re-use.
354            self.allocator.free_internal(self.stored);
355        }
356    }
357}
358
359/// Number of bits the org id is offset into a connection id.
360pub const ORG_ID_OFFSET: usize = 19;
361
362/// Max (inclusive) connection id that can be produced.
363pub const MAX_ORG_ID: u32 = (1 << ORG_ID_OFFSET) - 1;
364
365/// Extracts the lower 12 bits from an org id. These are later used as the [31, 20] bits of a
366/// connection id to help route cancellation requests.
367pub fn org_id_conn_bits(uuid: &Uuid) -> u32 {
368    let lower = uuid.as_u128();
369    let lower = (lower & 0xFFF) << ORG_ID_OFFSET;
370    let lower: u32 = lower.try_into().expect("must fit");
371    lower
372}
373
374/// Returns the portion of the org's UUID present in connection id.
375pub fn conn_id_org_uuid(conn_id: u32) -> String {
376    const UPPER: [char; 16] = [
377        '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F',
378    ];
379
380    // Extract UUID from conn_id: upper 12 bits excluding the first.
381    let orgid = usize::try_from((conn_id >> ORG_ID_OFFSET) & 0xFFF).expect("must cast");
382    // Convert the bits into a 3 char string and inject into the resolver template.
383    let mut dst = String::with_capacity(3);
384    dst.push(UPPER[(orgid >> 8) & 0xf]);
385    dst.push(UPPER[(orgid >> 4) & 0xf]);
386    dst.push(UPPER[orgid & 0xf]);
387    dst
388}
389
390/// Generate a random temporary ID.
391///
392/// Concretely we generate a UUIDv4 and return the last 12 characters for maximum uniqueness.
393///
394/// Note: the reason we use the last 12 characters is because the bits 6, 7, and 12 - 15
395/// are all hard coded <https://www.rfc-editor.org/rfc/rfc4122#section-4.4>.
396/// ```
397/// use mz_ore::id_gen::temp_id;
398///
399/// let temp = temp_id();
400/// assert_eq!(temp.len(), 12);
401/// assert!(temp.is_ascii());
402/// ```
403pub fn temp_id() -> String {
404    let temp_uuid = uuid::Uuid::new_v4().as_hyphenated().to_string();
405    temp_uuid.chars().rev().take_while(|c| *c != '-').collect()
406}
407
408#[cfg(test)]
409mod tests {
410    use std::collections::BTreeMap;
411
412    use crate::assert_none;
413
414    use super::*;
415
416    #[crate::test]
417    fn test_conn_org() {
418        let uuid = Uuid::parse_str("9e37ec59-56f4-450a-acbd-18ff14f10ca8").unwrap();
419        let lower = org_id_conn_bits(&uuid);
420        let org_lower_uuid = conn_id_org_uuid(lower);
421        assert_eq!(org_lower_uuid, "CA8");
422    }
423
424    #[crate::test]
425    fn test_id_gen() {
426        test_ad_allocator::<IdAllocatorInnerBitSet>();
427    }
428
429    // Test masks and maxs that intersect panic.
430    #[crate::test]
431    #[should_panic]
432    fn test_mask_intersect<A: IdAllocatorInner>() {
433        let env_lower = org_id_conn_bits(&uuid::Uuid::from_u128(u128::MAX));
434        let ida = IdAllocator::<IdAllocatorInnerBitSet>::new(
435            1 << ORG_ID_OFFSET,
436            1 << ORG_ID_OFFSET,
437            env_lower,
438        );
439        let id = ida.alloc().unwrap();
440        assert_eq!(id.unhandled(), (0xfff << ORG_ID_OFFSET) | MAX_ORG_ID);
441    }
442
443    fn test_ad_allocator<A: IdAllocatorInner>() {
444        test_id_alloc::<A>();
445        test_static_id_sorting::<A>();
446        test_id_reuse::<A>();
447        test_display::<A>();
448        test_map_lookup::<A>();
449        test_serialization::<A>();
450        test_mask::<A>();
451        test_mask_envd::<A>();
452    }
453
454    fn test_mask<A: IdAllocatorInner>() {
455        let ida = IdAllocator::<A>::new(1, 1, 0xfff << 20);
456        let id = ida.alloc().unwrap();
457        assert_eq!(id.unhandled(), (0xfff << 20) | 1);
458    }
459
460    // Test that the random conn id and and uuid each with all bits set don't intersect.
461    fn test_mask_envd<A: IdAllocatorInner>() {
462        let env_lower = org_id_conn_bits(&uuid::Uuid::from_u128(u128::MAX));
463        let ida = IdAllocator::<A>::new(MAX_ORG_ID, MAX_ORG_ID, env_lower);
464        let id = ida.alloc().unwrap();
465        assert_eq!(id.unhandled(), (0xfff << ORG_ID_OFFSET) | MAX_ORG_ID);
466    }
467
468    fn test_id_alloc<A: IdAllocatorInner>() {
469        let ida = IdAllocator::<A>::new(3, 5, 0);
470        let id3 = ida.alloc().unwrap();
471        let id4 = ida.alloc().unwrap();
472        let id5 = ida.alloc().unwrap();
473        assert_ne!(id3, id4);
474        assert_ne!(id3, id5);
475        assert_ne!(id4, id5);
476        drop(id4);
477        let _id4 = ida.alloc().unwrap();
478        drop(id5);
479        drop(id3);
480        let _id5 = ida.alloc().unwrap();
481        let _id3 = ida.alloc().unwrap();
482        match ida.alloc() {
483            Some(id) => panic!(
484                "id allocator returned {}, not expected id exhaustion error",
485                id
486            ),
487            None => (),
488        }
489    }
490
491    fn test_static_id_sorting<A: IdAllocatorInner>() {
492        let ida = IdAllocator::<A>::new(0, 0, 0);
493        let id0 = ida.alloc().unwrap();
494        let id1 = IdHandle::Static(1);
495        assert!(id0 < id1);
496
497        let ida = IdAllocator::<A>::new(1, 1, 0);
498        let id0 = IdHandle::Static(0);
499        let id1 = ida.alloc().unwrap();
500        assert!(id0 < id1);
501    }
502
503    fn test_id_reuse<A: IdAllocatorInner>() {
504        let allocator = IdAllocator::<A>::new(10, 11, 0);
505
506        let id_a = allocator.alloc().unwrap();
507        let a = id_a.unhandled();
508        let id_a_clone = id_a.clone();
509        // a should not get freed.
510        drop(id_a);
511
512        // There are only two slots, so trying to allocate 2 more should fail the second time.
513        let _id_b = allocator.alloc().unwrap();
514        assert_none!(allocator.alloc());
515
516        // a should get freed since all outstanding references have been dropped.
517        drop(id_a_clone);
518
519        // We should re-use a.
520        let id_c = allocator.alloc().unwrap();
521        assert_eq!(id_c.unhandled(), a);
522    }
523
524    fn test_display<A: IdAllocatorInner>() {
525        let allocator = IdAllocator::<A>::new(65_000, 65_000, 0);
526
527        let id_a = allocator.alloc().unwrap();
528        assert_eq!(id_a.unhandled(), 65_000);
529
530        // An IdHandle should use the inner type's Display impl.
531        let id_display = format!("{id_a}");
532        let val_display = format!("{}", id_a.unhandled());
533
534        assert_eq!(id_display, val_display);
535    }
536
537    fn test_map_lookup<A: IdAllocatorInner>() {
538        let allocator = IdAllocator::<A>::new(99, 101, 0);
539
540        let id_a = allocator.alloc().unwrap();
541        let a = id_a.unhandled();
542
543        let mut btree = BTreeMap::new();
544        btree.insert(id_a, "hello world");
545
546        // We should be able to lookup an IdHandle, based on just the value.
547        let entry = btree.remove(&a).unwrap();
548        assert_eq!(entry, "hello world");
549
550        assert!(btree.is_empty());
551    }
552
553    fn test_serialization<A: IdAllocatorInner>() {
554        let allocator = IdAllocator::<A>::new(42, 42, 0);
555
556        let id_a = allocator.alloc().unwrap();
557        assert_eq!(id_a.unhandled(), 42);
558
559        // An IdHandle should serialize the same as the inner value.
560        let id_json = serde_json::to_string(&id_a).unwrap();
561        let val_json = serde_json::to_string(&id_a.unhandled()).unwrap();
562
563        assert_eq!(id_json, val_json);
564    }
565}
mz_ore/id_gen.rs

mz_ore/
id_gen.rs
