|
| 1 | +from threading import Lock |
| 2 | +from time import sleep, time_ns |
| 3 | +from typing import Self, overload |
| 4 | + |
| 5 | +SONYFLAKE_EPOCH = 1409529600 # 2014-09-01 00:00:00 UTC |
| 6 | +SONYFLAKE_SEQUENCE_BITS = 8 |
| 7 | +SONYFLAKE_SEQUENCE_COUNT = 1 << SONYFLAKE_SEQUENCE_BITS |
| 8 | +SONYFLAKE_SEQUENCE_MAX = SONYFLAKE_SEQUENCE_COUNT - 1 |
| 9 | +SONYFLAKE_MACHINE_ID_BITS = 16 |
| 10 | +SONYFLAKE_MACHINE_ID_MAX = (1 << SONYFLAKE_MACHINE_ID_BITS) - 1 |
| 11 | +SONYFLAKE_MACHINE_ID_OFFSET = SONYFLAKE_SEQUENCE_BITS |
| 12 | +SONYFLAKE_TIME_OFFSET = SONYFLAKE_MACHINE_ID_BITS + SONYFLAKE_SEQUENCE_BITS |
| 13 | + |
| 14 | + |
| 15 | +def sf_to_ns(start_time: int, sf: int) -> int: |
| 16 | + return (start_time + sf) * 10_000_000 |
| 17 | + |
| 18 | + |
| 19 | +def ns_to_sf(start_time: int, ns: int) -> int: |
| 20 | + return ns // 10_000_000 - start_time |
| 21 | + |
| 22 | + |
| 23 | +def compose(i: int, elapsed: int, machine_ids: list[int]) -> int: |
| 24 | + t = elapsed << SONYFLAKE_TIME_OFFSET |
| 25 | + m = machine_ids[i >> SONYFLAKE_SEQUENCE_BITS] << SONYFLAKE_MACHINE_ID_OFFSET |
| 26 | + c = i & SONYFLAKE_SEQUENCE_MAX |
| 27 | + return t | m | c |
| 28 | + |
| 29 | + |
| 30 | +def diff(start_time: int, elapsed: int, current_ns: int) -> float: |
| 31 | + d = sf_to_ns(start_time, elapsed) - current_ns |
| 32 | + if d > 0: |
| 33 | + return d / 1_000_000_000 |
| 34 | + return 0 |
| 35 | + |
| 36 | + |
| 37 | +def machine_id_lcg(x: int) -> int: |
| 38 | + return (32309 * x + 13799) % 65536 |
| 39 | + |
| 40 | + |
| 41 | +class SonyFlake: |
| 42 | + __slots__ = ("_machine_ids", "_start_time", "_elapsed", "_max_i", "_i", "_lock") |
| 43 | + |
| 44 | + _start_time: int |
| 45 | + _elapsed: int |
| 46 | + _machine_ids: list[int] |
| 47 | + _max_i: int |
| 48 | + _i: int |
| 49 | + _lock: Lock |
| 50 | + |
| 51 | + def __init__(self, *machine_ids: int, start_time: int = SONYFLAKE_EPOCH): |
| 52 | + machine_ids_set = set(machine_ids) |
| 53 | + |
| 54 | + if len(machine_ids_set) == 0: |
| 55 | + raise ValueError("At least one machine ID must be provided") |
| 56 | + |
| 57 | + if len(machine_ids_set) > 65536: |
| 58 | + raise ValueError("Too many machine IDs, maximum is 65536") |
| 59 | + |
| 60 | + if len(machine_ids_set) != len(machine_ids): |
| 61 | + raise ValueError("Duplicate machine IDs are not allowed") |
| 62 | + |
| 63 | + if not all(map(lambda x: isinstance(x, int), machine_ids_set)): |
| 64 | + raise TypeError("Machine IDs must be integers") |
| 65 | + |
| 66 | + if min(machine_ids_set) < 0 or max(machine_ids_set) > 0xFFFF: |
| 67 | + raise ValueError("Machine IDs must be in range [0, 65535]") |
| 68 | + |
| 69 | + if not isinstance(start_time, int): |
| 70 | + raise TypeError("start_time must be an integer") |
| 71 | + |
| 72 | + self._machine_ids = sorted(machine_ids_set) |
| 73 | + self._max_i = len(machine_ids_set) * SONYFLAKE_SEQUENCE_COUNT |
| 74 | + self._start_time = start_time * 100 |
| 75 | + self._elapsed = 0 |
| 76 | + self._i = 0 |
| 77 | + self._lock = Lock() |
| 78 | + |
| 79 | + def __iter__(self) -> Self: |
| 80 | + return self |
| 81 | + |
| 82 | + def __next__(self) -> int: |
| 83 | + id_, to_sleep = self._next() |
| 84 | + if to_sleep > 0: |
| 85 | + sleep(to_sleep) |
| 86 | + return id_ |
| 87 | + |
| 88 | + def __call__(self, n: int) -> list[int]: |
| 89 | + ids, to_sleep = self._next_n(n) |
| 90 | + sleep(to_sleep) |
| 91 | + return ids |
| 92 | + |
| 93 | + def __repr__(self) -> str: |
| 94 | + cls = self.__class__.__name__ |
| 95 | + machine_ids = ", ".join(map(str, sorted(self._machine_ids))) |
| 96 | + return f"{cls}({machine_ids}, start_time={self._start_time // 100})" |
| 97 | + |
| 98 | + @overload |
| 99 | + def _raw(self, n: int) -> tuple[list[int], float]: ... |
| 100 | + @overload |
| 101 | + def _raw(self, n: None) -> tuple[int, float]: ... |
| 102 | + |
| 103 | + def _raw(self, n: int | None) -> tuple[int | list[int], float]: |
| 104 | + if n is None: |
| 105 | + return self._next() |
| 106 | + return self._next_n(n) |
| 107 | + |
| 108 | + def _next(self) -> tuple[int, float]: |
| 109 | + start_time = self._start_time |
| 110 | + machine_ids = self._machine_ids |
| 111 | + |
| 112 | + with self._lock: |
| 113 | + current_ns = time_ns() |
| 114 | + current = ns_to_sf(start_time, current_ns) |
| 115 | + elapsed = self._elapsed |
| 116 | + |
| 117 | + if elapsed < current: |
| 118 | + self._elapsed = elapsed = current |
| 119 | + self._i = i = 0 |
| 120 | + else: |
| 121 | + self._i = i = (self._i + 1) % self._max_i |
| 122 | + |
| 123 | + if i == 0: |
| 124 | + self._elapsed = elapsed = elapsed + 1 |
| 125 | + |
| 126 | + return ( |
| 127 | + compose(i, elapsed, machine_ids), |
| 128 | + diff(start_time, elapsed, current_ns), |
| 129 | + ) |
| 130 | + |
| 131 | + def _next_n(self, n: int) -> tuple[list[int], float]: |
| 132 | + if n < 1: |
| 133 | + return [], 0 |
| 134 | + |
| 135 | + start_time = self._start_time |
| 136 | + machine_ids = self._machine_ids |
| 137 | + max_i = self._max_i |
| 138 | + result: list[int] = [] |
| 139 | + |
| 140 | + with self._lock: |
| 141 | + current_ns = time_ns() |
| 142 | + current = ns_to_sf(start_time, current_ns) |
| 143 | + elapsed = self._elapsed |
| 144 | + |
| 145 | + if elapsed < current: |
| 146 | + elapsed = current |
| 147 | + i = 0 |
| 148 | + else: |
| 149 | + i = self._i |
| 150 | + |
| 151 | + while len(result) < n: |
| 152 | + result.append(compose(i, elapsed, machine_ids)) |
| 153 | + |
| 154 | + i = (i + 1) % max_i |
| 155 | + |
| 156 | + if i == 0: |
| 157 | + elapsed += 1 |
| 158 | + |
| 159 | + self._elapsed = elapsed |
| 160 | + self._i = i |
| 161 | + |
| 162 | + return result, diff(start_time, elapsed, current_ns) |
| 163 | + |
| 164 | + |
| 165 | +class MachineIDLCG: |
| 166 | + def __init__(self, seed: int = 0): |
| 167 | + self._seed = seed |
| 168 | + self._lock = Lock() |
| 169 | + |
| 170 | + def __next__(self) -> int: |
| 171 | + with self._lock: |
| 172 | + self._seed = x = machine_id_lcg(self._seed) |
| 173 | + return x |
| 174 | + |
| 175 | + def __repr__(self) -> str: |
| 176 | + return f"{self.__class__.__name__}({self._seed})" |
0 commit comments