futool-tiny-datahub/common/futool/core/fu_id.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/4/9 14:54
# @Author  : old tom
# @File    : fu_id.py
# @Project : futool-tiny-datahub
# @Desc    : ID 生成器
import time
import abc
from datahub.log_conf import Logger

logger = Logger().get_logger()


class IdGenerator(object):
    """
    ID 生成器
    """

    def __init__(self, id_type='snowflake'):
        """
        :param id_type: id类型,默认雪花算法
        """
        self.id_type = id_type
        self.snowflake = SnowFlakeId(datacenter_id=1, worker_id=1)
        self.uuid = UUID()

    def get_id(self):
        return {
            'snowflake': self.snowflake,
            'uuid': self.uuid
        }[self.id_type].get_id()


class AbsIdGenerator(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def get_id(self):
        pass


class SnowFlakeId(AbsIdGenerator):
    """
    雪花ID生成
    会生成一个64bit的整数，最终存到数据库就只占用8字节
    1bit: 一般是符号位，代表正负数的所以这一位不做处理
    41bit：这个部分用来记录时间戳，如果从1970-01-01 00:00:00来计算开始时间的话，它可以记录到2039年，足够我们用了，并且后续我们可以设置起始时间，这样就不用担心不够的问题， 这一个部分是保证我们生辰的id趋势递增的关键。
    10bit：这是用来记录机器id的， 默认情况下这10bit会分成两部分前5bit代表数据中心，后5bit代表某个数据中心的机器id，默认情况下计算大概可以支持32*32 - 1= 1023台机器。
    12bit：循环位，来对应1毫秒内产生的不同的id， 大概可以满足1毫秒并发生成2^12-1=4095次id的要求。
    """
    # 64位ID的划分
    WORKER_ID_BITS = 5
    DATACENTER_ID_BITS = 5
    SEQUENCE_BITS = 12
    # 最大取值计算
    MAX_WORKER_ID = -1 ^ (-1 << WORKER_ID_BITS)  # 2**5-1 0b11111
    MAX_DATACENTER_ID = -1 ^ (-1 << DATACENTER_ID_BITS)
    # 移位偏移计算
    WOKER_ID_SHIFT = SEQUENCE_BITS
    DATACENTER_ID_SHIFT = SEQUENCE_BITS + WORKER_ID_BITS
    TIMESTAMP_LEFT_SHIFT = SEQUENCE_BITS + WORKER_ID_BITS + DATACENTER_ID_BITS
    # 序号循环掩码
    SEQUENCE_MASK = -1 ^ (-1 << SEQUENCE_BITS)
    # Twitter元年时间戳
    TWEPOCH = 1288834974657

    def __init__(self, datacenter_id, worker_id, start_seq=0):
        """
        :param datacenter_id: 数据中心ID
        :param worker_id: 机器ID
        :param start_seq: 起始序号
        """
        # sanity check
        if worker_id > self.MAX_WORKER_ID or worker_id < 0:
            raise ValueError('worker_id值越界')
        if datacenter_id > self.MAX_DATACENTER_ID or datacenter_id < 0:
            raise ValueError('datacenter_id值越界')
        self.worker_id = worker_id
        self.datacenter_id = datacenter_id
        self.sequence = start_seq
        self.last_timestamp = -1  # 上次计算的时间戳

    @staticmethod
    def _gen_timestamp():
        """
        生成整数时间戳
        :return:int timestamp
        """
        return int(time.time() * 1000)

    def get_id(self):
        """
        获取新ID
        :return:
        """
        timestamp = self._gen_timestamp()
        # 时钟回拨
        if timestamp < self.last_timestamp:
            logger.error('clock is moving backwards. Rejecting requests until{}'.format(self.last_timestamp))
            raise Exception
        if timestamp == self.last_timestamp:
            self.sequence = (self.sequence + 1) & self.SEQUENCE_MASK
            if self.sequence == 0:
                timestamp = self._til_next_millis(self.last_timestamp)
        else:
            self.sequence = 0
        self.last_timestamp = timestamp
        new_id = ((timestamp - self.TWEPOCH) << self.TIMESTAMP_LEFT_SHIFT) | (
                self.datacenter_id << self.DATACENTER_ID_SHIFT) | \
                 (self.worker_id << self.WOKER_ID_SHIFT) | self.sequence
        return new_id

    def _til_next_millis(self, last_timestamp):
        """
        等到下一毫秒
        """
        timestamp = self._gen_timestamp()
        while timestamp <= last_timestamp:
            timestamp = self._gen_timestamp()
        return timestamp


class UUID(AbsIdGenerator):
    """
    UUID生成
    """

    def __init__(self):
        pass

    def get_id(self):
        pass


id_gen = IdGenerator()