dongjoon-hyun commented on code in PR #58763:
URL: https://github.com/apache/spark/pull/58763#discussion_r3998074889
##########
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:
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?
##########
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:
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?
##########
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:
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)?
##########
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:
`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?
##########
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:
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.
##########
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:
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()`?
##########
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 {
Review Comment:
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.
##########
core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala:
##########
@@ -206,7 +320,65 @@ private[spark] class UnifiedMemoryManager(
override def acquireStorageMemory(
blockId: BlockId,
numBytes: Long,
- memoryMode: MemoryMode): Boolean = synchronized {
+ memoryMode: MemoryMode): Boolean = {
+ if (isStorageMemoryRequestTooLarge(numBytes, memoryMode)) {
+ return synchronized { acquireStorageMemoryInternal(blockId, numBytes,
memoryMode) }
+ }
+ val gate = optionalAdmissionGate.readLock()
Review Comment:
`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`.
##########
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(
Review Comment:
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.
##########
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()
Review Comment:
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.
--
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]
