cloud-fan commented on code in PR #58763:
URL: https://github.com/apache/spark/pull/58763#discussion_r4048046585
##########
core/src/main/scala/org/apache/spark/memory/ExecutionMemoryPool.scala:
##########
@@ -131,7 +227,8 @@ private[memory] class ExecutionMemoryPool(
// How much we can grant this task; keep its share within 0 <= X <= 1 /
numActiveTasks
val maxToGrant = math.min(numBytes, math.max(0, maxMemoryPerTask -
curMem))
// Only give it as much memory as is free, which might be none if it
reached 1 / numTasks
- val toGrant = math.min(maxToGrant, memoryFree)
+ val available = if (computeMemoryFree == null) memoryFree else
computeMemoryFree()
Review Comment:
**Blocking (P1):** A failed optional reclaimer can make this wait permanent.
With a 1000-byte pool, 600 optional bytes still charged, and the only ordinary
task requesting 700, `computeMemoryFree` returns 400 while `minMemoryPerTask`
remains 500. This branch waits, but there is no ordinary holder that can
release memory or notify it; the reclamation pass has already finished.
**Recommended change:** Have the 1/2N wait test use the effective ordinary
execution capacity when a physical cap is active.
**Why this works:** Alongside free bytes, derive effective ordinary capacity
as current ordinary charges plus physically grantable execution capacity, and
use that value for the minimum-share floor. Keep the original pool-size floor
when no optional physical cap is active.
**Scope:** Unify the optional-enabled grant and liveness capacity domains
without changing ordinary allocator policy when the feature is disabled or no
optional bytes compete.
**Compatibility:** Ordinary 1/N ceilings, storage borrowing, and
disabled-feature allocation behavior remain unchanged.
**Risks:** An effective-capacity calculation that excludes current ordinary
charges would make the floor too small. The disabled/no-optional path must
retain existing Spark fairness semantics.
**Constraints:** Callbacks must remain outside the memory-manager monitor.
Nonfatal reclaimer failures must retain their outstanding charge.
**Success:** The sole ordinary task receives the physically available
partial grant after a failed optional drain instead of waiting forever.
Multiple ordinary tasks still wait only for a floor that can be reached by
ordinary releases. The no-optional allocator keeps its existing fair-share
behavior.
##########
core/src/main/scala/org/apache/spark/internal/config/package.scala:
##########
@@ -483,6 +483,17 @@ package object config {
.checkValue(_ >= 0, "The off-heap memory size must not be negative")
.createWithDefault(0)
+ private[spark] val MEMORY_OPTIONAL_ENABLED =
ConfigBuilder("spark.memory.optional.enabled")
+ .internal()
+ .doc("Enable optional execution-memory admission and reclamation. Set
before the memory " +
+ "manager is created; changing it later has no effect. Disabled managers
reject optional " +
+ "requests and skip admission coordination. Consumers opt in to
revocation of unused " +
+ "speculative work for ordinary allocations, including allocations from
other tasks.")
+ .version("5.0.0")
Review Comment:
**Non-blocking (P2):** This should use `4.4.0`. For a normal non-breaking PR
against master, Spark's branch policy uses the latest `branch-N.x` version
where the change first ships; the pinned `branch-4.x` version is
`4.4.0-SNAPSHOT`, not 5.0.0.
##########
core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala:
##########
@@ -193,20 +364,98 @@ private[spark] class UnifiedMemoryManager(
* when unmanaged components are consuming significant memory.
*/
def computeMaxExecutionPoolSize(): Long = {
- val unmanagedMemory = getUnmanagedMemoryUsed(memoryMode)
- val availableMemory = maxMemory - math.min(storagePool.memoryUsed,
storageRegionSize)
- // Reduce available memory by unmanaged memory usage to prevent
over-allocation
- math.max(0L, availableMemory - unmanagedMemory)
+ val unmanagedMemory =
unmanagedMemorySnapshot.getOrElse(getUnmanagedMemoryUsed(memoryMode))
+ maxExecutionMemory(storagePool, storageRegionSize, maxMemory,
unmanagedMemory)
}
+ // The gate prevents new optional grants, including while acquireMemory
waits. Avoid creating
+ // a capacity callback when no optional credit can compete with this
ordinary allocation.
+ val optionalCapacity = if (!optionalMemoryEnabled ||
+ executionPool.optionalMemoryUsed == 0L) {
+ null
+ } else {
+ () => {
+ val unmanaged =
unmanagedMemorySnapshot.getOrElse(getUnmanagedMemoryUsed(memoryMode))
+ availableExecutionMemory(memoryMode, unmanaged, borrowStorage = false)
+ }
+ }
executionPool.acquireMemory(
- numBytes, taskAttemptId, maybeGrowExecutionPool, () =>
computeMaxExecutionPoolSize())
+ numBytes, taskAttemptId, maybeGrowExecutionPool, () =>
computeMaxExecutionPoolSize(),
+ optionalCapacity)
}
override def acquireStorageMemory(
blockId: BlockId,
numBytes: Long,
- memoryMode: MemoryMode): Boolean = synchronized {
+ memoryMode: MemoryMode): Boolean = {
+ if (!optionalMemoryEnabled || isStorageMemoryRequestTooLarge(numBytes,
memoryMode)) {
+ return synchronized { acquireStorageMemoryInternal(blockId, numBytes,
memoryMode) }
+ }
+ withReclamation(memoryMode) {
+ val unmanaged = getUnmanagedMemoryUsed(memoryMode)
+ if (canAcquireStorageMemory(numBytes, memoryMode, unmanaged)) {
+ Some(acquireStorageMemoryInternal(blockId, numBytes, memoryMode,
Some(unmanaged)))
+ } else if (numBytes > maxStorageMemoryAfterReclamation(memoryMode,
unmanaged)) {
+ Some(acquireStorageMemoryInternal(blockId, numBytes, memoryMode,
Some(unmanaged)))
+ } else {
+ None
+ }
+ } {
+ acquireStorageMemoryInternal(blockId, numBytes, memoryMode)
+ }
+ }
+
+ private def maxStorageMemoryAfterReclamation(
+ memoryMode: MemoryMode,
+ unmanaged: Long): Long = {
+ val (executionPool, _, _, maxMemory) = poolsFor(memoryMode)
+ math.max(0L, maxMemory - executionPool.ordinaryMemoryUsed - unmanaged)
+ }
+
+ override private[spark] def withStorageMemoryReclamation[T](
+ numBytes: Long,
+ memoryMode: MemoryMode)(body: => T): T = {
+ require(numBytes >= 0L)
+ if (!optionalMemoryEnabled) return synchronized { body }
+ withReclamation(memoryMode) {
+ val unmanaged = getUnmanagedMemoryUsed(memoryMode)
+ if (canAcquireStorageMemory(numBytes, memoryMode, unmanaged) ||
+ numBytes > maxStorageMemoryAfterReclamation(memoryMode, unmanaged)) {
+ Some(body)
+ } else {
+ None
+ }
+ } {
+ body
+ }
+ }
+
+ /** Check storage capacity under the monitor without moving pool boundaries
or evicting blocks. */
+ private def canAcquireStorageMemory(
+ numBytes: Long,
+ memoryMode: MemoryMode,
+ unmanagedMemory: Long): Boolean = {
+ assert(Thread.holdsLock(this))
+ val (executionPool, storagePool, maxMemory) = memoryMode match {
+ case MemoryMode.ON_HEAP =>
+ (onHeapExecutionMemoryPool, onHeapStorageMemoryPool,
maxOnHeapStorageMemory)
+ case MemoryMode.OFF_HEAP =>
+ (offHeapExecutionMemoryPool, offHeapStorageMemoryPool,
maxOffHeapStorageMemory)
+ }
+ numBytes <= math.max(0L, maxMemory - unmanagedMemory) &&
Review Comment:
**Blocking (P1):** This preflight needs to enforce aggregate same-mode
occupancy, not just compare `numBytes` with the unmanaged-adjusted ceiling. For
example, with max memory 1000, existing storage 700, optional execution 100,
unmanaged usage 200, and a 100-byte request, both predicates pass; admission
leaves 800 storage + 100 optional + 200 unmanaged = 1100 bytes.
##########
core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala:
##########
@@ -122,6 +122,69 @@ private[spark] class UnifiedMemoryManager(
maxOffHeapMemory - offHeapExecutionMemoryPool.memoryUsed
}
+ override private[spark] def isStorageMemoryRequestTooLarge(
+ numBytes: Long,
+ memoryMode: MemoryMode): Boolean = memoryMode match {
+ case MemoryMode.ON_HEAP => numBytes > maxHeapMemory
+ case MemoryMode.OFF_HEAP => numBytes > maxOffHeapMemory
+ }
+
+ /** Pool selection and the fairness ceiling are shared by optional and
ordinary admission. */
+ private def poolsFor(memoryMode: MemoryMode):
+ (ExecutionMemoryPool, StorageMemoryPool, Long, Long) = memoryMode match {
+ case MemoryMode.ON_HEAP =>
+ (onHeapExecutionMemoryPool, onHeapStorageMemoryPool,
onHeapStorageRegionSize, maxHeapMemory)
+ case MemoryMode.OFF_HEAP =>
+ (offHeapExecutionMemoryPool, offHeapStorageMemoryPool,
offHeapStorageMemory, maxOffHeapMemory)
+ }
+
+ private def maxExecutionMemory(memoryMode: MemoryMode, unmanagedMemory:
Long): Long = {
+ val (_, storagePool, regionSize, maxMemory) = poolsFor(memoryMode)
+ maxExecutionMemory(storagePool, regionSize, maxMemory, unmanagedMemory)
+ }
+
+ private def maxExecutionMemory(
+ storagePool: StorageMemoryPool,
+ regionSize: Long,
+ maxMemory: Long,
+ unmanagedMemory: Long): Long = {
+ math.max(0L, maxMemory - math.min(storagePool.memoryUsed, regionSize) -
unmanagedMemory)
+ }
+
+ /**
+ * Reserve all requested bytes from existing free execution memory for
optional task work.
+ *
+ * Unlike ordinary admission, this does not borrow storage memory, evict
blocks, or wait for
+ * capacity. The fair-share ceiling is the same as ordinary admission,
including storage that
Review Comment:
**Non-blocking (P2):** The fair-share ceiling is not the same as ordinary
admission. `tryAcquireMemory` includes optional-only tasks in the divisor and
subtracts the requester's ordinary plus optional bytes; ordinary
`acquireMemory` considers only ordinary participants and ordinary bytes. Please
describe this as the intentionally more conservative optional ceiling instead
of claiming equality.
##########
core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java:
##########
@@ -187,6 +194,349 @@ public long spill(long size, MemoryConsumer trigger) {
}
}
+ /** One Comet-like consumer owns both ordinary and optional bytes through
the same task. */
+ private static final class OptionalMemoryConsumer extends MemoryConsumer {
+ private final AtomicLong optionalBytes = new AtomicLong();
+
+ OptionalMemoryConsumer(TaskMemoryManager manager) {
+ super(manager, 0L, MemoryMode.OFF_HEAP);
Review Comment:
**Nit (P3):** Because this fixture hard-codes `OFF_HEAP`, none of the
optional-memory tests exercises `TaskMemoryManager`'s new on-heap current/peak
accounting. Please parameterize the consumer by mode and assert current and
peak usage for mixed ordinary/optional acquisition and release in both modes.
##########
core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java:
##########
@@ -187,6 +194,349 @@ public long spill(long size, MemoryConsumer trigger) {
}
}
+ /** One Comet-like consumer owns both ordinary and optional bytes through
the same task. */
+ private static final class OptionalMemoryConsumer extends MemoryConsumer {
+ private final AtomicLong optionalBytes = new AtomicLong();
+
+ OptionalMemoryConsumer(TaskMemoryManager manager) {
+ super(manager, 0L, MemoryMode.OFF_HEAP);
+ }
+
+ synchronized long reserveOptional(long size) {
+ long got = taskMemoryManager.tryAcquireOptionalExecutionMemory(size,
this);
+ optionalBytes.addAndGet(got);
+ used.addAndGet(got);
+ return got;
+ }
+
+ synchronized long acquireOrdinary(long size) {
+ long got = taskMemoryManager.acquireExecutionMemory(size, this);
+ used.addAndGet(got);
+ return got;
+ }
+
+ synchronized void releaseOrdinary(long size) {
+ taskMemoryManager.releaseExecutionMemory(size, this);
+ used.addAndGet(-size);
+ }
+
+ synchronized void reclaimOptional() {
+ long size = optionalBytes.get();
+ if (size > 0) {
+ taskMemoryManager.releaseOptionalExecutionMemory(size, this);
+ optionalBytes.addAndGet(-size);
+ used.addAndGet(-size);
+ }
+ }
+
+ @Override
+ public long spill(long size, MemoryConsumer trigger) {
+ return 0L;
+ }
+ }
+
+ private static SparkConf optionalTestConf(boolean enabled) {
+ return new SparkConf(false)
+ .set("spark.memory.optional.enabled", Boolean.toString(enabled))
+ .set(package$.MODULE$.MEMORY_OFFHEAP_ENABLED(), true)
+ .set(package$.MODULE$.MEMORY_OFFHEAP_SIZE(), 1000L)
+ .set(package$.MODULE$.MEMORY_STORAGE_FRACTION(), 0.0);
+ }
+
+ private UnifiedMemoryManager optionalTestPool(boolean enabled) {
+ return new UnifiedMemoryManager(optionalTestConf(enabled), 1000L, 0L, 1);
+ }
+
+ @Test
+ public void optionalAdmissionCanReclaimAnotherTasksReadAhead() throws
Exception {
+ UnifiedMemoryManager pool = optionalTestPool(true);
+ TaskMemoryManager owner = new TaskMemoryManager(pool, 1);
+ TaskMemoryManager requester = new TaskMemoryManager(pool, 2);
+ OptionalMemoryConsumer readAhead = new OptionalMemoryConsumer(owner);
+ TestMemoryConsumer ordinary = new TestMemoryConsumer(requester,
MemoryMode.OFF_HEAP);
+
+ Assertions.assertEquals(0, readAhead.reserveOptional(1)); // Must register
first.
+ AutoCloseable registration =
+ owner.registerOptionalMemoryReclaimer(readAhead,
readAhead::reclaimOptional);
+ Assertions.assertNotNull(registration);
+ Assertions.assertEquals(400, readAhead.reserveOptional(400));
+ Assertions.assertEquals(400, owner.getMemoryConsumptionForThisTask());
+ String breakdown = owner.getMemoryConsumptionBreakdown();
+ Assertions.assertTrue(breakdown.contains(readAhead.toString() + ": "),
breakdown);
+ Assertions.assertFalse(breakdown.contains("(not attributed to a specific
consumer)"),
+ breakdown);
+
+ // The owner has no ordinary grant. Its optional bytes must not shrink the
peer's share.
+ ordinary.use(700);
+ Assertions.assertEquals(700, ordinary.getUsed());
+ Assertions.assertEquals(0, readAhead.getUsed());
+ Assertions.assertEquals(0, owner.getMemoryConsumptionForThisTask());
+ Assertions.assertEquals(700, pool.executionMemoryUsed());
+
+ registration.close();
+ registration.close();
+ ordinary.free(700);
+ Assertions.assertEquals(0, owner.cleanUpAllAllocatedMemory());
+ Assertions.assertEquals(0, requester.cleanUpAllAllocatedMemory());
+ }
+
+ @Test
+ public void optionalAndOrdinaryBytesOfOneConsumerHaveSeparateReleasePaths()
throws Exception {
+ UnifiedMemoryManager pool = optionalTestPool(true);
+ TaskMemoryManager owner = new TaskMemoryManager(pool, 1);
+ OptionalMemoryConsumer readAhead = new OptionalMemoryConsumer(owner);
+ AutoCloseable registration =
+ owner.registerOptionalMemoryReclaimer(readAhead,
readAhead::reclaimOptional);
+
+ Assertions.assertEquals(100, readAhead.reserveOptional(100));
+ Assertions.assertEquals(200, readAhead.acquireOrdinary(200));
+ Assertions.assertEquals(300, readAhead.getUsed());
+ readAhead.reclaimOptional();
+ Assertions.assertEquals(200, owner.getMemoryConsumptionForThisTask());
+ Assertions.assertEquals(200, readAhead.getUsed());
+ readAhead.releaseOrdinary(200);
+ registration.close();
+ Assertions.assertEquals(0, owner.cleanUpAllAllocatedMemory());
+ Assertions.assertEquals(300, owner.getPeakOffHeapExecutionMemory());
+ }
+
+ @Test
+ public void optionalReleaseCannotSpendAnotherConsumersCredit() throws
Exception {
+ UnifiedMemoryManager pool = optionalTestPool(true);
+ TaskMemoryManager owner = new TaskMemoryManager(pool, 1);
+ OptionalMemoryConsumer first = new OptionalMemoryConsumer(owner);
+ OptionalMemoryConsumer second = new OptionalMemoryConsumer(owner);
+ OptionalMemoryConsumer unregistered = new OptionalMemoryConsumer(owner);
+ AutoCloseable firstRegistration =
+ owner.registerOptionalMemoryReclaimer(first, first::reclaimOptional);
+ AutoCloseable secondRegistration =
+ owner.registerOptionalMemoryReclaimer(second, second::reclaimOptional);
+
+ Assertions.assertEquals(300, first.reserveOptional(300));
+ Assertions.assertEquals(300, second.reserveOptional(300));
+ // The task-level pool has 600 bytes; neither the 400-byte first release
nor an unrelated
+ // consumer's 100-byte release may spend credit owned by the second
consumer.
+ Assertions.assertThrows(AssertionError.class,
+ () -> owner.releaseOptionalExecutionMemory(400, first));
+ Assertions.assertThrows(AssertionError.class,
+ () -> owner.releaseOptionalExecutionMemory(100, unregistered));
+ Assertions.assertEquals(600, owner.getMemoryConsumptionForThisTask());
+ Assertions.assertEquals(300, first.optionalBytes.get());
+ Assertions.assertEquals(300, second.optionalBytes.get());
+
+ first.reclaimOptional();
+ Assertions.assertEquals(300, owner.getMemoryConsumptionForThisTask());
+ second.reclaimOptional();
+ firstRegistration.close();
+ secondRegistration.close();
+ Assertions.assertEquals(0, owner.cleanUpAllAllocatedMemory());
+ }
+
+ @Test
+ public void skipsDrainedOwnerWhenAnotherConsumerInTaskHoldsOptionalBytes()
throws Exception {
+ UnifiedMemoryManager pool = optionalTestPool(true);
+ TaskMemoryManager owner = new TaskMemoryManager(pool, 1);
+ TaskMemoryManager peer = new TaskMemoryManager(pool, 2);
+ OptionalMemoryConsumer idle = new OptionalMemoryConsumer(owner);
+ OptionalMemoryConsumer readAhead = new OptionalMemoryConsumer(owner);
+ TestMemoryConsumer ordinary = new TestMemoryConsumer(peer,
MemoryMode.OFF_HEAP);
+ AtomicInteger idleCalls = new AtomicInteger();
+ AutoCloseable idleRegistration =
+ owner.registerOptionalMemoryReclaimer(idle, () ->
idleCalls.incrementAndGet());
+ AutoCloseable activeRegistration =
+ owner.registerOptionalMemoryReclaimer(readAhead,
readAhead::reclaimOptional);
+
+ Assertions.assertEquals(400, readAhead.reserveOptional(400));
+ ordinary.use(700);
+ Assertions.assertEquals(700, ordinary.getUsed());
+ Assertions.assertEquals(0, idleCalls.get());
+ Assertions.assertEquals(0, readAhead.optionalBytes.get());
+
+ idleRegistration.close();
+ activeRegistration.close();
+ ordinary.free(700);
+ Assertions.assertEquals(0, owner.cleanUpAllAllocatedMemory());
+ Assertions.assertEquals(0, peer.cleanUpAllAllocatedMemory());
+ }
+
+ @Test
+ public void reclamationSeesGrantBeforeItsOwnerCreditIsPublished() throws
Exception {
+ CountDownLatch grantReady = new CountDownLatch(1);
+ CountDownLatch allowPublication = new CountDownLatch(1);
+ CountDownLatch callbackStarted = new CountDownLatch(1);
+ UnifiedMemoryManager pool = new
UnifiedMemoryManager(optionalTestConf(true), 1000L, 0L, 1) {
+ @Override
+ public long tryAcquireExecutionMemory(
+ long size, long taskAttemptId, MemoryMode mode) {
+ long granted = super.tryAcquireExecutionMemory(size, taskAttemptId,
mode);
+ if (taskAttemptId == 1L && granted > 0L) {
+ grantReady.countDown();
+ try {
+ if (!allowPublication.await(10, TimeUnit.SECONDS)) {
+ // Invariants must propagate even when JVM assertions are
disabled.
+ // checkstyle.off: RegexpSinglelineJava
+ throw new AssertionError("optional grant was not allowed to
publish its credit");
+ // checkstyle.on: RegexpSinglelineJava
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ // Invariants must propagate even when JVM assertions are disabled.
+ // checkstyle.off: RegexpSinglelineJava
+ throw new AssertionError(e);
+ // checkstyle.on: RegexpSinglelineJava
+ }
+ }
+ return granted;
+ }
+ };
+ TaskMemoryManager owner = new TaskMemoryManager(pool, 1);
+ TaskMemoryManager peer = new TaskMemoryManager(pool, 2);
+ OptionalMemoryConsumer readAhead = new OptionalMemoryConsumer(owner);
+ TestMemoryConsumer ordinary = new TestMemoryConsumer(peer,
MemoryMode.OFF_HEAP);
+ AutoCloseable registration =
owner.registerOptionalMemoryReclaimer(readAhead, () -> {
+ callbackStarted.countDown();
+ readAhead.reclaimOptional();
+ });
+ ExecutorService executor = Executors.newFixedThreadPool(2);
+ Future<?> peerAcquisition = null;
+ try {
+ Future<Long> optionalAcquisition = executor.submit(() ->
readAhead.reserveOptional(400));
+ Assertions.assertTrue(grantReady.await(10, TimeUnit.SECONDS));
+ Assertions.assertEquals(400, pool.executionMemoryUsed());
+ // The shared grant is visible, but the owner has not yet published its
local credit.
+ peerAcquisition = executor.submit(() -> ordinary.use(700));
+ Assertions.assertTrue(callbackStarted.await(10, TimeUnit.SECONDS));
+ Assertions.assertFalse(peerAcquisition.isDone());
+ allowPublication.countDown();
+ Assertions.assertEquals(400L, optionalAcquisition.get(10,
TimeUnit.SECONDS).longValue());
+ peerAcquisition.get(10, TimeUnit.SECONDS);
+ Assertions.assertEquals(700, ordinary.getUsed());
+ Assertions.assertEquals(0, readAhead.optionalBytes.get());
+ registration.close();
+ ordinary.free(700);
+ Assertions.assertEquals(0, owner.cleanUpAllAllocatedMemory());
+ Assertions.assertEquals(0, peer.cleanUpAllAllocatedMemory());
+ } finally {
+ allowPublication.countDown();
+ if (peerAcquisition != null) peerAcquisition.cancel(true);
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ public void incompleteOptionalDrainKeepsRegistrationAndCredit() throws
Exception {
+ UnifiedMemoryManager pool = optionalTestPool(true);
+ TaskMemoryManager owner = new TaskMemoryManager(pool, 1);
+ OptionalMemoryConsumer readAhead = new OptionalMemoryConsumer(owner);
+ AutoCloseable registration =
owner.registerOptionalMemoryReclaimer(readAhead, () -> { });
+
+ Assertions.assertEquals(400, readAhead.reserveOptional(400));
+ Assertions.assertThrows(AssertionError.class, registration::close);
+ Assertions.assertEquals(400, owner.getMemoryConsumptionForThisTask());
+ Assertions.assertEquals(0, readAhead.reserveOptional(1));
+ // Task cleanup must fail visibly rather than silently drop the unread
native reservation.
+ Assertions.assertThrows(AssertionError.class,
owner::cleanUpAllAllocatedMemory);
+ Assertions.assertEquals(400, owner.getMemoryConsumptionForThisTask());
+ Assertions.assertEquals(400, readAhead.optionalBytes.get());
+ readAhead.reclaimOptional();
+ registration.close();
+ Assertions.assertEquals(0, owner.cleanUpAllAllocatedMemory());
+ }
+
+ @Test
+ public void cleanupClosesOptionalAdmissionBeforeDrainingOwner() throws
Exception {
+ UnifiedMemoryManager pool = optionalTestPool(true);
+ TaskMemoryManager owner = new TaskMemoryManager(pool, 1);
+ OptionalMemoryConsumer readAhead = new OptionalMemoryConsumer(owner);
+ CountDownLatch callbackStarted = new CountDownLatch(1);
+ CountDownLatch allowRelease = new CountDownLatch(1);
+ AutoCloseable registration =
owner.registerOptionalMemoryReclaimer(readAhead, () -> {
+ Assertions.assertFalse(Thread.holdsLock(owner));
+ Assertions.assertFalse(Thread.holdsLock(pool));
+ callbackStarted.countDown();
+ try {
+ if (!allowRelease.await(5, TimeUnit.SECONDS)) {
+ // Invariants must propagate even when JVM assertions are disabled.
+ // checkstyle.off: RegexpSinglelineJava
+ throw new AssertionError("timed out before optional cleanup could
release memory");
+ // checkstyle.on: RegexpSinglelineJava
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ // Invariants must propagate even when JVM assertions are disabled.
+ // checkstyle.off: RegexpSinglelineJava
+ throw new AssertionError(e);
+ // checkstyle.on: RegexpSinglelineJava
+ }
+ readAhead.reclaimOptional();
+ });
+ Assertions.assertEquals(400, readAhead.reserveOptional(400));
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ Future<Long> cleanup = executor.submit(owner::cleanUpAllAllocatedMemory);
+ Assertions.assertTrue(callbackStarted.await(5, TimeUnit.SECONDS));
+ Assertions.assertEquals(400, pool.executionMemoryUsed());
+ // This call must fail immediately rather than wait behind the pending
callback.
+ Assertions.assertEquals(0, owner.tryAcquireOptionalExecutionMemory(1,
readAhead));
+ allowRelease.countDown();
+ Assertions.assertEquals(0L, cleanup.get(5,
TimeUnit.SECONDS).longValue());
+ Assertions.assertEquals(0, pool.executionMemoryUsed());
+ Assertions.assertNull(owner.registerOptionalMemoryReclaimer(readAhead,
() -> { }));
+ Assertions.assertEquals(0, readAhead.reserveOptional(1));
+ registration.close();
+ } finally {
+ allowRelease.countDown();
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ public void failedOwnerCleanupRetainsItsOptionalCharge() throws Exception {
+ UnifiedMemoryManager pool = optionalTestPool(true);
Review Comment:
**Nit (P3):** This failure test uses only one optional owner, so it cannot
detect regressions in the loop's two multi-owner guarantees. Please add a
failing owner followed by a healthy owner and verify the later reservation is
drained, plus a case with two nonfatal failures that checks the second is
suppressed onto the first.
##########
core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala:
##########
@@ -134,21 +197,129 @@ private[spark] class UnifiedMemoryManager(
override private[memory] def acquireExecutionMemory(
numBytes: Long,
taskAttemptId: Long,
- memoryMode: MemoryMode): Long = synchronized {
+ memoryMode: MemoryMode): Long = {
+ if (!optionalMemoryEnabled) {
+ return synchronized {
+ acquireExecutionMemoryInternal(numBytes, taskAttemptId, memoryMode)
+ }
+ }
+ require(!Thread.holdsLock(this),
+ "ordinary execution admission cannot start under the memory-manager
monitor")
+ withReclamation(memoryMode) {
+ tryAcquireExecutionMemoryWithoutReclamation(numBytes, taskAttemptId,
memoryMode)
+ } {
+ acquireExecutionMemoryInternal(numBytes, taskAttemptId, memoryMode)
+ }
+ }
+
+ /**
+ * Share the gate, atomic preflight, outside-monitor drain and retry for
both allocation paths.
+ * Nested storage calls reuse an outer boundary. Release-only markers never
trigger a drain.
+ */
+ private def withReclamation[T](memoryMode: MemoryMode)(
+ tryWithoutReclamation: => Option[T])(allocate: => T): T = {
+ if (!optionalMemoryEnabled) return synchronized { allocate }
+ val gate = acquireOptionalAdmissionReadLock()
+ try {
+ if (Thread.holdsLock(this)) return allocate
+ synchronized {
+ if (!hasOptionalExecutionMemory(memoryMode)) return allocate
+ val immediate = tryWithoutReclamation
+ if (immediate.isDefined) return immediate.get
+ }
+ var result: Option[T] = None
+ reclaimOptionalMemory(memoryMode, () => synchronized {
+ result = tryWithoutReclamation
+ result.isDefined
+ })
+ result.getOrElse(synchronized { allocate })
+ } finally {
+ gate.unlock()
+ }
+ }
+
+ /** Grant immediately when reclamation cannot improve the result, preserving
optional buffers. */
+ private def tryAcquireExecutionMemoryWithoutReclamation(
+ numBytes: Long,
+ taskAttemptId: Long,
+ memoryMode: MemoryMode): Option[Long] = {
+ assert(Thread.holdsLock(this))
+ val unmanaged = getUnmanagedMemoryUsed(memoryMode)
+ if (canAcquireExecutionMemory(numBytes, taskAttemptId, memoryMode,
unmanaged)) {
+ return Some(acquireExecutionMemoryInternal(
+ numBytes, taskAttemptId, memoryMode, Some(unmanaged)))
+ }
+ val (executionPool, storagePool, regionSize, maxMemory) =
poolsFor(memoryMode)
+ val storageProtected = storagePool.memoryUsed <= regionSize
+ // Whole-block eviction can raise the share ceiling. Only bypass it when
releasing optional
+ // bytes would make the original request fit without eviction, or storage
is protected.
+ val needsEviction = !storageProtected &&
+ numBytes > maxMemory - storagePool.memoryUsed -
executionPool.ordinaryMemoryUsed
+ val maxExecution = if (needsEviction) math.max(0L, maxMemory - unmanaged)
+ else maxExecutionMemory(memoryMode, unmanaged)
+ val headroom = executionPool.absoluteMemoryHeadroom(taskAttemptId,
maxExecution)
+ val available = availableExecutionMemory(memoryMode, unmanaged,
borrowStorage = true)
+ val poolAfterGrowth = executionPool.poolSize +
+ math.min(math.max(0L, numBytes - executionPool.memoryFree),
storagePool.memoryFree)
+ if (headroom >= 0L && headroom < numBytes && headroom <= available &&
!needsEviction &&
+ executionPool.canGrantWithoutWaiting(
+ taskAttemptId, numBytes, headroom, poolAfterGrowth)) {
+ // Keep normal protected-storage growth. With borrowed storage, avoid
evicting a block
+ // for bytes that cannot be granted even after every optional owner has
released.
+ val request = if (storageProtected) numBytes else headroom
+ if (request == 0L) Some(0L) else Some(acquireExecutionMemoryInternal(
+ request, taskAttemptId, memoryMode, Some(unmanaged)))
+ } else {
+ None
+ }
+ }
+
+ /**
+ * Test full ordinary admission against current free capacity and the
prospective task share.
+ * Called under this monitor before any task registration, pool growth,
eviction, or wait.
+ * Free storage can be borrowed without a drain; occupied storage is not
optimistically evicted.
+ */
+ private def canAcquireExecutionMemory(
+ numBytes: Long,
+ taskAttemptId: Long,
+ memoryMode: MemoryMode,
+ unmanagedMemory: Long): Boolean = {
+ assert(Thread.holdsLock(this))
+ val executionPool = executionPoolFor(memoryMode)
+ executionPool.canAcquireMemory(numBytes, taskAttemptId,
+ maxExecutionMemory(memoryMode, unmanagedMemory),
+ availableExecutionMemory(memoryMode, unmanagedMemory, borrowStorage =
true))
+ }
+
+ /**
+ * Optional bytes count against physical capacity after unmanaged usage,
independently of the
+ * ordinary share ceiling. Without optional credit, preserve the existing
allocator's policy.
+ */
+ private def availableExecutionMemory(
+ memoryMode: MemoryMode,
+ unmanaged: Long,
+ borrowStorage: Boolean): Long = {
+ val (executionPool, storagePool, _, maxMemory) = poolsFor(memoryMode)
+ val free = executionPool.memoryFree + (if (borrowStorage)
storagePool.memoryFree else 0L)
Review Comment:
**Non-blocking (P2):** With optional bytes present, this hot path traverses
all active-task accounting several times under the shared monitor: `memoryFree`
calls `memoryUsed`, this method calls `memoryUsed` again, the share-bound path
can call `ordinaryMemoryUsed`, and the capacity callback repeats the
calculation on wake-up. Please maintain the ordinary-byte and
optional-only-participant aggregates alongside the maps so these checks stay
O(1).
##########
core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java:
##########
@@ -118,6 +124,86 @@ public class TaskMemoryManager {
@GuardedBy("this")
private final HashSet<MemoryConsumer> consumers;
+ private final ReentrantLock optionalAdmissionLock = new ReentrantLock();
+
+ @GuardedBy("optionalAdmissionLock")
+ private boolean optionalAdmissionClosed;
+
+ @GuardedBy("optionalAdmissionLock")
+ private final Map<MemoryConsumer, OptionalMemoryReclaimerRegistration>
optionalReclaimers =
+ new IdentityHashMap<>();
+
+ // Diagnostics and typed releases cannot acquire the admission lock. A
closed registration is
+ // removed only after all of its owner's optional bytes have been released.
+ private final CopyOnWriteArrayList<OptionalMemoryReclaimerRegistration>
optionalRegistrations =
+ new CopyOnWriteArrayList<>();
+
+ /**
+ * A task-owned registration. Closing first drains its owner, then removes
the shared callback.
+ * A callback already captured by another allocator can run after close and
must be idempotent.
+ */
+ private final class OptionalMemoryReclaimerRegistration implements
AutoCloseable {
+ private final MemoryConsumer consumer;
+ private final Runnable reclaimer;
+ private final Runnable unregister;
+ private final AtomicLong optionalBytes = new AtomicLong();
+ private final AtomicBoolean admissionInProgress = new AtomicBoolean();
+ private final AtomicReference<Runnable> unregisterAfterReclaim = new
AtomicReference<>();
+
+ @GuardedBy("optionalAdmissionLock")
+ private boolean admissionClosed;
+
+ @GuardedBy("optionalAdmissionLock")
+ private boolean closed;
+
+ private OptionalMemoryReclaimerRegistration(MemoryConsumer consumer,
Runnable reclaimer) {
+ this.consumer = consumer;
+ this.reclaimer = reclaimer;
+ // The shared callback cannot capture this registration or take the
admission lock.
+ AtomicReference<Runnable> afterReclaim = unregisterAfterReclaim;
+ AtomicBoolean admitting = admissionInProgress;
+ AtomicLong credits = optionalBytes;
+ this.unregister = memoryManager.registerOptionalMemoryReclaimer(
+ taskAttemptId, consumer.getMode(), () -> {
+ // Read the flag first: a grant may precede credit publication, and
clearing the flag
+ // after publication makes the subsequent credit read see the
completed grant.
+ if (admitting.get() || credits.get() > 0L) {
+ reclaimer.run();
+ }
+ Runnable retire = afterReclaim.get();
+ if (retire != null && !admitting.get() && credits.get() == 0L)
retire.run();
+ });
+ }
+
+ @Override
+ public void close() {
+ optionalAdmissionLock.lock();
+ try {
+ if (!closed) {
+ admissionClosed = true;
+ // A failed or incomplete drain remains registered and charged for a
later retry.
+ reclaimer.run();
Review Comment:
**Non-blocking (P2):** If the callback releases its final optional
reservation and then throws a nonfatal error, this exits before
`unregister.run()` and before removing the local registration. Cleanup
propagates the error, while future shared reclamation skips this task because
its optional usage is now zero, so the deferred unregister can never run. The
zero-credit exceptional path should retire the registration while still
rethrowing the callback failure; positive-credit failures must remain
registered.
**Recommended change:** Retire a registration when its callback exits
exceptionally after reducing its exact optional credit to zero, while
preserving and rethrowing the callback failure.
**Why this works:** Capture callback failure, then under the registration's
existing synchronization check published admission and exact owner credit. If
fully drained, run unregister and remove the local registration exactly once
before rethrowing; if bytes remain, keep the current charged/retryable state.
**Scope:** Close the zero-credit exceptional lifecycle gap without
unregistering live buffers or swallowing cleanup failures.
**Compatibility:** Positive-credit failures remain visible, charged, and
available for a later retry.
**Risks:** Retiring before the final release is published could lose live
backing-memory ownership. The callback exception must remain the task-visible
failure.
**Constraints:** Incomplete or failed drains with positive credit stay
registered and charged. Retirement and local collection removal remain
idempotent under explicit-close and cleanup races.
**Success:** A release-then-throw callback leaves no global callback or
local consumer registration once its credit is zero. A callback that throws
with positive credit remains registered and retryable. The original callback
failure still propagates.
##########
core/src/main/scala/org/apache/spark/storage/memory/MemoryStore.scala:
##########
@@ -414,20 +440,52 @@ private[spark] class MemoryStore(
def freeMemoryEntry[T <: MemoryEntry[_]](entry: T): Unit = {
entry match {
case SerializedMemoryEntry(buffer, _, _) => buffer.dispose()
- case e: DeserializedMemoryEntry[_] => e.value.foreach {
- case o: AutoCloseable =>
- try {
- o.close()
- } catch {
- case NonFatal(e) =>
- logWarning("Fail to close a memory entry", e)
- }
- case _ =>
+ case e: DeserializedMemoryEntry[_] => freeValues(e.value.iterator)
+ }
+ }
+
+ private def disposeUnstoredEntry(entry: MemoryEntry[_]): Unit = entry match {
+ case SerializedMemoryEntry(buffer, _, _) => buffer.dispose()
+ // Deserialized values still belong to the caller until published in the
store.
+ case _: DeserializedMemoryEntry[_] => ()
+ }
+
+ private def freeValues(values: Iterator[_]): Unit = values.foreach {
+ case o: AutoCloseable =>
+ try {
+ o.close()
+ } catch {
+ case NonFatal(e) =>
+ logWarning("Fail to close a memory entry", e)
}
+ case _ =>
+ }
+
+ /** Dispose Spark's temporary serialized buffers after a failed unroll. */
+ private def discardUnstoredValues(valuesHolder: ValuesHolder[_]): Unit =
valuesHolder match {
+ // Deserialized values remain caller-owned when a put fails.
+ case _: DeserializedValuesHolder[_] => ()
+ case holder: SerializedValuesHolder[_] =>
Review Comment:
**Nit (P3):** The added exceptional-put regression uses
`putIteratorAsValues`, so it selects the deserialized branch and never
exercises this serialized cleanup. Please add a serialized put whose unroll
reservation throws and verify that MemoryStore closes the stream without
further flushing and disposes the temporary chunk buffers; the isolated
`dispose()` test cannot catch a missing handoff here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
