kotas · June 13, 2014 18:21 · Jun 13, 2014
diff --git a/id_generator.rb b/id_generator.rb
@@ -0,0 +1,90 @@
+# ID Generator for entities
+#
+# ## ID Generation
+# It uses Twitter's Snowflake algorithm for ID generation.
+# https://github.com/twitter/snowflake
+#
+# Our ID is an unsigned 64-bit integer that consists of four elements:
+#
+# - Timestamp:  41 bits, milliseconds from EPOCH_TIME
+# - Shard ID:   12 bits, logical shard ID
+# - Sequence:   11 bits, auto-incrementing sequence, modulus 2048.
+#
+class IdGenerator
+  EPOCH_TIME = 1388534400_000  # 2014-01-01 00:00:00 +00:00
+
+  TIMESTAMP_BITS    = 41
+  SHARD_ID_BITS     = 12
+  SEQUENCE_BITS     = 11
+
+  TIMESTAMP_SHIFT   = SEQUENCE_BITS + SHARD_ID_BITS
+  SHARD_ID_SHIFT    = SEQUENCE_BITS
+  SEQUENCE_SHIFT    = 0
+
+  TIMESTAMP_MASK    = (1 << TIMESTAMP_BITS) - 1
+  SHARD_ID_MASK     = (1 << SHARD_ID_BITS) - 1
+  SEQUENCE_MASK     = (1 << SEQUENCE_BITS) - 1
+
+  class << self
+    # Default shard ID for all generator instances
+    attr_accessor :shard_id
+  end
+
+  # @param  [Fixnum] shard_id     Shard ID. If not given, uses `IdGenerator.shard_id` or `0`
+  def initialize(shard_id = nil)
+    @shard_id       = shard_id || self.class.shard_id || 0
+    @sequence       = 0
+    @last_timestamp = 0
+    @mutex          = Mutex.new
+    raise ArgumentError, "shard_id is out of range"    if (@shard_id & ~SHARD_ID_MASK) != 0
+  end
+
+  # Generate a next ID.
+  #
+  # @return [Fixnum]        Generated ID.
+  # @raise  [RuntimeError]  if system clock moved backwards.
+  def next_id
+    @mutex.lock
+
+    timestamp = current_time
+    raise "System clock moved backwards" if timestamp < @last_timestamp
+
+    if timestamp == @last_timestamp
+      @sequence = (@sequence + 1) & SEQUENCE_MASK
+      timestamp = wait_til_next_tick(timestamp) if @sequence == 0
+    else
+      @sequence = 0
+    end
+
+    @last_timestamp = timestamp
+
+    (timestamp << TIMESTAMP_SHIFT) |
+      (@shard_id << SHARD_ID_SHIFT) |
+      @sequence
+  ensure
+    @mutex.unlock
+  end
+
+  # Reset internal sequence and timestamp.
+  #
+  # USE THIS METHOD ONLY FOR TESTING
+  #
+  def reset!
+    @sequence       = 0
+    @last_timestamp = 0
+  end
+
+  private
+
+  def wait_til_next_tick(last_time)
+    time = current_time
+    while time <= last_time
+      time = current_time
+    end
+    time
+  end
+
+  def current_time
+    (Time.now.to_f * 1000).to_i - EPOCH_TIME
+  end
+end