Add distributed deployment support for snowflake ID (#927)

* feat: Add support distributed deployment for Snowflake * Update the algorithm implementation * Remove duplicate codes and update error messages
2026-09-21 13:12:24 +00:00 · 2025-11-19 12:48:24 +08:00
parent 551dc51c52
commit 16bcaf963d
4 changed files with 190 additions and 64 deletions
@@ -70,6 +70,6 @@ class UploadUrl:
 class SnowflakeInfo:
    timestamp: int
    datetime: str
-    cluster_id: int
-    node_id: int
+    datacenter_id: int
+    worker_id: int
    sequence: int
@@ -52,6 +52,15 @@ class Settings(BaseSettings):
    # Redis
    REDIS_TIMEOUT: int = 5

+    # .env Snowflake
+    SNOWFLAKE_DATACENTER_ID: int | None = None
+    SNOWFLAKE_WORKER_ID: int | None = None
+
+    # Snowflake
+    SNOWFLAKE_REDIS_PREFIX: str = 'fba:snowflake'
+    SNOWFLAKE_HEARTBEAT_INTERVAL_SECONDS: int = 30
+    SNOWFLAKE_NODE_TTL_SECONDS: int = 60
+
    # .env Token
    TOKEN_SECRET_KEY: str  # 密钥 secrets.token_urlsafe(32)

@@ -34,6 +34,7 @@ from backend.utils.demo_site import demo_site
 from backend.utils.health_check import ensure_unique_route_names, http_limit_callback
 from backend.utils.openapi import simplify_operation_ids
 from backend.utils.serializers import MsgSpecJSONResponse
+from backend.utils.snowflake import snowflake


@asynccontextmanager
@@ -57,11 +58,17 @@ async def register_init(app: FastAPI) -> AsyncGenerator[None, None]:
        http_callback=http_limit_callback,
    )

+    # 初始化 snowflake 节点
+    await snowflake.init()
+
    # 创建操作日志任务
    create_task(OperaLogMiddleware.consumer())

    yield

+    # 释放 snowflake 节点
+    await snowflake.shutdown()
+
    # 关闭 redis 连接
    await redis_client.aclose()

@@ -1,15 +1,22 @@
+import asyncio
+import datetime
+import os
+import threading
 import time

 from dataclasses import dataclass

 from backend.common.dataclasses import SnowflakeInfo
 from backend.common.exception import errors
+from backend.common.log import log
 from backend.core.conf import settings
+from backend.database.redis import redis_client
+from backend.utils.timezone import timezone


@dataclass(frozen=True)
 class SnowflakeConfig:
-    """雪花算法配置类"""
+    """雪花算法配置类，采用 Twitter 原版 Snowflake 64 位 ID 位分配配置（通用标准）"""

    # 位分配
    WORKER_ID_BITS: int = 5
@@ -29,81 +36,184 @@ class SnowflakeConfig:
    # 元年时间戳
    EPOCH: int = 1262275200000

