Consolidate Streaming Clock class into common util Clock

Author: srowen

core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala

@@ -124,7 +124,7 @@ private[spark] class ExecutorAllocationManager(
   private val intervalMillis: Long = 100
 
   // Clock used to schedule when executors should be added and removed
-  private var clock: Clock = SystemClock
+  private var clock: Clock = new SystemClock()
 
   // Listener for Spark events that impact the allocation policy
   private val listener = new ExecutorAllocationListener

core/src/main/scala/org/apache/spark/deploy/worker/DriverRunner.scala

@@ -58,7 +58,7 @@ private[spark] class DriverRunner(
   // Decoupled for testing
   private[deploy] def setClock(_clock: Clock) = clock = _clock
   private[deploy] def setSleeper(_sleeper: Sleeper) = sleeper = _sleeper
-  private var clock: Clock = SystemClock
+  private var clock: Clock = new SystemClock()
   private var sleeper = new Sleeper {
     def sleep(seconds: Int): Unit = (0 until seconds).takeWhile(f => {Thread.sleep(1000); !killed})
   }

core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala

@@ -63,7 +63,7 @@ class DAGScheduler(
     mapOutputTracker: MapOutputTrackerMaster,
     blockManagerMaster: BlockManagerMaster,
     env: SparkEnv,
-    clock: org.apache.spark.util.Clock = SystemClock)
+    clock: Clock = new SystemClock())
   extends Logging {
 
   def this(sc: SparkContext, taskScheduler: TaskScheduler) = {

core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala

@@ -51,7 +51,7 @@ private[spark] class TaskSetManager(
     sched: TaskSchedulerImpl,
     val taskSet: TaskSet,
     val maxTaskFailures: Int,
-    clock: Clock = SystemClock)
+    clock: Clock = new SystemClock())
   extends Schedulable with Logging {
 
   val conf = sched.sc.conf

core/src/main/scala/org/apache/spark/util/Clock.scala

@@ -22,8 +22,35 @@ package org.apache.spark.util
  */
 private[spark] trait Clock {
   def getTime(): Long
+  def waitTillTime(targetTime: Long): Long
 }
 
-private[spark] object SystemClock extends Clock {
+private[spark] class SystemClock extends Clock {
+
+  val minPollTime = 25L
+
   def getTime(): Long = System.currentTimeMillis()
+
+  def waitTillTime(targetTime: Long): Long = {
+    var currentTime = 0L
+    currentTime = System.currentTimeMillis()
+
+    var waitTime = targetTime - currentTime
+    if (waitTime <= 0) {
+      return currentTime
+    }
+
+    val pollTime = math.max(waitTime / 10.0, minPollTime).toLong
+
+    while (true) {
+      currentTime = System.currentTimeMillis()
+      waitTime = targetTime - currentTime
+      if (waitTime <= 0) {
+        return currentTime
+      }
+      val sleepTime = math.min(waitTime, pollTime)
+      Thread.sleep(sleepTime)
+    }
+    -1
+  }
 }

core/src/test/scala/org/apache/spark/util/FakeClock.scala renamed to core/src/main/scala/org/apache/spark/util/FakeClock.scala

@@ -17,9 +17,33 @@
 
 package org.apache.spark.util
 
-class FakeClock(private var time: Long = 0L) extends Clock {
+class FakeClock(private var time: Long) extends Clock {
 
-  def advance(millis: Long): Unit = time += millis
+  def this() = this(0L)
+
+  def getTime(): Long =
+    synchronized {
+      time
+    }
+
+  def setTime(timeToSet: Long) =
+    synchronized {
+      time = timeToSet
+      notifyAll()
+    }
+
+  def advance(timeToAdd: Long) =
+    this.synchronized {
+      time += timeToAdd
+      notifyAll()
+    }
+
+  def waitTillTime(targetTime: Long): Long =
+    synchronized {
+      while (time < targetTime) {
+        wait(100)
+      }
+      getTime()
+    }
 
-  def getTime(): Long = time
 }

external/flume/src/test/java/org/apache/spark/streaming/LocalJavaStreamingContext.java

@@ -31,7 +31,7 @@ public void setUp() {
         SparkConf conf = new SparkConf()
             .setMaster("local[2]")
             .setAppName("test")
-            .set("spark.streaming.clock", "org.apache.spark.streaming.util.ManualClock");
+            .set("spark.streaming.clock", "org.apache.spark.util.FakeClock");
         ssc = new JavaStreamingContext(conf, new Duration(1000));
         ssc.checkpoint("checkpoint");
     }

external/flume/src/test/scala/org/apache/spark/streaming/flume/FlumePollingStreamSuite.scala

@@ -34,10 +34,9 @@ import org.scalatest.{BeforeAndAfter, FunSuite}
 import org.apache.spark.{SparkConf, Logging}
 import org.apache.spark.storage.StorageLevel
 import org.apache.spark.streaming.dstream.ReceiverInputDStream
-import org.apache.spark.streaming.util.ManualClock
 import org.apache.spark.streaming.{Seconds, TestOutputStream, StreamingContext}
 import org.apache.spark.streaming.flume.sink._
-import org.apache.spark.util.Utils
+import org.apache.spark.util.{FakeClock, Utils}
 
 class FlumePollingStreamSuite extends FunSuite with BeforeAndAfter with Logging {
 
@@ -54,7 +53,7 @@ class FlumePollingStreamSuite extends FunSuite with BeforeAndAfter with Logging
 
   def beforeFunction() {
     logInfo("Using manual clock")
-    conf.set("spark.streaming.clock", "org.apache.spark.streaming.util.ManualClock")
+    conf.set("spark.streaming.clock", "org.apache.spark.util.FakeClock")
   }
 
   before(beforeFunction())
@@ -171,7 +170,7 @@ class FlumePollingStreamSuite extends FunSuite with BeforeAndAfter with Logging
 
   def writeAndVerify(channels: Seq[MemoryChannel], ssc: StreamingContext,
     outputBuffer: ArrayBuffer[Seq[SparkFlumeEvent]]) {
-    val clock = ssc.scheduler.clock.asInstanceOf[ManualClock]
+    val clock = ssc.scheduler.clock.asInstanceOf[FakeClock]
     val executor = Executors.newCachedThreadPool()
     val executorCompletion = new ExecutorCompletionService[Void](executor)
     channels.map(channel => {
@@ -221,7 +220,7 @@ class FlumePollingStreamSuite extends FunSuite with BeforeAndAfter with Logging
     assert(m.invoke(queueRemaining.get(channel)).asInstanceOf[Int] === 5000)
   }
 
-  private class TxnSubmitter(channel: MemoryChannel, clock: ManualClock) extends Callable[Void] {
+  private class TxnSubmitter(channel: MemoryChannel, clock: FakeClock) extends Callable[Void] {
     override def call(): Void = {
       var t = 0
       for (i <- 0 until batchCount) {
@@ -236,7 +235,7 @@ class FlumePollingStreamSuite extends FunSuite with BeforeAndAfter with Logging
         tx.commit()
         tx.close()
         Thread.sleep(500) // Allow some time for the events to reach
-        clock.addToTime(batchDuration.milliseconds)
+        clock.advance(batchDuration.milliseconds)
       }
       null
     }

external/mqtt/src/test/java/org/apache/spark/streaming/LocalJavaStreamingContext.java

@@ -31,7 +31,7 @@ public void setUp() {
         SparkConf conf = new SparkConf()
             .setMaster("local[2]")
             .setAppName("test")
-            .set("spark.streaming.clock", "org.apache.spark.streaming.util.ManualClock");
+            .set("spark.streaming.clock", "org.apache.spark.util.FakeClock");
         ssc = new JavaStreamingContext(conf, new Duration(1000));
         ssc.checkpoint("checkpoint");
     }

external/twitter/src/test/java/org/apache/spark/streaming/LocalJavaStreamingContext.java

@@ -31,7 +31,7 @@ public void setUp() {
         SparkConf conf = new SparkConf()
             .setMaster("local[2]")
             .setAppName("test")
-            .set("spark.streaming.clock", "org.apache.spark.streaming.util.ManualClock");
+            .set("spark.streaming.clock", "org.apache.spark.util.FakeClock");
         ssc = new JavaStreamingContext(conf, new Duration(1000));
         ssc.checkpoint("checkpoint");
     }