[SPARK-24225][CORE] Support closing AutoClosable objects in MemoryStore

core/src/main/scala/org/apache/spark/storage/memory/MemoryStore.scala

@@ -21,6 +21,7 @@ import java.io.OutputStream
 import java.nio.ByteBuffer
 import java.util.LinkedHashMap
 
+import scala.collection.JavaConverters._
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 import scala.reflect.ClassTag
@@ -384,15 +385,36 @@ private[spark] class MemoryStore(
     }
   }
 
+  private def maybeReleaseResources(entry: MemoryEntry[_]): Unit = {
+    entry match {
+      case SerializedMemoryEntry(buffer, _, _) => buffer.dispose()
+      case DeserializedMemoryEntry(objs: Array[Any], _, _) => maybeCloseValues(objs)
+      case _ =>
+    }
+  }
+
+  private def maybeCloseValues(objs: Array[Any]): Unit = {
+    objs.foreach {
+        case closable: AutoCloseable =>
+          safelyCloseValue(closable)
+        case _ =>
+      }
+  }
+
+  private def safelyCloseValue(closable: AutoCloseable) = {
+    try {
+      closable.close()
+    } catch {
+      case ex: Exception => logWarning(s"Failed to close AutoClosable $closable", ex)
+    }
+  }
+
   def remove(blockId: BlockId): Boolean = memoryManager.synchronized {
     val entry = entries.synchronized {
       entries.remove(blockId)
     }
     if (entry != null) {
-      entry match {
-        case SerializedMemoryEntry(buffer, _, _) => buffer.dispose()
-        case _ =>
-      }
+      maybeReleaseResources(entry)
       memoryManager.releaseStorageMemory(entry.size, entry.memoryMode)
       logDebug(s"Block $blockId of size ${entry.size} dropped " +
         s"from memory (free ${maxMemory - blocksMemoryUsed})")
@@ -404,6 +426,7 @@ private[spark] class MemoryStore(
 
   def clear(): Unit = memoryManager.synchronized {
     entries.synchronized {
+      entries.values().asScala.foreach(maybeReleaseResources)
       entries.clear()
     }
     onHeapUnrollMemoryMap.clear()

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

@@ -526,4 +526,84 @@ class MemoryStoreSuite
       }
     }
   }
+
+  test("[SPARK-24225]: remove should close AutoCloseable object") {
+
+    val (store, _) = makeMemoryStore(12000)
+
+    val id = BroadcastBlockId(0)
+    val tracker = new CloseTracker()
+    store.putIteratorAsValues(id, Iterator(tracker), ClassTag.Any)
+    assert(store.contains(id))
+    store.remove(id)
+    assert(tracker.getClosed())
+  }
+
+  test("[SPARK-24225]: remove should close AutoCloseable objects even if they throw exceptions") {
+
+    val (store, _) = makeMemoryStore(12000)
+
+    val id = BroadcastBlockId(0)
+    val tracker = new CloseTracker(true)
+    store.putIteratorAsValues(id, Iterator(tracker), ClassTag.Any)
+    assert(store.contains(id))
+    store.remove(id)
+    assert(tracker.getClosed())
+  }
+
+  test("[SPARK-24225]: clear should close AutoCloseable objects") {
+
+    val (store, _) = makeMemoryStore(12000)
+
+    val id = BroadcastBlockId(0)
+    val tracker = new CloseTracker
+    store.putIteratorAsValues(id, Iterator(tracker), ClassTag.Any)
+    assert(store.contains(id))
+    store.clear()
+    assert(tracker.getClosed())
+  }
+
+  test("[SPARK-24225]: clear should close all AutoCloseable objects put together in an iterator") {
+
+    val (store, _) = makeMemoryStore(12000)
+
+    val id1 = BroadcastBlockId(1)
+    val tracker2 = new CloseTracker
+    val tracker1 = new CloseTracker
+    store.putIteratorAsValues(id1, Iterator(tracker1, tracker2), ClassTag.Any)
+    assert(store.contains(id1))
+    store.clear()
+    assert(tracker1.getClosed())
+    assert(tracker2.getClosed())
+  }
+
+  test("[SPARK-24225]: clear should close AutoCloseable objects even if they throw exceptions") {
+
+    val (store, _) = makeMemoryStore(12000)
+
+    val id1 = BroadcastBlockId(1)
+    val id2 = BroadcastBlockId(2)
+    val tracker2 = new CloseTracker(true)
+    val tracker1 = new CloseTracker(true)
+    store.putIteratorAsValues(id1, Iterator(tracker1), ClassTag.Any)
+    store.putIteratorAsValues(id2, Iterator(tracker2), ClassTag.Any)
+    assert(store.contains(id1))
+    assert(store.contains(id2))
+    store.clear()
+    assert(tracker1.getClosed())
+    assert(tracker2.getClosed())
+  }
+}
+
+private class CloseTracker (val throwsOnClosed: Boolean = false) extends AutoCloseable {
+  var closed = false
+  override def close(): Unit = {
+    closed = true
+    if (throwsOnClosed) {
+      throw new RuntimeException("Throwing")
+    }
+  }
+  def getClosed(): Boolean = {
+    closed
+  }
 }