-    # 默认值
-    DEFAULT_DATACENTER_ID: int = 1
-    DEFAULT_WORKER_ID: int = 0
-    DEFAULT_SEQUENCE: int = 0
+    # 时钟回拨容忍阈值，应对 NTP 自动同步引起的正常回跳（非标准）
+    CLOCK_BACKWARD_TOLERANCE_MS: int = 10_000
+
+
+class SnowflakeNodeManager:
+    """雪花算法节点管理器，负责从 Redis 分配和管理节点 ID"""
+
+    def __init__(self) -> None:
+        """初始化节点管理器"""
+        self.datacenter_id: int | None = None
+        self.worker_id: int | None = None
+        self.node_redis_prefix: str = f'{settings.SNOWFLAKE_REDIS_PREFIX}:nodes'
+        self._heartbeat_task: asyncio.Task | None = None
+
+    async def acquire_node_id(self) -> tuple[int, int]:
+        """从 Redis 获取可用的 datacenter_id 和 worker_id"""
+        occupied_nodes = set()
+        async for key in redis_client.scan_iter(match=f'{self.node_redis_prefix}:*'):
+            parts = key.split(':')
+            if len(parts) >= 5:
+                try:
+                    datacenter_id = int(parts[-2])
+                    worker_id = int(parts[-1])
+                    occupied_nodes.add((datacenter_id, worker_id))
+                except ValueError:
+                    continue
+
+        # 顺序查找第一个可用的 ID 组合
+        for datacenter_id in range(SnowflakeConfig.MAX_DATACENTER_ID + 1):
+            for worker_id in range(SnowflakeConfig.MAX_WORKER_ID + 1):
+                if (datacenter_id, worker_id) not in occupied_nodes and await self._register(datacenter_id, worker_id):
+                    return datacenter_id, worker_id
+
+        raise errors.ServerError(msg='无可用的雪花算法节点，节点已耗尽')
+
+    async def _register(self, datacenter_id: int, worker_id: int) -> bool:
+        key = f'{self.node_redis_prefix}:{datacenter_id}:{worker_id}'
+        value = f'pid:{os.getpid()}-ts:{timezone.now().timestamp()}'
+        return await redis_client.set(key, value, nx=True, ex=settings.SNOWFLAKE_NODE_TTL_SECONDS)
+
+    async def start_heartbeat(self, datacenter_id: int, worker_id: int) -> None:
+        """启动节点心跳"""
+        self.datacenter_id = datacenter_id
+        self.worker_id = worker_id
+
+        async def heartbeat() -> None:
+            key = f'{self.node_redis_prefix}:{datacenter_id}:{worker_id}'
+            while True:
+                await asyncio.sleep(settings.SNOWFLAKE_HEARTBEAT_INTERVAL_SECONDS)
+                try:
+                    await redis_client.expire(key, settings.SNOWFLAKE_NODE_TTL_SECONDS)
+                    log.debug(f'雪花算法节点心跳任务开始：datacenter_id={datacenter_id}, worker_id={worker_id}')
+                except Exception as e:
+                    log.error(f'雪花算法节点心跳任务失败：{e}')
+
+        self._heartbeat_task = asyncio.create_task(heartbeat())
+
+    async def release(self) -> None:
+        """释放节点"""
+        if self._heartbeat_task:
+            self._heartbeat_task.cancel()
+            try:
+                await self._heartbeat_task
+            except asyncio.CancelledError:
+                log.debug(f'雪花算法节点心跳任务释放：datacenter_id={self.datacenter_id}, worker_id={self.worker_id}')
+
+        if self.datacenter_id is not None and self.worker_id is not None:
+            key = f'{self.node_redis_prefix}:{self.datacenter_id}:{self.worker_id}'
+            await redis_client.delete(key)


 class Snowflake:
    """雪花算法类"""

-    def __init__(
-        self,
-        cluster_id: int = SnowflakeConfig.DEFAULT_DATACENTER_ID,
-        node_id: int = SnowflakeConfig.DEFAULT_WORKER_ID,
-        sequence: int = SnowflakeConfig.DEFAULT_SEQUENCE,
-    ) -> None:
-        """
-        初始化雪花算法生成器
+    def __init__(self) -> None:
+        """初始化雪花算法"""
+        self.datacenter_id: int | None = None
+        self.worker_id: int | None = None
+        self.sequence: int = 0
+        self.last_timestamp: int = -1

-        :param cluster_id: 集群 ID (0-31)
-        :param node_id: 节点 ID (0-31)
-        :param sequence: 起始序列号
-        """
-        if cluster_id < 0 or cluster_id > SnowflakeConfig.MAX_DATACENTER_ID:
-            raise errors.RequestError(msg=f'集群编号必须在 0-{SnowflakeConfig.MAX_DATACENTER_ID} 之间')
-        if node_id < 0 or node_id > SnowflakeConfig.MAX_WORKER_ID:
-            raise errors.RequestError(msg=f'节点编号必须在 0-{SnowflakeConfig.MAX_WORKER_ID} 之间')
+        self._lock = threading.Lock()
+        self._initialized = False
+        self._node_manager: SnowflakeNodeManager | None = None
+        self._auto_allocated = False  # 标记是否由 Redis 自动分配 ID

