mirror of
https://github.com/xinnan-tech/xiaozhi-esp32-server.git
synced 2026-07-22 15:13:55 +08:00
57 lines
2.1 KiB
Python
57 lines
2.1 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import socket
|
|
import threading
|
|
import time
|
|
|
|
|
|
class SnowflakeIdGenerator:
|
|
"""MyBatis-Plus compatible 41/5/5/12-bit Snowflake identifier generator."""
|
|
|
|
EPOCH = 1288834974657
|
|
SEQUENCE_BITS = 12
|
|
WORKER_BITS = 5
|
|
DATACENTER_BITS = 5
|
|
MAX_SEQUENCE = (1 << SEQUENCE_BITS) - 1
|
|
WORKER_SHIFT = SEQUENCE_BITS
|
|
DATACENTER_SHIFT = SEQUENCE_BITS + WORKER_BITS
|
|
TIMESTAMP_SHIFT = SEQUENCE_BITS + WORKER_BITS + DATACENTER_BITS
|
|
|
|
def __init__(self, worker_id: int | None = None, datacenter_id: int | None = None):
|
|
host_hash = sum(socket.gethostname().encode("utf-8"))
|
|
self.worker_id = worker_id if worker_id is not None else (host_hash ^ os.getpid()) & 31
|
|
self.datacenter_id = datacenter_id if datacenter_id is not None else host_hash & 31
|
|
if not 0 <= self.worker_id <= 31 or not 0 <= self.datacenter_id <= 31:
|
|
raise ValueError("worker_id and datacenter_id must be in [0, 31]")
|
|
self._sequence = 0
|
|
self._last_timestamp = -1
|
|
self._lock = threading.Lock()
|
|
|
|
@staticmethod
|
|
def _milliseconds() -> int:
|
|
return time.time_ns() // 1_000_000
|
|
|
|
def next_id(self) -> int:
|
|
with self._lock:
|
|
timestamp = self._milliseconds()
|
|
if timestamp < self._last_timestamp:
|
|
raise RuntimeError("clock moved backwards; refusing to generate a duplicate Snowflake ID")
|
|
if timestamp == self._last_timestamp:
|
|
self._sequence = (self._sequence + 1) & self.MAX_SEQUENCE
|
|
if self._sequence == 0:
|
|
while timestamp <= self._last_timestamp:
|
|
timestamp = self._milliseconds()
|
|
else:
|
|
self._sequence = 0
|
|
self._last_timestamp = timestamp
|
|
return (
|
|
((timestamp - self.EPOCH) << self.TIMESTAMP_SHIFT)
|
|
| (self.datacenter_id << self.DATACENTER_SHIFT)
|
|
| (self.worker_id << self.WORKER_SHIFT)
|
|
| self._sequence
|
|
)
|
|
|
|
|
|
snowflake = SnowflakeIdGenerator()
|