diff --git a/backend/common/utils/snowflake.py b/backend/common/utils/snowflake.py new file mode 100644 index 0000000..5f75e8f --- /dev/null +++ b/backend/common/utils/snowflake.py @@ -0,0 +1,57 @@ +import time +import threading +class SnowflakeGenerator: + def __init__(self, worker_id=0, datacenter_id=0): + self.worker_id = worker_id + self.datacenter_id = datacenter_id + self.sequence = 0 + self.last_timestamp = -1 + self.lock = threading.Lock() + + 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) + + if self.worker_id > self.max_worker_id or self.worker_id < 0: + raise ValueError(f"worker ID can't be greater than {self.max_worker_id} or less than 0") + if self.datacenter_id > self.max_datacenter_id or self.datacenter_id < 0: + raise ValueError(f"datacenter ID can't be greater than {self.max_datacenter_id} or less than 0") + + def _current_time(self): + return int(time.time() * 1000) + + def _wait_next_millis(self, last_timestamp): + timestamp = self._current_time() + while timestamp <= last_timestamp: + timestamp = self._current_time() + return timestamp + + def generate_id(self): + with self.lock: + timestamp = self._current_time() + + if timestamp < self.last_timestamp: + raise ValueError("Clock moved backwards. Refusing to generate ID") + + if timestamp == self.last_timestamp: + self.sequence = (self.sequence + 1) & self.sequence_mask + if self.sequence == 0: + timestamp = self._wait_next_millis(self.last_timestamp) + else: + self.sequence = 0 + + self.last_timestamp = timestamp + + return ((timestamp << self.timestamp_left_shift) | + (self.datacenter_id << self.datacenter_id_shift) | + (self.worker_id << self.worker_id_shift) | + self.sequence) + +snowflake = SnowflakeGenerator(worker_id=1) \ No newline at end of file