snowflake_me 0.5.0

// Copyright 2022 houseme
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use crate::builder::Builder;
use crate::error::*;
use base64::Engine;
use base64::engine::general_purpose;
use chrono::prelude::*;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

///SharedSnowflake stores the state and configurations shared by all instances of Snowflake Builder.
/// It is designed to be lockless for high concurrency performance.
pub struct SharedSnowflake {
    /// state atomic variables, packaged with 'elapsed_time' and 'sequence'.
    /// High storage time, low storage serial number.
    pub(crate) state: AtomicU64,
    /// The start timestamp (in milliseconds) generated by the ID.
    pub(crate) start_time: i64,
    /// Data center ID.
    pub(crate) data_center_id: u16,
    /// Machine ID.
    pub(crate) machine_id: u16,
    /// The bit length of the timestamp section.
    pub(crate) bit_len_time: u8,
    /// The bit length of the serial number part.
    pub(crate) bit_len_sequence: u8,
    /// The bit length of the data center ID section.
    pub(crate) bit_len_data_center_id: u8,
    /// The bit length of the machine ID section.
    pub(crate) bit_len_machine_id: u8,
}

/// Snowflake is a high-performance, distributed, unique ID generator.
/// It can be securely cloned and shared across multiple threads.
pub struct Snowflake(pub(crate) Arc<SharedSnowflake>);

impl Snowflake {
    /// Create a new Snowflake builder with the default configuration.
    ///
    /// # Errors
    ///
    /// If the default machine ID or data center ID cannot be obtained (e.g., when the 'ip-fallback' feature is disabled), an error will be returned.
    pub fn new() -> Result<Self, Error> {
        Builder::new().finalize()
    }

    /// Create a new 'Builder' to build the Snowflake Builder in a customized way.
    pub fn builder<'a>() -> Builder<'a> {
        Builder::new()
    }

    /// (Internal method) Create a Snowflake instance using shared state.
    pub(crate) fn new_inner(shared: Arc<SharedSnowflake>) -> Self {
        Self(shared)
    }

    /// Generate the next unique ID.
    ///
    /// This method is lockless and thread-safe. It uses atomic operations to guarantee performance and uniqueness at high concurrency.
    ///
    /// # Errors
    ///
    /// If the timestamp exceeds the maximum value that its bit length can represent, 'Error::OverTimeLimit' is returned.
    pub fn next_id(&self) -> Result<u64, Error> {
        let sequence_mask = (1u64 << self.0.bit_len_sequence) - 1;
        let time_shift = self.0.bit_len_sequence;
        let time_max = (1u64 << self.0.bit_len_time) - 1;

        loop {
            let current_state = self.0.state.load(Ordering::Relaxed);
            let last_time = current_state >> time_shift;

            let elapsed_time = current_elapsed_time(self.0.start_time) as u64;

            let (next_time, next_sequence) = if elapsed_time == last_time {
                // In the same millisecond, the serial number is incremented
                let sequence = (current_state & sequence_mask) + 1;
                if sequence > sequence_mask {
                    // The serial number has run out, busy waiting until the next millisecond
                    til_next_millis(self.0.start_time + last_time as i64);
                    continue; // Restart the loop to get a new timestamp
                }
                (last_time, sequence)
            } else {
                // new milliseconds, the serial number resets to 0
                (elapsed_time, 0)
            };

            if next_time > time_max {
                return Err(Error::OverTimeLimit);
            }

            // Pack the new time and serial number into a new state
            let new_state = (next_time << time_shift) | next_sequence;

            // Use CAS (Compare-And-Swap) to update status atomically
            // 'compare_exchange_weak' performs better at high concurrency because it allows for spurious failures,
            // It is safe to use in cycles.
            if self
                .0
                .state
                .compare_exchange_weak(
                    current_state,
                    new_state,
                    Ordering::AcqRel,
                    Ordering::Relaxed,
                )
                .is_ok()
            {
                // CAS succeeds, assembles, and returns the final ID
                let id = (next_time
                    << (self.0.bit_len_data_center_id
                        + self.0.bit_len_machine_id
                        + self.0.bit_len_sequence))
                    | ((self.0.data_center_id as u64)
                        << (self.0.bit_len_machine_id + self.0.bit_len_sequence))
                    | ((self.0.machine_id as u64) << self.0.bit_len_sequence)
                    | next_sequence;
                return Ok(id);
            }
            // CAS failure means that another thread has modified its state and the loop will be retried
        }
    }

    /// Decomposes a Snowflake ID into its constituent parts using the generator's configuration.
    ///
    /// # Arguments
    ///
    /// * `id` - The Snowflake ID to decompose.
    ///
    /// # Returns
    ///
    /// A `DecomposedSnowflake` struct containing the parts of the ID.
    pub fn decompose(&self, id: u64) -> DecomposedSnowflake {
        DecomposedSnowflake::decompose(
            id,
            self.0.bit_len_time,
            self.0.bit_len_sequence,
            self.0.bit_len_data_center_id,
            self.0.bit_len_machine_id,
        )
    }
}

impl Clone for Snowflake {
    fn clone(&self) -> Self {
        Self(self.0.clone())
    }
}

const SNOWFLAKE_TIME_UNIT: i64 = 1_000_000; // Nanosecond, i.e. 1 millisecond

