[SPARK-59461][CORE] Reclaim optional execution memory before active query allocations - #58763

This changes behavior even with no optional reclaimer registered. freeUnrolledValues (and freeMemoryEntry(entry) in the transfer catch below) close every already-unrolled AutoCloseable value, but MemoryStore only takes ownership of values on a successful put (PartiallyUnrolledIterator.close never closes values; freeMemoryEntry is only called on entries already in entries).

This throw path is reachable today: reserveUnrollMemoryForThisTask -> acquireStorageMemory -> evictBlocksToFreeSpace -> dropFromMemory -> diskStore.put throwing an IOException. For example, TorrentBroadcast.writeBlocks does putSingle(broadcastId, value, MEMORY_AND_DISK) with the user's live object, so sc.broadcast(autoCloseable) failing on a disk error now hands the caller an exception and a closed object. The catch is also on Throwable, so an InterruptedException closes values as well.

Could we limit the close to the reclaimer-failure case (or drop it, matching the existing ownership semantics) rather than closing on every throw?

A NonFatal failure from another task's reclaimer propagates out of acquireExecutionMemory / acquireStorageMemory / reserveUnrollMemoryForThisTask of the requesting task, and none of those paths fall back to ordinary admission (spill / evict / partial grant). So a bug in one optional owner fails unrelated tasks that would have succeeded before this PR; in putIterator it additionally closes the unrolled values.

I see this is intentional and tested, but a failed reclaimer never "invents" credit: its bytes stay charged, so ordinary admission just sees less free memory. Could this be treated like a failed spill() of another consumer, i.e. log and continue into ordinary admission, instead of failing the requester?

Unlike acquireExecutionMemory / acquireStorageMemory, which preflight and only reclaim on a miss, this drains every optional owner of both memory modes at any outermost MemoryStore entry with no capacity check:

• reserveUnrollMemoryForThisTask drains on the initial reservation and on every growth step of every cache put, even when storage is nearly empty (the nested acquireStorageMemory takes the enteredWithMonitor short-circuit and skips canAcquireStorageMemory).
• The unroll -> storage transfer drains although it can never need new memory (it releases >= entry.size and re-acquires entry.size under one monitor hold).
• remove / clear drain although they free memory, so BlockManager.removeRdd / removeBroadcast (ContextCleaner traffic) drain all owners once per block.
• An ON_HEAP put drains OFF_HEAP owners.

With any registered owner, optional memory becomes unusable on an executor that ever caches or uncaches a block. Could the unroll path use the same preflight-under-monitor-then-drain pattern as acquireStorageMemory, pass Some(memoryMode), and have release-only paths take the marker without draining? The read gate alone already makes tryAcquireExecutionMemory decline while they run.

This preflight requires the full request to fit within the caller's fair share, so a request that merely exceeds the share (with plenty of free memory) fails the preflight and drains every optional owner of that mode. But draining cannot raise maxPoolSize / numActiveTasks: computeMaxExecutionPoolSize depends only on storage usage and unmanaged memory, and the task count only shrinks when an owner's entry drops to zero. Meanwhile ordinary acquireMemory would have granted the partial maxToGrant immediately without waiting.

Example: pool 1000, tasks A and B; A holds 400, B holds 100 ordinary + 100 optional. A requests 200: share cap = 500 - 400 = 100 < 200, preflight fails, B's optional 100 is reclaimed, retry grants A 100 -- exactly what it would have gotten without the drain. A task already at its cap drains all owners on every request and still gets 0. Since TaskMemoryManager routinely issues page-sized requests and accepts got < required, this looks like the common case.

Could the preflight distinguish "capacity short" (drain helps) from "share-bound" (drain cannot help)?

The read gate is held across the whole call, including ExecutionMemoryPool.acquireMemory's lock.wait(), which releases the monitor but not the RRWL read hold. Since tryAcquireExecutionMemory uses writeLock().tryLock(), optional admission is denied executor-wide for as long as any task is parked in the 1/2N fairness wait (potentially minutes), even when hundreds of MB of execution memory are free. The no-owner fast path below also holds the read lock during ordinary grants, so a steady allocation stream causes tryLock misses too.

The scaladoc mentions "eviction or a capacity wait", but the PR description doesn't call out the capacity-wait starvation and there's no test with a task in lock.wait() alongside an optional request. Is this intended? If so, could it be documented and tested; if not, could the marker be released around lock.wait()?