-        self.node_id = node_id
-        self.cluster_id = cluster_id
-        self.sequence = sequence
-        self.last_timestamp = -1
+    async def init(self) -> None:
+        """初始化雪花算法"""
+        if self._initialized:
+            return
+
+        with self._lock:
+            # 环境变量固定分配
+            if settings.SNOWFLAKE_DATACENTER_ID is not None and settings.SNOWFLAKE_WORKER_ID is not None:
+                self.datacenter_id = settings.SNOWFLAKE_DATACENTER_ID
+                self.worker_id = settings.SNOWFLAKE_WORKER_ID
+                log.debug(
+                    f'雪花算法使用环境变量固定节点：datacenter_id={self.datacenter_id}, worker_id={self.worker_id}'
+                )
+            elif (settings.SNOWFLAKE_DATACENTER_ID is not None and settings.SNOWFLAKE_WORKER_ID is None) or (
+                settings.SNOWFLAKE_DATACENTER_ID is None and settings.SNOWFLAKE_WORKER_ID is not None
+            ):
+                log.error('雪花算法 datacenter_id 和 worker_id 配置错误，只允许同时非 None 或同时为 None')
+                raise errors.ServerError(msg='雪花算法配置失败，请联系系统管理员')
+            else:
+                # Redis 动态分配
+                self._node_manager = SnowflakeNodeManager()
+                self.datacenter_id, self.worker_id = await self._node_manager.acquire_node_id()
+                self._auto_allocated = True
+                await self._node_manager.start_heartbeat(self.datacenter_id, self.worker_id)
+                log.debug(
+                    f'雪花算法使用 Redis 动态分配节点：datacenter_id={self.datacenter_id}, worker_id={self.worker_id}'
+                )
+
+            # 严格校验范围
+            if not (0 <= self.datacenter_id <= SnowflakeConfig.MAX_DATACENTER_ID):
+                log.error(f'雪花算法 datacenter_id 配置失败，必须在 0~{SnowflakeConfig.MAX_DATACENTER_ID} 之间')
+                raise errors.ServerError(msg='雪花算法数据中心配置失败，请联系系统管理员')
+            if not (0 <= self.worker_id <= SnowflakeConfig.MAX_WORKER_ID):
+                log.error(f'雪花算法 worker_id 配置失败，必须在 0~{SnowflakeConfig.MAX_WORKER_ID} 之间')
+                raise errors.ServerError(msg='雪花算法工作机器配置失败，请联系系统管理员')
+
+            self._initialized = True
+
+    async def shutdown(self) -> None:
+        """释放 Redis 节点"""
+        if self._node_manager and self._auto_allocated:
+            await self._node_manager.release()

    @staticmethod
-    def _current_millis() -> int:
-        """返回当前毫秒时间戳"""
-        return int(time.time() * 1000)
+    def _current_ms() -> int:
+        return int(timezone.now().timestamp() * 1000)

-    def _next_millis(self, last_timestamp: int) -> int:
-        """
-        等待至下一毫秒
-
-        :param last_timestamp: 上次生成 ID 的时间戳
-        :return:
-        """
-        timestamp = self._current_millis()
-        while timestamp <= last_timestamp:
-            time.sleep((last_timestamp - timestamp + 1) / 1000.0)
-            timestamp = self._current_millis()
-        return timestamp
+    def _till_next_ms(self, last_timestamp: int) -> int:
+        """等待直到下一毫秒"""
+        ts = self._current_ms()
+        while ts <= last_timestamp:
+            time.sleep(0.0001)
+            ts = self._current_ms()
+        return ts

    def generate(self) -> int:
        """生成雪花 ID"""
-        timestamp = self._current_millis()
+        if not self._initialized:
+            raise errors.ServerError(msg='雪花 ID 生成失败，雪花算法未初始化')

