sunchao commented on code in PR #5998: URL: https://github.com/apache/datafusion-comet/pull/5998#discussion_r4042175831
########## spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala: ########## @@ -0,0 +1,278 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.comet + +import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong} + +import scala.util.control.NonFatal + +import org.apache.arrow.memory.AllocationListener +import org.apache.spark.internal.Logging +import org.apache.spark.memory.{MemoryConsumer, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager} + +import org.apache.comet.CometConf + +/** + * Accounts one task's JVM-side Arrow allocations against Spark's off-heap execution pool. + * + * `CometArrowAllocator` is a process-wide `RootAllocator` with no limit, so until now the + * off-heap bytes it hands out were counted by nobody: not Spark's `TaskMemoryManager`, and not + * Comet's native memory pool. They are still resident in the container, which makes them a blind + * spot when an executor is killed for exceeding its memory limit. This closes the reporting half + * of that gap: the bytes appear in `TaskMemoryManager.showMemoryUsage` and are arbitrated against + * Spark's other off-heap consumers. + * + * '''Ownership.''' One instance is created per task and attached to that task's Arrow allocator + * by [[CometTaskArrowAllocator]]. Arrow reports an allocation and its matching release to the + * listener of the allocator that '''owns''' the buffer, on whichever thread happens to drop the + * last reference, and `AllocationListener` is handed nothing but a size. Binding the listener to + * an allocator is therefore the only way to attribute a release, and reading `TaskContext` inside + * the callbacks would get it wrong: a shuffle-read batch handed on to a native operator is pinned + * by native and dropped later from a Tokio worker with no task context installed. That release + * would be lost, leaving the task charged for memory it had already freed, batch after batch. + * + * '''Reporting only.''' A short grant is logged and the allocation proceeds, because Arrow + * allocation on these paths cannot fail today and making it fail is a behavioural change that + * belongs in its own commit. Enforcement belongs in `onPreAllocation`, the only callback + * permitted to throw, and in `onFailedAllocation`, not here. See + * [[https://github.com/apache/datafusion-comet/issues/5997]]. + * + * '''Neither callback may throw.''' Arrow's `AllocationListener` documents that, and + * `BaseAllocator.buffer` marks the allocation successful before calling `onAllocation`, so + * throwing from here loses the buffer Arrow has already created and never hands back. Spark's + * acquisition is fallible in three ways, and only the first is caught by `NonFatal`: it runs + * other consumers' `spill`, which turns a task interrupt into a `RuntimeException` and an I/O + * failure into a `SparkOutOfMemoryError`, and the execution pool itself parks in `lock.wait()`, + * so killing a task can raise a plain `InterruptedException` here. Every call into the memory + * manager is wrapped and reported rather than propagated, and an interrupt additionally re-arms + * the thread's flag so the cancellation is not swallowed. A failed acquisition can also leave the + * task charged for bytes Spark never reported back; see [[acquire]]. + * + * '''Lock order.''' [[getUsed]] and [[spill]] must stay lock-free, because Spark calls both while + * holding the `TaskMemoryManager` monitor, and [[adjust]] holds this listener's monitor across + * `acquireExecutionMemory`, which takes that monitor. Were the snapshot to take this monitor + * instead, a native reservation arriving through `CometTaskMemoryManager` on a Comet Tokio thread + * could hold Spark's monitor and wait for ours while an Arrow allocation on the same task held + * ours and waited for Spark's. + */ +private[comet] class CometArrowAllocationListener(taskMemoryManager: TaskMemoryManager) + extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) + with AllocationListener { + + import CometArrowAllocationListener._ + + /** + * Bytes Arrow currently holds on this task's behalf. An atomic rather than a guarded field so + * that [[getUsed]] can read it without taking this listener's monitor; see the lock order note + * above. + */ + private val live = new AtomicLong(0L) + + /** Bytes currently reserved with Spark. Guarded by this listener's monitor. */ + private var reserved = 0L + + /** Set once the owning task has finished. Volatile so [[getUsed]] can read it lock-free. */ + @volatile private var completed = false + + override def onAllocation(size: Long): Unit = { + live.addAndGet(size) + adjustQuietly() + } + + override def onRelease(size: Long): Unit = { + live.addAndGet(-size) + adjustQuietly() + } + + /** + * Reports our own tally. Spark reads this for spill-victim ordering, `showMemoryUsage` and + * end-of-task leak reporting. The inherited `used` counter stays at zero because this consumer + * never calls `acquireMemory` or `allocatePage`; Arrow has already obtained the memory and we + * are only accounting for it. + * + * Reports zero once the task has finished, so that buffers deliberately allowed to outlive + * their task are not reported by `cleanUpAllAllocatedMemory` as a Spark memory leak. + */ + override def getUsed: Long = if (completed) 0L else math.max(0L, live.get()) + + /** Comet's native operators cannot be made to spill from here. See issue #5997. */ + override def spill(size: Long, trigger: MemoryConsumer): Long = 0L + + /** + * Drops the whole reservation and stops accounting. + * + * Called from the owning task's completion listener. Anything still alive afterwards is a + * buffer that outlives its task, which the process-wide allocator exists to allow; those + * releases are ignored rather than charged to whichever task happens to be running by then. + */ + private[comet] def taskCompleted(): Unit = { + try { + synchronized { + completed = true + if (reserved > 0L) { + taskMemoryManager.releaseExecutionMemory(reserved, this) + reserved = 0L + } + } + } catch { + case e: InterruptedException => reportAndReinterrupt(e) + case NonFatal(e) => warnOnMemoryManagerFailure(e) + case e: SparkOutOfMemoryError => warnOnMemoryManagerFailure(e) + } + } + + /** Bytes Arrow currently holds on this task's behalf. Visible for testing. */ + private[comet] def liveBytes: Long = live.get() + + /** Bytes currently reserved with Spark on this task's behalf. Visible for testing. */ + private[comet] def reservedBytes: Long = synchronized(reserved) + + private def adjustQuietly(): Unit = { + try { + adjust() + } catch { + // Growth handles its own failures in `acquire`, so this is the net for the release path and + // for anything unforeseen. All three are reachable from the memory manager: + // `acquireExecutionMemory` runs other consumers' `spill`, `TaskMemoryManager` turns an + // interrupted spill into a RuntimeException and an IOException into a SparkOutOfMemoryError, + // and the execution pool itself parks in `lock.wait()`. The last two slip past NonFatal, + // which excludes Errors and InterruptedException. + case e: InterruptedException => reportAndReinterrupt(e) + case NonFatal(e) => warnOnMemoryManagerFailure(e) + case e: SparkOutOfMemoryError => warnOnMemoryManagerFailure(e) + } + } + + private def adjust(): Unit = synchronized { + if (!completed) { + val liveBytes = math.max(0L, live.get()) + if (reserved < liveBytes) { + // Round up so `reserved` stays a block multiple and growth always leaves headroom. + // Requesting the bare deficit would land exactly on `liveBytes` for any buffer at or above + // the block size, sending the very next allocation straight back into Spark's lock. + val request = roundUpToBlock(liveBytes - reserved) + val granted = acquire(request) + reserved += granted + if (granted < request) { + warnOnShortGrant(request, granted) + } + } else { + // Returned in one call rather than one per block: `releaseExecutionMemory` synchronizes on + // the executor-wide pool, so a per-block loop would take that lock once per megabyte freed. + val excess = ((reserved - liveBytes) / BLOCK_SIZE) * BLOCK_SIZE + if (excess > 0L) { + taskMemoryManager.releaseExecutionMemory(excess, this) + reserved -= excess + } + } + } + } + + /** + * Asks Spark for `request` bytes and returns what this consumer ends up holding, which is not + * always what Spark returns. + * + * `acquireExecutionMemory` takes its first grant from the pool and only then asks other + * consumers to spill, so when a spill throws it has already charged the task for bytes it never + * reports back. Nothing would release them: [[taskCompleted]] only knows about `reserved`, and + * Spark itself only reclaims them in `cleanUpAllAllocatedMemory` at the very end of the task, + * so until then they are headroom nobody can use. They are adopted here instead, measured as + * the change in what the pool says this task holds. + * + * That measurement is an estimate, but a safe one. Another consumer in the same task cannot + * acquire concurrently, because `acquireExecutionMemory` holds the `TaskMemoryManager` monitor + * throughout, so the only interference is a concurrent release, which makes the figure too + * small rather than too large; and `request` bounds it from above either way. Too small + * degrades to what would have happened anyway. + */ + private def acquire(request: Long): Long = { + val heldBefore = taskMemoryManager.getMemoryConsumptionForThisTask + try { + taskMemoryManager.acquireExecutionMemory(request, this) Review Comment: ### Correctness [P2] Keep both usage snapshots inside the acquisition transaction `heldBefore` is read before `acquireExecutionMemory` takes the task-manager monitor, and `adoptOrphanedGrant` reads again after that monitor has been released on the exception path. Another consumer can therefore acquire between either snapshot and the protected call. This lets the listener adopt and later release someone else's reservation, rather than only underestimate its own grant. With the exact current listener and real Spark `UnifiedMemoryManager`/`TaskMemoryManager`, I paused after the real before-snapshot, let another OFF_HEAP consumer take the 1 MiB pool, and then let Arrow's zero-grant acquisition trigger that consumer's failing spill. The listener adopted 1 MiB despite receiving no grant. Closing the Arrow buffer reduced Spark's charge to zero while the other consumer still owned 1 MiB. A control that acquired before the snapshot retained the correct 1 MiB charge. Could failure recovery use an atomic, consumer-specific accounting boundary and add this interleaving as a regression? Bounding the result by `request` does not establish that those bytes belong to this listener. -- 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]
