sunchao commented on code in PR #58763:
URL: https://github.com/apache/spark/pull/58763#discussion_r4010944081
##########
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:
##########
@@ -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:
Updated in `79b3464d7889fcc66282e47804706178f6a2ec42`.
Added typed optional-byte accounting per mode/task and exact owner
accounting in TaskMemoryManager. The zero-outstanding path avoids
snapshotting/draining; drained registrations skip owner callbacks, and
reclamation stops once admission can proceed. Two new consumer tests cover an
empty owner beside an active owner in the same task and the
grant-before-credit-publication race, so skipping an empty owner cannot lose a
newly admitted reservation. Optional admission can still occur between
independent ordinary operations; sustained-pressure prefetch backoff remains a
consumer scheduling concern.
##########
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:
I measured the ordinary on-heap allocation paths with no optional owners
registered. The comparison uses the same JDK, dependency classpath, and harness
in separate JVMs, changing only the allocator classes and the candidate's
startup setting. Each operation acquires and releases 64 or 128 bytes; every
grant and final memory balance is checked.
Revisions:
- Before this PR: `5d388249c3cf28f5ba11cdd5275d1fe5ae9cc1c1`.
- Original PR: `9d965e4be4be060c7c99bd68569cc44d7616d7cd`.
- Updated PR: `79b3464d7889fcc66282e47804706178f6a2ec42`; tested with
`spark.memory.optional.enabled=false` and `true`.
Environment: Linux x86-64; AMD EPYC-Milan (32 available CPUs); OpenJDK
17.0.20.1+1; Scala 2.13.18; `-Xms512m -Xmx512m`. Three JVM forks per case with
variant order rotated across forks; three warm-up rounds and seven measured
rounds. The single-thread case uses 500,000 pairs per round; four-thread cases
use 50,000 pairs per worker. Timing covers the concurrent round, so four-thread
ns/pair expresses aggregate throughput rather than individual request latency.
Allocation counts come from the JVM's per-thread allocated-byte counters around
the worker loop; these are cumulative allocations, not retained heap size.
Median ns per acquire/release pair; brackets show the minimum–maximum across
all 21 measured rounds:
| Path / threads | Pre-PR | Original PR | Updated, disabled | Updated,
enabled |
|---|---:|---:|---:|---:|
| Execution / 1 | 205.3 [193.9–324.9] | 211.3 [196.0–732.8] | 206.9
[192.2–448.7] | 289.0 [243.5–593.4] |
| Execution / 4 | 514.7 [275.8–1251.5] | 1226.3 [757.6–2025.7] | 996.6
[549.7–1755.8] | 1440.0 [1053.7–1997.8] |
| Storage / 1 | 110.1 [68.9–590.6] | 89.3 [75.4–132.0] | 114.2 [72.9–116.3]
| 151.4 [110.3–199.3] |
| Storage / 4 | 335.6 [292.6–405.0] | 542.2 [430.8–1371.7] | 321.2
[261.7–894.1] | 666.0 [460.1–879.5] |
Median allocated bytes per pair:
| Path / threads | Pre-PR | Original PR | Updated, disabled | Updated,
enabled |
|---|---:|---:|---:|---:|
| Execution / 1 | 256.0 | 264.0 | 264.0 | 346.3 |
| Execution / 4 | 245.2 | 260.5 | 240.8 | 357.5 |
| Storage / 1 | 48.0 | 48.0 | 48.0 | 48.0 |
| Storage / 4 | 48.0 | 75.3 | 48.0 | 78.6 |
All 48 JVM runs passed grant, class-origin, and final-accounting checks. The
updated disabled execution path overlaps the controls. The enabled path has
measurable coordination/allocation cost; disabled storage was about 25 ns/pair
above the original PR in this run. Four-thread timing varied substantially, and
the pre-PR single-thread storage fork medians were 180.9, 73.8, and 84.0
ns/pair, so these data do not establish precise small differences or contention
speedups.
The startup setting `spark.memory.optional.enabled` defaults to `false`. The
disabled cases measure the ordinary path with that feature gate off; enabled
cases include its coordination cost even with no optional owners.
This is a bounded allocator microbenchmark, not a SQL query, storage
eviction, optional-owner contention, or I/O-performance result. It measures
ordinary allocations with ample free capacity and no reclaimers. The variants
share unchanged runtime dependencies; the pre-PR and original allocator
implementations are compiled unchanged into separate directories and placed
first on their JVM classpaths. The harness prints the origins of the three
allocator classes; the run script checks they match the requested variant.
<details>
<summary>Full standalone harness and reproduction commands</summary>
Save the following as `bench/NoOwnerAllocationBenchmark.java`:
```java
package org.apache.spark.memory;
import java.lang.management.ManagementFactory;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.spark.SparkConf;
import org.apache.spark.storage.RDDBlockId;
/** Isolated no-owner allocation probe. One operation is an acquire/release
pair. */
public final class NoOwnerAllocationBenchmark {
private static volatile long sink;
private static final com.sun.management.ThreadMXBean ALLOC =
(com.sun.management.ThreadMXBean) ManagementFactory.getThreadMXBean();
public static void main(String[] args) throws Exception {
String label = args[0];
boolean enabled = Boolean.parseBoolean(args[1]);
String scenario = args[2];
int threads = Integer.parseInt(args[3]);
int iterations = Integer.parseInt(args[4]);
int warmups = Integer.parseInt(args[5]);
int measurements = Integer.parseInt(args[6]);
SparkConf conf = new
SparkConf(false).set("spark.memory.optional.enabled", "" + enabled);
UnifiedMemoryManager manager = new UnifiedMemoryManager(conf, 67108864L,
33554432L, threads);
if (ALLOC.isThreadAllocatedMemorySupported() &&
!ALLOC.isThreadAllocatedMemoryEnabled()) {
ALLOC.setThreadAllocatedMemoryEnabled(true);
}
System.out.println("ORIGIN " +
MemoryManager.class.getProtectionDomain().getCodeSource().getLocation());
System.out.println("ORIGIN " +
UnifiedMemoryManager.class.getProtectionDomain().getCodeSource().getLocation());
System.out.println("ORIGIN " +
ExecutionMemoryPool.class.getProtectionDomain().getCodeSource().getLocation());
ExecutorService executor = Executors.newFixedThreadPool(threads);
try {
for (int round = -warmups; round < measurements; round++) {
CountDownLatch ready = new CountDownLatch(threads);
CountDownLatch start = new CountDownLatch(1);
CountDownLatch done = new CountDownLatch(threads);
AtomicReference<Throwable> failure = new AtomicReference<>();
AtomicLong allocated = new AtomicLong();
AtomicLong checksum = new AtomicLong();
for (int worker = 0; worker < threads; worker++) {
final long taskId = worker + 1;
executor.submit(() -> {
boolean signaled = false;
try {
RDDBlockId block = new RDDBlockId(1, (int) taskId);
if (scenario.equals("execution") &&
manager.acquireExecutionMemory(64L, taskId,
MemoryMode.ON_HEAP) != 64L) {
throw new AssertionError("anchor refused");
}
ready.countDown();
signaled = true;
start.await();
long threadId = Thread.currentThread().getId();
long before = ALLOC.getThreadAllocatedBytes(threadId);
long sum = 0L;
for (int i = 0; i < iterations; i++) {
long bytes = (i & 1) == 0 ? 64L : 128L;
if (scenario.equals("execution")) {
long got = manager.acquireExecutionMemory(bytes, taskId,
MemoryMode.ON_HEAP);
if (got != bytes) throw new AssertionError("partial
ordinary grant " + got);
manager.releaseExecutionMemory(got, taskId,
MemoryMode.ON_HEAP);
sum += got;
} else {
if (!manager.acquireStorageMemory(block, bytes,
MemoryMode.ON_HEAP)) {
throw new AssertionError("storage grant refused");
}
manager.releaseStorageMemory(bytes, MemoryMode.ON_HEAP);
sum += bytes;
}
}
allocated.addAndGet(ALLOC.getThreadAllocatedBytes(threadId) -
before);
checksum.addAndGet(sum);
if (scenario.equals("execution")) {
manager.releaseExecutionMemory(64L, taskId,
MemoryMode.ON_HEAP);
}
} catch (Throwable error) {
failure.compareAndSet(null, error);
} finally {
if (!signaled) ready.countDown();
done.countDown();
}
});
}
if (!ready.await(10L, TimeUnit.SECONDS)) throw new
AssertionError("worker startup timeout");
long begin = System.nanoTime();
start.countDown();
if (!done.await(30L, TimeUnit.SECONDS)) throw new
AssertionError("round timeout");
long elapsed = System.nanoTime() - begin;
if (failure.get() != null) throw new AssertionError("worker failed",
failure.get());
if (manager.onHeapExecutionMemoryUsed() != 0L ||
manager.onHeapStorageMemoryUsed() != 0L) {
throw new AssertionError("memory credit leaked");
}
sink = checksum.get();
if (round >= 0) {
long operations = (long) threads * iterations;
System.out.printf(java.util.Locale.ROOT,
"RESULT,%s,%s,%d,%d,%d,%d,%.3f,%.3f,%d%n", label, scenario,
threads, round,
operations, elapsed, (double) elapsed / operations,
(double) allocated.get() / operations, sink);
}
}
} finally {
executor.shutdownNow();
}
}
}
```
From the updated PR checkout, build its core test classpath once. This
recipe requires JDK 17 on `PATH`; run from the checkout root in a Unix shell.
```sh
mkdir -p bench/harness-classes bench/base-src bench/base-classes \
bench/original-src bench/original-classes bench/results
build/sbt 'core/Test/compile' 'export core/Test/fullClasspath' >
bench/classpath-export.log 2>&1
python3 - <<'PYCP'
from pathlib import Path
lines = Path('bench/classpath-export.log').read_text().splitlines()
cp = [line for line in lines if line.startswith('/') and '/core/target/' in
line and ':/' in line][-1]
Path('bench/candidate-classpath.txt').write_text(cp + '\n')
PYCP
for variant in base original; do
if [ "$variant" = base ]; then
revision=5d388249c3cf28f5ba11cdd5275d1fe5ae9cc1c1
else
revision=9d965e4be4be060c7c99bd68569cc44d7616d7cd
fi
for name in MemoryManager UnifiedMemoryManager ExecutionMemoryPool; do
git show
"$revision:core/src/main/scala/org/apache/spark/memory/$name.scala" >
"bench/$variant-src/$name.scala"
done
done
python3 - <<'PYBUILD'
from pathlib import Path
import subprocess
cp = Path('bench/candidate-classpath.txt').read_text().strip()
compiler = ':'.join(path for path in cp.split(':') if any(
key in path for key in ('/scala-compiler/', '/scala-library/',
'/scala-reflect/')))
for variant in ('base', 'original'):
subprocess.run(['java', '-Xmx2g', '-cp', compiler,
'scala.tools.nsc.Main',
'-classpath', cp, '-d', f'bench/{variant}-classes'] +
[f'bench/{variant}-src/{name}.scala' for name in
('MemoryManager', 'UnifiedMemoryManager',
'ExecutionMemoryPool')], check=True)
subprocess.run(['javac', '-cp', cp, '-d', 'bench/harness-classes',
'bench/NoOwnerAllocationBenchmark.java'], check=True)
PYBUILD
python3 - <<'PYRUN'
from pathlib import Path
import subprocess
cp = Path('bench/candidate-classpath.txt').read_text().strip()
variants = [('base', False), ('original', False),
('candidate_disabled', False), ('candidate_enabled', True)]
for fork in range(3):
order = variants[fork:] + variants[:fork]
for scenario in ('execution', 'storage'):
for threads in (1, 4):
for label, enabled in order:
prefix = ['bench/harness-classes']
if label in ('base', 'original'):
prefix.append(f'bench/{label}-classes')
classpath = ':'.join(prefix + [cp])
iterations = 500000 if threads == 1 else 50000
args = ['java', '-Xms512m', '-Xmx512m',
'--add-opens=java.base/sun.nio.ch=ALL-UNNAMED',
'-cp', classpath,
'org.apache.spark.memory.NoOwnerAllocationBenchmark',
label, str(enabled).lower(), scenario, str(threads),
str(iterations), '3', '7']
log =
Path(f'bench/results/{fork}-{scenario}-{threads}-{label}.log')
with log.open('w') as output:
subprocess.run(args, stdout=output,
stderr=subprocess.STDOUT,
timeout=120, check=True)
origins = [line for line in log.read_text().splitlines()
if line.startswith('ORIGIN ')]
expected = str(Path(f'bench/{label}-classes').resolve()) \
if label in ('base', 'original') else str(Path.cwd())
assert len(origins) == 3 and all(expected in line for line
in origins)
PYRUN
```
Each `RESULT` line reports variant, path, threads, round, pair count,
elapsed nanoseconds, ns/pair, allocated bytes/pair, and checksum. Compare
medians across the 21 measured rounds per case; retain individual fork results
to see JVM and host variation.
</details>
##########
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:
Updated in `79b3464d7889fcc66282e47804706178f6a2ec42`.
Factored gate/preflight/outside-monitor-reclaim/retry into one
withReclamation helper, with shared pool selection and execution-ceiling
helpers. Removed the unsupported entered-with-monitor execution path and added
the ordering assertion. Serialized failure disposal remains a small helper with
the same redirect-before-close contract as PartiallySerializedBlock.discard; it
owns pre-builder state rather than a constructed PartiallySerializedBlock.
--
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]
