sunchao commented on code in PR #58763:
URL: https://github.com/apache/spark/pull/58763#discussion_r4010944819
##########
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:
Updated in `79b3464d7889fcc66282e47804706178f6a2ec42`.
The revision now includes the task/consumer ownership surface: registration
requires a MemoryConsumer, typed optional release checks the exact owner, task
diagnostics include it, and task cleanup closes admission and drains before
releasing credits. I kept the executor-wide callback boundary outside the
manager monitor: invoking a native owner from ExecutionMemoryPool's
synchronized loop would introduce the lock inversion you and Peter describe,
while existing spill is task-local. The exact API has also passed the adopting
Comet JNI fixture (12 scenarios), and the in-tree TaskMemoryManager tests
include a representative optional consumer. A production reader is not included
in Apache Spark here, so I am leaving this design discussion open for review
rather than treating the fixture as a shipped consumer. Owner cost/tail-latency
measurements are still needed; the startup setting defaults false.
##########
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:
Updated in `79b3464d7889fcc66282e47804706178f6a2ec42`.
Fixed both unroll and transfer failure cleanup. Deserialized values remain
caller-owned until the put succeeds; only this operation's unroll credits and
Spark-created serialized buffers are released on failure. The regression
injects a real eviction IOException with no reclaimer and checks that the
consumed AutoCloseable remains open, the old block remains present, and unroll
accounting is restored.
##########
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:
Updated in `79b3464d7889fcc66282e47804706178f6a2ec42`.
Changed nonfatal callback failures to log task/mode attribution and retain
the outstanding charge while continuing ordinary admission. Other owners can
still release enough capacity for the requester. Fatal errors and accounting
AssertionErrors propagate; allocation and lock-order failures are outside the
callback catch. Execution and real unroll tests cover both a successful grant
after another owner's failure and honest denial when capacity remains
unavailable.
##########
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:
Updated in `79b3464d7889fcc66282e47804706178f6a2ec42`.
Separated the admission marker from reclamation. Unroll now uses a
same-mode, capacity-aware boundary before the manager monitor; transfers,
remove and clear take only the marker. MemoryStore/BlockManager tests cover
ample-capacity puts, cross-mode isolation, real unroll pressure,
removal/clearing and transfer without discarding optional buffers.
##########
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:
Updated in `79b3464d7889fcc66282e47804706178f6a2ec42`.
Added a no-reclamation path for an immediately available share-bound partial
grant, including your 400 / (100 ordinary + 100 optional) example: the request
returns 100, then zero at the cap, without invoking the owner. Ordinary and
optional ledgers are now separate, so optional-only participants and the
requester's own optional bytes do not lower ordinary shares; their bytes still
count against physical capacity. The preflight remains conservative where
whole-block storage eviction can raise the grant.
##########
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:
Updated in `79b3464d7889fcc66282e47804706178f6a2ec42`.
Yes, retained deliberately as a conservative opt-in policy. The PR
description and gate documentation now state that a fairness/capacity waiter
denies optional admission executor-wide, including the other memory mode. A
regression places a thread in ExecutionMemoryPool's actual wait loop, verifies
same-mode and cross-mode denial, then verifies ordinary progress and optional
admission after release. The startup setting defaults false and bypasses the
gate when disabled; this change does not promise optional progress during a
sustained ordinary allocation stream.
##########
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:
Updated in `79b3464d7889fcc66282e47804706178f6a2ec42`.
Moved marker-before-monitor validation ahead of gate acquisition, outside
the callback-failure catch. Correctly marked nested storage operations remain
allowed; unmarked monitor-holding calls fail immediately. Added assertions for
invalid ordering and propagation of an accounting AssertionError, and updated
MemoryStore's class-level lock-order documentation.
##########
core/src/main/scala/org/apache/spark/memory/MemoryManager.scala:
##########
@@ -128,6 +258,24 @@ private[spark] abstract class MemoryManager(
taskAttemptId: Long,
memoryMode: MemoryMode): Long
+ /**
+ * Reserve all `numBytes` for optional task work without waiting for
capacity or reclaiming
+ * memory.
+ *
+ * Managers must opt in to this policy; the default rejects the request
without changing their
+ * accounting. Implementations may still contend on bookkeeping locks. A
successful reservation
+ * must be task-attributed and released through the existing
execution-memory release methods.
+ *
+ * @return `numBytes` on success, or zero with no reservation on denial
+ */
+ private[memory] def tryAcquireExecutionMemory(
Review Comment:
Updated in `79b3464d7889fcc66282e47804706178f6a2ec42`.
Included task-completion handling and consumer attribution in this PR.
Optional registration now belongs to TaskMemoryManager and a specific
MemoryConsumer. Cleanup closes admission, drains owners, then performs ordinary
leak detection/release; failed or incomplete drains stay registered and charged
rather than returning credit for live buffers. Diagnostics include optional
consumers. Tests cover same/peer-task reclamation, mixed ordinary/optional
accounting, cleanup racing a native-style worker, failed cleanup,
wrong-consumer release and a no-op drain; the real Comet JNI fixture also
passes.
##########
core/src/main/scala/org/apache/spark/storage/memory/MemoryStore.scala:
##########
@@ -414,20 +438,47 @@ 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)
}
}
- def remove(blockId: BlockId): Boolean = memoryManager.synchronized {
+ 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 only values consumed by this unroll operation. */
+ private def freeUnrolledValues(valuesHolder: ValuesHolder[_]): Unit =
valuesHolder match {
+ case holder: DeserializedValuesHolder[_] =>
+ // Final sizing has already moved the vector into arrayValues; do not
build it again.
+ freeValues(if (holder.vector != null) holder.vector.iterator else
holder.arrayValues.iterator)
+ case holder: SerializedValuesHolder[_] =>
+ Utils.tryWithSafeFinally {
+ // As in PartiallySerializedBlock.discard, closing must not allocate
or flush more data.
+
holder.redirectableStream.setOutputStream(OutputStream.nullOutputStream())
+ holder.serializationStream.close()
+ } {
+ holder.bbos.dispose()
+ }
+ }
+
+ /**
+ * Remove a block and release its storage charge; optional admission skips
object close callbacks.
+ */
+ def remove(blockId: BlockId): Boolean = {
Review Comment:
Updated in `79b3464d7889fcc66282e47804706178f6a2ec42`.
Updated MemoryStore's class comment to require admission marker first,
manager monitor second, and to distinguish withStorageMemoryReclamation for
capacity-changing unroll admission from withMemoryReclamation for
transfers/removal. Runtime ordering checks now run before attempting the gate.
##########
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()
+ try {
+ // A first registration can race this check, but cannot admit optional
bytes while this
+ // read gate is held. Keep the no-owner path free of reentrant-monitor
checks/preflight.
+ if (!hasOptionalMemoryReclaimers(memoryMode)) {
+ return synchronized {
+ acquireExecutionMemoryInternal(numBytes, taskAttemptId, memoryMode)
+ }
+ }
+ val enteredWithMonitor = Thread.holdsLock(this)
Review Comment:
Updated in `79b3464d7889fcc66282e47804706178f6a2ec42`.
Removed the unsupported execution enteredWithMonitor path. With optional
coordination enabled, acquireExecutionMemory requires that the caller does not
already hold the memory-manager monitor. Added the corresponding rejection
test; correctly marked nested storage admission still uses its explicit
boundary.
##########
core/src/main/scala/org/apache/spark/storage/memory/MemoryStore.scala:
##########
@@ -589,17 +659,23 @@ private[spark] class MemoryStore(
blockId: BlockId,
memory: Long,
memoryMode: MemoryMode): Boolean = {
- memoryManager.synchronized {
- val success = memoryManager.acquireUnrollMemory(blockId, memory,
memoryMode)
- if (success) {
- val taskAttemptId = currentTaskAttemptId()
- val unrollMemoryMap = memoryMode match {
- case MemoryMode.ON_HEAP => onHeapUnrollMemoryMap
- case MemoryMode.OFF_HEAP => offHeapUnrollMemoryMap
+ if (memoryManager.isStorageMemoryRequestTooLarge(memory, memoryMode)) {
Review Comment:
Updated in `79b3464d7889fcc66282e47804706178f6a2ec42`.
Restored the impossible-unroll diagnostic before returning false, including
block id, requested bytes and the mode's memory limit. The impossible-request
tests still verify no optional reclamation.
##########
core/src/main/scala/org/apache/spark/storage/memory/MemoryStore.scala:
##########
@@ -414,20 +438,47 @@ 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)
}
}
- def remove(blockId: BlockId): Boolean = memoryManager.synchronized {
+ 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 only values consumed by this unroll operation. */
+ private def freeUnrolledValues(valuesHolder: ValuesHolder[_]): Unit =
valuesHolder match {
Review Comment:
Updated in `79b3464d7889fcc66282e47804706178f6a2ec42`.
Sealed ValuesHolder; both implementations remain in this file, and temporary
cleanup now explicitly preserves deserialized values while disposing serialized
state.
##########
core/src/main/scala/org/apache/spark/util/io/ChunkedByteBufferOutputStream.scala:
##########
@@ -53,6 +53,16 @@ private[spark] class ChunkedByteBufferOutputStream(
def size: Long = _size
+ /** Release untransferred chunks without allocating a compact final chunk. */
+ def dispose(): Unit = {
Review Comment:
Updated in `79b3464d7889fcc66282e47804706178f6a2ec42`.
Added a separate disposed flag. A conversion after dispose now reports
'cannot call toChunkedByteBuffer() after dispose()'; a second conversion
retains the original one-conversion diagnostic. The tests assert both messages.
--
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]
