Uncomment BlockManagerSuite test and get it to compile again.

Author: JoshRosen

core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala

@@ -842,71 +842,87 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with BeforeAndAfterE
     assert(Arrays.equals(notMappedAsArray, bytes))
   }
 
-//  test("updated block statuses") {
-//    store = makeBlockManager(12000)
-//    val list = List.fill(2)(new Array[Byte](2000))
-//    val bigList = List.fill(8)(new Array[Byte](2000))
-//
-//    // 1 updated block (i.e. list1)
-//    val updatedBlocks1 =
-//      store.putIterator("list1", list.iterator, StorageLevel.MEMORY_ONLY, tellMaster = true)
-//    assert(updatedBlocks1.size === 1)
-//    assert(updatedBlocks1.head._1 === TestBlockId("list1"))
-//    assert(updatedBlocks1.head._2.storageLevel === StorageLevel.MEMORY_ONLY)
-//
-//    // 1 updated block (i.e. list2)
-//    val updatedBlocks2 =
-//      store.putIterator("list2", list.iterator, StorageLevel.MEMORY_AND_DISK, tellMaster = true)
-//    assert(updatedBlocks2.size === 1)
-//    assert(updatedBlocks2.head._1 === TestBlockId("list2"))
-//    assert(updatedBlocks2.head._2.storageLevel === StorageLevel.MEMORY_ONLY)
-//
-//    // 2 updated blocks - list1 is kicked out of memory while list3 is added
-//    val updatedBlocks3 =
-//      store.putIterator("list3", list.iterator, StorageLevel.MEMORY_ONLY, tellMaster = true)
-//    assert(updatedBlocks3.size === 2)
-//    updatedBlocks3.foreach { case (id, status) =>
-//      id match {
-//        case TestBlockId("list1") => assert(status.storageLevel === StorageLevel.NONE)
-//        case TestBlockId("list3") => assert(status.storageLevel === StorageLevel.MEMORY_ONLY)
-//        case _ => fail("Updated block is neither list1 nor list3")
-//      }
-//    }
-//    assert(store.memoryStore.contains("list3"), "list3 was not in memory store")
-//
-//    // 2 updated blocks - list2 is kicked out of memory (but put on disk) while list4 is added
-//    val updatedBlocks4 =
-//      store.putIterator("list4", list.iterator, StorageLevel.MEMORY_ONLY, tellMaster = true)
-//    assert(updatedBlocks4.size === 2)
-//    updatedBlocks4.foreach { case (id, status) =>
-//      id match {
-//        case TestBlockId("list2") => assert(status.storageLevel === StorageLevel.DISK_ONLY)
-//        case TestBlockId("list4") => assert(status.storageLevel === StorageLevel.MEMORY_ONLY)
-//        case _ => fail("Updated block is neither list2 nor list4")
-//      }
-//    }
-//    assert(store.diskStore.contains("list2"), "list2 was not in disk store")
-//    assert(store.memoryStore.contains("list4"), "list4 was not in memory store")
-//
-//    // No updated blocks - list5 is too big to fit in store and nothing is kicked out
-//    val updatedBlocks5 =
-//      store.putIterator("list5", bigList.iterator, StorageLevel.MEMORY_ONLY, tellMaster = true)
-//    assert(updatedBlocks5.size === 0)
-//
-//    // memory store contains only list3 and list4
-//    assert(!store.memoryStore.contains("list1"), "list1 was in memory store")
-//    assert(!store.memoryStore.contains("list2"), "list2 was in memory store")
-//    assert(store.memoryStore.contains("list3"), "list3 was not in memory store")
-//    assert(store.memoryStore.contains("list4"), "list4 was not in memory store")
-//    assert(!store.memoryStore.contains("list5"), "list5 was in memory store")
-//
-//    // disk store contains only list2
-//    assert(!store.diskStore.contains("list1"), "list1 was in disk store")
-//    assert(store.diskStore.contains("list2"), "list2 was not in disk store")
-//    assert(!store.diskStore.contains("list3"), "list3 was in disk store")
-//    assert(!store.diskStore.contains("list4"), "list4 was in disk store")
-//    assert(!store.diskStore.contains("list5"), "list5 was in disk store")
-//  }
+  test("updated block statuses") {
+    store = makeBlockManager(12000)
+    val list = List.fill(2)(new Array[Byte](2000))
+    val bigList = List.fill(8)(new Array[Byte](2000))
+
+    def getUpdatedBlocks(task: => Unit): Seq[(BlockId, BlockStatus)] = {
+      val context = TaskContext.empty()
+      try {
+        TaskContext.setTaskContext(context)
+        task
+      } finally {
+        TaskContext.unset()
+      }
+      context.taskMetrics.updatedBlocks.getOrElse(Seq.empty)
+    }
+
+    // 1 updated block (i.e. list1)
+    val updatedBlocks1 = getUpdatedBlocks {
+      store.putIterator("list1", list.iterator, StorageLevel.MEMORY_ONLY, tellMaster = true)
+    }
+    assert(updatedBlocks1.size === 1)
+    assert(updatedBlocks1.head._1 === TestBlockId("list1"))
+    assert(updatedBlocks1.head._2.storageLevel === StorageLevel.MEMORY_ONLY)
+
+    // 1 updated block (i.e. list2)
+    val updatedBlocks2 = getUpdatedBlocks {
+      store.putIterator("list2", list.iterator, StorageLevel.MEMORY_AND_DISK, tellMaster = true)
+    }
+    assert(updatedBlocks2.size === 1)
+    assert(updatedBlocks2.head._1 === TestBlockId("list2"))
+    assert(updatedBlocks2.head._2.storageLevel === StorageLevel.MEMORY_ONLY)
+
+    // 2 updated blocks - list1 is kicked out of memory while list3 is added
+    val updatedBlocks3 = getUpdatedBlocks {
+      store.putIterator("list3", list.iterator, StorageLevel.MEMORY_ONLY, tellMaster = true)
+    }
+    assert(updatedBlocks3.size === 2)
+    updatedBlocks3.foreach { case (id, status) =>
+      id match {
+        case TestBlockId("list1") => assert(status.storageLevel === StorageLevel.NONE)
+        case TestBlockId("list3") => assert(status.storageLevel === StorageLevel.MEMORY_ONLY)
+        case _ => fail("Updated block is neither list1 nor list3")
+      }
+    }
+    assert(store.memoryStore.contains("list3"), "list3 was not in memory store")
+
+    // 2 updated blocks - list2 is kicked out of memory (but put on disk) while list4 is added
+    val updatedBlocks4 = getUpdatedBlocks {
+      store.putIterator("list4", list.iterator, StorageLevel.MEMORY_ONLY, tellMaster = true)
+    }
+    assert(updatedBlocks4.size === 2)
+    updatedBlocks4.foreach { case (id, status) =>
+      id match {
+        case TestBlockId("list2") => assert(status.storageLevel === StorageLevel.DISK_ONLY)
+        case TestBlockId("list4") => assert(status.storageLevel === StorageLevel.MEMORY_ONLY)
+        case _ => fail("Updated block is neither list2 nor list4")
+      }
+    }
+    assert(store.diskStore.contains("list2"), "list2 was not in disk store")
+    assert(store.memoryStore.contains("list4"), "list4 was not in memory store")
+
+    // No updated blocks - list5 is too big to fit in store and nothing is kicked out
+    val updatedBlocks5 = getUpdatedBlocks {
+      store.putIterator("list5", bigList.iterator, StorageLevel.MEMORY_ONLY, tellMaster = true)
+    }
+    assert(updatedBlocks5.size === 0)
+
+    // memory store contains only list3 and list4
+    assert(!store.memoryStore.contains("list1"), "list1 was in memory store")
+    assert(!store.memoryStore.contains("list2"), "list2 was in memory store")
+    assert(store.memoryStore.contains("list3"), "list3 was not in memory store")
+    assert(store.memoryStore.contains("list4"), "list4 was not in memory store")
+    assert(!store.memoryStore.contains("list5"), "list5 was in memory store")
+
+    // disk store contains only list2
+    assert(!store.diskStore.contains("list1"), "list1 was in disk store")
+    assert(store.diskStore.contains("list2"), "list2 was not in disk store")
+    assert(!store.diskStore.contains("list3"), "list3 was in disk store")
+    assert(!store.diskStore.contains("list4"), "list4 was in disk store")
+    assert(!store.diskStore.contains("list5"), "list5 was in disk store")
+  }
 
   test("query block statuses") {
     store = makeBlockManager(12000)