shanxia/crm/py-backend/app/utils/__init__.py

import time
import threading


class SnowflakeIdWorker:
    def __init__(self, worker_id: int = 1, datacenter_id: int = 1, sequence: int = 0):
        self.worker_id = worker_id
        self.datacenter_id = datacenter_id
        self.sequence = sequence
        self.twepoch = 1288834974657
        self.worker_id_bits = 5
        self.datacenter_id_bits = 5
        self.max_worker_id = -1 ^ (-1 << self.worker_id_bits)
        self.max_datacenter_id = -1 ^ (-1 << self.datacenter_id_bits)
        self.sequence_bits = 12
        self.worker_id_shift = self.sequence_bits
        self.datacenter_id_shift = self.sequence_bits + self.worker_id_bits
        self.timestamp_left_shift = (
            self.sequence_bits + self.worker_id_bits + self.datacenter_id_bits
        )
        self.sequence_mask = -1 ^ (-1 << self.sequence_bits)
        self.last_timestamp = -1
        self._lock = threading.Lock()

    def _gen_timestamp(self) -> int:
        return int(time.time() * 1000)

    def _til_next_millis(self, last_timestamp: int) -> int:
        ts = self._gen_timestamp()
        while ts <= last_timestamp:
            ts = self._gen_timestamp()
        return ts

    def get_id(self) -> int:
        with self._lock:
            ts = self._gen_timestamp()
            if ts < self.last_timestamp:
                raise RuntimeError(f"Clock moved backwards. Refusing to generate id for {self.last_timestamp - ts} milliseconds")
            if ts == self.last_timestamp:
                self.sequence = (self.sequence + 1) & self.sequence_mask
                if self.sequence == 0:
                    ts = self._til_next_millis(self.last_timestamp)
            else:
                self.sequence = 0
            self.last_timestamp = ts
            return (
                ((ts - self.twepoch) << self.timestamp_left_shift)
                | (self.datacenter_id << self.datacenter_id_shift)
                | (self.worker_id << self.worker_id_shift)
                | self.sequence
            )


snowflake = SnowflakeIdWorker()


def next_id() -> int:
    return snowflake.get_id()