Merge branch 'cached_sanity_check'

src/main/scala/com/twitter/collections/FIFOHashSet.scala

@@ -0,0 +1,28 @@
+/** Copyright 2010 Twitter, Inc.*/
+package com.twitter.collections
+
+import scala.collection.mutable.LinkedHashSet
+
+/**
+ * A size-limited HashSet. Eviction is FIFO
+ */
+class FIFOHashSet[T](sizeLimit: Int) extends LinkedHashSet[T] {
+
+  def this() {
+    this(Integer.MAX_VALUE)
+  }
+
+  override def add(t: T): Boolean = {
+    if (size >= sizeLimit) {
+      trim
+    }
+    super.add(t)
+  }
+
+  protected def trim() {
+    val iter = elements
+    while (size >= sizeLimit) {
+      remove(iter.next)
+    }
+  }
+}

src/main/scala/com/twitter/service/snowflake/IdWorker.scala

@@ -5,6 +5,7 @@ import com.twitter.ostrich.Stats
 import com.twitter.service.snowflake.gen._
 import net.lag.logging.Logger
 import java.util.Random
+import com.twitter.collections.FIFOHashSet
 
 /**
  * An object that generates IDs.
@@ -17,6 +18,9 @@ class IdWorker(workerId: Long, datacenterId: Long) extends Snowflake.Iface {
   val genCounter = Stats.getCounter("ids_generated")
   val reporter = new Reporter
   val rand = new Random
+  val idsCache = new FIFOHashSet[Long](100000) // TODO make this configurable
+
+  Stats.makeGauge("ids_cache_size") { idsCache.size }
 
   val twepoch = 1288834974657L
 
@@ -54,6 +58,10 @@ class IdWorker(workerId: Long, datacenterId: Long) extends Snowflake.Iface {
     val id = nextId()
 
     reporter.report(new AuditLogEntry(id, useragent, rand.nextLong))
+    if (idsCache.contains(id)) {
+      throw new SnowflakeIdCollisionError("id collision error on the server")
+    }
+    idsCache.add(id)
     id
   }
 

src/main/thrift/Snowflake.thrift

@@ -9,6 +9,10 @@ exception InvalidUserAgentError {
   1: string message,
 }
 
+exception SnowflakeIdCollisionError {
+  1: string message,
+}
+
 service Snowflake {
   i64 get_worker_id()
   i64 get_timestamp()

src/test/scala/com/twitter/collections/FIFOHashSetSpec.scala

@@ -0,0 +1,34 @@
+package com.twitter.collections
+import org.specs._
+
+class FIFOHashSetSpec extends Specification {
+  "FIFOHashSet" should {
+    "store and retrieve an item" in {
+      val set = new FIFOHashSet[Int]()
+      set.add(1) must be_==(true)
+      set.contains(1) must be_==(true)
+    }
+
+    "not grow over the size limit" in {
+      val set = new FIFOHashSet[Int](1)
+      set.add(1) must be_==(true)
+      set.add(2) must be_==(true)
+      set.size must be_==(1)
+    }
+
+    "evict the oldest item" in {
+      val set = new FIFOHashSet[Int](10)
+      for (i <- 1 to 10) {
+        set.add(i)
+      }
+      for (i <- 11 to 100) {
+        set.add(i)
+        set.size must be_==(10)
+        set.contains(i-10) must be_==(false)
+        for (j <- i-9 to i) {
+          set.contains(j) must be_==(true)
+        }
+      }
+    }
+  }
+}

src/test/scala/com/twitter/service/snowflake/ReporterSpec.scala

@@ -18,5 +18,4 @@ class ReporterSpec extends Specification {
       true
     }
   }
-
 }