qubit_id/snowflake/

snowflake_generator.rs

/*******************************************************************************
 *
 *    Copyright (c) 2026 Haixing Hu.
 *
 *    SPDX-License-Identifier: Apache-2.0
 *
 *    Licensed under the Apache License, Version 2.0.
 *
 ******************************************************************************/
//! Classic 41/10/12 Snowflake generator.

use std::sync::{
    Arc,
    Mutex,
};
use std::thread;
use std::time::{
    Duration,
    SystemTime,
    UNIX_EPOCH,
};

use super::constants::DEFAULT_QUBIT_EPOCH_MILLIS;
use super::time_slice::TimeSlice;
use crate::{
    IdError,
    IdGenerator,
};

const TIMESTAMP_BITS: u8 = 41;
const NODE_BITS: u8 = 10;
const SEQUENCE_BITS: u8 = 12;
const MAX_NODE_ID: u64 = (1_u64 << NODE_BITS) - 1;

/// Classic Snowflake generator using 41 timestamp, 10 node, and 12 sequence bits.
pub struct SnowflakeGenerator {
    node_id: u64,
    epoch: SystemTime,
    clock: Arc<dyn Fn() -> SystemTime + Send + Sync>,
    state: Mutex<TimeSlice>,
}

impl SnowflakeGenerator {
    /// Creates a classic Snowflake generator with the default Qubit epoch.
    ///
    /// # Parameters
    /// - `node_id`: Node identifier in `0..=1023`.
    ///
    /// # Returns
    /// A configured generator.
    ///
    /// # Errors
    /// Returns [`IdError::NodeOutOfRange`] when `node_id` does not fit in 10 bits.
    pub fn new(node_id: u64) -> Result<Self, IdError> {
        Self::with_epoch(
            node_id,
            UNIX_EPOCH + Duration::from_millis(DEFAULT_QUBIT_EPOCH_MILLIS),
        )
    }

    /// Creates a classic Snowflake generator with an explicit epoch.
    ///
    /// # Parameters
    /// - `node_id`: Node identifier in `0..=1023`.
    /// - `epoch`: Timestamp origin.
    ///
    /// # Returns
    /// A configured generator using the system clock.
    ///
    /// # Errors
    /// Returns [`IdError::NodeOutOfRange`] when `node_id` does not fit in 10 bits.
    pub fn with_epoch(node_id: u64, epoch: SystemTime) -> Result<Self, IdError> {
        Self::with_clock(node_id, epoch, SystemTime::now)
    }

    /// Creates a classic Snowflake generator with an explicit clock.
    ///
    /// # Parameters
    /// - `node_id`: Node identifier in `0..=1023`.
    /// - `epoch`: Timestamp origin.
    /// - `clock`: Function returning the current time.
    ///
    /// # Returns
    /// A configured generator.
    ///
    /// # Errors
    /// Returns [`IdError::NodeOutOfRange`] when `node_id` does not fit in 10 bits.
    pub fn with_clock<F>(node_id: u64, epoch: SystemTime, clock: F) -> Result<Self, IdError>
    where
        F: Fn() -> SystemTime + Send + Sync + 'static,
    {
        if node_id > MAX_NODE_ID {
            return Err(IdError::NodeOutOfRange {
                node_id,
                max: MAX_NODE_ID,
            });
        }
        Ok(Self {
            node_id,
            epoch,
            clock: Arc::new(clock),
            state: Mutex::new(TimeSlice::new(0)),
        })
    }

    /// Returns the configured node identifier.
    ///
    /// # Returns
    /// Node identifier.
    pub const fn node_id(&self) -> u64 {
        self.node_id
    }

    /// Returns the configured epoch.
    ///
    /// # Returns
    /// Timestamp origin.
    pub const fn epoch(&self) -> SystemTime {
        self.epoch
    }

    /// Returns the maximum representable timestamp.
    ///
    /// # Returns
    /// Maximum timestamp in milliseconds since the epoch.
    pub const fn max_timestamp(&self) -> u64 {
        (1_u64 << TIMESTAMP_BITS) - 1
    }

    /// Returns the maximum representable sequence.
    ///
    /// # Returns
    /// Maximum sequence number.
    pub const fn max_sequence(&self) -> u64 {
        (1_u64 << SEQUENCE_BITS) - 1
    }