require(!Thread.holdsLock(this), ...) in reclaimOptionalMemory throws IllegalArgumentException, which is NonFatal, so on the releaseOnly path (MemoryStore.remove / clear) a marker-before-monitor violation is caught here, logged as a transient "Failed to reclaim optional memory" warning, and cleanup proceeds without draining. No current caller violates it (all monitor-holding callers reach remove via evictBlocksToFreeSpace with the gate already held), but StorageMemoryPool.acquireMemory / freeSpaceToShrinkPool rely purely on convention, and a future memoryManager.synchronized { memoryStore.remove(...) } would either hit this swallowed assertion or, if a writer already won tryLock and is blocked on synchronized, deadlock on readLock().lock().

Could the lock-order check be moved outside the releaseOnly catch (or throw an AssertionError) so it always propagates, keeping the catch for callback failures only?

Every preflight miss snapshots the registry and runs every callback (each taking its owner lock), and because optional bytes are released through the undifferentiated releaseExecutionMemory path there's no "outstanding optional bytes" state, so already-drained owners are re-invoked as no-ops on every later miss. Under memory pressure a miss is the normal case: TaskMemoryManager.trySpillAndAcquire calls back into acquireExecutionMemory after each consumer spill, so each spill iteration pays monitor + registry lock + O(R) list + R callbacks + monitor with nothing left to reclaim. Between iterations the gate is released, so another task's owner can re-admit via tryLock, start prefetch I/O, and be cancelled by the next drain.

A per-mode "optional bytes admitted since last drain" flag (set in tryAcquireMemory, cleared after a completed drain) would let the preflight/drain be skipped when nothing is reclaimable.

The read gate is taken before the hasOptionalMemoryReclaimers check (same in acquireStorageMemory and withMemoryReclamation), so every TaskMemoryManager page allocation and every storage/unroll/remove/evict op now pays a ReentrantReadWriteLock shared acquire/release (CAS on one global state word plus hold-count bookkeeping) on top of the existing monitor -- in the no-reclaimer case, which is 100% of production since this PR ships no consumer. All task threads contend on that one word.

I realize checking the counter first isn't trivially safe (a waiter that skipped the gate could be starved by a later-registered owner). Do we have a no-owner allocator overhead measurement for the hot path? The description says the overhead "is not claimed to be free" but doesn't quantify it.

acquireExecutionMemory and acquireStorageMemory repeat the same ~30-line skeleton (read lock, no-owner fast path, Thread.holdsLock(this), synchronized { nested-or-no-owner return; unmanaged snapshot; preflight return }, reclaimOptionalMemory(Some(mode)), synchronized { internal }, finally unlock), and withMemoryReclamation is a third variant. The pool-selection 4-tuple match now appears three times (and the 3-tuple twice), and the computeMaxExecutionPoolSize formula is retyped in tryAcquireExecutionMemory and canAcquireExecutionMemory. Also, the enteredWithMonitor branch in acquireExecutionMemory looks dead: its only production caller, TaskMemoryManager, never holds the manager monitor.

Since the comments call this choreography correctness-critical, could it be factored into one private helper (plus a poolsFor(memoryMode) helper and a hoisted maxExecutionPoolSize) so a fix to e.g. the "first registration can race this check" window lands in one place? Relatedly, freeUnrolledValues' serialized branch duplicates PartiallySerializedBlock.discard.

Stepping back: this adds a second cross-task reclamation mechanism (a Runnable registry with its own lock, a prose lock-order contract, and its own failure semantics) alongside the existing MemoryConsumer.spill / TaskMemoryManager mechanism, and it lands ~700 lines (gate, preflights, five MemoryStore wrappers, ~90 lines of failure cleanup, ChunkedByteBufferOutputStream.dispose) with no production caller of registerOptionalMemoryReclaimer or tryAcquireExecutionMemory.

The two mechanisms are invisible to each other: optional reservations don't show up in TaskMemoryManager.showMemoryUsage / getMemoryConsumptionBreakdown, the spill-priority ordering can't see that a cheaper cross-task drain exists, and deadlock-freedom rests on every future callback obeying "Never hold a lock needed by a reclaimer while requesting ordinary memory", which a read-ahead holding a buffer-pool lock while decoding will naturally violate.

