teaql_runtime/
id.rs

1use std::sync::atomic::{AtomicU64, Ordering};
2use std::sync::{Mutex, OnceLock};
3use std::time::{Duration, SystemTime, UNIX_EPOCH};
4
5use crate::RuntimeError;
6
7pub trait InternalIdGenerator: Send + Sync {
8    fn generate_id(&self, entity: &str) -> Result<u64, RuntimeError>;
9}
10
11// ---------------------------------------------------------------------------
12// AtomicCounterIdGenerator — process-level counter, suitable for in-memory use
13// ---------------------------------------------------------------------------
14
15/// A simple atomic counter that produces sequential IDs starting from a
16/// configurable base value (default 1000).
17///
18/// Suitable for in-memory / test / single-process scenarios where readable,
19/// compact IDs are preferred over globally unique snowflake IDs.
20#[derive(Debug)]
21pub struct AtomicCounterIdGenerator {
22    counter: AtomicU64,
23}
24
25impl Default for AtomicCounterIdGenerator {
26    fn default() -> Self {
27        Self::new(1000)
28    }
29}
30
31impl AtomicCounterIdGenerator {
32    /// Create a new counter starting from `start`.
33    /// The first call to `generate_id` will return `start + 1`.
34    pub fn new(start: u64) -> Self {
35        Self {
36            counter: AtomicU64::new(start),
37        }
38    }
39}
40
41impl InternalIdGenerator for AtomicCounterIdGenerator {
42    fn generate_id(&self, _entity: &str) -> Result<u64, RuntimeError> {
43        Ok(self.counter.fetch_add(1, Ordering::Relaxed) + 1)
44    }
45}
46
47// ---------------------------------------------------------------------------
48// SnowflakeIdGenerator — distributed-friendly, timestamp-based
49// ---------------------------------------------------------------------------
50
51#[derive(Debug)]
52pub struct SnowflakeIdGenerator {
53    epoch_millis: u64,
54    worker_id: u64,
55    datacenter_id: u64,
56    state: Mutex<SnowflakeState>,
57}
58
59#[derive(Debug, Default)]
60struct SnowflakeState {
61    last_timestamp: u64,
62    sequence: u64,
63}
64
65impl Default for SnowflakeIdGenerator {
66    fn default() -> Self {
67        Self::new(0, 0)
68    }
69}
70
71impl SnowflakeIdGenerator {
72    const DEFAULT_EPOCH_MILLIS: u64 = 1_288_834_974_657;
73    const WORKER_ID_BITS: u64 = 5;
74    const DATACENTER_ID_BITS: u64 = 5;
75    const SEQUENCE_BITS: u64 = 12;
76    const MAX_WORKER_ID: u64 = (1 << Self::WORKER_ID_BITS) - 1;
77    const MAX_DATACENTER_ID: u64 = (1 << Self::DATACENTER_ID_BITS) - 1;
78    const SEQUENCE_MASK: u64 = (1 << Self::SEQUENCE_BITS) - 1;
79    const WORKER_ID_SHIFT: u64 = Self::SEQUENCE_BITS;
80    const DATACENTER_ID_SHIFT: u64 = Self::SEQUENCE_BITS + Self::WORKER_ID_BITS;
81    const TIMESTAMP_SHIFT: u64 =
82        Self::SEQUENCE_BITS + Self::WORKER_ID_BITS + Self::DATACENTER_ID_BITS;
83
84    pub fn new(worker_id: u64, datacenter_id: u64) -> Self {
85        assert!(worker_id <= Self::MAX_WORKER_ID, "worker id out of range");
86        assert!(
87            datacenter_id <= Self::MAX_DATACENTER_ID,
88            "datacenter id out of range"
89        );
90
91        Self {
92            epoch_millis: Self::DEFAULT_EPOCH_MILLIS,
93            worker_id,
94            datacenter_id,
95            state: Mutex::new(SnowflakeState::default()),
96        }
97    }
98
99    fn current_millis() -> Result<u64, RuntimeError> {
100        let now = SystemTime::now()
101            .duration_since(UNIX_EPOCH)
102            .map_err(|err| RuntimeError::IdGeneration(err.to_string()))?;
103        Ok(now.as_millis() as u64)
104    }
105
106    fn wait_until_next_millis(last_timestamp: u64) -> Result<u64, RuntimeError> {
107        loop {
108            let timestamp = Self::current_millis()?;
109            if timestamp > last_timestamp {
110                return Ok(timestamp);
111            }
112            std::thread::sleep(Duration::from_millis(1));
113        }
114    }
115}
116
117impl InternalIdGenerator for SnowflakeIdGenerator {
118    fn generate_id(&self, _entity: &str) -> Result<u64, RuntimeError> {
119        let mut state = self
120            .state
121            .lock()
122            .map_err(|_| RuntimeError::IdGeneration("snowflake state poisoned".to_owned()))?;
123        let mut timestamp = Self::current_millis()?;
124
125        if timestamp < state.last_timestamp {
126            timestamp = Self::wait_until_next_millis(state.last_timestamp)?;
127        }
128
129        match timestamp == state.last_timestamp {
130            true => {
131                state.sequence = (state.sequence + 1) & Self::SEQUENCE_MASK;
132                if state.sequence == 0 {
133                    timestamp = Self::wait_until_next_millis(state.last_timestamp)?;
134                }
135            }
136            false => state.sequence = 0,
137        }
138
139        state.last_timestamp = timestamp;
140
141        let relative_timestamp = timestamp.checked_sub(self.epoch_millis).ok_or_else(|| {
142            RuntimeError::IdGeneration("system clock is before snowflake epoch".to_owned())
143        })?;
144
145        Ok((relative_timestamp << Self::TIMESTAMP_SHIFT)
146            | (self.datacenter_id << Self::DATACENTER_ID_SHIFT)
147            | (self.worker_id << Self::WORKER_ID_SHIFT)
148            | state.sequence)
149    }
150}
151
152pub(crate) fn local_id_generator() -> &'static AtomicCounterIdGenerator {
153    static LOCAL_ID_GENERATOR: OnceLock<AtomicCounterIdGenerator> = OnceLock::new();
154    LOCAL_ID_GENERATOR.get_or_init(AtomicCounterIdGenerator::default)
155}
teaql_runtime/id.rs

teaql_runtime/
id.rs
