Shell指令碼實現Twitter的Snowflake演算法的ID生成器

大部分時候，需要透過shell指令碼批次處理一些資料，在分散式環境下，資料庫表的主鍵儲存的都是分佈id，透過Java程式碼生成。

shell指令碼都是透過mysql命令生成insert語句，以前生成insert語句時，我都是先select MAX(id) from table賦值到MAX_ID,然後拼接,類似於

max_id_sql="select MAX(id) from table";
MAX_ID="$(query ${max_id_sql})";

echo "insert into table (id, col1,col2) values" > insert.sql

count=0;

while read -r line
do
    ((count++))
    col1="$(echo ${line} | cut -f1)";
    col2="$(echo ${line} | cut -f2)";
    echo "(${MAX_ID} + ${count}, ${col1}, ${col2})" >> insert.sql
done < datas;

這樣的弊端是要自己計算，且id不符合規律，程式碼量還不小，同時寫多張表難管理，抽空完成了一個shell版的Snowflake-id生成指令碼

#!/bin/bash

# 定義起始時間戳（2024-01-01 00:00:00 UTC）時間越小，生成的id越大
readonly __BEGIN_EPOCH=$(date -d "2000-01-01 00:00:00" +%s%3N)

# 從MAC地址中提取MACHINE_ID
readonly __MACHINE_ID=$(ip link show | grep -Po 'link/ether \K[0-9a-f:]{17}' | head -n 1 | tr -d ':' | cut -c1-8 | xxd -r -p | od -An -t u4 | tr -d ' ' | awk '{print $1 % 1024}')

# 從IP地址中提取DATACENTER_ID
readonly __DATACENTER_ID=$(ip -4 addr show scope global | grep -oP '(?<=inet\s)\d+(\.\d+){3}' | head -n 1 | cut -d'.' -f3 | awk '{print $1 % 1024}')

# 序列號初始化為0
_SN_SEQUENCE=0

# 上一次的時間戳
_SN_LAST_TIMESTAMP=0

# 🔒檔案
readonly _SN_LOCKFILE="/tmp/snowflake-id.lock"

# 生成Snowflake ID的函式
next-id() {
    exec 200>> "${_SN_LOCKFILE}";
    flock -w 1 200

    # 獲取當前時間戳（毫秒）
    local current_timestamp=$(date +%s%N | cut -b1-13)

    # 如果當前時間戳小於上次的時間戳，則報錯
    if [ "${current_timestamp}" -lt "${_SN_LAST_TIMESTAMP}" ]; then
        echo "Clock error : offset : $((_SN_LAST_TIMESTAMP - current_timestamp)) ms" >&2
        flock -u 200;
        return 1
    fi

    # 如果時間戳相同，則遞增序列號, 重複時使用
    # local timestamp_diff=$((current_timestamp - _SN_LAST_TIMESTAMP))
       # if [ "${timestamp_diff}" -le 10 ]; then
    if [ "${current_timestamp}" -eq "${_SN_LAST_TIMESTAMP}" ]; then
        _SN_SEQUENCE=$(( (_SN_SEQUENCE + 1) & 4095 ))  # 12位，最大4095
        if [ "$_SN_SEQUENCE" -eq 0 ]; then
            # 序列號溢位，等待下一毫秒
            sleep 0.001
            current_timestamp=$(date +%s%N | cut -b1-13)
        fi
    else
        _SN_SEQUENCE=0
    fi

    # 更新上次的時間戳
    _SN_LAST_TIMESTAMP=$current_timestamp

    # 計算偏移時間戳
    local time_diff=$((current_timestamp - __BEGIN_EPOCH))

    # 檢查時間戳是否超過41位的限制
    local max_time_diff=$(( (1 << 41) - 1 ))
    if [ "${time_diff}" -gt "${max_time_diff}" ]; then
        echo "Time over flow : ${time_diff} > ${max_time_diff}" >&2
        flock -u 200;
        return 2
    fi

    local lastGenerateId=$(((time_diff << 22) | (__DATACENTER_ID << 17) | (__MACHINE_ID << 7) | _SN_SEQUENCE));
    flock -u 200;
    echo "${lastGenerateId}";
}


# # 測試生成多個Snowflake ID
# for i in {1..10}; do
#    # next-id 
#     # echo "${lastGenerateId}";
#     echo "$(next-id)"
# done

使用時直接呼叫即可

#!/bin/bash

. snowflake-id.sh

id="$(next-id)"

現在機器上測試，透過後再使用。

可能會出現重複的問題，重複時，將next-id函式的最後兩行修改

lastGenerateId=$(((time_diff << 22) | (__DATACENTER_ID << 17) | (__MACHINE_ID << 7) | _SN_SEQUENCE));
flock -u 200;

在呼叫時先呼叫函式，再從全域性變數lastGenerateId獲取,這樣的好處是不會重複，但只能在一個shell程序中使用

next-id;
echo "${lastGenerateId}";

如果仍然重複，可將前面的註釋放開

# 如果時間戳相同，則遞增序列號, 重複時使用
local timestamp_diff=$((current_timestamp - _SN_LAST_TIMESTAMP))
if [ "${timestamp_diff}" -le 10 ]; then
# if [ "${current_timestamp}" -eq "${_SN_LAST_TIMESTAMP}" ]; then
    _SN_SEQUENCE=$(( (_SN_SEQUENCE + 1) & 4095 ))  # 12位，最大4095
    if [ "$_SN_SEQUENCE" -eq 0 ]; then
        # 序列號溢位，等待下一毫秒
        sleep 0.001
        current_timestamp=$(date +%s%N | cut -b1-13)
    fi
else
    _SN_SEQUENCE=0
fi

Shell指令碼實現Twitter的Snowflake演算法的ID生成器

相關文章