Have you considered modelling the optional owner as a MemoryConsumer subtype (discardable flag, zero-cost spill) registered executor-wide by TaskMemoryManager, and plumbing the reclaim hook into ExecutionMemoryPool.acquireMemory's existing loop (next to maybeGrowPool / computeMaxPoolSize)? That would remove the preflight duplication and the RW gate, and keep MemoryStore untouched. Either way, I'd prefer to review this API surface together with its first consumer rather than freezing it beforehand.

I agree this changes ownership semantics on an existing failure path. Consuming a value from the iterator does not transfer ownership of that value to MemoryStore before the put succeeds.

For example, TorrentBroadcast.writeBlocks passes the caller's object through putSingle. If a later unroll reservation triggers eviction and the disk write throws, this catch now closes the original object even when no optional reclaimer exists. The transfer catch has the same ownership concern.

Could we retain cleanup of this operation's unroll credits and internally allocated serialized buffers, while leaving uncommitted deserialized values open? Restricting the close to reclaimer failures would not itself establish ownership either. The regression test should inject an eviction failure without a reclaimer and verify that the caller's object remains open.

Throwing here prevents the requesting task from reaching ordinary admission, even if the other callbacks have already released enough memory. It also means a registered owner holding zero bytes can still fail unrelated tasks when its callback throws.

For a non-fatal callback failure that leaves accounting valid, could we retain the outstanding charge, log the failing owner/task/mode, and continue into ordinary admission? The allocator can then determine whether the request can proceed with the capacity actually available. Fatal failures or broken accounting invariants should remain separate cases.

This differs from existing task-local spill failure handling primarily in its failure scope: a problem in optional work owned by task A now becomes a failure of required work in task B.

Could we separate holding the admission gate from actually reclaiming optional memory?

This unconditional, both-mode drain is reached by unroll reservations with ample free capacity, the unroll-to-storage accounting transfer, and release-only operations such as remove and clear. An on-heap cache operation consequently discards off-heap optional buffers too.

The nested storage preflight cannot prevent this: reserveUnrollMemoryForThisTask has already drained before entering the manager monitor, and its nested admission takes the enteredWithMonitor branch.

Release-only operations should not need reclamation. For allocations, could the outer boundary preflight under the monitor, reclaim only the relevant mode outside it when necessary, and then recheck? Please cover this through the actual MemoryStore paths; the direct acquireStorageMemory tests do not exercise this behavior.

A failed full-request check does not necessarily mean reclamation can improve the grant.

With a 1000-byte execution pool, suppose A holds 400 ordinary bytes and B holds 100 ordinary plus 100 optional bytes. If A requests 200, its remaining share is 100. Reclaiming B's optional bytes leaves B active, so A still receives exactly 100—the same immediate partial grant available before reclamation.

Could the preflight distinguish capacity pressure from a share-bound request where draining cannot help? It should still account for cases where reclaiming removes an optional-only task or releases the requester's own optional bytes. A regression test should verify that the example above preserves B's optional reservation.

This read hold survives ExecutionMemoryPool.acquireMemory calling lock.wait(): the wait releases the manager monitor, but not the admission gate. Consequently, a fairness waiter disables optional admission executor-wide, including admission in the other memory mode.

Is that deliberately part of the policy? If so, please document the cross-mode effect and add a test that places an allocator in the actual capacity-wait loop before attempting optional admission. The existing marked-operation and stalled-eviction tests do not establish that behavior.

If the gate is instead released around the wait, the wake-up path would need to re-establish the reclamation boundary before granting memory.

The require(!Thread.holdsLock(this)) failure is also NonFatal, so releaseOnly can swallow a lock-order violation as though it were a callback failure.

Could we validate marker-before-monitor ordering before attempting to acquire the gate, allowing correctly marked nested calls, and keep that validation outside this catch? Checking only inside reclaimOptionalMemory is too late if the caller already holds the manager monitor while an optional admission holds the write gate and is waiting for that monitor.

I have not found a current production caller violating the ordering; this is about making the correctness-critical contract fail reliably rather than depending on convention.

Finding 1. An optional reservation is charged to the task in ExecutionMemoryPool.memoryForTask but is owned by no MemoryConsumer, so TaskMemoryManager cannot see it. Two consequences, both reachable by the first consumer:

