added test coverage for user timeouts

twitter-archive · Sep 1, 2010 · 7819d82 · 7819d82
1 parent f71ae86
commit 7819d82
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 2 deletions.
diff --git a/src/main/scala/com/twitter/flockdb/queries/IntersectionQuery.scala b/src/main/scala/com/twitter/flockdb/queries/IntersectionQuery.scala
@@ -44,8 +44,12 @@ class IntersectionQuery(query1: Query, query2: Query, val userTimeoutMS: Int) ex
       val guessedPageSize = (count / config("edges.average_intersection_proportion").toDouble).toInt
       val internalPageSize = guessedPageSize min config("edges.intersection_page_size_max").toInt
 
+      val now = System.currentTimeMillis
+
       var resultWindow = pageIntersection(smallerQuery, largerQuery, internalPageSize, count, cursor)
-      while (resultWindow.page.size < count && resultWindow.continueCursor != Cursor.End) {
+      while (resultWindow.page.size < count &&
+             resultWindow.continueCursor != Cursor.End &&
+             (userTimeoutMS <= 0 || System.currentTimeMillis - now < userTimeoutMS)) {
         resultWindow = resultWindow ++ pageIntersection(smallerQuery, largerQuery, internalPageSize, count, resultWindow.continueCursor)
       }
       resultWindow

diff --git a/src/test/scala/com/twitter/flockdb/unit/IntersectionQuerySpec.scala b/src/test/scala/com/twitter/flockdb/unit/IntersectionQuerySpec.scala
@@ -24,17 +24,32 @@ import org.specs.mock.JMocker
 import conversions.Results._
 import thrift.{Results, Page}
 
+class SlowQuery(s: Seq[Long], userTimeoutMS: Int) extends queries.SeqQuery(s: Seq[Long], userTimeoutMS: Int) {
+  override def selectPageByDestinationId(count: Int, cursor: Cursor) = {
+    Thread.sleep(userTimeoutMS * 2)
+    super.selectPageByDestinationId(count, cursor)
+  }
+}
 
 object IntersectionQuerySpec extends ConfiguredSpecification with JMocker {
-  val timeout = 0
+  val timeout = 100
   "IntersectionQuery" should {
     val query1 = new queries.SeqQuery(List(1,2,3,4,5,6,7,8,9,10), timeout)
     val query2 = new queries.SeqQuery(List(1,2,3,4,11), timeout)
+    val slow = new SlowQuery(List(1,2,3,4,11), timeout)
 
     doBefore {
       config("edges.average_intersection_proportion") = "1.0"
     }
 
+    "slow queries should make progress" in {
+      config("edges.intersection_page_size_max") = "3"
+      val intersectionQuery = new queries.IntersectionQuery(query1, slow, timeout)
+      val window = intersectionQuery.selectPage(5, Cursor.Start)
+      window.data.length mustEqual 2
+      window.continueCursor mustEqual new Cursor(3)
+    }
+
     "sizeEstimate" in {
       val intersectionQuery = new queries.IntersectionQuery(query1, query2, timeout)
       intersectionQuery.sizeEstimate() mustEqual 5