diff --git a/snowflake_generator.py b/snowflake_generator.py new file mode 100644 index 0000000..2861677 --- /dev/null +++ b/snowflake_generator.py @@ -0,0 +1,66 @@ +import time +import threading +EPOCH_JAVA_COMMON = 1288834974657 +class Snowflake: + def __init__(self, datacenter_id: int = 0, worker_id: int = 0, epoch: int = 1480166465631): + # 机器和数据中心配置 + self.worker_id_bits = 5 + self.datacenter_id_bits = 5 + self.sequence_bits = 12 + + self.max_worker_id = -1 ^ (-1 << self.worker_id_bits) + self.max_datacenter_id = -1 ^ (-1 << self.datacenter_id_bits) + + if worker_id > self.max_worker_id or worker_id < 0: + raise ValueError(f"worker_id 超出范围 (0 ~ {self.max_worker_id})") + if datacenter_id > self.max_datacenter_id or datacenter_id < 0: + raise ValueError(f"datacenter_id 超出范围 (0 ~ {self.max_datacenter_id})") + + self.worker_id = worker_id + self.datacenter_id = datacenter_id + self.epoch = epoch + + self.sequence = 0 + self.last_timestamp = -1 + + 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.lock = threading.Lock() + + def _timestamp(self): + return int(time.time() * 1000) + + def _til_next_millis(self, last_timestamp): + timestamp = self._timestamp() + while timestamp <= last_timestamp: + timestamp = self._timestamp() + return timestamp + + def next_id(self) -> int: + with self.lock: + timestamp = self._timestamp() + if timestamp < self.last_timestamp: + raise Exception("时钟回拨,拒绝生成ID") + + 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 + + return ((timestamp - self.epoch) << self.timestamp_left_shift) | \ + (self.datacenter_id << self.datacenter_id_shift) | \ + (self.worker_id << self.worker_id_shift) | \ + self.sequence + + +if __name__ == "__main__": + snowflake = Snowflake(datacenter_id=0, worker_id=0, epoch=EPOCH_JAVA_COMMON) + for _ in range(10): + print(str(snowflake.next_id()))