snowflake uid generator

geolinks.snowflake.clj

(ns geolinks.snowflake
  (:import java.util.Random))

(defprotocol ICounter
  (counter-name [this])
  (counter-count [this])
  (increase-count [this]))

(defrecord Counter [name count-atom]
  ICounter
  (counter-name [this] name)
  (counter-count [this] @count-atom)
  (increase-count [this] (swap! count-atom inc)))

(defn- counter [name]
  (new Counter name (atom (long 0))))

(defprotocol ISnowflake
  (next-uid [this]))

(def ^{:private true} twepoch (long 1288834974657))
(def ^{:private true} worker-id-bits (long 5))
(def ^{:private true} datacenter-id-bits (long 5))
(def ^{:private true} max-worker-id (long (bit-xor -1 (bit-shift-left -1 worker-id-bits))))
(def ^{:private true} max-datacenter-id (long (bit-xor -1 (bit-shift-left -1 datacenter-id-bits))))
(def ^{:private true} sequence-bits (long 12))
(def ^{:private true} worker-id-shift sequence-bits)
(def ^{:private true} datacenter-id-shift (+ sequence-bits worker-id-bits))
(def ^{:private true} timestamp-left-shift (+ sequence-bits worker-id-bits datacenter-id-bits))
(def ^{:private true} sequence-mask (long (bit-xor -1 (bit-shift-left -1 sequence-bits))))

(def ^{:private true} useragent-re #"([a-zA-Z][a-zA-Z\-0-9]*)")

(defn- ^Boolean valid-useragent? [^String useragent]
  (boolean (re-matches useragent-re useragent)))

(defn- ^Long time-gen [] (System/currentTimeMillis))

(defn- ^Long til-next-millis [^Long last-timestamp]
  (loop [timestamp (time-gen)]
    (cond (> timestamp last-timestamp) timestamp
          :else (recur (time-gen)))))

(defrecord Snowflake [worker-id, datacenter-id, timestamp,
                      gen-counter, exception-counter, random,
                      last-timestamp-atom, sequence-atom]
  ISnowflake
  (next-uid
   [this]
   (with-local-vars [timestamp-var (time-gen)]
     (cond (= @timestamp-var @last-timestamp-atom)
           (do (swap! sequence-atom #(bit-and % (inc %)))
               (when (zero? @sequence-atom)
                 (var-set timestamp-var (til-next-millis @last-timestamp-atom))))

           :else (reset! sequence-atom 0))

     (when (< @timestamp-var @last-timestamp-atom)
       (increase-count exception-counter)
       (throw (new Exception (str "Clock moved backwards.  Refusing to generate id for "
                                  (- @last-timestamp-atom @timestamp-var) " milliseconds"))))

     (reset! last-timestamp-atom @timestamp-var)
     (increase-count gen-counter)
     (reduce bit-or [(bit-shift-left (- @timestamp-var twepoch) timestamp-left-shift)
                     (bit-shift-left datacenter-id datacenter-id-shift)
                     (bit-shift-left worker-id worker-id-shift)
                     @sequence-atom]))))

(defn snowflake
  "An object that generates IDs.
   This is broken into a separate class in case we ever want to
   support multiple worker threads per process."

  [^Long worker-id, ^Long datacenter-id]

  (when (or (> worker-id max-worker-id)
            (neg? worker-id))
    (throw (new IllegalArgumentException (str "worker Id can't be greater than " max-worker-id " or less than 0"))))

  (when (or (> datacenter-id max-datacenter-id)
            (neg? datacenter-id))
    (throw (new IllegalArgumentException (str "datacenter Id can't be greater than " max-datacenter-id " or less than 0"))))

  ;; (print (str "worker starting. timestamp left shift " timestamp-left-shift
  ;;             ", datacenter id bits " datacenter-id-bits
  ;;             ", worker id bits " worker-id-bits
  ;;             ", sequence bits" sequence-bits
  ;;             ", workerid " worker-id))
  
  (new Snowflake
       worker-id
       datacenter-id
       (System/currentTimeMillis)
       (counter "ids_generated")
       (counter "exceptions")
       (new Random)
       (atom (long -1))
       (atom (long 0))))