-        if timestamp < self.last_timestamp:
-            raise errors.ServerError(msg=f'系统时间倒退，拒绝生成 ID 直到 {self.last_timestamp}')
+        with self._lock:
+            timestamp = self._current_ms()

-        if timestamp == self.last_timestamp:
-            self.sequence = (self.sequence + 1) & SnowflakeConfig.SEQUENCE_MASK
-            if self.sequence == 0:
-                timestamp = self._next_millis(self.last_timestamp)
-        else:
-            self.sequence = 0
+            # 时钟回拨处理
+            if timestamp < self.last_timestamp:
+                back_ms = self.last_timestamp - timestamp
+                if back_ms <= SnowflakeConfig.CLOCK_BACKWARD_TOLERANCE_MS:
+                    log.warning(f'检测到时钟回拨 {back_ms} ms，等待恢复...')
+                    timestamp = self._till_next_ms(self.last_timestamp)
+                else:
+                    raise errors.ServerError(msg=f'雪花 ID 生成失败，时钟回拨超过 {back_ms} ms，请立即联系系统管理员')

-        self.last_timestamp = timestamp
+            # 同毫秒内序列号递增
+            if timestamp == self.last_timestamp:
+                self.sequence = (self.sequence + 1) & SnowflakeConfig.SEQUENCE_MASK
+                if self.sequence == 0:
+                    timestamp = self._till_next_ms(self.last_timestamp)
+            else:
+                self.sequence = 0

-        return (
-            ((timestamp - SnowflakeConfig.EPOCH) << SnowflakeConfig.TIMESTAMP_LEFT_SHIFT)
-            | (self.cluster_id << SnowflakeConfig.DATACENTER_ID_SHIFT)
-            | (self.node_id << SnowflakeConfig.WORKER_ID_SHIFT)
-            | self.sequence
-        )
+            self.last_timestamp = timestamp
+
+            # 组合 64 位 ID
+            return (
+                ((timestamp - SnowflakeConfig.EPOCH) << SnowflakeConfig.TIMESTAMP_LEFT_SHIFT)
+                | (self.datacenter_id << SnowflakeConfig.DATACENTER_ID_SHIFT)
+                | (self.worker_id << SnowflakeConfig.WORKER_ID_SHIFT)
+                | self.sequence
+            )

    @staticmethod
-    def parse_id(snowflake_id: int) -> SnowflakeInfo:
+    def parse(snowflake_id: int) -> SnowflakeInfo:
        """
        解析雪花 ID，获取其包含的详细信息

@@ -111,15 +221,15 @@ class Snowflake:
        :return:
        """
        timestamp = (snowflake_id >> SnowflakeConfig.TIMESTAMP_LEFT_SHIFT) + SnowflakeConfig.EPOCH
-        cluster_id = (snowflake_id >> SnowflakeConfig.DATACENTER_ID_SHIFT) & SnowflakeConfig.MAX_DATACENTER_ID
-        node_id = (snowflake_id >> SnowflakeConfig.WORKER_ID_SHIFT) & SnowflakeConfig.MAX_WORKER_ID
+        datacenter_id = (snowflake_id >> SnowflakeConfig.DATACENTER_ID_SHIFT) & SnowflakeConfig.MAX_DATACENTER_ID
+        worker_id = (snowflake_id >> SnowflakeConfig.WORKER_ID_SHIFT) & SnowflakeConfig.MAX_WORKER_ID
        sequence = snowflake_id & SnowflakeConfig.SEQUENCE_MASK

        return SnowflakeInfo(
            timestamp=timestamp,
-            datetime=time.strftime(settings.DATETIME_FORMAT, time.localtime(timestamp / 1000)),
-            cluster_id=cluster_id,
-            node_id=node_id,
+            datetime=timezone.to_str(datetime.datetime.fromtimestamp(timestamp / 1000, timezone.tz_info)),
+            datacenter_id=datacenter_id,
+            worker_id=worker_id,
            sequence=sequence,
        )