    /// Composes an ID from timestamp and sequence parts.
    ///
    /// # Parameters
    /// - `timestamp`: Milliseconds elapsed since the epoch.
    /// - `sequence`: Sequence value inside the timestamp millisecond.
    ///
    /// # Returns
    /// Encoded ID.
    ///
    /// # Errors
    /// Returns [`IdError::TimestampOverflow`] or [`IdError::SequenceOverflow`]
    /// when a part does not fit the classic Snowflake layout.
    pub fn compose(&self, timestamp: u64, sequence: u64) -> Result<u64, IdError> {
        if timestamp > self.max_timestamp() {
            return Err(IdError::TimestampOverflow {
                timestamp,
                max: self.max_timestamp(),
            });
        }
        if sequence > self.max_sequence() {
            return Err(IdError::SequenceOverflow {
                sequence,
                max: self.max_sequence(),
            });
        }
        Ok((timestamp << (NODE_BITS + SEQUENCE_BITS)) | (self.node_id << SEQUENCE_BITS) | sequence)
    }

    /// Extracts the timestamp part from an ID.
    ///
    /// # Parameters
    /// - `id`: ID generated by this layout.
    ///
    /// # Returns
    /// Milliseconds elapsed since the epoch.
    pub const fn extract_timestamp(&self, id: u64) -> u64 {
        id >> (NODE_BITS + SEQUENCE_BITS)
    }

    /// Extracts the node identifier from an ID.
    ///
    /// # Parameters
    /// - `id`: ID generated by this layout.
    ///
    /// # Returns
    /// Node identifier.
    pub const fn extract_node_id(&self, id: u64) -> u64 {
        (id >> SEQUENCE_BITS) & MAX_NODE_ID
    }

    /// Extracts the sequence number from an ID.
    ///
    /// # Parameters
    /// - `id`: ID generated by this layout.
    ///
    /// # Returns
    /// Sequence number.
    pub const fn extract_sequence(&self, id: u64) -> u64 {
        id & ((1_u64 << SEQUENCE_BITS) - 1)
    }

    /// Converts a clock time into milliseconds since the epoch.
    ///
    /// # Parameters
    /// - `time`: Time to convert.
    ///
    /// # Returns
    /// Milliseconds elapsed since the epoch.
    ///
    /// # Errors
    /// Returns [`IdError::TimeBeforeEpoch`] when `time` is before the epoch.
    fn timestamp_for(&self, time: SystemTime) -> Result<u64, IdError> {
        let elapsed = time
            .duration_since(self.epoch)
            .map_err(|_| IdError::TimeBeforeEpoch)?;
        let timestamp = elapsed.as_millis();
        if timestamp > u128::from(self.max_timestamp()) {
            return Err(IdError::TimestampOverflow {
                timestamp: u64::try_from(timestamp).unwrap_or(u64::MAX),
                max: self.max_timestamp(),
            });
        }
        Ok(timestamp as u64)
    }

    /// Reads the current timestamp from the configured clock.
    ///
    /// # Returns
    /// Current timestamp in milliseconds.
    ///
    /// # Errors
    /// Returns [`IdError::TimeBeforeEpoch`] when the clock is before the epoch.
    fn current_timestamp(&self) -> Result<u64, IdError> {
        self.timestamp_for((self.clock)())
    }

    /// Waits until the clock reaches a later millisecond.
    ///
    /// # Parameters
    /// - `last_timestamp`: Millisecond that exhausted its sequence range.
    ///
    /// # Returns
    /// First observed millisecond greater than `last_timestamp`.
    ///
    /// # Errors
    /// Returns [`IdError::TimeBeforeEpoch`] when the clock is before the epoch.
    fn wait_for_next_timestamp(&self, last_timestamp: u64) -> Result<u64, IdError> {
        let mut timestamp = self.current_timestamp()?;
        while timestamp <= last_timestamp {
            thread::sleep(Duration::from_millis(1));
            timestamp = self.current_timestamp()?;
        }
        Ok(timestamp)
    }
}

impl IdGenerator<u64> for SnowflakeGenerator {
    type Error = IdError;

    /// Generates the next classic Snowflake ID.
    fn next_id(&self) -> Result<u64, Self::Error> {
        let mut state = self
            .state
            .lock()
            .expect("generator state mutex should not be poisoned");
        let mut timestamp = self.current_timestamp()?;

        if state.timestamp > timestamp {
            return Err(IdError::ClockMovedBackwards {
                last_timestamp: state.timestamp,
                current_timestamp: timestamp,
                skew_millis: state.timestamp - timestamp,
                max_skew_millis: 0,
            });
        }

        let sequence = if timestamp == state.timestamp {
            let next_sequence = state.sequence + 1;
            if next_sequence > self.max_sequence() {
                drop(state);
                timestamp = self.wait_for_next_timestamp(timestamp)?;
                let mut state = self
                    .state
                    .lock()
                    .expect("generator state mutex should not be poisoned");
                state.timestamp = timestamp;
                state.sequence = 0;
                return self.compose(timestamp, 0);
            }
            next_sequence
        } else {
            0
        };

        state.timestamp = timestamp;
        state.sequence = sequence;
        self.compose(timestamp, sequence)
    }
}