Changed class names, made allocation more explicit and added cleanup

Author: tdas

streaming/src/main/scala/org/apache/spark/streaming/dstream/ReceiverInputDStream.scala

@@ -95,6 +95,6 @@ abstract class ReceiverInputDStream[T: ClassTag](@transient ssc_ : StreamingCont
    */
   private[streaming] override def clearMetadata(time: Time) {
     super.clearMetadata(time)
-    ssc.scheduler.receiverTracker.cleanupOldInfo(time - rememberDuration)
+    ssc.scheduler.receiverTracker.cleanupOldMetadata(time - rememberDuration)
   }
 }

streaming/src/main/scala/org/apache/spark/streaming/scheduler/JobGenerator.scala

@@ -112,7 +112,7 @@ class JobGenerator(jobScheduler: JobScheduler) extends Logging {
       // Wait until all the received blocks in the network input tracker has
       // been consumed by network input DStreams, and jobs have been generated with them
       logInfo("Waiting for all received blocks to be consumed for job generation")
-      while(!hasTimedOut && jobScheduler.receiverTracker.hasMoreReceivedBlockIds) {
+      while(!hasTimedOut && jobScheduler.receiverTracker.hasUnallocatedBlocks) {
         Thread.sleep(pollTime)
       }
       logInfo("Waited for all received blocks to be consumed for job generation")
@@ -218,7 +218,10 @@ class JobGenerator(jobScheduler: JobScheduler) extends Logging {
   /** Generate jobs and perform checkpoint for the given `time`.  */
   private def generateJobs(time: Time) {
     SparkEnv.set(ssc.env)
-    Try(graph.generateJobs(time)) match {
+    Try {
+      jobScheduler.receiverTracker.allocateBlocksToBatch(time) // allocate received blocks to batch
+      graph.generateJobs(time) // generate jobs using allocated block
+    } match {
       case Success(jobs) =>
         val receivedBlockInfo = graph.getReceiverInputStreams.map { stream =>
           val streamId = stream.id
@@ -235,6 +238,7 @@ class JobGenerator(jobScheduler: JobScheduler) extends Logging {
   /** Clear DStream metadata for the given `time`. */
   private def clearMetadata(time: Time) {
     ssc.graph.clearMetadata(time)
+    jobScheduler.receiverTracker.cleanupOldMetadata(time - graph.batchDuration)
 
     // If checkpointing is enabled, then checkpoint,
     // else mark batch to be fully processed

streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceivedBlockTracker.scala

@@ -26,32 +26,31 @@ import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 
 import org.apache.spark.{Logging, SparkConf}
-import org.apache.spark.storage.StreamBlockId
 import org.apache.spark.streaming.Time
 import org.apache.spark.streaming.util.{Clock, WriteAheadLogManager}
 import org.apache.spark.util.Utils
 
-/** Trait representing any action done in the ReceivedBlockTracker */
-private[streaming] sealed trait ReceivedBlockTrackerAction
+/** Trait representing any event in the ReceivedBlockTracker that updates its state. */
+private[streaming] sealed trait ReceivedBlockTrackerLogEvent
 
-private[streaming] case class BlockAddition(receivedBlockInfo: ReceivedBlockInfo)
-  extends ReceivedBlockTrackerAction
-private[streaming] case class BatchAllocations(time: Time, allocatedBlocks: AllocatedBlocks)
-  extends ReceivedBlockTrackerAction
-private[streaming] case class BatchCleanup(times: Seq[Time])
-  extends ReceivedBlockTrackerAction
+private[streaming] case class BlockAdditionEvent(receivedBlockInfo: ReceivedBlockInfo)
+  extends ReceivedBlockTrackerLogEvent
+private[streaming] case class BatchAllocationEvent(time: Time, allocatedBlocks: AllocatedBlocks)
+  extends ReceivedBlockTrackerLogEvent
+private[streaming] case class BatchCleanupEvent(times: Seq[Time])
+  extends ReceivedBlockTrackerLogEvent
 
 
 /** Class representing the blocks of all the streams allocated to a batch */
 case class AllocatedBlocks(streamIdToAllocatedBlocks: Map[Int, Seq[ReceivedBlockInfo]]) {
-  def apply(streamId: Int) = streamIdToAllocatedBlocks(streamId)
+  def getBlockForStream(streamId: Int) = streamIdToAllocatedBlocks(streamId)
 }
 
 /**
  * Class that keep track of all the received blocks, and allocate them to batches
- * when required. All actions taken by this class can be saved to a write ahead log,
- * so that the state of the tracker (received blocks and block-to-batch allocations)
- * can be recovered after driver failure.
+ * when required. All actions taken by this class can be saved to a write ahead log
+ * (if a checkpoint directory has been provided), so that the state of the tracker
+ * (received blocks and block-to-batch allocations) can be recovered after driver failure.
  */
 private[streaming]
 class ReceivedBlockTracker(
@@ -60,8 +59,8 @@ class ReceivedBlockTracker(
 
   private type ReceivedBlockQueue = mutable.Queue[ReceivedBlockInfo]
 
-  private val streamIdToUnallocatedBlockInfo = new mutable.HashMap[Int, ReceivedBlockQueue]
-  private val timeToAllocatedBlockInfo = new mutable.HashMap[Time, AllocatedBlocks]
+  private val streamIdToUnallocatedBlockQueues = new mutable.HashMap[Int, ReceivedBlockQueue]
+  private val timeToAllocatedBlocks = new mutable.HashMap[Time, AllocatedBlocks]
 
   private val logManagerRollingIntervalSecs = conf.getInt(
     "spark.streaming.receivedBlockTracker.writeAheadLog.rotationIntervalSecs", 60)
@@ -78,70 +77,69 @@ class ReceivedBlockTracker(
   // Recover block information from write ahead logs
   recoverFromWriteAheadLogs()
 
-  /** Add received block */
+  /** Add received block. This event will get written to the write ahead log (if enabled). */
   def addBlock(receivedBlockInfo: ReceivedBlockInfo): Boolean = synchronized {
     try {
-      writeToLog(BlockAddition(receivedBlockInfo))
+      writeToLog(BlockAdditionEvent(receivedBlockInfo))
       getReceivedBlockQueue(receivedBlockInfo.streamId) += receivedBlockInfo
       logDebug(s"Stream ${receivedBlockInfo.streamId} received " +
         s"block ${receivedBlockInfo.blockStoreResult.blockId}")
       true
     } catch {
       case e: Exception =>
-        logError("Error adding block " + receivedBlockInfo, e)
+        logError(s"Error adding block $receivedBlockInfo", e)
         false
     }
   }
 
-  /** Get blocks that have been added but not yet allocated to any batch */
+  /**
+   * Allocate all unallocated blocks to the given batch.
+   * This event will get written to the write ahead log (if enabled).
+   */
+  def allocateBlocksToBatch(batchTime: Time): Unit = synchronized {
+    val allocatedBlocks = AllocatedBlocks(streamIds.map { streamId =>
+        (streamId, getReceivedBlockQueue(streamId).dequeueAll(x => true))
+    }.toMap)
+    writeToLog(BatchAllocationEvent(batchTime, allocatedBlocks))
+    timeToAllocatedBlocks(batchTime) = allocatedBlocks
+    allocatedBlocks
+  }
+
+  /** Get blocks that have been added but not yet allocated to any batch. */
   def getUnallocatedBlocks(streamId: Int): Seq[ReceivedBlockInfo] = synchronized {
     getReceivedBlockQueue(streamId).toSeq
   } 
 
-  /** Get the blocks allocated to a batch, or allocate blocks to the batch and then get them */
-  def getOrAllocateBlocksToBatch(batchTime: Time, streamId: Int): Seq[ReceivedBlockInfo] = {
-    synchronized {
-      if (!timeToAllocatedBlockInfo.contains(batchTime)) {
-        allocateAllUnallocatedBlocksToBatch(batchTime)
-      }
-      timeToAllocatedBlockInfo(batchTime)(streamId)
-    }
+  /** Get the blocks allocated to a batch, or allocate blocks to the batch and then get them. */
+  def getBlocksOfBatch(batchTime: Time, streamId: Int): Seq[ReceivedBlockInfo] = synchronized {
+    timeToAllocatedBlocks.get(batchTime).map { _.getBlockForStream(streamId) }.getOrElse(Seq.empty)
   }
 
-  /** Check if any blocks are left to be allocated to batches */
-  def hasUnallocatedReceivedBlocks(): Boolean = synchronized {
-    !streamIdToUnallocatedBlockInfo.values.forall(_.isEmpty)
+  /** Check if any blocks are left to be allocated to batches. */
+  def hasUnallocatedReceivedBlocks: Boolean = synchronized {
+    !streamIdToUnallocatedBlockQueues.values.forall(_.isEmpty)
   }
 
-  /** Clean up block information of old batches */
+  /** Clean up block information of old batches. */
   def cleanupOldBatches(cleanupThreshTime: Time): Unit = synchronized {
     assert(cleanupThreshTime.milliseconds < clock.currentTime())
-    val timesToCleanup = timeToAllocatedBlockInfo.keys.filter { _ < cleanupThreshTime }.toSeq
+    val timesToCleanup = timeToAllocatedBlocks.keys.filter { _ < cleanupThreshTime }.toSeq
     logInfo("Deleting batches " + timesToCleanup)
-    writeToLog(BatchCleanup(timesToCleanup))
-    timeToAllocatedBlockInfo --= timesToCleanup
+    writeToLog(BatchCleanupEvent(timesToCleanup))
+    timeToAllocatedBlocks --= timesToCleanup
     logManagerOption.foreach(_.cleanupOldLogs(cleanupThreshTime.milliseconds))
     log
   }
 
-  /** Stop the block tracker */
+  /** Stop the block tracker. */
   def stop() {
     logManagerOption.foreach { _.stop() }
   }
 
-  /** Allocate all unallocated blocks to the given batch */
-  private def allocateAllUnallocatedBlocksToBatch(batchTime: Time): AllocatedBlocks = synchronized {
-    val allocatedBlockInfos = AllocatedBlocks(streamIds.map { streamId =>
-      (streamId, getReceivedBlockQueue(streamId).dequeueAll(x => true))
-    }.toMap)
-    writeToLog(BatchAllocations(batchTime, allocatedBlockInfos))
-    timeToAllocatedBlockInfo(batchTime) = allocatedBlockInfos
-    allocatedBlockInfos
-  }
 
   /**
    * Recover all the tracker actions from the write ahead logs to recover the state (unallocated
-   * and allocated block info) prior to failure
+   * and allocated block info) prior to failure.
    */
   private def recoverFromWriteAheadLogs(): Unit = synchronized {
     logInfo("Recovering from checkpoint")
@@ -160,34 +158,34 @@ class ReceivedBlockTracker(
         s"${allocatedBlocks.streamIdToAllocatedBlocks}")
       // println(s"Recovery: Inserting allocated batch for time $time to " +
       // s"${allocatedBlocks.streamIdToAllocatedBlocks}")
-      streamIdToUnallocatedBlockInfo.values.foreach { _.clear() }
-      timeToAllocatedBlockInfo.put(time, allocatedBlocks)
+      streamIdToUnallocatedBlockQueues.values.foreach { _.clear() }
+      timeToAllocatedBlocks.put(time, allocatedBlocks)
     }
 
     // Cleanup the batch allocations
     def cleanupBatches(batchTimes: Seq[Time]) {
       logTrace(s"Recovery: Cleaning up batches $batchTimes")
       // println(s"Recovery: Cleaning up batches ${batchTimes}")
-      timeToAllocatedBlockInfo --= batchTimes
+      timeToAllocatedBlocks --= batchTimes
     }
 
     logManagerOption.foreach { logManager =>
       logManager.readFromLog().foreach { byteBuffer =>
         logTrace("Recovering record " + byteBuffer)
-        Utils.deserialize[ReceivedBlockTrackerAction](byteBuffer.array) match {
-          case BlockAddition(receivedBlockInfo) =>
+        Utils.deserialize[ReceivedBlockTrackerLogEvent](byteBuffer.array) match {
+          case BlockAdditionEvent(receivedBlockInfo) =>
             insertAddedBlock(receivedBlockInfo)
-          case BatchAllocations(time, allocatedBlocks) =>
+          case BatchAllocationEvent(time, allocatedBlocks) =>
             insertAllocatedBatch(time, allocatedBlocks)
-          case BatchCleanup(batchTimes) =>
+          case BatchCleanupEvent(batchTimes) =>
             cleanupBatches(batchTimes)
         }
       }
     }
   }
 
   /** Write an update to the tracker to the write ahead log */
-  private def writeToLog(record: ReceivedBlockTrackerAction) {
+  private def writeToLog(record: ReceivedBlockTrackerLogEvent) {
     logDebug(s"Writing to log $record")
     logManagerOption.foreach { logManager =>
         logManager.writeToLog(ByteBuffer.wrap(Utils.serialize(record)))
@@ -196,7 +194,7 @@ class ReceivedBlockTracker(
 
   /** Get the queue of received blocks belonging to a particular stream */
   private def getReceivedBlockQueue(streamId: Int): ReceivedBlockQueue = {
-    streamIdToUnallocatedBlockInfo.getOrElseUpdate(streamId, new ReceivedBlockQueue)
+    streamIdToUnallocatedBlockQueues.getOrElseUpdate(streamId, new ReceivedBlockQueue)
   }
 }
 

streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverTracker.scala

@@ -100,12 +100,20 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
     }
   }
 
-  /** Return all the blocks received from a receiver. */
+  /** Allocate all unallocated blocks to the given batch. */
+  def allocateBlocksToBatch(batchTime: Time): Unit = {
+    if (receiverInputStreams.nonEmpty) {
+      receivedBlockTracker.allocateBlocksToBatch(batchTime)
+    }
+  }
+
+  /** Get all the block for batch time . */
   def getReceivedBlocks(batchTime: Time, streamId: Int): Seq[ReceivedBlockInfo] = {
-    receivedBlockTracker.getOrAllocateBlocksToBatch(batchTime, streamId)
+    receivedBlockTracker.getBlocksOfBatch(batchTime, streamId)
   }
 
-  def cleanupOldInfo(cleanupThreshTime: Time) {
+  /** Clean up metadata older than the given threshold time */
+  def cleanupOldMetadata(cleanupThreshTime: Time) {
     receivedBlockTracker.cleanupOldBatches(cleanupThreshTime)
   }
 
@@ -170,8 +178,8 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
   }
 
   /** Check if any blocks are left to be processed */
-  def hasMoreReceivedBlockIds: Boolean = {
-    receivedBlockTracker.hasUnallocatedReceivedBlocks()
+  def hasUnallocatedBlocks: Boolean = {
+    receivedBlockTracker.hasUnallocatedReceivedBlocks
   }
 
   /** Actor to receive messages from the receivers. */