teaql_runtime/

id.rs

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Mutex, OnceLock};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use crate::RuntimeError;

pub trait InternalIdGenerator: Send + Sync {
    fn generate_id(&self, entity: &str) -> Result<u64, RuntimeError>;

    fn ensure_floor(&self, entity: &str, floor: u64) -> Result<(), RuntimeError> {
        Err(RuntimeError::IdGeneration(format!(
            "ID generator cannot reserve fixed bootstrap ID floor {floor} for {entity}"
        )))
    }
}

/// Normalize generated Rust type names and model entity names to the same
/// stable key used by generated `ENTITY_NAME` constants.
pub fn canonical_id_space_entity(entity: &str) -> String {
    let mut result = String::with_capacity(entity.len() + 4);
    for (index, character) in entity.chars().enumerate() {
        if character.is_ascii_uppercase() {
            if index > 0 {
                result.push('_');
            }
            result.push(character.to_ascii_lowercase());
        } else {
            result.push(character);
        }
    }
    result
}

// ---------------------------------------------------------------------------
// AtomicCounterIdGenerator — process-level counter, suitable for in-memory use
// ---------------------------------------------------------------------------

/// A simple atomic counter that produces sequential IDs starting from a
/// configurable base value (default 1000).
///
/// Suitable for in-memory / test / single-process scenarios where readable,
/// compact IDs are preferred over globally unique snowflake IDs.
#[derive(Debug)]
pub struct AtomicCounterIdGenerator {
    counter: AtomicU64,
}

impl Default for AtomicCounterIdGenerator {
    fn default() -> Self {
        Self::new(1000)
    }
}

impl AtomicCounterIdGenerator {
    /// Create a new counter starting from `start`.
    /// The first call to `generate_id` will return `start + 1`.
    pub fn new(start: u64) -> Self {
        Self {
            counter: AtomicU64::new(start),
        }
    }
}

impl InternalIdGenerator for AtomicCounterIdGenerator {
    fn generate_id(&self, _entity: &str) -> Result<u64, RuntimeError> {
        Ok(self.counter.fetch_add(1, Ordering::Relaxed) + 1)
    }

    fn ensure_floor(&self, _entity: &str, floor: u64) -> Result<(), RuntimeError> {
        self.counter.fetch_max(floor, Ordering::Relaxed);
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// SnowflakeIdGenerator — distributed-friendly, timestamp-based
// ---------------------------------------------------------------------------

#[derive(Debug)]
pub struct SnowflakeIdGenerator {
    epoch_millis: u64,
    worker_id: u64,
    datacenter_id: u64,
    state: Mutex<SnowflakeState>,
}

#[derive(Debug, Default)]
struct SnowflakeState {
    last_timestamp: u64,
    sequence: u64,
}

impl Default for SnowflakeIdGenerator {
    fn default() -> Self {
        Self::new(0, 0)
    }
}

impl SnowflakeIdGenerator {
    const DEFAULT_EPOCH_MILLIS: u64 = 1_288_834_974_657;
    const WORKER_ID_BITS: u64 = 5;
    const DATACENTER_ID_BITS: u64 = 5;
    const SEQUENCE_BITS: u64 = 12;
    const MAX_WORKER_ID: u64 = (1 << Self::WORKER_ID_BITS) - 1;
    const MAX_DATACENTER_ID: u64 = (1 << Self::DATACENTER_ID_BITS) - 1;
    const SEQUENCE_MASK: u64 = (1 << Self::SEQUENCE_BITS) - 1;
    const WORKER_ID_SHIFT: u64 = Self::SEQUENCE_BITS;
    const DATACENTER_ID_SHIFT: u64 = Self::SEQUENCE_BITS + Self::WORKER_ID_BITS;
    const TIMESTAMP_SHIFT: u64 =
        Self::SEQUENCE_BITS + Self::WORKER_ID_BITS + Self::DATACENTER_ID_BITS;

    pub fn new(worker_id: u64, datacenter_id: u64) -> Self {
        assert!(worker_id <= Self::MAX_WORKER_ID, "worker id out of range");
        assert!(
            datacenter_id <= Self::MAX_DATACENTER_ID,
            "datacenter id out of range"
        );

        Self {
            epoch_millis: Self::DEFAULT_EPOCH_MILLIS,
            worker_id,
            datacenter_id,
            state: Mutex::new(SnowflakeState::default()),
        }
    }

    fn current_millis() -> Result<u64, RuntimeError> {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map_err(|err| RuntimeError::IdGeneration(err.to_string()))?;
        Ok(now.as_millis() as u64)
    }

    fn wait_until_next_millis(last_timestamp: u64) -> Result<u64, RuntimeError> {
        loop {
            let timestamp = Self::current_millis()?;
            if timestamp > last_timestamp {
                return Ok(timestamp);
            }
            std::thread::sleep(Duration::from_millis(1));
        }
    }
}

impl InternalIdGenerator for SnowflakeIdGenerator {
    fn generate_id(&self, _entity: &str) -> Result<u64, RuntimeError> {
        let mut state = self
            .state
            .lock()
            .map_err(|_| RuntimeError::IdGeneration("snowflake state poisoned".to_owned()))?;
        let mut timestamp = Self::current_millis()?;

        if timestamp < state.last_timestamp {
            timestamp = Self::wait_until_next_millis(state.last_timestamp)?;
        }

        match timestamp == state.last_timestamp {
            true => {
                state.sequence = (state.sequence + 1) & Self::SEQUENCE_MASK;
                if state.sequence == 0 {
                    timestamp = Self::wait_until_next_millis(state.last_timestamp)?;
                }
            }
            false => state.sequence = 0,
        }

        state.last_timestamp = timestamp;

        let relative_timestamp = timestamp.checked_sub(self.epoch_millis).ok_or_else(|| {
            RuntimeError::IdGeneration("system clock is before snowflake epoch".to_owned())
        })?;

        Ok((relative_timestamp << Self::TIMESTAMP_SHIFT)
            | (self.datacenter_id << Self::DATACENTER_ID_SHIFT)
            | (self.worker_id << Self::WORKER_ID_SHIFT)
            | state.sequence)
    }

    fn ensure_floor(&self, _entity: &str, _floor: u64) -> Result<(), RuntimeError> {
        Ok(())
    }
}

pub(crate) fn local_id_generator() -> &'static AtomicCounterIdGenerator {
    static LOCAL_ID_GENERATOR: OnceLock<AtomicCounterIdGenerator> = OnceLock::new();
    LOCAL_ID_GENERATOR.get_or_init(AtomicCounterIdGenerator::default)
}