/// (internal method) <Utc>converts 'DateTime' to millisecond timestamps.
pub(crate) fn to_snowflake_time(time: DateTime<Utc>) -> i64 {
    time.timestamp_nanos_opt().unwrap_or(0) / SNOWFLAKE_TIME_UNIT
}

/// (Internal method) Gets the elapsed time (in milliseconds) from the start time to the present.
fn current_elapsed_time(start_time: i64) -> i64 {
    to_snowflake_time(Utc::now()) - start_time
}

/// (Internal method) Busy waiting (Spin-Wait) until the next millisecond。
/// `last_timestamp` is the full timestamp of the previous ID (start_time + elapsed_time).
fn til_next_millis(last_timestamp: i64) {
    let mut now = to_snowflake_time(Utc::now());
    while now <= last_timestamp {
        // You can consider adding thread::yield_now() or short sleep to reduce CPU usage
        // But for low-latency scenarios, pure spin may be the fastest
        now = to_snowflake_time(Utc::now());
    }
}

///DecomposedSnowflake contains all components of a Snowflake ID.
///
/// Created by calling the `DecomposedSnowflake::decompose` method and providing the bit length configuration used when generating the ID.
pub struct DecomposedSnowflake {
    /// Original u64 format ID
    pub id: u64,
    /// The time stamp part from the epoch (milliseconds)
    pub time: u64,
    /// Serial number part
    pub sequence: u64,
    /// Data Center ID section
    pub data_center_id: u64,
    /// Machine ID section
    pub machine_id: u64,
}

impl DecomposedSnowflake {
    /// Decompose a Snowflake ID into its components based on the provided bit length.
    pub fn decompose(
        id: u64,
        bit_len_time: u8,
        bit_len_sequence: u8,
        bit_len_data_center_id: u8,
        bit_len_machine_id: u8,
    ) -> Self {
        let total_bits = bit_len_time as u32
            + bit_len_sequence as u32
            + bit_len_data_center_id as u32
            + bit_len_machine_id as u32;
        assert_eq!(total_bits, 63, "Total bit length must be 63");

        // 修正位移计算
        let sequence_shift = 0;
        let machine_id_shift = sequence_shift + bit_len_sequence;
        let data_center_id_shift = machine_id_shift + bit_len_machine_id;
        let time_shift = data_center_id_shift + bit_len_data_center_id;

        let sequence_mask = (1u64 << bit_len_sequence) - 1;
        let machine_id_mask = (1u64 << bit_len_machine_id) - 1;
        let data_center_id_mask = (1u64 << bit_len_data_center_id) - 1;

        Self {
            id,
            time: id >> time_shift,
            data_center_id: (id >> data_center_id_shift) & data_center_id_mask,
            machine_id: (id >> machine_id_shift) & machine_id_mask,
            sequence: (id >> sequence_shift) & sequence_mask,
        }
    }

    /// Returns the timestamp in nanoseconds without an epoch.
    pub fn nanos_time(&self) -> i64 {
        (self.time as i64) * SNOWFLAKE_TIME_UNIT
    }

    /// Returns the timestamp in milliseconds since the epoch.
    pub fn int64(&self) -> i64 {
        self.id as i64
    }

    /// Returns the string representation of the Snowflake ID.
    pub fn string(&self) -> String {
        self.id.to_string()
    }

    /// Returns the base2 encoded string.
    pub fn base2(&self) -> String {
        format!("{:b}", self.id)
    }

    /// Returns the base32 encoded string.
    pub fn base32(&self) -> String {
        const ENCODE_BASE32_MAP: &str = "ybndrfg8ejkmcpqxot1uwisza345h769";
        let mut id = self.id;
        if id < 32 {
            return ENCODE_BASE32_MAP
                .chars()
                .nth(id as usize)
                .unwrap()
                .to_string();
        }

        let mut b = Vec::new();
        while id >= 32 {
            b.push(ENCODE_BASE32_MAP.chars().nth((id % 32) as usize).unwrap());
            id /= 32;
        }
        b.push(ENCODE_BASE32_MAP.chars().nth(id as usize).unwrap());

        b.reverse();
        b.into_iter().collect()
    }

    /// Returns the base36 encoded string.
    pub fn base36(&self) -> String {
        format!("{:x}", self.id)
    }

    /// Returns the base58 encoded string.
    pub fn base58(&self) -> String {
        const ENCODE_BASE58_MAP: &str =
            "123456789abcdefghijkmnopqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ";
        let mut id = self.id;
        if id < 58 {
            return ENCODE_BASE58_MAP
                .chars()
                .nth(id as usize)
                .unwrap()
                .to_string();
        }

        let mut b = Vec::new();
        while id >= 58 {
            b.push(ENCODE_BASE58_MAP.chars().nth((id % 58) as usize).unwrap());
            id /= 58;
        }
        b.push(ENCODE_BASE58_MAP.chars().nth(id as usize).unwrap());

        b.reverse();
        b.into_iter().collect()
    }

    /// Returns the base64 encoded string.
    pub fn base64(&self) -> String {
        general_purpose::STANDARD.encode(self.id.to_be_bytes())
    }
    /// Returns the bytes of the Snowflake ID.
    pub fn bytes(&self) -> Vec<u8> {
        self.id.to_string().into_bytes()
    }
    /// Returns the bytes of the Snowflake ID.
    pub fn int_bytes(&self) -> [u8; 8] {
        self.id.to_be_bytes()
    }
    /// Returns the timestamp in seconds since the epoch.
    pub fn time(&self) -> i64 {
        self.time as i64
    }
}