Github user andrewor14 commented on a diff in the pull request:
https://github.com/apache/spark/pull/10170#discussion_r47012724
--- Diff:
core/src/test/scala/org/apache/spark/memory/MemoryManagerSuite.scala ---
@@ -36,105 +37,98 @@ import org.apache.spark.storage.{BlockId, BlockStatus,
MemoryStore, StorageLevel
/**
* Helper trait for sharing code among [[MemoryManager]] tests.
*/
-private[memory] trait MemoryManagerSuite extends SparkFunSuite {
+private[memory] trait MemoryManagerSuite extends SparkFunSuite with
BeforeAndAfterEach {
- import MemoryManagerSuite.DEFAULT_ENSURE_FREE_SPACE_CALLED
+ protected val evictedBlocks = new mutable.ArrayBuffer[(BlockId,
BlockStatus)]
+
+ import MemoryManagerSuite.DEFAULT_EVICT_BLOCKS_TO_FREE_SPACE_CALLED
// Note: Mockito's verify mechanism does not provide a way to reset
method call counts
// without also resetting stubbed methods. Since our test code relies on
the latter,
- // we need to use our own variable to track invocations of
`ensureFreeSpace`.
+ // we need to use our own variable to track invocations of
`evictBlocksToFreeSpace`.
/**
- * The amount of free space requested in the last call to
[[MemoryStore.ensureFreeSpace]]
+ * The amount of space requested in the last call to
[[MemoryStore.evictBlocksToFreeSpace]].
*
- * This set whenever [[MemoryStore.ensureFreeSpace]] is called, and
cleared when the test
- * code makes explicit assertions on this variable through
[[assertEnsureFreeSpaceCalled]].
+ * This set whenever [[MemoryStore.evictBlocksToFreeSpace]] is called,
and cleared when the test
+ * code makes explicit assertions on this variable through
+ * [[assertEvictBlocksToFreeSpaceCalled]].
*/
- private val ensureFreeSpaceCalled = new
AtomicLong(DEFAULT_ENSURE_FREE_SPACE_CALLED)
+ private val evictBlocksToFreeSpaceCalled = new AtomicLong(0)
+
+ override def beforeEach(): Unit = {
+ super.beforeEach()
+ evictedBlocks.clear()
+
evictBlocksToFreeSpaceCalled.set(DEFAULT_EVICT_BLOCKS_TO_FREE_SPACE_CALLED)
+ }
/**
- * Make a mocked [[MemoryStore]] whose [[MemoryStore.ensureFreeSpace]]
method is stubbed.
+ * Make a mocked [[MemoryStore]] whose
[[MemoryStore.evictBlocksToFreeSpace]] method is stubbed.
*
- * This allows our test code to release storage memory when
[[MemoryStore.ensureFreeSpace]]
- * is called without relying on
[[org.apache.spark.storage.BlockManager]] and all of its
- * dependencies.
+ * This allows our test code to release storage memory when these
methods are called
+ * without relying on [[org.apache.spark.storage.BlockManager]] and all
of its dependencies.
*/
protected def makeMemoryStore(mm: MemoryManager): MemoryStore = {
- val ms = mock(classOf[MemoryStore])
- when(ms.ensureFreeSpace(anyLong(),
any())).thenAnswer(ensureFreeSpaceAnswer(mm, 0))
- when(ms.ensureFreeSpace(any(), anyLong(),
any())).thenAnswer(ensureFreeSpaceAnswer(mm, 1))
+ val ms = mock(classOf[MemoryStore], RETURNS_SMART_NULLS)
+ when(ms.evictBlocksToFreeSpace(any(), anyLong(), any()))
+ .thenAnswer(evictBlocksToFreeSpaceAnswer(mm))
mm.setMemoryStore(ms)
ms
}
/**
- * Make an [[Answer]] that stubs [[MemoryStore.ensureFreeSpace]] with
the right arguments.
- */
- private def ensureFreeSpaceAnswer(mm: MemoryManager, numBytesPos: Int):
Answer[Boolean] = {
+ * Simulate the part of [[MemoryStore.evictBlocksToFreeSpace]] that
releases storage memory.
+ *
+ * This is a significant simplification of the real method, which
actually drops existing
+ * blocks based on the size of each block. Instead, here we simply
release as many bytes
+ * as needed to ensure the requested amount of free space. This allows
us to set up the
+ * test without relying on the
[[org.apache.spark.storage.BlockManager]], which brings in
+ * many other dependencies.
+ *
+ * Every call to this method will set a global variable,
[[evictBlocksToFreeSpaceCalled]], that
+ * records the number of bytes this is called with. This variable is
expected to be cleared
+ * by the test code later through
[[assertEvictBlocksToFreeSpaceCalled]].
+ */
+ private def evictBlocksToFreeSpaceAnswer(mm: MemoryManager):
Answer[Boolean] = {
new Answer[Boolean] {
override def answer(invocation: InvocationOnMock): Boolean = {
val args = invocation.getArguments
- require(args.size > numBytesPos, s"bad test: expected
>$numBytesPos arguments " +
- s"in ensureFreeSpace, found ${args.size}")
- require(args(numBytesPos).isInstanceOf[Long], s"bad test: expected
ensureFreeSpace " +
- s"argument at index $numBytesPos to be a Long:
${args.mkString(", ")}")
- val numBytes = args(numBytesPos).asInstanceOf[Long]
- val success = mockEnsureFreeSpace(mm, numBytes)
- if (success) {
+ val numBytesToFree = args(1).asInstanceOf[Long]
+ assert(numBytesToFree > 0)
+ require(evictBlocksToFreeSpaceCalled.get() ===
DEFAULT_EVICT_BLOCKS_TO_FREE_SPACE_CALLED,
+ "bad test: evictBlocksToFreeSpace() variable was not reset")
+ evictBlocksToFreeSpaceCalled.set(numBytesToFree)
+ if (numBytesToFree <= mm.storageMemoryUsed) {
+ // We can evict enough blocks to fulfill the request for space
+ mm.releaseStorageMemory(numBytesToFree)
args.last.asInstanceOf[mutable.Buffer[(BlockId,
BlockStatus)]].append(
- (null, BlockStatus(StorageLevel.MEMORY_ONLY, numBytes, 0L,
0L)))
+ (null, BlockStatus(StorageLevel.MEMORY_ONLY, numBytesToFree,
0L, 0L)))
+ evictedBlocks.append(
+ (null, BlockStatus(StorageLevel.MEMORY_ONLY, numBytesToFree,
0L, 0L)))
--- End diff --
as mentioned offline, this will likely add the same block to
`evictedBlocks` twice
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastruct...@apache.org or file a JIRA ticket
with INFRA.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
