Merge pull request #282 from gnes-ai/snowflake

feat(snoflake-uuid): add snowflake uuid generator

gnes/helper.py

@@ -19,6 +19,7 @@
 import os
 import sys
 import time
+import threading
 from copy import copy
 from functools import wraps
 from itertools import islice
@@ -41,7 +42,24 @@
            'profile_logger', 'load_contrib_module',
            'parse_arg', 'profiling', 'FileLock',
            'train_required', 'get_first_available_gpu',
-           'PathImporter', 'progressbar']
+           'PathImporter', 'progressbar', 'Singleton']
+
+
+class Singleton:
+    """
+    Make your class singeton
+    """
+    def __init__(self, cls):
+        self.__instance = None
+        self.__cls = cls
+        self._lock = threading.Lock()
+
+    def __call__(self, *args, **kwargs):
+        self._lock.acquire()
+        if self.__instance is None:
+            self.__instance = self.__cls(*args, **kwargs)
+        self._lock.release()
+        return self.__instance
 
 
 def progressbar(i, prefix="", suffix="", count=100, size=60):

gnes/uuid.py

@@ -0,0 +1,89 @@
+import threading
+import time
+from datetime import datetime
+
+from . import helper
+
+
+@helper.Singleton
+class BaseIDGenerator(object):
+    """
+    Thread-safe (auto incremental) uuid generator
+    """
+
+    def __init__(self, start_id: int = 0, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+        self._lock = threading.Lock()
+        self._next_id = start_id
+
+    def reset(self, start_id: int = 0):
+        with self._lock:
+            self._next_id = start_id
+
+    def next(self) -> int:
+        with self._lock:
+            temp = self._next_id
+            self._next_id += 1
+            return temp
+
+
+@helper.Singleton
+class SnowflakeIDGenerator(object):
+
+    def __init__(self,
+                 machine_id: int = 0,
+                 datacenter_id: int = 0,
+                 *args,
+                 **kwargs):
+        self._lock = threading.Lock()
+        self._next_id = 0
+
+        self.machine_id = machine_id
+        self.datacenter_id = datacenter_id
+
+        self.machine_bits = 5
+        self.datacenter_bits = 5
+        self.max_machine_id = -1 ^ -1 << self.machine_bits
+        self.max_datacenter_id = -1 ^ (-1 << self.datacenter_bits)
+
+        self.counter_bits = 12
+        self.max_counter_mask = -1 ^ -1 << self.counter_bits
+
+        self.machine_shift = self.counter_bits
+        self.datacenter_shift = self.counter_bits + self.machine_shift
+        self.timestamp_shift = self.counter_bits + self.machine_shift + self.datacenter_shift
+
+        self.twepoch = int(time.mktime(time.strptime('2019-01-01 00:00:00', "%Y-%m-%d %H:%M:%S")))
+        self.last_timestamp = -1
+        self.current_timestamp = lambda: int(datetime.now().timestamp() * 1000)
+
+    # def _get_timestamp(self) -> int:
+    #     return int(datetime.now().timestamp() * 1000)
+
+    def _get_next_timestamp(self, last_timestamp) -> int:
+        timestamp = self.current_timestamp()
+        while timestamp <= last_timestamp:
+            timestamp = self.current_timestamp()
+        return timestamp
+
+    def next(self) -> int:
+        with self._lock:
+            timestamp = self.current_timestamp()
+            if self.last_timestamp == timestamp:
+                self._next_id = (self._next_id + 1) & self.max_counter_mask
+                if self._next_id == 0:
+                    timestamp = self._get_next_timestamp(self.last_timestamp)
+            else:
+                self._next_id = 0
+
+            if timestamp < self.last_timestamp:
+                raise ValueError(
+                    'the current timestamp is smaller than the last timestamp')
+
+            self.last_timestamp = timestamp
+            uuid = ((timestamp - self.twepoch) << self.timestamp_shift) \
+                    | (self.datacenter_id << self.datacenter_shift) \
+                    | (self.machine_id << self.machine_shift) \
+                    | self._next_id
+            return uuid

tests/test_uuid.py

@@ -0,0 +1,21 @@
+import unittest
+
+from gnes.uuid import BaseIDGenerator, SnowflakeIDGenerator
+
+class TestUUID(unittest.TestCase):
+    def test_base_uuid(self):
+        uuid_generator = BaseIDGenerator()
+        last = -1
+        for _ in range(10000):
+            nid = uuid_generator.next()
+            self.assertGreater(nid, last)
+            last = nid
+
+
+    def test_snoflake(self):
+        uuid_generator = SnowflakeIDGenerator()
+        last = -1
+        for _ in range(10000):
+            nid = uuid_generator.next()
+            self.assertGreater(nid, last)
+            last = nid