Refactored ReceivedBlockTracker API a bit to make things a little cleaner for users of the tracker.

Author: tdas

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

@@ -54,11 +54,22 @@ abstract class ReceiverInputDStream[T: ClassTag](@transient ssc_ : StreamingCont
 
   def stop() {}
 
-  /** Ask ReceiverInputTracker for received data blocks and generates RDDs with them. */
+  /**
+   * Generates RDDs with blocks received by the receiver of this stream. */
   override def compute(validTime: Time): Option[RDD[T]] = {
     val blockRDD = {
-      if (validTime >= graph.startTime) {
-        val blockStoreResults = getReceivedBlockInfo(validTime).map { _.blockStoreResult }
+
+      if (validTime < graph.startTime) {
+        // If this is called for any time before the start time of the context,
+        // then this returns an empty RDD. This may happen when recovering from a
+        // driver failure without any write ahead log to recover pre-failure data.
+        new BlockRDD[T](ssc.sc, Array.empty)
+      } else {
+        // Otherwise, ask the tracker for all the blocks that have been allocated to this stream
+        // for this batch
+        val blockInfos =
+          ssc.scheduler.receiverTracker.getBlocksOfBatch(validTime).get(id).getOrElse(Seq.empty)
+        val blockStoreResults = blockInfos.map { _.blockStoreResult }
         val blockIds = blockStoreResults.map { _.blockId.asInstanceOf[BlockId] }.toArray
         val isWriteAheadLogBased = blockStoreResults.forall {
           _.isInstanceOf[WriteAheadLogBasedStoreResult]
@@ -72,21 +83,11 @@ abstract class ReceiverInputDStream[T: ClassTag](@transient ssc_ : StreamingCont
         } else {
           new BlockRDD[T](ssc.sc, blockIds)
         }
-      } else {
-        // If this is called for any time before the start time of the context,
-        // then this returns an empty RDD. This may happen when recovering from a
-        // driver failure, a
-        new BlockRDD[T](ssc.sc, Array.empty)
       }
     }
     Some(blockRDD)
   }
 
-  /** Get information on received blocks. */
-  private[streaming] def getReceivedBlockInfo(time: Time): Seq[ReceivedBlockInfo] = {
-    ssc.scheduler.receiverTracker.getReceivedBlocks(time, id)
-  }
-
   /**
    * Clear metadata that are older than `rememberDuration` of this DStream.
    * This is an internal method that should not be called directly. This

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

@@ -223,12 +223,9 @@ class JobGenerator(jobScheduler: JobScheduler) extends Logging {
       graph.generateJobs(time) // generate jobs using allocated block
     } match {
       case Success(jobs) =>
-        val receivedBlockInfo = graph.getReceiverInputStreams.map { stream =>
-          val streamId = stream.id
-          val receivedBlockInfo = stream.getReceivedBlockInfo(time).toArray
-          (streamId, receivedBlockInfo)
-        }.toMap
-        jobScheduler.submitJobSet(JobSet(time, jobs, receivedBlockInfo.toMap))
+        val receivedBlockInfos =
+          jobScheduler.receiverTracker.getBlocksOfBatch(time).mapValues { _.toArray }
+        jobScheduler.submitJobSet(JobSet(time, jobs, receivedBlockInfos))
       case Failure(e) =>
         jobScheduler.reportError("Error generating jobs for time " + time, e)
     }

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

@@ -25,7 +25,7 @@ import scala.language.implicitConversions
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 
-import org.apache.spark.{Logging, SparkConf}
+import org.apache.spark.{SparkException, Logging, SparkConf}
 import org.apache.spark.streaming.Time
 import org.apache.spark.streaming.util.{Clock, WriteAheadLogManager}
 import org.apache.spark.util.Utils
@@ -42,8 +42,11 @@ private[streaming] case class BatchCleanupEvent(times: Seq[Time])
 
 
 /** Class representing the blocks of all the streams allocated to a batch */
-private[streaming] case class AllocatedBlocks(streamIdToAllocatedBlocks: Map[Int, Seq[ReceivedBlockInfo]]) {
-  def getBlockForStream(streamId: Int) = streamIdToAllocatedBlocks(streamId)
+private[streaming]
+case class AllocatedBlocks(streamIdToAllocatedBlocks: Map[Int, Seq[ReceivedBlockInfo]]) {
+  def getBlocksOfStream(streamId: Int): Seq[ReceivedBlockInfo] = {
+    streamIdToAllocatedBlocks.get(streamId).getOrElse(Seq.empty)
+  }
 }
 
 /**
@@ -80,6 +83,8 @@ private[streaming] class ReceivedBlockTracker(
     )
   }
 
+  private var lastAllocatedBatchTime: Time = null
+
   // Recover block information from write ahead logs
   recoverFromWriteAheadLogs()
 
@@ -103,22 +108,40 @@ private[streaming] class ReceivedBlockTracker(
    * 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
+    if (lastAllocatedBatchTime == null || batchTime > lastAllocatedBatchTime) {
+      val allocatedBlocks = {
+        val streamIdToBlocks = streamIds.map { streamId =>
+            (streamId, getReceivedBlockQueue(streamId).dequeueAll(x => true))
+        }.toMap
+        AllocatedBlocks(streamIdToBlocks)
+      }
+      writeToLog(BatchAllocationEvent(batchTime, allocatedBlocks))
+      timeToAllocatedBlocks(batchTime) = allocatedBlocks
+      lastAllocatedBatchTime = batchTime
+      allocatedBlocks
+    } else {
+      throw new SparkException(s"Unexpected allocation of blocks, " +
+        s"last batch = $lastAllocatedBatchTime, batch time to allocate = $batchTime  ")
+    }
   }
 
   /** 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 the given batch. */
+  def getBlocksOfBatch(batchTime: Time): Map[Int, Seq[ReceivedBlockInfo]] = synchronized {
+    timeToAllocatedBlocks.get(batchTime).map { _.streamIdToAllocatedBlocks }.getOrElse(Map.empty)
+  }
 
-  /** 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)
+  /** Get the blocks allocated to the given batch and stream. */
+  def getBlocksOfBatchAndStream(batchTime: Time, streamId: Int): Seq[ReceivedBlockInfo] = {
+    synchronized {
+      timeToAllocatedBlocks.get(batchTime).map {
+        _.getBlocksOfStream(streamId)
+      }.getOrElse(Seq.empty)
+    }
   }
 
   /** Check if any blocks are left to be allocated to batches. */
@@ -155,11 +178,12 @@ private[streaming] class ReceivedBlockTracker(
 
     // Insert the recovered block-to-batch allocations and clear the queue of received blocks
     // (when the blocks were originally allocated to the batch, the queue must have been cleared).
-    def insertAllocatedBatch(time: Time, allocatedBlocks: AllocatedBlocks) {
-      logTrace(s"Recovery: Inserting allocated batch for time $time to " +
+    def insertAllocatedBatch(batchTime: Time, allocatedBlocks: AllocatedBlocks) {
+      logTrace(s"Recovery: Inserting allocated batch for time $batchTime to " +
         s"${allocatedBlocks.streamIdToAllocatedBlocks}")
       streamIdToUnallocatedBlockQueues.values.foreach { _.clear() }
-      timeToAllocatedBlocks.put(time, allocatedBlocks)
+      lastAllocatedBatchTime = batchTime
+      timeToAllocatedBlocks.put(batchTime, allocatedBlocks)
     }
 
     // Cleanup the batch allocations

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

@@ -56,13 +56,13 @@ private[streaming]
 class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false) extends Logging {
 
   private val receiverInputStreams = ssc.graph.getReceiverInputStreams()
-  private val receiverInputStreamMap = Map(receiverInputStreams.map(x => (x.id, x)): _*)
+  private val receiverInputStreamIds = receiverInputStreams.map { _.id }
   private val receiverExecutor = new ReceiverLauncher()
   private val receiverInfo = new HashMap[Int, ReceiverInfo] with SynchronizedMap[Int, ReceiverInfo]
   private val receivedBlockTracker = new ReceivedBlockTracker(
     ssc.sparkContext.conf,
     ssc.sparkContext.hadoopConfiguration,
-    receiverInputStreams.map { _.id },
+    receiverInputStreamIds,
     ssc.scheduler.clock,
     Option(ssc.checkpointDir)
   )
@@ -107,12 +107,19 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
     }
   }
 
-  /** Get all the block for batch time . */
-  def getReceivedBlocks(batchTime: Time, streamId: Int): Seq[ReceivedBlockInfo] = {
-    receivedBlockTracker.getBlocksOfBatch(batchTime, streamId)
+  /** Get the blocks for the given batch and all input streams. */
+  def getBlocksOfBatch(batchTime: Time): Map[Int, Seq[ReceivedBlockInfo]] = {
+    receivedBlockTracker.getBlocksOfBatch(batchTime)
   }
 
-  /** Clean up metadata older than the given threshold time */
+  /** Get the blocks allocated to the given batch and stream. */
+  def getBlocksOfBatchAndStream(batchTime: Time, streamId: Int): Seq[ReceivedBlockInfo] = {
+    synchronized {
+      receivedBlockTracker.getBlocksOfBatchAndStream(batchTime, streamId)
+    }
+  }
+
+    /** Clean up metadata older than the given threshold time */
   def cleanupOldMetadata(cleanupThreshTime: Time) {
     receivedBlockTracker.cleanupOldBatches(cleanupThreshTime)
   }
@@ -125,8 +132,8 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
       receiverActor: ActorRef,
       sender: ActorRef
     ) {
-    if (!receiverInputStreamMap.contains(streamId)) {
-      throw new Exception("Register received for unexpected id " + streamId)
+    if (!receiverInputStreamIds.contains(streamId)) {
+      throw new SparkException("Register received for unexpected id " + streamId)
     }
     receiverInfo(streamId) = ReceiverInfo(
       streamId, s"${typ}-${streamId}", receiverActor, true, host)

streaming/src/test/scala/org/apache/spark/streaming/ReceivedBlockTrackerSuite.scala

@@ -28,7 +28,7 @@ import scala.util.Random
 import com.google.common.io.Files
 import org.apache.commons.io.FileUtils
 import org.apache.hadoop.conf.Configuration
-import org.apache.spark.{Logging, SparkConf}
+import org.apache.spark.{SparkException, Logging, SparkConf}
 import org.apache.spark.storage.StreamBlockId
 import org.apache.spark.streaming.scheduler._
 import org.apache.spark.streaming.util.WriteAheadLogSuite._
@@ -73,12 +73,27 @@ class ReceivedBlockTrackerSuite
     val blockInfos = generateBlockInfos()
     blockInfos.map(receivedBlockTracker.addBlock)
 
+    // Verify added blocks are unallocated blocks
     receivedBlockTracker.getUnallocatedBlocks(streamId) shouldEqual blockInfos
+
+    // Allocate the blocks to a batch and verify that all of them have been allocated
     receivedBlockTracker.allocateBlocksToBatch(1)
-    receivedBlockTracker.getBlocksOfBatch(1, streamId) shouldEqual blockInfos
-    receivedBlockTracker.getUnallocatedBlocks(streamId) should have size 0
+    receivedBlockTracker.getBlocksOfBatchAndStream(1, streamId) shouldEqual blockInfos
+    receivedBlockTracker.getUnallocatedBlocks(streamId) shouldBe empty
+
+    // Allocate no blocks to another batch
     receivedBlockTracker.allocateBlocksToBatch(2)
-    receivedBlockTracker.getBlocksOfBatch(2, streamId) should have size 0
+    receivedBlockTracker.getBlocksOfBatchAndStream(2, streamId) shouldBe empty
+
+    // Verify that batch 2 cannot be allocated again
+    intercept[SparkException] {
+      receivedBlockTracker.allocateBlocksToBatch(2)
+    }
+
+    // Verify that older batches cannot be allocated again
+    intercept[SparkException] {
+      receivedBlockTracker.allocateBlocksToBatch(1)
+    }
   }
 
   test("block addition, block to batch allocation and cleanup with write ahead log") {
@@ -126,14 +141,14 @@ class ReceivedBlockTrackerSuite
     // Allocate blocks to batch and verify whether the unallocated blocks got allocated
     val batchTime1 = manualClock.currentTime
     tracker2.allocateBlocksToBatch(batchTime1)
-    tracker2.getBlocksOfBatch(batchTime1, streamId) shouldEqual blockInfos1
+    tracker2.getBlocksOfBatchAndStream(batchTime1, streamId) shouldEqual blockInfos1
 
     // Add more blocks and allocate to another batch
     incrementTime()
     val batchTime2 = manualClock.currentTime
     val blockInfos2 = addBlockInfos(tracker2)
     tracker2.allocateBlocksToBatch(batchTime2)
-    tracker2.getBlocksOfBatch(batchTime2, streamId) shouldEqual blockInfos2
+    tracker2.getBlocksOfBatchAndStream(batchTime2, streamId) shouldEqual blockInfos2
 
     // Verify whether log has correct contents
     val expectedWrittenData2 = expectedWrittenData1 ++
@@ -145,8 +160,8 @@ class ReceivedBlockTrackerSuite
     // Restart tracker and verify recovered state
     incrementTime()
     val tracker3 = createTracker(enableCheckpoint = true, clock = manualClock)
-    tracker3.getBlocksOfBatch(batchTime1, streamId) shouldEqual blockInfos1
-    tracker3.getBlocksOfBatch(batchTime2, streamId) shouldEqual blockInfos2
+    tracker3.getBlocksOfBatchAndStream(batchTime1, streamId) shouldEqual blockInfos1
+    tracker3.getBlocksOfBatchAndStream(batchTime2, streamId) shouldEqual blockInfos2
     tracker3.getUnallocatedBlocks(streamId) shouldBe empty
 
     // Cleanup first batch but not second batch
@@ -155,7 +170,7 @@ class ReceivedBlockTrackerSuite
     tracker3.cleanupOldBatches(batchTime2)
 
     // Verify that the batch allocations have been cleaned, and the act has been written to log
-    tracker3.getBlocksOfBatch(batchTime1, streamId) shouldEqual Seq.empty
+    tracker3.getBlocksOfBatchAndStream(batchTime1, streamId) shouldEqual Seq.empty
     getWrittenLogData(getWriteAheadLogFiles().last) should contain(createBatchCleanup(batchTime1))
 
     // Verify that at least one log file gets deleted
@@ -169,8 +184,8 @@ class ReceivedBlockTrackerSuite
     incrementTime()
     val tracker4 = createTracker(enableCheckpoint = true, clock = manualClock)
     tracker4.getUnallocatedBlocks(streamId) shouldBe empty
-    tracker4.getBlocksOfBatch(batchTime1, streamId) shouldBe empty  // should be cleaned
-    tracker4.getBlocksOfBatch(batchTime2, streamId) shouldEqual blockInfos2
+    tracker4.getBlocksOfBatchAndStream(batchTime1, streamId) shouldBe empty  // should be cleaned
+    tracker4.getBlocksOfBatchAndStream(batchTime2, streamId) shouldEqual blockInfos2
   }
 
   /**