-- 創(chuàng)建配置表
CREATE TABLE SnowflakeConfig (
    MachineId BIGINT NOT NULL,
    DatacenterId BIGINT NOT NULL,
    LastTimestamp BIGINT NOT NULL,
    Sequence BIGINT NOT NULL,
    CONSTRAINT PK_SnowflakeConfig PRIMARY KEY (MachineId, DatacenterId)
);

-- 初始化配置 (機器ID和數(shù)據(jù)中心ID需要在每個節(jié)點上配置不同)
INSERT INTO SnowflakeConfig (MachineId, DatacenterId, LastTimestamp, Sequence)
VALUES (1, 1, 0, 0);

-- 創(chuàng)建獲取當前時間戳的函數(shù)
CREATE FUNCTION GetCurrentTimestamp()
RETURNS BIGINT
AS
BEGIN
    DECLARE @epoch DATETIME2 = '1970-01-01 00:00:00';
    DECLARE @now DATETIME2 = SYSUTCDATETIME();
    RETURN CAST(DATEDIFF_BIG(MILLISECOND, @epoch, @now) AS BIGINT);
END;

-- 創(chuàng)建等待下一毫秒的函數(shù)
CREATE FUNCTION TilNextMillis(@lastTimestamp BIGINT)
RETURNS BIGINT
AS
BEGIN
    DECLARE @timestamp BIGINT;
    SET @timestamp = dbo.GetCurrentTimestamp();
    
    WHILE @timestamp <= @lastTimestamp
    BEGIN
        SET @timestamp = dbo.GetCurrentTimestamp();
    END
    
    RETURN @timestamp;
END;
GO

-- 創(chuàng)建計算冪的函數(shù)（替代位移操作）
CREATE FUNCTION PowerOfTwo(@exponent BIGINT)
RETURNS BIGINT
AS
BEGIN
    RETURN CAST(POWER(CAST(2 AS FLOAT), @exponent) AS BIGINT);
END;
GO

-- 創(chuàng)建生成雪花ID的存儲過程
CREATE PROCEDURE GenerateSnowflakeId
    @MachineId BIGINT = 1,
    @DatacenterId BIGINT = 1,
    @SnowflakeId BIGINT OUTPUT
AS
BEGIN
    SET NOCOUNT ON;
    
    -- 常量定義
    DECLARE @Twepoch BIGINT = 1700058600000;
    DECLARE @MachineIdBits BIGINT = 5;
    DECLARE @DatacenterIdBits BIGINT = 5;
    DECLARE @SequenceBits BIGINT = 12;
    
    -- 使用POWER計算替代位移
    DECLARE @MaxMachineId BIGINT = dbo.PowerOfTwo(@MachineIdBits) - 1;
    DECLARE @MaxDatacenterId BIGINT = dbo.PowerOfTwo(@DatacenterIdBits) - 1;
    DECLARE @SequenceMask BIGINT = dbo.PowerOfTwo(@SequenceBits) - 1;
    
    DECLARE @MachineIdShift BIGINT = @SequenceBits;
    DECLARE @DatacenterIdShift BIGINT = @SequenceBits + @MachineIdBits;
    DECLARE @TimestampLeftShift BIGINT = @SequenceBits + @MachineIdBits + @DatacenterIdBits;
    
    -- 驗證參數(shù)
    IF @MachineId > @MaxMachineId OR @MachineId < 0
    BEGIN
        THROW 50000, 'MachineId can''t be greater than maxMachineId or less than 0', 1;
        RETURN;
    END
    
    IF @DatacenterId > @MaxDatacenterId OR @DatacenterId < 0
    BEGIN
        THROW 50000, 'DatacenterId can''t be greater than maxDatacenterId or less than 0', 1;
        RETURN;
    END
    
    -- 使用事務確保原子性
    BEGIN TRANSACTION;
    
    BEGIN TRY
        DECLARE @LastTimestamp BIGINT;
        DECLARE @Sequence BIGINT;
        DECLARE @Timestamp BIGINT;
        
        -- 獲取當前狀態(tài)
        SELECT @LastTimestamp = LastTimestamp, @Sequence = Sequence
        FROM SnowflakeConfig WITH (UPDLOCK)
        WHERE MachineId = @MachineId AND DatacenterId = @DatacenterId;
        
        -- 獲取當前時間戳
        SET @Timestamp = dbo.GetCurrentTimestamp();
        
        -- 檢查時鐘回撥
        IF @Timestamp < @LastTimestamp
        BEGIN
            ROLLBACK TRANSACTION;
            THROW 50000, 'Clock moved backwards. Refusing to generate id', 1;
            RETURN;
        END
        
        -- 同一毫秒內(nèi)生成多個ID
        IF @LastTimestamp = @Timestamp
        BEGIN
            SET @Sequence = (@Sequence + 1) & @SequenceMask;
            IF @Sequence = 0
            BEGIN
                -- 序列耗盡，等待下一毫秒
                SET @Timestamp = dbo.TilNextMillis(@LastTimestamp);
            END
        END
        ELSE
        BEGIN
            SET @Sequence = 0;
        END
        
        -- 更新狀態(tài)
        UPDATE SnowflakeConfig
        SET LastTimestamp = @Timestamp,
            Sequence = @Sequence
        WHERE MachineId = @MachineId AND DatacenterId = @DatacenterId;
        
        -- 生成ID (使用乘法替代位移)
        SET @SnowflakeId = 
            (@Timestamp - @Twepoch) * dbo.PowerOfTwo(@TimestampLeftShift) +
            @DatacenterId * dbo.PowerOfTwo(@DatacenterIdShift) +
            @MachineId * dbo.PowerOfTwo(@MachineIdShift) +
            @Sequence;
        
        COMMIT TRANSACTION;
    END TRY
    BEGIN CATCH
        ROLLBACK TRANSACTION;
        THROW;
    END CATCH
END;
GO

-- 使用存儲過程版本
DECLARE @Id BIGINT;
EXEC GenerateSnowflakeId @MachineId = 1, @DatacenterId = 1, @SnowflakeId = @Id OUTPUT;
SELECT @Id AS SnowflakeId;