• cleanUpAllAllocatedMemory() returns the leftover bytes, so Executor reports Managed memory leak detected (core/src/main/scala/org/apache/spark/executor/Executor.scala:917-925) for any reservation still held when the task body returns. spark.unsafe.exceptionOnMemoryLeak turns that warning into a thrown SparkException. It defaults to false, but every Spark test run sets it to true (project/SparkBuild.scala:2038, pom.xml:2903). So a consumer whose prefetch outlives the task body only warns in production and fails outright in CI.
• snapshotMemoryUsage() computes memoryNotAccountedFor as the task's pool charge minus the sum over consumers (core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java:490). Optional bytes therefore land in the unattributed bucket of the UNABLE_TO_ACQUIRE_MEMORY breakdown, which is the diagnostic someone reads when a task OOMs next to a prefetcher.

The description defers task-completion integration to a follow-up. The contract as written makes it a prerequisite instead. Nothing ties the registration or the reservation to task completion, so a task that dies before its finally leaves a live registration. Its counter then keeps every later allocation on the preflight-and-drain path for the life of the executor, invoking a callback into a dead reader.

Two ways out. Add the completion-listener integration here, releasing the outstanding bytes and unregistering. Or model the owner as a MemoryConsumer, as suggested at r3998074909, which gets the lifecycle, the breakdown and showMemoryUsage for free. One caveat on that second option: ExecutionMemoryPool.acquireMemory is lock.synchronized, so hooking the reclaim into its loop would run callbacks under the manager monitor. That is the thing this design spends the RW gate to avoid.

Finding 2. The class comment at MemoryStore.scala:91-92 still reads "all changes to memory allocations, notably putting blocks, evicting blocks, and acquiring or releasing unroll memory, must be synchronized on memoryManager". After this PR that recipe hangs.

I ran it on this head. One reclaimer registered, thread A in mm.synchronized { mm.acquireStorageMemory(block, 100, ON_HEAP) }, thread B in mm.tryAcquireExecutionMemory(100, 2, ON_HEAP) once A holds the monitor. Result: threadA=WAITING threadB=BLOCKED, neither returns in 15 seconds. That is the deadlock described at r3998074901.

The new rule is the opposite of the comment: gate first, monitor second. It is stated only in the scaladoc of withMemoryReclamation, which a MemoryStore contributor is less likely to read than the comment at the top of the class they are editing.

Please rewrite lines 91-92 to state the ordering and name the methods that establish it. The runtime check asked for on that thread catches the accident. This catches the person writing the next MemoryStore method.

Finding 3. No in-tree caller reaches acquireExecutionMemory with this monitor held. The only production caller is TaskMemoryManager.acquireExecutionMemory, which holds its own monitor, not this one, and a reclaimer callback is forbidden from allocating. The storage twin at line 335 is covered by "outer storage marker drains before inherited monitor and excludes new optional admission", but nothing covers this one.

The branch is load-bearing where it does fire: without it, a nested call whose full request exceeds the caller's fair share falls through to reclaimOptionalMemory, whose require(!Thread.holdsLock(this)) then fails. So the test needs a registered owner that the drain does not fully release, plus a second task holding ordinary bytes so the task count stays put, and then an execution request under the monitor that must come back as a partial grant rather than an exception.

If there is no caller and none planned, dropping the branch and asserting !Thread.holdsLock(this) here is smaller than keeping an untested one.

Finding 4. This early return drops the diagnostic that used to fire for an impossible unroll request. Before, the call reached acquireStorageMemoryInternal and logged "Will not store <block> as the required space (N bytes) exceeds our memory limit (M bytes)" on its way to false (UnifiedMemoryManager.scala:401). Same result now, one fewer line explaining why nothing got cached. A logInfo here, or reusing that message, keeps it.

Finding 5. This match is not exhaustive and ValuesHolder is not sealed (MemoryStore.scala:760). A third implementation would fail with a MatchError at the worst moment: inside the cleanup of an already-failing unroll, replacing the original error with a confusing one. The trait and both implementations live in this file, so private sealed trait ValuesHolder[T] turns that into a compile error. MemoryEntry next door is already declared that way (MemoryStore.scala:44), and freeMemoryEntry relies on it.

Finding 6. dispose() reuses toChunkedByteBufferWasCalled to mean "the chunks are gone", so a later toChunkedByteBuffer fails on require(!toChunkedByteBufferWasCalled, "toChunkedByteBuffer() can only be called once") with a message saying the opposite of what happened. The new test asserts the IllegalArgumentException without checking its text, so the wrong message is locked in. A separate disposed flag with its own require(!disposed, "cannot call toChunkedByteBuffer() after dispose()") keeps both messages honest for the price of one field.