zwf
127 lines
3.2 KiB
Python
127 lines
3.2 KiB
Python
"""
|
|
雪花 ID 生成程序。
|
|
"""
|
|
|
|
import time
|
|
import logging
|
|
|
|
|
|
# 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)
|
|
|
|
# 移位偏移计算
|
|
WORKER_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元年时间戳
|
|
TW_EPOCH = 1288834974657
|
|
|
|
ID_WORKER = None
|
|
|
|
|
|
class InvalidSystemClock(Exception):
|
|
"""
|
|
时钟回拨异常
|
|
"""
|
|
pass
|
|
|
|
|
|
class IdWorker(object):
|
|
"""
|
|
用于生成 Snow ID。
|
|
"""
|
|
|
|
@classmethod
|
|
def get_id_worker(cls, datacenter_id=1, worker_id=1, sequence=0):
|
|
"""
|
|
创建 Snow ID 对象。
|
|
|
|
:param datacenter_id: 数据中心（机器区域）ID
|
|
:param worker_id: 机器ID
|
|
:param sequence: 起始序号
|
|
"""
|
|
global ID_WORKER
|
|
if ID_WORKER is None:
|
|
ID_WORKER = IdWorker(datacenter_id, worker_id, sequence)
|
|
return ID_WORKER
|
|
|
|
def __init__(self, datacenter_id, worker_id, sequence=0):
|
|
"""
|
|
初始化。
|
|
|
|
:param datacenter_id: 数据中心（机器区域）ID
|
|
:param worker_id: 机器ID
|
|
:param sequence: 起始序号
|
|
"""
|
|
# sanity check
|
|
if worker_id > MAX_WORKER_ID or worker_id < 0:
|
|
raise ValueError('worker_id值越界')
|
|
|
|
if datacenter_id > MAX_DATACENTER_ID or datacenter_id < 0:
|
|
raise ValueError('datacenter_id值越界')
|
|
|
|
self.worker_id = worker_id
|
|
self.datacenter_id = datacenter_id
|
|
self.sequence = sequence
|
|
|
|
self.last_timestamp = -1 # 上次计算的时间戳
|
|
|
|
@staticmethod
|
|
def _gen_timestamp():
|
|
"""
|
|
生成整数时间戳。
|
|
|
|
:return:int timestamp
|
|
"""
|
|
return int(time.time() * 1000)
|
|
|
|
def get_id(self):
|
|
"""
|
|
获取新ID。
|
|
|
|
:return: 新的 Snow ID
|
|
"""
|
|
timestamp = self._gen_timestamp()
|
|
|
|
# 时钟回拨
|
|
if timestamp < self.last_timestamp:
|
|
logging.error(f"时钟正在向后倒转。拒绝请求直至 {self.last_timestamp}.")
|
|
raise InvalidSystemClock
|
|
|
|
if timestamp == self.last_timestamp:
|
|
self.sequence = (self.sequence + 1) & 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 - TW_EPOCH) << TIMESTAMP_LEFT_SHIFT) | (self.datacenter_id << DATACENTER_ID_SHIFT) | \
|
|
(self.worker_id << WORKER_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
|
|
|
|
|
|
if __name__ == '__main__':
|
|
worker = IdWorker(1, 1, 0)
|
|
print(worker.get_id())
|