Port tests from ShuffleMemoryManagerSuite

core/src/test/scala/org/apache/spark/memory/MemoryManagerSuite.scala

@@ -19,10 +19,14 @@ package org.apache.spark.memory
 
 import java.util.concurrent.atomic.AtomicLong
 
+import scala.concurrent.duration.Duration
+import scala.concurrent.{Await, ExecutionContext, Future}
+
 import org.mockito.Matchers.{any, anyLong}
 import org.mockito.Mockito.{mock, when}
 import org.mockito.invocation.InvocationOnMock
 import org.mockito.stubbing.Answer
+import org.scalatest.time.SpanSugar._
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.storage.MemoryStore
@@ -126,6 +130,136 @@ private[memory] trait MemoryManagerSuite extends SparkFunSuite {
     assert(ensureFreeSpaceCalled.get() === DEFAULT_ENSURE_FREE_SPACE_CALLED,
       "ensure free space should not have been called!")
   }
+
+  /**
+   * Create a MemoryManager with the specified execution memory limit and no storage memory.
+   */
+  protected def createMemoryManager(maxExecutionMemory: Long): MemoryManager
+
+  // -- Tests of sharing of execution memory between tasks ----------------------------------------
+  // Prior to Spark 1.6, these tests were part of ShuffleMemoryManagerSuite.
+
+  implicit val ec = ExecutionContext.global
+
+  test("single task requesting execution memory") {
+    val manager = createMemoryManager(1000L)
+    val taskMemoryManager = new TaskMemoryManager(manager, 0)
+
+    assert(taskMemoryManager.acquireExecutionMemory(100L) === 100L)
+    assert(taskMemoryManager.acquireExecutionMemory(400L) === 400L)
+    assert(taskMemoryManager.acquireExecutionMemory(400L) === 400L)
+    assert(taskMemoryManager.acquireExecutionMemory(200L) === 100L)
+    assert(taskMemoryManager.acquireExecutionMemory(100L) === 0L)
+    assert(taskMemoryManager.acquireExecutionMemory(100L) === 0L)
+
+    taskMemoryManager.releaseExecutionMemory(500L)
+    assert(taskMemoryManager.acquireExecutionMemory(300L) === 300L)
+    assert(taskMemoryManager.acquireExecutionMemory(300L) === 200L)
+
+    taskMemoryManager.cleanUpAllAllocatedMemory()
+    assert(taskMemoryManager.acquireExecutionMemory(1000L) === 1000L)
+    assert(taskMemoryManager.acquireExecutionMemory(100L) === 0L)
+  }
+
+  test("two tasks requesting full execution memory") {
+    val memoryManager = createMemoryManager(1000L)
+    val t1MemManager = new TaskMemoryManager(memoryManager, 1)
+    val t2MemManager = new TaskMemoryManager(memoryManager, 2)
+    val futureTimeout: Duration = 20.seconds
+
+    // Have both tasks request 500 bytes, then wait until both requests have been granted:
+    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(500L) }
+    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(500L) }
+    assert(Await.result(t1Result1, futureTimeout) === 500L)
+    assert(Await.result(t2Result1, futureTimeout) === 500L)
+
+    // Have both tasks each request 500 bytes more; both should immediately return 0 as they are
+    // both now at 1 / N
+    val t1Result2 = Future { t1MemManager.acquireExecutionMemory(500L) }
+    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(500L) }
+    assert(Await.result(t1Result2, 200.millis) === 0L)
+    assert(Await.result(t2Result2, 200.millis) === 0L)
+  }
+
+  test("two tasks cannot grow past 1 / N of execution memory") {
+    val memoryManager = createMemoryManager(1000L)
+    val t1MemManager = new TaskMemoryManager(memoryManager, 1)
+    val t2MemManager = new TaskMemoryManager(memoryManager, 2)
+    val futureTimeout: Duration = 20.seconds
+
+    // Have both tasks request 250 bytes, then wait until both requests have been granted:
+    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(250L) }
+    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(250L) }
+    assert(Await.result(t1Result1, futureTimeout) === 250L)
+    assert(Await.result(t2Result1, futureTimeout) === 250L)
+
+    // Have both tasks each request 500 bytes more.
+    // We should only grant 250 bytes to each of them on this second request
+    val t1Result2 = Future { t1MemManager.acquireExecutionMemory(500L) }
+    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(500L) }
+    assert(Await.result(t1Result2, futureTimeout) === 250L)
+    assert(Await.result(t2Result2, futureTimeout) === 250L)
+  }
+
+  test("tasks can block to get at least 1 / 2N of execution memory") {
+    val memoryManager = createMemoryManager(1000L)
+    val t1MemManager = new TaskMemoryManager(memoryManager, 1)
+    val t2MemManager = new TaskMemoryManager(memoryManager, 2)
+    val futureTimeout: Duration = 20.seconds
+
+    // t1 grabs 1000 bytes and then waits until t2 is ready to make a request.
+    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(1000L) }
+    assert(Await.result(t1Result1, futureTimeout) === 1000L)
+    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(250L) }
+    // Make sure that t2 didn't grab the memory right away. This is hacky but it would be difficult
+    // to make sure the other thread blocks for some time otherwise.
+    Thread.sleep(300)
+    t1MemManager.releaseExecutionMemory(250L)
+    // The memory freed from t1 should now be granted to t2.
+    assert(Await.result(t2Result1, futureTimeout) === 250L)
+    // Further requests by t2 should be denied immediately because it now has 1 / 2N of the memory.
+    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(100L) }
+    assert(Await.result(t2Result2, 200.millis) === 0L)
+  }
+
+  test("TaskMemoryManager.cleanUpAllAllocatedMemory") {
+    val memoryManager = createMemoryManager(1000L)
+    val t1MemManager = new TaskMemoryManager(memoryManager, 1)
+    val t2MemManager = new TaskMemoryManager(memoryManager, 2)
+    val futureTimeout: Duration = 20.seconds
+
+    // t1 grabs 1000 bytes and then waits until t2 is ready to make a request.
+    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(1000L) }
+    assert(Await.result(t1Result1, futureTimeout) === 1000L)
+    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(500L) }
+    // Make sure that t2 didn't grab the memory right away. This is hacky but it would be difficult
+    // to make sure the other thread blocks for some time otherwise.
+    Thread.sleep(300)
+    // t1 releases all of its memory, so t2 should be able to grab all of the memory
+    t1MemManager.cleanUpAllAllocatedMemory()
+    assert(Await.result(t2Result1, futureTimeout) === 500L)
+    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(500L) }
+    assert(Await.result(t2Result2, futureTimeout) === 500L)
+    val t2Result3 = Future { t2MemManager.acquireExecutionMemory(500L) }
+    assert(Await.result(t2Result3, 200.millis) === 0L)
+  }
+
+  test("tasks should not be granted a negative amount of execution memory") {
+    // This is a regression test for SPARK-4715.
+    val memoryManager = createMemoryManager(1000L)
+    val t1MemManager = new TaskMemoryManager(memoryManager, 1)
+    val t2MemManager = new TaskMemoryManager(memoryManager, 2)
+    val futureTimeout: Duration = 20.seconds
+
+    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(700L) }
+    assert(Await.result(t1Result1, futureTimeout) === 700L)
+
+    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(300L) }
+    assert(Await.result(t2Result1, futureTimeout) === 300L)
+
+    val t1Result2 = Future { t1MemManager.acquireExecutionMemory(300L) }
+    assert(Await.result(t1Result2, 200.millis) === 0L)
+  }
 }
 
 private object MemoryManagerSuite {

core/src/test/scala/org/apache/spark/memory/StaticMemoryManagerSuite.scala

@@ -41,6 +41,14 @@ class StaticMemoryManagerSuite extends MemoryManagerSuite {
     (mm, ms)
   }
 
+  override protected def createMemoryManager(maxMemory: Long): MemoryManager = {
+    new StaticMemoryManager(
+      conf,
+      maxExecutionMemory = maxMemory,
+      maxStorageMemory = 0,
+      numCores = 1)
+  }
+
   test("basic execution memory") {
     val maxExecutionMem = 1000L
     val (mm, _) = makeThings(maxExecutionMem, Long.MaxValue)

core/src/test/scala/org/apache/spark/memory/UnifiedMemoryManagerSuite.scala

@@ -39,6 +39,10 @@ class UnifiedMemoryManagerSuite extends MemoryManagerSuite with PrivateMethodTes
     (mm, ms)
   }
 
+  override protected def createMemoryManager(maxMemory: Long): MemoryManager = {
+    new UnifiedMemoryManager(conf, maxMemory, numCores = 1)
+  }
+
   private def getStorageRegionSize(mm: UnifiedMemoryManager): Long = {
     mm invokePrivate PrivateMethod[Long]('storageRegionSize)()
   }