id_generator.rb

# 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