viirya commented on code in PR #58763:
URL: https://github.com/apache/spark/pull/58763#discussion_r3998224486
##########
core/src/main/scala/org/apache/spark/storage/memory/MemoryStore.scala:
##########
@@ -211,9 +211,22 @@ private[spark] class MemoryStore(
// Keep track of unroll memory used by this particular block /
putIterator() operation
var unrollMemoryUsedByThisBlock = 0L
+ def reserveUnrollMemory(memory: Long): Boolean = {
+ try {
+ reserveUnrollMemoryForThisTask(blockId, memory, memoryMode)
+ } catch {
+ case error: Throwable =>
+ // No entry or partial iterator can own these values when
reclamation throws.
+ // A normal denial must retain them for the returned partial
iterator.
+ Utils.tryWithSafeFinally { throw error } {
+ releaseUnrollMemoryForThisTask(memoryMode,
unrollMemoryUsedByThisBlock)
+ freeUnrolledValues(valuesHolder)
Review Comment:
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.
##########
core/src/main/scala/org/apache/spark/memory/MemoryManager.scala:
##########
@@ -44,6 +48,127 @@ private[spark] abstract class MemoryManager(
require(onHeapExecutionMemory > 0, "onHeapExecutionMemory must be > 0")
+ // Acquire the marker before this manager's monitor. Shared ownership never
excludes ordinary
+ // operations, including capacity waiters; optional admission only tries the
exclusive side.
+ protected val optionalAdmissionGate = new ReentrantReadWriteLock()
+
+ // This lock protects only registrations. Callbacks never run while it or
this manager's
+ // monitor is held, and registration does not wait behind ordinary capacity
waiters.
+ private val optionalReclaimers = new mutable.LinkedHashMap[Runnable, (Long,
MemoryMode)]()
+ @volatile private var onHeapOptionalReclaimers = 0
+ @volatile private var offHeapOptionalReclaimers = 0
+
+ /**
+ * Register a task-owned, release-only callback before its first optional
admission.
+ * Returns an idempotent unregister action; callers must drain the owner
before unregistering.
+ * Callbacks may run concurrently, repeatedly, or after unregistering and
must release each
+ * reservation exactly once. They may take a short owner-state lock, but
must not acquire a
+ * TaskMemoryManager monitor, allocate execution memory, or wait for I/O or
task cleanup.
+ *
+ * Never hold a lock needed by a reclaimer while requesting ordinary memory
or invoking another
+ * operation that may reclaim optional memory, including storage cleanup.
Otherwise two tasks
+ * can hold their own owner locks while reclaiming each other. Optional
admission and release
+ * may use that lock: neither invokes reclamation nor acquires a
TaskMemoryManager monitor.
+ */
+ private[memory] final def registerOptionalMemoryReclaimer(
+ taskAttemptId: Long,
+ memoryMode: MemoryMode,
+ reclaimer: Runnable): Runnable = {
+ // A distinct forwarding object gives each registration identity even if
callbacks are reused.
+ val registered = new Runnable {
+ /** Release this owner's optional bytes without allocating or destroying
a whole reader. */
+ override def run(): Unit = reclaimer.run()
+ }
+ optionalReclaimers.synchronized {
+ optionalReclaimers(registered) = (taskAttemptId, memoryMode)
+ memoryMode match {
+ case MemoryMode.ON_HEAP => onHeapOptionalReclaimers += 1
+ case MemoryMode.OFF_HEAP => offHeapOptionalReclaimers += 1
+ }
+ }
+ new Runnable {
+ /** Remove this registration only; an already-captured callback remains
safe to invoke. */
+ override def run(): Unit = optionalReclaimers.synchronized {
+ if (optionalReclaimers.remove(registered).isDefined) {
+ memoryMode match {
+ case MemoryMode.ON_HEAP => onHeapOptionalReclaimers -= 1
+ case MemoryMode.OFF_HEAP => offHeapOptionalReclaimers -= 1
+ }
+ }
+ }
+ }
+ }
+
+ /** Check for eligible owners without invoking callbacks or inspecting
native state. */
+ protected final def hasOptionalMemoryReclaimers(memoryMode: MemoryMode):
Boolean = {
+ memoryMode match {
+ case MemoryMode.ON_HEAP => onHeapOptionalReclaimers != 0
+ case MemoryMode.OFF_HEAP => offHeapOptionalReclaimers != 0
+ }
+ }
+
+ /**
+ * Drain a snapshot of matching owners under ordinary admission's shared
gate, outside all
+ * this manager's and the registry's monitors. Owners may run under the
requesting task's monitor
+ * and must follow the registration's lock-order contract. They must
synchronously cancel pure
+ * I/O and release exact credits, without dropping readers/sessions or
awaiting task cleanup.
+ * Continue draining other owners after a non-fatal failure, then propagate
it without inventing
+ * freed credit. Registrations remain live so a failed drain may be retried
safely.
+ */
+ protected final def reclaimOptionalMemory(memoryMode: Option[MemoryMode]):
Unit = {
+ require(!Thread.holdsLock(this), "optional callbacks cannot run under the
memory manager")
+ val callbacks = optionalReclaimers.synchronized {
+ optionalReclaimers.iterator.collect {
+ case (callback, (_, mode)) if memoryMode.forall(_ == mode) => callback
+ }.toList
+ }
+ var failure: Throwable = null
+ callbacks.foreach { callback =>
+ try {
+ callback.run()
+ } catch {
+ case NonFatal(error) =>
+ if (failure == null) failure = error else if (failure ne error) {
+ failure.addSuppressed(error)
+ }
+ }
+ }
+ if (failure != null) throw failure
Review Comment:
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.
##########
core/src/main/scala/org/apache/spark/memory/ExecutionMemoryPool.scala:
##########
@@ -67,6 +67,54 @@ private[memory] class ExecutionMemoryPool(
memoryForTask.getOrElse(taskAttemptId, 0L)
}
+ /**
+ * Check a prospective ordinary request without registering a task or
changing its charge.
+ * `availableMemory` includes free storage the caller can borrow without
eviction; `maxPoolSize`
+ * uses the same potential fair-share ceiling as ordinary acquisition. False
asks the caller to
+ * drain optional owners before any grant, eviction, or capacity wait.
+ */
+ private[memory] def canAcquireMemory(
+ numBytes: Long,
+ taskAttemptId: Long,
+ maxPoolSize: Long,
+ availableMemory: Long): Boolean = lock.synchronized {
+ val tasks = memoryForTask.size + (if
(memoryForTask.contains(taskAttemptId)) 0 else 1)
+ val current = memoryForTask.getOrElse(taskAttemptId, 0L)
+ numBytes <= availableMemory && numBytes <= math.max(0L, maxPoolSize /
tasks - current)
Review Comment:
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.
##########
core/src/main/scala/org/apache/spark/memory/MemoryManager.scala:
##########
@@ -44,6 +48,127 @@ private[spark] abstract class MemoryManager(
require(onHeapExecutionMemory > 0, "onHeapExecutionMemory must be > 0")
+ // Acquire the marker before this manager's monitor. Shared ownership never
excludes ordinary
+ // operations, including capacity waiters; optional admission only tries the
exclusive side.
+ protected val optionalAdmissionGate = new ReentrantReadWriteLock()
+
+ // This lock protects only registrations. Callbacks never run while it or
this manager's
+ // monitor is held, and registration does not wait behind ordinary capacity
waiters.
+ private val optionalReclaimers = new mutable.LinkedHashMap[Runnable, (Long,
MemoryMode)]()
+ @volatile private var onHeapOptionalReclaimers = 0
+ @volatile private var offHeapOptionalReclaimers = 0
+
+ /**
+ * Register a task-owned, release-only callback before its first optional
admission.
+ * Returns an idempotent unregister action; callers must drain the owner
before unregistering.
+ * Callbacks may run concurrently, repeatedly, or after unregistering and
must release each
+ * reservation exactly once. They may take a short owner-state lock, but
must not acquire a
+ * TaskMemoryManager monitor, allocate execution memory, or wait for I/O or
task cleanup.
+ *
+ * Never hold a lock needed by a reclaimer while requesting ordinary memory
or invoking another
+ * operation that may reclaim optional memory, including storage cleanup.
Otherwise two tasks
+ * can hold their own owner locks while reclaiming each other. Optional
admission and release
+ * may use that lock: neither invokes reclamation nor acquires a
TaskMemoryManager monitor.
+ */
+ private[memory] final def registerOptionalMemoryReclaimer(
+ taskAttemptId: Long,
+ memoryMode: MemoryMode,
+ reclaimer: Runnable): Runnable = {
+ // A distinct forwarding object gives each registration identity even if
callbacks are reused.
+ val registered = new Runnable {
+ /** Release this owner's optional bytes without allocating or destroying
a whole reader. */
+ override def run(): Unit = reclaimer.run()
+ }
+ optionalReclaimers.synchronized {
+ optionalReclaimers(registered) = (taskAttemptId, memoryMode)
+ memoryMode match {
+ case MemoryMode.ON_HEAP => onHeapOptionalReclaimers += 1
+ case MemoryMode.OFF_HEAP => offHeapOptionalReclaimers += 1
+ }
+ }
+ new Runnable {
+ /** Remove this registration only; an already-captured callback remains
safe to invoke. */
+ override def run(): Unit = optionalReclaimers.synchronized {
+ if (optionalReclaimers.remove(registered).isDefined) {
+ memoryMode match {
+ case MemoryMode.ON_HEAP => onHeapOptionalReclaimers -= 1
+ case MemoryMode.OFF_HEAP => offHeapOptionalReclaimers -= 1
+ }
+ }
+ }
+ }
+ }
+
+ /** Check for eligible owners without invoking callbacks or inspecting
native state. */
+ protected final def hasOptionalMemoryReclaimers(memoryMode: MemoryMode):
Boolean = {
+ memoryMode match {
+ case MemoryMode.ON_HEAP => onHeapOptionalReclaimers != 0
+ case MemoryMode.OFF_HEAP => offHeapOptionalReclaimers != 0
+ }
+ }
+
+ /**
+ * Drain a snapshot of matching owners under ordinary admission's shared
gate, outside all
+ * this manager's and the registry's monitors. Owners may run under the
requesting task's monitor
+ * and must follow the registration's lock-order contract. They must
synchronously cancel pure
+ * I/O and release exact credits, without dropping readers/sessions or
awaiting task cleanup.
+ * Continue draining other owners after a non-fatal failure, then propagate
it without inventing
+ * freed credit. Registrations remain live so a failed drain may be retried
safely.
+ */
+ protected final def reclaimOptionalMemory(memoryMode: Option[MemoryMode]):
Unit = {
+ require(!Thread.holdsLock(this), "optional callbacks cannot run under the
memory manager")
+ val callbacks = optionalReclaimers.synchronized {
+ optionalReclaimers.iterator.collect {
+ case (callback, (_, mode)) if memoryMode.forall(_ == mode) => callback
+ }.toList
+ }
+ var failure: Throwable = null
+ callbacks.foreach { callback =>
+ try {
+ callback.run()
+ } catch {
+ case NonFatal(error) =>
+ if (failure == null) failure = error else if (failure ne error) {
+ failure.addSuppressed(error)
+ }
+ }
+ }
+ if (failure != null) throw failure
+ }
+
+ /**
+ * Mark an operation that can hold this monitor while evicting blocks or
waiting for capacity.
+ *
+ * MemoryStore uses this before its atomic unroll/storage transfers take the
monitor, preserving
+ * marker-before-monitor ordering when they call back into ordinary
allocation. An outermost
+ * MemoryStore operation drains optional owners before taking the monitor: a
preflight outside
+ * that monitor could otherwise race ordinary allocations and require a
callback inside it.
+ * Nested operations reuse the outer drain. Failures propagate and the
marker is always released.
+ * Release-only storage cleanup logs non-fatal drain failures and continues:
aborting removal
+ * could leave a cached entry behind after BlockManager deletes its
metadata. Failed owners keep
+ * their memory charges and registrations; failures from the cleanup body
still propagate.
+ */
+ private[spark] final def withMemoryReclamation[T](
+ body: => T,
+ releaseOnly: Boolean = false): T = {
+ val gate = optionalAdmissionGate.readLock()
+ gate.lock()
+ try {
+ if ((onHeapOptionalReclaimers != 0 || offHeapOptionalReclaimers != 0) &&
+ optionalAdmissionGate.getReadHoldCount == 1) {
+ try {
+ reclaimOptionalMemory(None)
+ } catch {
+ case NonFatal(error) if releaseOnly =>
Review Comment:
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.
##########
core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala:
##########
@@ -134,7 +180,75 @@ private[spark] class UnifiedMemoryManager(
override private[memory] def acquireExecutionMemory(
numBytes: Long,
taskAttemptId: Long,
- memoryMode: MemoryMode): Long = synchronized {
+ memoryMode: MemoryMode): Long = {
+ val gate = optionalAdmissionGate.readLock()
+ gate.lock()
Review Comment:
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.
##########
core/src/main/scala/org/apache/spark/memory/MemoryManager.scala:
##########
@@ -44,6 +48,127 @@ private[spark] abstract class MemoryManager(
require(onHeapExecutionMemory > 0, "onHeapExecutionMemory must be > 0")
+ // Acquire the marker before this manager's monitor. Shared ownership never
excludes ordinary
+ // operations, including capacity waiters; optional admission only tries the
exclusive side.
+ protected val optionalAdmissionGate = new ReentrantReadWriteLock()
+
+ // This lock protects only registrations. Callbacks never run while it or
this manager's
+ // monitor is held, and registration does not wait behind ordinary capacity
waiters.
+ private val optionalReclaimers = new mutable.LinkedHashMap[Runnable, (Long,
MemoryMode)]()
+ @volatile private var onHeapOptionalReclaimers = 0
+ @volatile private var offHeapOptionalReclaimers = 0
+
+ /**
+ * Register a task-owned, release-only callback before its first optional
admission.
+ * Returns an idempotent unregister action; callers must drain the owner
before unregistering.
+ * Callbacks may run concurrently, repeatedly, or after unregistering and
must release each
+ * reservation exactly once. They may take a short owner-state lock, but
must not acquire a
+ * TaskMemoryManager monitor, allocate execution memory, or wait for I/O or
task cleanup.
+ *
+ * Never hold a lock needed by a reclaimer while requesting ordinary memory
or invoking another
+ * operation that may reclaim optional memory, including storage cleanup.
Otherwise two tasks
+ * can hold their own owner locks while reclaiming each other. Optional
admission and release
+ * may use that lock: neither invokes reclamation nor acquires a
TaskMemoryManager monitor.
+ */
+ private[memory] final def registerOptionalMemoryReclaimer(
+ taskAttemptId: Long,
+ memoryMode: MemoryMode,
+ reclaimer: Runnable): Runnable = {
+ // A distinct forwarding object gives each registration identity even if
callbacks are reused.
+ val registered = new Runnable {
+ /** Release this owner's optional bytes without allocating or destroying
a whole reader. */
+ override def run(): Unit = reclaimer.run()
+ }
+ optionalReclaimers.synchronized {
+ optionalReclaimers(registered) = (taskAttemptId, memoryMode)
+ memoryMode match {
+ case MemoryMode.ON_HEAP => onHeapOptionalReclaimers += 1
+ case MemoryMode.OFF_HEAP => offHeapOptionalReclaimers += 1
+ }
+ }
+ new Runnable {
+ /** Remove this registration only; an already-captured callback remains
safe to invoke. */
+ override def run(): Unit = optionalReclaimers.synchronized {
+ if (optionalReclaimers.remove(registered).isDefined) {
+ memoryMode match {
+ case MemoryMode.ON_HEAP => onHeapOptionalReclaimers -= 1
+ case MemoryMode.OFF_HEAP => offHeapOptionalReclaimers -= 1
+ }
+ }
+ }
+ }
+ }
+
+ /** Check for eligible owners without invoking callbacks or inspecting
native state. */
+ protected final def hasOptionalMemoryReclaimers(memoryMode: MemoryMode):
Boolean = {
+ memoryMode match {
+ case MemoryMode.ON_HEAP => onHeapOptionalReclaimers != 0
+ case MemoryMode.OFF_HEAP => offHeapOptionalReclaimers != 0
+ }
+ }
+
+ /**
+ * Drain a snapshot of matching owners under ordinary admission's shared
gate, outside all
+ * this manager's and the registry's monitors. Owners may run under the
requesting task's monitor
+ * and must follow the registration's lock-order contract. They must
synchronously cancel pure
+ * I/O and release exact credits, without dropping readers/sessions or
awaiting task cleanup.
+ * Continue draining other owners after a non-fatal failure, then propagate
it without inventing
+ * freed credit. Registrations remain live so a failed drain may be retried
safely.
+ */
+ protected final def reclaimOptionalMemory(memoryMode: Option[MemoryMode]):
Unit = {
+ require(!Thread.holdsLock(this), "optional callbacks cannot run under the
memory manager")
+ val callbacks = optionalReclaimers.synchronized {
+ optionalReclaimers.iterator.collect {
+ case (callback, (_, mode)) if memoryMode.forall(_ == mode) => callback
+ }.toList
+ }
+ var failure: Throwable = null
+ callbacks.foreach { callback =>
+ try {
+ callback.run()
+ } catch {
+ case NonFatal(error) =>
+ if (failure == null) failure = error else if (failure ne error) {
+ failure.addSuppressed(error)
+ }
+ }
+ }
+ if (failure != null) throw failure
+ }
+
+ /**
+ * Mark an operation that can hold this monitor while evicting blocks or
waiting for capacity.
+ *
+ * MemoryStore uses this before its atomic unroll/storage transfers take the
monitor, preserving
+ * marker-before-monitor ordering when they call back into ordinary
allocation. An outermost
+ * MemoryStore operation drains optional owners before taking the monitor: a
preflight outside
+ * that monitor could otherwise race ordinary allocations and require a
callback inside it.
+ * Nested operations reuse the outer drain. Failures propagate and the
marker is always released.
+ * Release-only storage cleanup logs non-fatal drain failures and continues:
aborting removal
+ * could leave a cached entry behind after BlockManager deletes its
metadata. Failed owners keep
+ * their memory charges and registrations; failures from the cleanup body
still propagate.
+ */
+ private[spark] final def withMemoryReclamation[T](
+ body: => T,
+ releaseOnly: Boolean = false): T = {
+ val gate = optionalAdmissionGate.readLock()
+ gate.lock()
+ try {
+ if ((onHeapOptionalReclaimers != 0 || offHeapOptionalReclaimers != 0) &&
+ optionalAdmissionGate.getReadHoldCount == 1) {
+ try {
+ reclaimOptionalMemory(None)
Review Comment:
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.
--
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]
