【算法】自创的冰雹ID算法(Rust版本)

这个帖子里所记录的算法跟另一个帖子里的Go版本算法功能一样,容量也一样。只是这个Rust版本算法实际上建立在先,Go版本的建立在后。

Rust版本的时间戳生成功能没有独立出去,所以代码都在一起。

use core::time;
use std::sync::{Arc, Mutex};

use chrono::NaiveDateTime;
use once_cell::sync::{Lazy, OnceCell};
use thiserror::Error;

const HAIL_PERIOD_START: Lazy<i64> = Lazy::new(|| {
    crate::time::date(2022, 2, 22)
        .map(|d| d.and_hms_opt(22, 22, 22))
        .flatten()
        .map(|dt| crate::time::attach_asia_shanghai(dt))
        .map(|dt| dt.timestamp())
        .unwrap_or_else(|| NaiveDateTime::MIN.timestamp())
});
type TimestampValidator = fn(i64) -> bool;
type TimestampGenerator = fn() -> i64;

static INSTANCE: OnceCell<HailSerialCodeAlgorithm> = OnceCell::new();

#[derive(Debug, Error)]
pub enum HailSerialCodeAlgorithmError {
    #[error("Algorithm is already initialized")]
    AlgorithmAlreadyInitialized,
}

/// 冰雹序列ID算法。
/// 缩减了时间戳的位数,相比雪花算法可以额外支持近40年。
pub struct HailSerialCodeAlgorithm {
    validator: Option<TimestampValidator>,
    generator: Option<TimestampGenerator>,
    host_id: i64,
    last_timestamp: Arc<Mutex<i64>>,
    counter: Arc<Mutex<i64>>,
}

impl HailSerialCodeAlgorithm {
    /// 获取一个算法实例用于获取序列ID。
    pub fn get() -> &'static Self {
        INSTANCE.get().unwrap()
    }

    /// 初始化整个序列ID算法。
    /// ! 注意,如果选择使用内置的主机独立时间戳生成器和验证器,那么将不能保证多主机状态下的序列ID一致性。可能会存在个别主机时间回拨现象。
    ///
    /// - `host_id`:主机ID,取值范围为0~65535。
    /// - `timestamp_generator`:时间戳生成器,用于生成时间戳。如果不提供,则使用算法内置的主机独立时间戳生成器。
    /// - `timestamp_validator`:时间戳验证器,用于验证时间戳是否有效。如果不提供,则使用算法内置的主机独立时间戳验证器。
    pub fn initialize_algorithm(
        host_id: i64,
        timestamp_generator: Option<TimestampGenerator>,
        timestamp_validator: Option<TimestampValidator>,
    ) -> Result<(), HailSerialCodeAlgorithmError> {
        let algorithm = HailSerialCodeAlgorithm {
            validator: timestamp_validator,
            generator: timestamp_generator,
            host_id,
            last_timestamp: Arc::new(Mutex::new(0)),
            counter: Arc::new(Mutex::new(0)),
        };
        INSTANCE
            .set(algorithm)
            .map_err(|_| HailSerialCodeAlgorithmError::AlgorithmAlreadyInitialized)
    }

    /// 生成一个自计时起点以来的时间戳。
    fn generate_timestamp(&self) -> i64 {
        let current_time = crate::time::now_asia_shanghai().timestamp();
        current_time - *HAIL_PERIOD_START
    }

    /// 生成一个64位长整型序列ID。
    pub fn generate_serial(&self) -> i64 {
        let last_timestamp = self.last_timestamp.clone();
        let mut last_timestamp = last_timestamp.lock().unwrap();
        let counter = self.counter.clone();
        let mut counter = counter.lock().unwrap();
        loop {
            let timestamp = if let Some(generator) = self.generator {
                generator()
            } else {
                self.generate_timestamp()
            };
            if let Some(validator) = self.validator {
                if !validator(timestamp) {
                    std::thread::sleep(time::Duration::from_secs(1));
                    continue;
                }
            } else if timestamp < *last_timestamp {
                std::thread::sleep(time::Duration::from_secs(1));
                continue;
            }
            if *last_timestamp < timestamp {
                // 对齐时间戳并重置序列计数器
                *last_timestamp = timestamp;
                *counter = 0;
            }
            *counter += 1;
            return (timestamp << 20) | ((self.host_id & 0xFFFF) << 16) | (*counter & 0xFFFF_FFFF);
        }
    }

    /// 生成一个17位长前补零的序列ID字符串。
    pub fn generate_string_serial(&self) -> String {
        let serial = self.generate_serial();
        format!("{:017}", serial)
    }

    /// 生成一个带字符串前缀17位长前补零的序列ID字符串。
    pub fn generate_prefixed_string_serial(&self, prefix: &str) -> String {
        let serial = self.generate_serial();
        format!("{}{:017}", prefix, serial)
    }
}

这个算法设计使用的时候一定要使用提供的.get()方法获取算法实例,否则可能还需要自己unwrap()

4 Likes

分布式环境下要先保证时间一致,不然会有问题

分布式下可以借助另外定义的TimestampGenerator来保证。

From #dev to 开发调优

1 Like