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

teaql_runtime/
id.rs
