milli_core/update/
concurrent_available_ids.rs

1use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
2
3use roaring::RoaringBitmap;
4
5/// A concurrent ID generate that will never return the same ID twice.
6#[derive(Debug)]
7pub struct ConcurrentAvailableIds {
8    /// The current tree node ID we should use if there is no other IDs available.
9    current: AtomicU32,
10    /// The total number of tree node IDs used.
11    used: AtomicU64,
12
13    /// A list of IDs to exhaust before picking IDs from `current`.
14    available: RoaringBitmap,
15    /// The current Nth ID to select in the bitmap.
16    select_in_bitmap: AtomicU32,
17    /// Tells if you should look in the roaring bitmap or if all the IDs are already exhausted.
18    look_into_bitmap: AtomicBool,
19}
20
21impl ConcurrentAvailableIds {
22    /// Creates an ID generator returning unique IDs, avoiding the specified used IDs.
23    pub fn new(used: RoaringBitmap) -> ConcurrentAvailableIds {
24        let last_id = used.max().map_or(0, |id| id + 1);
25        let used_ids = used.len();
26        let available = RoaringBitmap::from_sorted_iter(0..last_id).unwrap() - used;
27
28        ConcurrentAvailableIds {
29            current: AtomicU32::new(last_id),
30            used: AtomicU64::new(used_ids),
31            select_in_bitmap: AtomicU32::new(0),
32            look_into_bitmap: AtomicBool::new(!available.is_empty()),
33            available,
34        }
35    }
36
37    /// Returns a new unique ID and increase the count of IDs used.
38    pub fn next(&self) -> Option<u32> {
39        if self.used.fetch_add(1, Ordering::Relaxed) > u32::MAX as u64 {
40            None
41        } else if self.look_into_bitmap.load(Ordering::Relaxed) {
42            let current = self.select_in_bitmap.fetch_add(1, Ordering::Relaxed);
43            match self.available.select(current) {
44                Some(id) => Some(id),
45                None => {
46                    self.look_into_bitmap.store(false, Ordering::Relaxed);
47                    Some(self.current.fetch_add(1, Ordering::Relaxed))
48                }
49            }
50        } else {
51            Some(self.current.fetch_add(1, Ordering::Relaxed))
52        }
53    }
54
55    /// Returns the number of used ids in total.
56    pub fn used(&self) -> u64 {
57        self.used.load(Ordering::Relaxed)
58    }
59}
milli_core/update/concurrent_available_ids.rs

milli_core/update/
concurrent_available_ids.rs
