Fix the issue that blockIntervalTimer won't call updateCurrentBuffer when stopping

Author: zsxwing

streaming/src/main/scala/org/apache/spark/streaming/receiver/BlockGenerator.scala

@@ -140,6 +140,16 @@ private[streaming] class BlockGenerator(
       }
     }
 
+    // When we arrive here, no data will be added to `currentBuffer`. However, `currentBuffer` may
+    // not be empty. If so, we should wait until all data in `currentBuffer` is consumed, because
+    // `blockIntervalTimer.stop(interruptTimer = false)` doesn't guarantee calling
+    // `updateCurrentBuffer` for us.
+    var isCurrentBufferEmpty = synchronized { currentBuffer.isEmpty }
+    while(!isCurrentBufferEmpty) {
+      Thread.sleep(blockIntervalMs)
+      isCurrentBufferEmpty = synchronized { currentBuffer.isEmpty }
+    }
+
     // Stop generating blocks and set the state for block pushing thread to start draining the queue
     logInfo("Stopping BlockGenerator")
     blockIntervalTimer.stop(interruptTimer = false)

streaming/src/test/scala/org/apache/spark/streaming/receiver/BlockGeneratorSuite.scala

@@ -184,9 +184,10 @@ class BlockGeneratorSuite extends SparkFunSuite with BeforeAndAfter {
     // Verify that the final data is present in the final generated block and
     // pushed before complete stop
     assert(blockGenerator.isStopped() === false) // generator has not stopped yet
-    clock.advance(blockIntervalMs)   // force block generation
-    failAfter(1 second) {
-      thread.join()
+    eventually(timeout(10 seconds), interval(10 milliseconds)) {
+      // Keep calling `advance` to avoid blocking forever in `clock.waitTillTime`
+      clock.advance(blockIntervalMs)
+      assert(thread.isAlive === false)
     }
     assert(blockGenerator.isStopped() === true) // generator has finally been completely stopped
     assert(listener.pushedData === data, "All data not pushed by stop()")