diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index a6a494e8f90..27185d1fd31 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -537,6 +537,7 @@ jobs: org.apache.spark.CometPluginsSuite org.apache.spark.CometRuntimeShutdownSuite org.apache.spark.CometTaskMemoryManagerSuite + org.apache.spark.comet.CometArrowAllocationListenerSuite org.apache.spark.CometExecIteratorLifecycleSuite org.apache.spark.CometPluginsDefaultSuite org.apache.spark.CometPluginsNonOverrideSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 3faca44dd75..7e67c71b609 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -185,6 +185,7 @@ jobs: org.apache.spark.CometPluginsSuite org.apache.spark.CometRuntimeShutdownSuite org.apache.spark.CometTaskMemoryManagerSuite + org.apache.spark.comet.CometArrowAllocationListenerSuite org.apache.spark.CometExecIteratorLifecycleSuite org.apache.spark.CometPluginsDefaultSuite org.apache.spark.CometPluginsNonOverrideSuite diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index c252afedc45..ed963070d4f 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -64,23 +64,63 @@ Enabling Comet does not add one new memory consumer, it adds several, and they a accounted by the same party. This inventory is worth internalizing before reading the rest of the page: -| Allocator | Lives in | Bounded by | Visible to Spark? | -| --------------------------------------- | ----------- | ------------------------------------------------------------- | ----------------- | -| Spark execution + storage (on-heap) | JVM heap | `spark.executor.memory` and the unified memory manager | Yes | -| Spark Tungsten (off-heap) | Off-heap | `spark.memory.offHeap.size` via `TaskMemoryManager` | Yes | -| Comet native (Rust global allocator) | Native heap | `memory_limit` (see below), enforced only via the memory pool | No | -| Comet JVM Arrow (`CometArrowAllocator`) | Off-heap | **Nothing**: a `RootAllocator(Long.MaxValue)` | No | -| Comet JVM shuffle pages | Off-heap | `spark.memory.offHeap.size` via `TaskMemoryManager` | Yes | - -Two observations follow. - -**Comet's JVM-side Arrow allocator is unbounded and accounted by nobody.** `CometArrowAllocator` -(`spark/src/main/scala/org/apache/comet/package.scala`) is a single process-wide -`new RootAllocator(Long.MaxValue)`. Child allocators are cut from it for FFI stream export -(`CometNativeArrowSource`), broadcast coalescing, and `CometSparkToColumnarExec`. These are real -off-heap bytes in container RSS that neither Spark's `TaskMemoryManager` nor Comet's native memory -pool sees. In practice the volume is modest, a batch at a time per stream, but there is no -ceiling and no backpressure. +| Allocator | Lives in | Bounded by | Visible to Spark? | +| ----------------------------------------- | ----------- | ------------------------------------------------------------- | ----------------- | +| Spark execution + storage (on-heap) | JVM heap | `spark.executor.memory` and the unified memory manager | Yes | +| Spark Tungsten (off-heap) | Off-heap | `spark.memory.offHeap.size` via `TaskMemoryManager` | Yes | +| Comet native (Rust global allocator) | Native heap | `memory_limit` (see below), enforced only via the memory pool | No | +| Comet JVM Arrow that stays in the JVM | Off-heap | **Nothing**, but reported to `TaskMemoryManager` | Yes | +| Comet JVM Arrow crossing the FFI boundary | Off-heap | **Nothing**: a `RootAllocator(Long.MaxValue)` | No | +| Comet JVM shuffle pages | Off-heap | `spark.memory.offHeap.size` via `TaskMemoryManager` | Yes | + +Several observations follow. + +**Comet's JVM-side Arrow allocations inside a task are reported to Spark, but still unbounded.** +`CometArrowAllocator` (`spark/src/main/scala/org/apache/comet/package.scala`) is a single +process-wide `new RootAllocator(Long.MaxValue)`. Child allocators are cut from it for FFI stream +export (`CometNativeArrowSource`), broadcast coalescing, and `CometSparkToColumnarExec`. JVM-owned +allocations do not use the root directly: `CometTaskArrowAllocator.forCurrentTask()` hands each +task a child carrying its own `CometArrowAllocationListener`, which charges the bytes to a +`MemoryConsumer` for that task in whole blocks, so they appear in `showMemoryUsage` and are +arbitrated against Spark's other off-heap consumers. The limit is still `Long.MaxValue`: the +listener reports without enforcing, so an allocation here cannot fail, though the bytes do consume +the off-heap pool and other consumers see correspondingly less headroom. Set +`spark.comet.arrowAllocator.accounting.enabled=false` to stop reporting them. + +**The owner is the allocator, not the calling thread.** Arrow hands `AllocationListener` nothing +but a size, and it reports a release on whichever thread drops the last reference, which for +anything exported over the C Data Interface is a Tokio worker with no task context installed. +Attributing by `TaskContext` would therefore drop those releases and leave a task charged for +memory it had already freed. Binding one listener to one task's allocator is what makes the release +land on the task that allocated. Children cut from a task allocator inherit its listener, so the +paths that make their own children are covered without knowing about any of this. + +**Anything crossing the FFI boundary uses the root, and is not accounted.** Both directions, for +the same reason: native's pool is the authority for bytes native holds. Coming in, Arrow's +`wrapForeignAllocation` reports an imported buffer to the allocator's listener at full capacity +even though no JVM-side allocation happened. Going out, whichever DataFusion operator retains the +batch reserves those buffers itself -- `ExternalSorter` through `get_reserved_bytes_for_record_batch`, +the hash join build side through `get_record_batch_memory_size`, both of which read sizes off the +`ArrayData` and so count imported buffers -- and through a unified pool that charges the same Spark +task. Reporting either direction on the JVM side would reserve the same memory twice and could +reject an allocation that fits. So `NativeUtil`, the JVM UDF result in `CometUdfBridge` +together with the codegen output vector it exports, and `CometNativeArrowSource.stream` all use +the listener-less root; IPC reads, the cached batch serializer and +`CometNativeArrowSource.readerBatchIter` use the task allocator. The +driver, broadcast coalescing and the cached batch serializer also fall back to the root when there +is no task to charge, as does Comet's on-heap mode. + +**The split is by allocation site, so it is not exact.** A buffer that is used in the JVM and only +later handed to native -- a shuffle-read batch feeding a native operator is the common shape -- is +allocated somewhere that cannot know its future, so it stays charged on the JVM side while native +may also reserve it. Closing that needs reservation ownership to be handed over at the boundary, +which is a change on both sides of it. + +**A task allocator is closed only once it has been drained.** The process-wide allocator exists +because Arrow buffers can outlive the task that created them, and Arrow treats closing an allocator +that still owns bytes as a leak. At task completion the Spark reservation is dropped and the +allocator is closed if it is empty; otherwise it is parked and closed by a later task once the +stragglers are released, since `BaseAllocator` keeps every child in a map until it closes. **The JVM shuffle allocator is an ordinary Spark consumer.** `CometShuffleMemoryAllocator.getInstance` returns `CometUnifiedShuffleMemoryAllocator`, a Spark `MemoryConsumer` drawing from @@ -205,13 +245,13 @@ accounting (if any) saw the bytes appear. Whether the buffer is _reserved_ in Co separate decision, made by whichever operator holds it, and that operator neither knows nor cares which side of the boundary the bytes came from. -**JVM → native (`ScanExec`).** The JVM allocates the Arrow buffers from a child of -`CometArrowAllocator` and exports the whole per-partition iterator once as an `ArrowArrayStream`. +**JVM → native (`ScanExec`).** The JVM allocates the Arrow buffers from a child of the +unaccounted root and exports the whole per-partition iterator once as an `ArrowArrayStream`. `ScanExec` imports each batch through `AlignedArrowStreamReader` with `CopyMode::UnpackOrClone`: dictionary columns are unpacked into new native arrays, everything else is an `Arc` clone of the imported buffers. Those bytes stay where Java Arrow put them and are pinned for as long as any native -reference survives. They are invisible to Spark's `TaskMemoryManager`, and `CometArrowAllocator` is -unbounded, so nobody charged for them at allocation time. Whether they are charged _later_ depends +reference survives. They are invisible to Spark's `TaskMemoryManager` at allocation time, deliberately: +this is the export path, so nobody charged for them yet. Whether they are charged _later_ depends on who holds them. DataFusion's `ExternalSorter` reserves `get_reserved_bytes_for_record_batch` for every batch it retains, and the hash join build side reserves `get_record_batch_memory_size` for each incoming batch; both read buffer sizes off the `ArrayData` and apply equally to imported @@ -286,7 +326,7 @@ hard ceiling on the sum of everything in the container. That cgroup counts, amon - JVM non-heap: metaspace, code cache, thread stacks, GC structures, Netty direct buffers, - Spark's own off-heap allocations, - **all of Comet's native allocations**, -- Comet's JVM-side Arrow buffers (`CometArrowAllocator`), +- Comet's JVM-side Arrow buffers (`CometArrowAllocator` and the per-task allocators cut from it), - page cache charged to the cgroup by the container's file I/O, including spill files. Only the first and a portion of the third are visible to Spark's accounting. When the total crosses @@ -321,7 +361,16 @@ much they matter: the fact. There is no runtime value that an operator, a metric, or a policy could read. - **`spark.comet.exec.memoryPool.fraction` is a manual proxy for the gap.** It asks operators to guess a per-workload haircut rather than measuring anything. -- **`CometArrowAllocator` is unbounded** and participates in no budget. +- **`CometArrowAllocator` is unbounded.** Allocations that live and die inside a task are now + reported to Spark's memory manager, so they are no longer invisible, but nothing caps them: the + listener reports without enforcing, and enforcing would mean failing allocations on paths that + cannot fail today. Allocations made off a task, and anything crossing the FFI boundary in either + direction, are not reported at all. +- **Reservation ownership is not handed over at the FFI boundary.** Which side accounts for a + buffer is decided by where it was allocated, not by who holds it, so a buffer allocated for JVM + use and later handed to native is charged on the JVM side while a native operator that retains it + charges the same task again. Buffers allocated for export dodge this only because their + allocation site knows where they are going. - **Buffer and reservation lifetimes are independent across the FFI boundary.** A batch can be resident on either side with no reservation covering it, because reservations are made and withdrawn by individual operators while the bytes outlive them. diff --git a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java index d8dea731353..5c70bf370ca 100644 --- a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java +++ b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java @@ -209,6 +209,10 @@ private static void evaluateInternal( }); assert udf != null : "reflective instantiation returned null for " + udfClassName; + // The unaccounted root on both sides. The inputs wrap native memory that the native side owns + // and frees, and the result below is exported straight back to native, where whichever + // operator retains the batch reserves those same buffers in Comet's pool and so charges the + // same Spark task. Reporting either direction to Spark would count native memory a second time. BufferAllocator allocator = org.apache.comet.package$.MODULE$.CometArrowAllocator(); ValueVector[] inputs = new ValueVector[inputArrayPtrs.length]; diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 4c4e4c668f8..4f7eb4e044b 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -852,6 +852,21 @@ object CometConf extends ShimCometConf { .doubleConf .createWithDefault(1.0) + val COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.arrowAllocator.accounting.enabled") + .category(CATEGORY_TUNING) + .doc( + "When enabled, JVM-side Arrow allocations made inside a Spark task are reported to " + + "Spark's memory manager, so they are visible to Spark's off-heap accounting rather " + + "than to no budget at all. Reporting only: this setting never fails an allocation, " + + "but the bytes do consume the off-heap pool, so other consumers see correspondingly " + + "less headroom. Buffers imported over the Arrow C Data Interface are never reported, " + + "because the memory belongs to the native side. Disable to restore the previous " + + "behaviour of not accounting for these allocations. " + + s"$TUNING_GUIDE.") + .booleanConf + .createWithDefault(true) + val COMET_NATIVE_LOAD_REQUIRED: ConfigEntry[Boolean] = conf("spark.comet.nativeLoadRequired") .category(CATEGORY_EXEC) .doc( diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala index 33e6c0c0355..e277f00e4e7 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala @@ -87,21 +87,27 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim { * Closes the vector on any failure so a partially-initialized tree doesn't leak buffers. */ def allocateOutput(field: Field, numRows: Int, estimatedBytes: Int): FieldVector = { + // The unaccounted root, because this vector is allocated to be handed to native. Its only + // caller is `CometScalaUDFCodegen.evaluate`, whose result `CometUdfBridge` exports over the C + // Data Interface before closing its own reference. Whichever DataFusion operator retains the + // batch reserves these same buffers through Comet's unified pool, which charges the same Spark + // task, so reporting them here as well would reserve the same memory twice. + val allocator: BufferAllocator = CometArrowAllocator val vec: FieldVector = field.getType match { case _: ArrowType.List | _: ArrowType.LargeList | _: ArrowType.FixedSizeList => - val v = new RenamedListVector(field, CometArrowAllocator) + val v = new RenamedListVector(field, allocator) v.initializeChildrenFromFields(field.getChildren) v case _: ArrowType.Map => - val v = new RenamedMapVector(field, CometArrowAllocator) + val v = new RenamedMapVector(field, allocator) v.initializeChildrenFromFields(field.getChildren) v case _: ArrowType.Struct => - val v = new RenamedStructVector(field, CometArrowAllocator) + val v = new RenamedStructVector(field, allocator) v.initializeChildrenFromFields(field.getChildren) v case _ => - field.createVector(CometArrowAllocator).asInstanceOf[FieldVector] + field.createVector(allocator).asInstanceOf[FieldVector] } try { vec.setInitialCapacity(numRows) diff --git a/spark/src/main/scala/org/apache/comet/package.scala b/spark/src/main/scala/org/apache/comet/package.scala index 0eb65c9ba6b..b67aa301364 100644 --- a/spark/src/main/scala/org/apache/comet/package.scala +++ b/spark/src/main/scala/org/apache/comet/package.scala @@ -32,6 +32,19 @@ package object comet { * Until the reference count is zero, the memory will not be released. If the consumer side is * finished later than the close of the allocator, the allocator will think the memory is * leaked. To avoid this, we use a single allocator for the whole execution process. + * + * It carries no allocation listener, so allocating from it directly is not reported to Spark's + * memory manager. That is what memory on either side of the C Data Interface wants. Imported + * buffers wrap memory the native side owns and frees, already charged to Comet's native pool, + * and Arrow's `wrapForeignAllocation` would otherwise report the full buffer capacity as though + * a JVM-side allocation had happened. Buffers allocated to be exported are the mirror image: + * whichever native operator retains the batch reserves them through Comet's unified pool, which + * charges the same Spark task, so reporting them here too would reserve the same memory twice. + * + * Allocations that live and die in the JVM should go through + * `CometTaskArrowAllocator.forCurrentTask()` instead, which cuts a per-task child whose + * listener reports the bytes to Spark. Off a task it hands back this allocator, so the + * driver-side paths are unchanged. */ val CometArrowAllocator = new RootAllocator(Long.MaxValue) diff --git a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala index 173086d2fd7..9ce5a6ed18e 100644 --- a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala +++ b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala @@ -47,7 +47,12 @@ import org.apache.comet.CometArrowAllocator class NativeUtil extends AutoCloseable { import Utils._ - /** Use the global allocator */ + // The unaccounted root, on both sides of the boundary this class straddles. Imported buffers + // wrap memory the native side owns and frees, already charged to Comet's native pool. Everything + // allocated here goes the other way -- the FFI structs, the materialised constant vectors, the + // exported vectors -- and whichever native operator retains a batch reserves those same buffers + // in Comet's pool, which charges the same Spark task. Reporting either direction to Spark would + // count native memory a second time. private val allocator = CometArrowAllocator /** ArrowImporter does not hold any state and does not need to be closed */ diff --git a/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala b/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala index 805eae988e8..8a47fe44557 100644 --- a/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala +++ b/spark/src/main/scala/org/apache/comet/vector/StreamReader.scala @@ -26,17 +26,19 @@ import scala.util.control.NonFatal import org.apache.arrow.vector.VectorSchemaRoot import org.apache.arrow.vector.ipc.{ArrowStreamReader, ReadChannel} import org.apache.arrow.vector.ipc.message.MessageChannelReader +import org.apache.spark.comet.CometTaskArrowAllocator import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.comet.CometArrowAllocator - /** * A reader that consumes Arrow data from an input channel, and produces Comet batches. */ case class StreamReader(channel: ReadableByteChannel, source: String) extends AutoCloseable { + // Decoded into JVM-owned buffers, so it is accounted to the task doing the read. Held in a val + // because the reader and the channel reader must share one allocator. + private val allocator = CometTaskArrowAllocator.forCurrentTask() private val channelReader = - new MessageChannelReader(new ReadChannel(channel), CometArrowAllocator) - private var arrowReader = new ArrowStreamReader(channelReader, CometArrowAllocator) + new MessageChannelReader(new ReadChannel(channel), allocator) + private var arrowReader = new ArrowStreamReader(channelReader, allocator) // Reading the schema allocates the root's vectors, so it can fail with buffers already taken. // No caller holds this reader until its constructor returns, so close it here or nothing will. diff --git a/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala new file mode 100644 index 00000000000..de8e88a7404 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/comet/CometArrowAllocationListener.scala @@ -0,0 +1,288 @@ +/* + * 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.''' This listener's monitor is taken before Spark's and never the other way + * round: [[adjust]] holds ours across `acquireExecutionMemory`, and [[acquire]] additionally + * takes the `TaskMemoryManager` monitor itself, so that the two usage snapshots either side of + * that call cannot be split by another consumer. [[getUsed]] and [[spill]] must therefore stay + * lock-free, because Spark calls both while holding its own monitor: were either to take ours, 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. For the same reason, nothing reachable from a `spill` callback may allocate + * JVM Arrow memory. + */ +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. + * + * Spark reports that figure per task rather than per consumer, so it only measures '''our''' + * grant if nothing else in the task can move it while we are looking. Both snapshots and the + * acquisition therefore run as one transaction under the `TaskMemoryManager` monitor. That is + * the same monitor `acquireExecutionMemory` takes and holds for its whole duration, spills + * included, and it is reentrant, so taking it here only widens that window to cover the two + * reads. Every acquisition in the task funnels through that method, so with it held no other + * consumer can take memory between a snapshot and the call and have it adopted here. + * + * What the monitor does not cover is a release, which reaches the pool without it. Another + * consumer returning memory, or a spill that frees some bytes before throwing, makes the figure + * too small, which is the safe direction: too small degrades to what would have happened + * anyway. `request` bounds it from above. + */ + private def acquire(request: Long): Long = taskMemoryManager.synchronized { + val heldBefore = taskMemoryManager.getMemoryConsumptionForThisTask + try { + taskMemoryManager.acquireExecutionMemory(request, this) + } catch { + case e: InterruptedException => + reportAndReinterrupt(e) + adoptOrphanedGrant(heldBefore, request) + case NonFatal(e) => + warnOnMemoryManagerFailure(e) + adoptOrphanedGrant(heldBefore, request) + case e: SparkOutOfMemoryError => + warnOnMemoryManagerFailure(e) + adoptOrphanedGrant(heldBefore, request) + } + } + + private def adoptOrphanedGrant(heldBefore: Long, request: Long): Long = { + val orphaned = taskMemoryManager.getMemoryConsumptionForThisTask - heldBefore + math.max(0L, math.min(orphaned, request)) + } + + /** + * An interrupt cannot be allowed out of an Arrow callback any more than anything else can, but + * swallowing the cancellation would be wrong too. Spark's execution pool parks in `lock.wait()` + * when a task is below its fair share, so killing a task lands here, and `NonFatal` + * deliberately excludes `InterruptedException`. Re-arming the flag leaves the cancellation for + * the task to observe at its next interruptible point, which is the only place it can act on it + * anyway. + */ + private def reportAndReinterrupt(e: InterruptedException): Unit = { + Thread.currentThread().interrupt() + warnOnMemoryManagerFailure(e) + } +} + +object CometArrowAllocationListener extends Logging { + + /** + * Batching granularity for reservations. Arrow allocates per buffer and + * `acquireExecutionMemory` takes an executor-wide lock, so the reservation is grown and shrunk + * in whole blocks and only block-crossing changes reach Spark. Deliberately not configurable: + * it trades lock chatter against reservation slack and has no plausible per-workload tuning. + */ + private[comet] val BLOCK_SIZE = 1024L * 1024L + + private val shortGrantLogged = new AtomicBoolean(false) + private val memoryManagerFailureLogged = new AtomicBoolean(false) + + private def roundUpToBlock(bytes: Long): Long = + ((bytes + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE + + private def warnOnShortGrant(requested: Long, granted: Long): Unit = { + if (shortGrantLogged.compareAndSet(false, true)) { + logWarning( + s"Spark granted $granted of $requested bytes requested for JVM Arrow allocations. " + + "The allocation proceeds regardless, so this is a reporting gap rather than a failure. " + + s"Set ${CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.key}=false to stop " + + "reporting these allocations to Spark.") + } + } + + private def warnOnMemoryManagerFailure(e: Throwable): Unit = { + if (memoryManagerFailureLogged.compareAndSet(false, true)) { + logWarning( + "Failed to report a JVM Arrow allocation to Spark's memory manager. The allocation " + + "itself is unaffected, so this is a reporting gap rather than a failure, but Spark's " + + "view of these bytes will be short until the task ends. " + + s"Set ${CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.key}=false to stop " + + "reporting these allocations to Spark.", + e) + } + } +} diff --git a/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala new file mode 100644 index 00000000000..fc5074e5a11 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/comet/CometTaskArrowAllocator.scala @@ -0,0 +1,204 @@ +/* + * 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.{ConcurrentHashMap, ConcurrentLinkedQueue} + +import scala.util.control.NonFatal + +import org.apache.arrow.memory.BufferAllocator +import org.apache.spark.{SparkEnv, TaskContext} +import org.apache.spark.internal.Logging +import org.apache.spark.memory.MemoryMode + +import org.apache.comet.{CometArrowAllocator, CometConf} + +/** + * Hands out the Arrow allocator that JVM-owned allocations should use, one per Spark task. + * + * Each task gets a child of `CometArrowAllocator` carrying its own + * [[CometArrowAllocationListener]], so the bytes it hands out are reported to that task's + * `TaskMemoryManager`. The allocator, not the calling thread, is what identifies the owner: + * Arrow's `AllocationListener` is given only a size, and a buffer is released on whichever thread + * drops the last reference, which for anything that reaches native over the C Data Interface is a + * Comet Tokio worker with no task context installed. Child allocators cut from the returned + * allocator inherit its listener, so the paths that make their own children are covered too. + * + * '''This is for buffers the JVM owns.''' Anything allocated to be handed straight to native -- + * `NativeUtil`, the JVM UDF result, `CometNativeArrowSource.stream` -- uses the unaccounted root + * instead, and so does anything imported from native. Native's pool is the authority for bytes + * native holds: whichever DataFusion operator retains the batch reserves those buffers through + * Comet's unified pool, which charges the same Spark task, so reporting them here as well would + * reserve the same memory twice and could reject an allocation that fits. The same fallback + * covers callers with no task to charge -- the driver, broadcast coalescing, the cached batch + * serializer + * -- and Comet's on-heap mode, where charging an off-heap consumer would be wrong. + * + * What this does not fix is a buffer that is used in the JVM and only later handed to native, a + * shuffle-read batch feeding a native operator being the common shape. Its allocation site cannot + * know, so it stays charged here while native may also reserve it. Coordinating reservation + * ownership across the FFI boundary needs changes on both sides; see + * [[https://github.com/apache/datafusion-comet/issues/5997]]. + * + * '''Lifetime.''' The task allocator cannot simply be closed when the task ends. The process-wide + * allocator exists precisely because Arrow buffers can outlive the task that created them, and + * Arrow treats closing an allocator that still owns bytes as a leak. So at task completion the + * Spark reservation is dropped and the allocator is closed only if it has been drained; otherwise + * it is parked and closed by a later task once the stragglers are released. Leaving it open + * indefinitely is not an option: `BaseAllocator` keeps every child in a map until it closes. + */ +object CometTaskArrowAllocator extends Logging { + + private class TaskAllocator( + val allocator: BufferAllocator, + val listener: CometArrowAllocationListener) + + private val perTask = new ConcurrentHashMap[Long, TaskAllocator]() + + /** Allocators whose task has ended but which still own bytes. Guarded by [[closeLock]]. */ + private val awaitingClose = new ConcurrentLinkedQueue[BufferAllocator]() + + private val closeLock = new Object + + /** + * Resolved once per JVM. Read from the `SparkConf` rather than `SQLConf`, because this is + * reached from executor threads where `SQLConf` does not carry Comet's settings. The `Option` + * guard covers tests that install a task context without a `SparkEnv`. + */ + private lazy val accountingEnabled: Boolean = Option(SparkEnv.get).forall { env => + env.conf.getBoolean( + CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.key, + CometConf.COMET_ARROW_ALLOCATOR_ACCOUNTING_ENABLED.defaultValue.get) + } + + /** + * The allocator to use for JVM-owned Arrow buffers on the calling thread. Never null, and never + * an allocator belonging to a task other than the current one. + */ + def forCurrentTask(): BufferAllocator = { + // Cheapest check first, and the one that eliminates the most callers: the driver, broadcast + // coalescing and the cached batch serializer all allocate with no task in scope. + val taskContext = TaskContext.get() + if (taskContext == null) { + CometArrowAllocator + } else { + val existing = perTask.get(taskContext.taskAttemptId()) + if (existing != null) existing.allocator else create(taskContext) + } + } + + private def create(taskContext: TaskContext): BufferAllocator = { + val taskMemoryManager = taskContext.taskMemoryManager() + // Comet's on-heap mode exists so the Spark SQL suite can run without off-heap memory + // configured, and charging an off-heap consumer there would be wrong. Note this differs from + // CometUnifiedShuffleMemoryAllocator, which throws in that situation; throwing here would + // break those tests. + if (!accountingEnabled || taskMemoryManager == null || + taskMemoryManager.getTungstenMemoryMode != MemoryMode.OFF_HEAP) { + return CometArrowAllocator + } + + val taskAttemptId = taskContext.taskAttemptId() + closeDrained() + val listener = new CometArrowAllocationListener(taskMemoryManager) + val allocator = CometArrowAllocator + .newChildAllocator(s"comet-task-$taskAttemptId", listener, 0L, Long.MaxValue) + val created = new TaskAllocator(allocator, listener) + val previous = perTask.putIfAbsent(taskAttemptId, created) + if (previous != null) { + // Lost the race with another thread in the same task; the loser's allocator is untouched. + closeQuietly(allocator) + previous.allocator + } else { + // Deliberately outside the map operation: `addTaskCompletionListener` runs the callback + // inline if the task has already completed, and that callback removes from this same map, + // which would be a recursive update inside a mapping function. + taskContext.addTaskCompletionListener[Unit](_ => taskCompleted(taskAttemptId)) + // ...and if it did run inline, the allocator just created is already closed, so hand back + // the root rather than something the caller cannot allocate from. + if (perTask.containsKey(taskAttemptId)) allocator else CometArrowAllocator + } + } + + private def taskCompleted(taskAttemptId: Long): Unit = { + val finished = perTask.remove(taskAttemptId) + if (finished != null) { + finished.listener.taskCompleted() + closeLock.synchronized { + if (!tryClose(finished.allocator)) { + awaitingClose.add(finished.allocator) + } + } + } + closeDrained() + } + + /** Closes any parked allocator whose stragglers have since been released. */ + private def closeDrained(): Unit = { + if (!awaitingClose.isEmpty) { + closeLock.synchronized { + val parked = awaitingClose.iterator() + while (parked.hasNext) { + if (tryClose(parked.next())) { + parked.remove() + } + } + } + } + } + + /** Closes the allocator if it has been drained. Returns false if it must stay open. */ + private def tryClose(allocator: BufferAllocator): Boolean = { + if (allocator.getAllocatedMemory != 0L) { + false + } else { + closeQuietly(allocator) + true + } + } + + private def closeQuietly(allocator: BufferAllocator): Unit = { + try { + allocator.close() + } catch { + case NonFatal(e) => + // Closing is bookkeeping: the bytes are already gone and the Spark reservation is already + // released, so a failure here must not propagate into a task completion listener. + logWarning(s"Failed to close Arrow allocator ${allocator.getName}", e) + } + } + + /** Number of tasks currently holding an accounted allocator. Visible for testing. */ + private[comet] def trackedTaskCount: Int = perTask.size() + + /** Number of finished tasks whose allocator is still draining. Visible for testing. */ + private[comet] def awaitingCloseCount: Int = awaitingClose.size() + + /** The listener accounting for the given task, if it has one. Visible for testing. */ + private[comet] def listenerForTask( + taskAttemptId: Long): Option[CometArrowAllocationListener] = { + Option(perTask.get(taskAttemptId)).map(_.listener) + } + + /** Bytes currently reserved with Spark on behalf of the given task. Visible for testing. */ + private[comet] def reservedBytesForTask(taskAttemptId: Long): Long = { + listenerForTask(taskAttemptId).map(_.reservedBytes).getOrElse(0L) + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala index d0e7d87911a..a9ef74f5e18 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala @@ -25,6 +25,7 @@ import scala.collection.JavaConverters._ import scala.util.control.NonFatal import org.apache.spark.TaskContext +import org.apache.spark.comet.CometTaskArrowAllocator import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, GenericInternalRow, IsNotNull, IsNull, UnsafeProjection} @@ -39,7 +40,7 @@ import org.apache.spark.storage.StorageLevel import org.apache.spark.unsafe.types.UTF8String import org.apache.spark.util.io.ChunkedByteBuffer -import org.apache.comet.{CometArrowAllocator, DataTypeSupport} +import org.apache.comet.DataTypeSupport /** * Cached batch format used when Comet writes Spark in-memory cache data. @@ -356,7 +357,10 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { Utils.serializeBatchColumns(batch) } else { val arrowBatch = - CometArrowConverters.columnarBatchToArrowBatch(batch, arrowSchema, CometArrowAllocator) + CometArrowConverters.columnarBatchToArrowBatch( + batch, + arrowSchema, + CometTaskArrowAllocator.forCurrentTask()) try Utils.serializeBatchColumns(arrowBatch) finally arrowBatch.close() } @@ -624,7 +628,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // the Unix epoch regardless of session timezone, so no values are converted. It also // matches Comet's native schema, avoiding a cast at the native boundary. CometArrowStream.NATIVE_TIMEZONE, - CometArrowAllocator) + CometTaskArrowAllocator.forCurrentTask()) encodeBatches(iter, schema) } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala index a60e2811ecb..d5bc9c41a33 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala @@ -27,6 +27,7 @@ import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.vector.ipc.ArrowReader import org.apache.arrow.vector.types.pojo.{Field, FieldType, Schema} import org.apache.spark.TaskContext +import org.apache.spark.comet.CometTaskArrowAllocator import org.apache.spark.internal.Logging import org.apache.spark.rdd.RDD import org.apache.spark.sql.comet.util.Utils @@ -230,6 +231,10 @@ object CometArrowStream extends Logging { name: String, readerFactory: BufferAllocator => ArrowReader): Iterator[ArrowArrayStream] = { val context = TaskContext.get() + // A child of the unaccounted root, not of the task allocator: every byte allocated here is + // exported to native and exists only for native to consume. Whichever native operator retains + // the batch reserves those same buffers in Comet's pool, which charges the same Spark task, so + // reporting them on this side too would reserve the same memory twice. val allocator = CometArrowAllocator.newChildAllocator(name, 0, Long.MaxValue) var reader: ArrowReader = null var arrowStream: ArrowArrayStream = null @@ -274,7 +279,10 @@ object CometArrowStream extends Logging { name: String, readerFactory: BufferAllocator => ArrowReader): Iterator[ColumnarBatch] = { val context = TaskContext.get() - val allocator = CometArrowAllocator.newChildAllocator(name, 0, Long.MaxValue) + // Accounted, unlike `stream`: these batches are consumed in the JVM, so nothing on the native + // side reserves them. + val allocator = + CometTaskArrowAllocator.forCurrentTask().newChildAllocator(name, 0, Long.MaxValue) val reader = try readerFactory(allocator) catch { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index d70fdab35e6..ff95daef47c 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -373,7 +373,8 @@ object Utils extends CometTypeShim with Logging { return (Array.empty, 0L, 0L) } - val allocator = org.apache.comet.CometArrowAllocator + val allocator = org.apache.spark.comet.CometTaskArrowAllocator + .forCurrentTask() .newChildAllocator("broadcast-coalesce", 0, Long.MaxValue) try { var targetRoot: VectorSchemaRoot = null @@ -567,7 +568,7 @@ object Utils extends CometTypeShim with Logging { cv.dataType(), rows, s"_const_$index", - org.apache.comet.CometArrowAllocator, + org.apache.spark.comet.CometTaskArrowAllocator.forCurrentTask(), "UTC") (materialized, None) diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala index fecf781baac..3229f11d7c3 100644 --- a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala +++ b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala @@ -34,13 +34,13 @@ import org.apache.arrow.vector.ipc.message.{ArrowFieldNode, ArrowRecordBatch, Me import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType} import org.apache.spark.{SparkEnv, TaskContext} import org.apache.spark.api.python.{BasePythonRunner, PythonRDD, PythonWorker, SpecialLengths} +import org.apache.spark.comet.CometTaskArrowAllocator import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StructType import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} -import org.apache.comet.CometArrowAllocator import org.apache.comet.vector.{CometDecodedVector, CometDictionaryVector, CometVector, CometVectorUtils} /** @@ -124,7 +124,9 @@ private[python] trait CometArrowPythonRunnerBase new Writer(env, worker, inputIterator, partitionIndex, context) { private val allocator = - CometArrowAllocator.newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue) + CometTaskArrowAllocator + .forCurrentTask() + .newChildAllocator(s"stdout writer for $pythonExec", 0, Long.MaxValue) private var batches = inputIterator.flatten // Upstream owns this batch. Even hasNext may close it and reuse its native buffers, so // leave the upstream iterators untouched until all ranges have been serialized. @@ -264,7 +266,9 @@ private[python] trait CometArrowPythonRunnerBase new ReaderIterator(stream, writer, startTime, env, worker, pid, releasedOrClosed, context) { private val allocator = - CometArrowAllocator.newChildAllocator(s"stdin reader for $pythonExec", 0, Long.MaxValue) + CometTaskArrowAllocator + .forCurrentTask() + .newChildAllocator(s"stdin reader for $pythonExec", 0, Long.MaxValue) private var reader: ArrowStreamReader = _ private var root: VectorSchemaRoot = _ private var batchLoaded = true diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala new file mode 100644 index 00000000000..07e2896f2f5 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerBenchmark.scala @@ -0,0 +1,312 @@ +/* + * 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.Properties +import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong} + +import org.apache.arrow.memory.{ArrowBuf, BufferAllocator, RootAllocator} +import org.apache.spark.{SparkConf, TaskContext, TaskContextImpl} +import org.apache.spark.benchmark.{Benchmark, BenchmarkBase} +import org.apache.spark.executor.TaskMetrics +import org.apache.spark.memory.{MemoryConsumer, MemoryManager, MemoryMode, TaskMemoryManager, TestMemoryManager} + +/** + * Measures what reporting JVM Arrow allocations to Spark costs, since it is on by default. + * + * Each case runs the same allocate/release loop twice: once against a plain `RootAllocator`, + * which is what Comet did before [[CometArrowAllocationListener]] existed and what + * `spark.comet.arrowAllocator.accounting.enabled=false` restores, and once against a task + * allocator carrying the listener. Reservation call counts are printed under each table, because + * the interesting variable is not the per-buffer bookkeeping but how often a buffer size crosses + * a block boundary and has to go into `TaskMemoryManager` at all. + * + * To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.comet.CometArrowAllocationListenerBenchmark + * }}} + */ +object CometArrowAllocationListenerBenchmark extends BenchmarkBase { + + private val blockSize = 1024L * 1024L + private val poolBytes = 1024L * 1024L * 1024L + + override def runBenchmarkSuite(mainArgs: Array[String]): Unit = { + runBenchmark("JVM Arrow allocations reported to Spark") { + // Many sub-block buffers, the shape a codegen output vector's validity and offset buffers + // take. None of them comes close to a block, so the listener should never reach Spark after + // the first one. + allocateAndRelease("small buffers", bufferSize = 128L, buffersPerIteration = 512) + // A wide batch: many medium buffers alive at once, one block boundary crossed per few + // buffers on the way up and the same on the way down. + allocateAndRelease( + "wide batch buffers", + bufferSize = 64L * 1024L, + buffersPerIteration = 128) + // Worst case for the block batching: every allocation crosses a boundary, so every one of + // them takes the executor-wide lock in `acquireExecutionMemory`. + allocateAndRelease("block-sized buffers", bufferSize = blockSize, buffersPerIteration = 8) + // Same, with a second thread reserving from the same constrained pool the way Comet's native + // side does through CometTaskMemoryManager. + allocateUnderNativePressure() + // Per call site rather than per buffer, but it is the cost the root allocator `val` did not + // have: a TaskContext lookup and a concurrent map read. + allocatorLookup() + } + } + + private def allocateAndRelease( + name: String, + bufferSize: Long, + buffersPerIteration: Int): Unit = { + val benchmark = + new Benchmark( + s"$name (${buffersPerIteration}x$bufferSize)", + buffersPerIteration, + output = output) + + // Both allocators are built once, outside the timed body, so what is measured is the + // steady-state cost of allocating and releasing rather than the cost of standing a task up. + val root = new RootAllocator(Long.MaxValue) + try { + withTaskAllocator() { (accounted, memory) => + benchmark.addCase("not accounted") { _ => + churn(root, bufferSize, buffersPerIteration) + } + benchmark.addCase("accounted") { _ => + churn(accounted, bufferSize, buffersPerIteration) + } + benchmark.run() + + // One more round with the counters zeroed, to report how often a single iteration reaches + // the memory manager. That, rather than the per-buffer bookkeeping, is the cost that + // scales with buffer size. + memory.reset() + churn(accounted, bufferSize, buffersPerIteration) + writeLine(s" accounted: ${memory.summary(buffersPerIteration)} per iteration") + } + } finally { + root.close() + } + } + + private def allocateUnderNativePressure(): Unit = { + val buffersPerIteration = 8 + val benchmark = new Benchmark( + s"block-sized buffers under native pressure (${buffersPerIteration}x$blockSize)", + buffersPerIteration, + output = output) + + val root = new RootAllocator(Long.MaxValue) + try { + // Two blocks of pool for both arms, with the pressure thread running throughout. In the + // unaccounted arm only that thread touches the pool, which is the point: the delta is what + // the Arrow side adds once it competes for the same budget. + withTaskAllocator(poolBytes = blockSize * 2) { (accounted, memory) => + withNativePressure(memory) { + benchmark.addCase("not accounted") { _ => + churn(root, blockSize, buffersPerIteration) + } + benchmark.addCase("accounted") { _ => + churn(accounted, blockSize, buffersPerIteration) + } + benchmark.run() + + memory.reset() + churn(accounted, blockSize, buffersPerIteration) + writeLine(s" accounted: ${memory.summary(buffersPerIteration)} per iteration") + } + } + } finally { + root.close() + } + } + + private def allocatorLookup(): Unit = { + val lookupsPerIteration = 100000 + val benchmark = + new Benchmark("allocator lookup", lookupsPerIteration.toLong, output = output) + + // The "before" shape: call sites read a package-object `val`, which the JIT folds away + // entirely. The interesting number is therefore the absolute cost of the second case. + benchmark.addCase("process-wide val") { _ => + var i = 0 + var sink = 0 + while (i < lookupsPerIteration) { + sink += System.identityHashCode(org.apache.comet.CometArrowAllocator) + i += 1 + } + assert(sink != Int.MinValue) + } + benchmark.addCase("forCurrentTask()") { _ => + withTaskAllocator() { (_, _) => + var i = 0 + var sink = 0 + while (i < lookupsPerIteration) { + sink += System.identityHashCode(CometTaskArrowAllocator.forCurrentTask()) + i += 1 + } + assert(sink != Int.MinValue) + } + } + + benchmark.run() + } + + /** + * Allocates the whole set, then releases it, so the peak is what the reservation has to cover. + */ + private def churn(allocator: BufferAllocator, bufferSize: Long, count: Int): Unit = { + val buffers = new Array[ArrowBuf](count) + var i = 0 + while (i < count) { + buffers(i) = allocator.buffer(bufferSize) + i += 1 + } + i = 0 + while (i < count) { + buffers(i).close() + i += 1 + } + } + + /** + * Runs the body against a task allocator, torn down afterwards, so that repeated iterations do + * not accumulate reservations or allocators. + */ + private def withTaskAllocator[T](poolBytes: Long = poolBytes)( + f: (BufferAllocator, CountingTaskMemoryManager) => T): T = { + val memory = newTaskMemoryManager(poolBytes) + val context = newTaskContext(memory) + val previous = TaskContext.get() + TaskContext.setTaskContext(context) + try { + f(CometTaskArrowAllocator.forCurrentTask(), memory) + } finally { + try { + context.markTaskCompleted(None) + memory.cleanUpAllAllocatedMemory() + } finally { + if (previous == null) TaskContext.unset() else TaskContext.setTaskContext(previous) + } + } + } + + /** Hammers the same pool from another thread, the way native reservations do. */ + private def withNativePressure[T](memory: TaskMemoryManager)(f: => T): T = { + val stop = new AtomicBoolean(false) + val consumer = new NativeLikeConsumer(memory) + val thread = new Thread(() => { + while (!stop.get()) { + val granted = consumer.reserve(blockSize) + consumer.release(granted) + } + }) + thread.setDaemon(true) + thread.setName("native-reservations") + thread.start() + try f + finally { + stop.set(true) + thread.join() + } + } + + private val nextTaskAttemptId = new AtomicLong(1L) + + private def newTaskMemoryManager(poolBytes: Long): CountingTaskMemoryManager = { + val conf = new SparkConf(false) + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", poolBytes.toString) + val memoryManager = new TestMemoryManager(conf) + memoryManager.limit(poolBytes) + new CountingTaskMemoryManager(memoryManager, nextTaskAttemptId.getAndIncrement()) + } + + private def newTaskContext(memory: CountingTaskMemoryManager): TaskContextImpl = { + new TaskContextImpl( + stageId = 0, + stageAttemptNumber = 0, + partitionId = 0, + numPartitions = 1, + taskAttemptId = memory.getTaskAttemptId, + attemptNumber = 0, + taskMemoryManager = memory, + localProperties = new Properties, + metricsSystem = null, + taskMetrics = TaskMetrics.empty, + cpus = 1, + resources = Map.empty) + } + + private def writeLine(line: String): Unit = { + // scalastyle:off println + println(line) + // scalastyle:on println + output.foreach(_.write(s"$line\n".getBytes("UTF-8"))) + } + + private class CountingTaskMemoryManager(memoryManager: MemoryManager, taskAttemptId: Long) + extends TaskMemoryManager(memoryManager, taskAttemptId) { + private val acquires = new AtomicLong(0L) + private val releases = new AtomicLong(0L) + + def getTaskAttemptId: Long = taskAttemptId + + // Only the Arrow listener's own calls are counted, so the pressure thread's traffic does not + // land in the reported figure. + override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = { + if (consumer.isInstanceOf[CometArrowAllocationListener]) acquires.incrementAndGet() + super.acquireExecutionMemory(required, consumer) + } + + override def releaseExecutionMemory(size: Long, consumer: MemoryConsumer): Unit = { + if (consumer.isInstanceOf[CometArrowAllocationListener]) releases.incrementAndGet() + super.releaseExecutionMemory(size, consumer) + } + + def reset(): Unit = { + acquires.set(0L) + releases.set(0L) + } + + def summary(buffers: Int): String = + s"${acquires.get()} acquire and ${releases.get()} release calls for $buffers buffers" + } + + /** Stands in for `CometTaskMemoryManager`: reserves from Spark directly and never spills. */ + private class NativeLikeConsumer(memory: TaskMemoryManager) + extends MemoryConsumer(memory, 0L, MemoryMode.OFF_HEAP) { + private val reserved = new AtomicLong(0L) + override def spill(size: Long, trigger: MemoryConsumer): Long = 0L + override def getUsed: Long = reserved.get() + def reserve(bytes: Long): Long = { + val granted = memory.acquireExecutionMemory(bytes, this) + reserved.addAndGet(granted) + granted + } + def release(bytes: Long): Unit = { + if (bytes > 0L) { + reserved.addAndGet(-bytes) + memory.releaseExecutionMemory(bytes, this) + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala new file mode 100644 index 00000000000..c5bbe055b3d --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/comet/CometArrowAllocationListenerSuite.scala @@ -0,0 +1,625 @@ +/* + * 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.io.{InterruptedIOException, IOException} +import java.util.Properties +import java.util.concurrent.atomic.{AtomicLong, AtomicReference} + +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.arrow.memory.BufferAllocator +import org.apache.spark.{SparkConf, TaskContext, TaskContextImpl} +import org.apache.spark.executor.TaskMetrics +import org.apache.spark.memory.{MemoryConsumer, MemoryManager, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager, TestMemoryManager} + +import org.apache.comet.CometArrowAllocator + +/** + * Tests that JVM Arrow allocations are reported to Spark, that they are reported against the task + * that made them rather than whichever task happens to be on the releasing thread, and, just as + * importantly, that the paths where they cannot be reported fail quietly rather than throwing. + * Arrow allocation on these paths cannot fail today and this listener must not change that. + */ +class CometArrowAllocationListenerSuite extends AnyFunSuite { + + private val blockSize = CometArrowAllocationListener.BLOCK_SIZE + private val poolBytes = 64L * 1024 * 1024 + + /** Task attempt ids are keys in a process-wide map, so no two tests may share one. */ + private val nextTaskAttemptId = new AtomicLong(1000L) + + // --------------------------------------------------------------------------------------------- + // Reservation arithmetic. Driven through the listener directly, since Arrow's rounding policy + // would otherwise decide the sizes under test. + // --------------------------------------------------------------------------------------------- + + test("allocations are charged to the current task in whole blocks") { + withTask() { task => + val listener = new CometArrowAllocationListener(task.taskMemoryManager) + + // Far smaller than a block, so the reservation should round up to exactly one block. + listener.onAllocation(128L) + assert(listener.reservedBytes == blockSize) + + // Still inside the first block, so Spark is not asked again. + listener.onAllocation(1024L) + assert(listener.reservedBytes == blockSize) + + listener.taskCompleted() + } + } + + test("a request larger than a block rounds up to a block multiple") { + withTask() { task => + val listener = new CometArrowAllocationListener(task.taskMemoryManager) + listener.onAllocation(blockSize * 3 + 7L) + // Rounded up rather than sized to the exact deficit, so growth leaves headroom and the next + // small allocation does not go straight back into Spark. + assert(listener.reservedBytes == blockSize * 4) + listener.taskCompleted() + } + } + + test("releasing returns whole blocks to Spark") { + withTask() { task => + val listener = new CometArrowAllocationListener(task.taskMemoryManager) + listener.onAllocation(blockSize * 2) + assert(listener.reservedBytes == blockSize * 2) + + listener.onRelease(blockSize * 2) + assert(listener.reservedBytes == 0L) + listener.taskCompleted() + } + } + + // --------------------------------------------------------------------------------------------- + // Which allocator is handed out, and what it charges. + // --------------------------------------------------------------------------------------------- + + test("a real Arrow allocation from the task allocator is charged to that task") { + withTask() { task => + val allocator = CometTaskArrowAllocator.forCurrentTask() + assert(allocator ne CometArrowAllocator) + val buf = allocator.buffer(blockSize) + try { + assert(reservedFor(task) == blockSize) + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize) + } finally { + buf.close() + } + assert(reservedFor(task) == 0L) + } + } + + test("a child of the task allocator is charged to the same task") { + withTask() { task => + // The Python runner and the native Arrow source cut their own children. Arrow passes the + // parent's listener down, so they are accounted without knowing anything about it. + val child = + CometTaskArrowAllocator.forCurrentTask().newChildAllocator("probe", 0L, Long.MaxValue) + try { + val buf = child.buffer(blockSize) + try { + assert(reservedFor(task) == blockSize) + } finally { + buf.close() + } + } finally { + child.close() + } + assert(reservedFor(task) == 0L) + } + } + + test("the process-wide root is not accounted, which is what the FFI paths rely on") { + withTask() { task => + // Establish the task allocator first, so this asserts "not charged" rather than "no task". + CometTaskArrowAllocator.forCurrentTask() + // Both directions across the C Data Interface use the listener-less root: imported buffers + // wrap memory the native side owns, and buffers allocated for export are reserved again by + // whichever native operator retains the batch, through a pool that charges the same task. + val buf = CometArrowAllocator.buffer(blockSize) + try { + assert(reservedFor(task) == 0L) + } finally { + buf.close() + } + } + } + + test("no active task uses the unaccounted root allocator") { + TaskContext.unset() + // Broadcast coalescing and the cached batch serializer can allocate off a task thread. + assert(CometTaskArrowAllocator.forCurrentTask() eq CometArrowAllocator) + } + + test("on-heap mode uses the unaccounted root allocator") { + withTask(offHeap = false) { task => + // Comet's on-heap mode exists so the Spark SQL suite can run without off-heap memory. + // Charging an off-heap consumer there would be wrong. + assert(CometTaskArrowAllocator.forCurrentTask() eq CometArrowAllocator) + assert(CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).isEmpty) + } + } + + // --------------------------------------------------------------------------------------------- + // Ownership: a release is attributed to the task that allocated, not to the releasing thread. + // --------------------------------------------------------------------------------------------- + + test("a release on a thread with no task context is charged to the allocating task") { + withTask() { task => + val buf = CometTaskArrowAllocator.forCurrentTask().buffer(blockSize) + assert(reservedFor(task) == blockSize) + + // This is what happens when a shuffle-read batch is handed on to a native operator: native + // pins it and drops it later from a Tokio worker with no task context installed. Reading + // TaskContext in onRelease would ignore this release and leave the task charged for memory + // it had already freed. + onDetachedThread(buf.close()) + + assert(reservedFor(task) == 0L) + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + } + } + + test("a release under a different task does not touch that task's accounting") { + withTask() { taskA => + val bufA = CometTaskArrowAllocator.forCurrentTask().buffer(blockSize) + withTask() { taskB => + val bufB = CometTaskArrowAllocator.forCurrentTask().buffer(blockSize) + assert(reservedFor(taskB) == blockSize) + + // A's buffer, released while B's context is on the thread. B must not pay for it. + bufA.close() + + assert(reservedFor(taskB) == blockSize) + assert(reservedFor(taskA) == 0L) + bufB.close() + } + } + } + + // --------------------------------------------------------------------------------------------- + // Task completion, and buffers that outlive their task. + // --------------------------------------------------------------------------------------------- + + test("task completion releases the whole reservation and closes the allocator") { + val task = newTask() + val allocatorName = withInstalledTask(task) { + val allocator = CometTaskArrowAllocator.forCurrentTask() + val buf = allocator.buffer(blockSize) + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize) + buf.close() + + task.context.markTaskCompleted(None) + + assert(CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).isEmpty) + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + allocator.getName + } + assert(!rootChildNames().contains(allocatorName)) + } + + test("a buffer outliving its task parks the allocator until it is released") { + val task = newTask() + withInstalledTask(task) { + val allocator = CometTaskArrowAllocator.forCurrentTask() + val buf = allocator.buffer(blockSize) + + task.context.markTaskCompleted(None) + + // The reservation goes back to Spark even though the buffer is still alive: the task is + // over, and leaving it charged would be reported as a Spark memory leak. + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + assert(CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).isEmpty) + // Arrow treats closing an allocator that still owns bytes as a leak, so it has to stay open. + assert(rootChildNames().contains(allocator.getName)) + + // A late release is ignored rather than charged to whoever is running by then... + buf.close() + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + + // ...and the drained allocator is reaped by the next task, so the root does not accumulate + // one child per task attempt. + withTask() { _ => CometTaskArrowAllocator.forCurrentTask() } + assert(!rootChildNames().contains(allocator.getName)) + } + } + + // --------------------------------------------------------------------------------------------- + // Failure containment: Spark's acquisition is fallible, Arrow's callbacks are not allowed to be. + // --------------------------------------------------------------------------------------------- + + for ((label, failure) <- Seq( + "an I/O failure" -> new IOException("spill failed"), + "an interrupt" -> new InterruptedIOException("task killed"))) { + test(s"$label while spilling does not fail or leak the Arrow allocation") { + // Exactly one block of budget, already taken by a consumer that refuses to spill, so the + // listener's acquisition has to go through Spark's spill path and comes back throwing. + withTask(pool = blockSize) { task => + val hostile = new FailingSpillConsumer(task.taskMemoryManager, failure) + assert(hostile.take(blockSize) == blockSize) + + val allocator = CometTaskArrowAllocator.forCurrentTask() + // `BaseAllocator.buffer` marks the allocation successful before calling `onAllocation`, so + // throwing from the listener would lose this buffer: Arrow neither returns nor frees it. + val buf = allocator.buffer(blockSize) + try { + assert(allocator.getAllocatedMemory == blockSize) + // Nothing was reserved, which is what says the acquisition really did go down the spill + // path and throw rather than quietly succeeding and making this test vacuous. + assert(reservedFor(task) == 0L) + } finally { + buf.close() + } + // Zero here is the leak check: a buffer Arrow created but never handed back would still + // be counted. + assert(allocator.getAllocatedMemory == 0L) + } + } + } + + test("an interrupt while spilling is re-armed rather than thrown or swallowed") { + // Spark's execution pool parks in `lock.wait()` when a task is below its fair share, so a task + // kill raises a plain InterruptedException out of `acquireExecutionMemory`. `NonFatal` excludes + // it, so before this it escaped `onAllocation` and Arrow lost the buffer it had just created. + // TestMemoryManager never parks, so the interrupt is injected through a failing spill instead. + withTask(pool = blockSize) { task => + val hostile = + new FailingSpillConsumer(task.taskMemoryManager, new InterruptedException("task killed")) + assert(hostile.take(blockSize) == blockSize) + + val allocator = CometTaskArrowAllocator.forCurrentTask() + val buf = allocator.buffer(blockSize) + try { + assert(allocator.getAllocatedMemory == blockSize) + assert(reservedFor(task) == 0L) + } finally { + buf.close() + } + assert(allocator.getAllocatedMemory == 0L) + // Cleared here as well as asserted, so the flag does not leak into the next test. + assert(Thread.interrupted(), "the interrupt was swallowed instead of being re-armed") + } + } + + test("a partial grant lost to a failing spill is adopted rather than stranded") { + // One block already taken, one still in the pool, and a two-block request: Spark hands over the + // block it has and only then asks the other consumer to spill, which throws. It never reports + // the block it already took, so nothing would release it before the task ended. + withTask(pool = blockSize * 2) { task => + val hostile = + new FailingSpillConsumer(task.taskMemoryManager, new IOException("spill failed")) + assert(hostile.take(blockSize) == blockSize) + + val allocator = CometTaskArrowAllocator.forCurrentTask() + val buf = allocator.buffer(blockSize * 2) + try { + assert(reservedFor(task) == blockSize) + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize * 2) + } finally { + buf.close() + } + assert(reservedFor(task) == 0L) + // Only the other consumer's block is left. Without adopting the orphan this would still be + // two blocks, with one of them charged to the task and owned by nobody. + assert(task.taskMemoryManager.getMemoryConsumptionForThisTask == blockSize) + } + } + + test("a concurrent acquisition is not adopted as this listener's lost grant") { + // The orphan is measured as a change in the task's total consumption, which Spark reports per + // task rather than per consumer. If another consumer could acquire between the snapshot taken + // before the acquisition and the acquisition itself, its bytes would be adopted here and handed + // back when this listener next shrank, leaving two consumers holding the same bytes between + // them. Both snapshots and the call therefore run as one transaction under the + // TaskMemoryManager monitor. This forces the interleaving that transaction exists to exclude. + val spare = 1024L + val task = newTask(pool = blockSize * 2 + spare) + withInstalledTask(task) { + val hostile = + new FailingSpillConsumer(task.taskMemoryManager, new IOException("spill failed")) + assert(hostile.take(blockSize) == blockSize) + + val interloper = new PlainConsumer(task.taskMemoryManager) + // Once the acquisition has taken what was left, the pool is empty and Spark answers the + // interloper by asking the hostile consumer to spill, which throws. That is a legitimate + // outcome for the interloper and not what is under test here; what it ends up holding is. + val interloperThread = daemonThread("interloper") { + try interloper.take(spare) + catch { case _: SparkOutOfMemoryError => } + } + + // Fires once, in place of the snapshot taken before the acquisition: exactly the window the + // transaction has to close. The interloper either runs to completion here, which is the bug, + // or blocks on the monitor the acquisition is holding, which is the fix. + task.snapshotHook.set(() => { + interloperThread.start() + awaitBlockedOrFinished(interloperThread) + }) + + val allocator = CometTaskArrowAllocator.forCurrentTask() + val buf = allocator.buffer(blockSize * 2) + try { + interloperThread.join(30000L) + assert( + !interloperThread.isAlive, + "the interloper never finished; the transaction deadlocked") + // Nothing is claimed twice: what the listener adopted has to fit alongside what the other + // two consumers hold. Without the transaction the listener adopts the interloper's bytes on + // top of its own grant, and this sum comes out over what the task actually holds. + assert( + listenerFor(task).reservedBytes + hostile.getUsed + interloper.getUsed == + task.taskMemoryManager.getMemoryConsumptionForThisTask, + "the listener adopted bytes belonging to another consumer") + } finally { + buf.close() + } + } + } + + // --------------------------------------------------------------------------------------------- + // Lock order. Spark calls getUsed and spill while holding the TaskMemoryManager monitor, and the + // listener holds its own monitor while waiting for that one. + // --------------------------------------------------------------------------------------------- + + test("the usage snapshot does not take the reservation monitor") { + withTask() { task => + val buf = CometTaskArrowAllocator.forCurrentTask().buffer(blockSize) + try { + val listener = CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).get + val used = new AtomicLong(-1L) + val spilled = new AtomicLong(-1L) + listener.synchronized { + // A native reservation arriving through CometTaskMemoryManager on a Tokio thread holds + // Spark's monitor here. If either call waited on this one, it would deadlock against an + // Arrow allocation on the same task that already holds this monitor and wants Spark's. + val probe = new Thread(() => { + used.set(listener.getUsed) + spilled.set(listener.spill(blockSize, listener)) + }) + probe.setDaemon(true) + probe.setName("lock-order-probe") + probe.start() + probe.join(30000L) + assert(!probe.isAlive, "getUsed or spill blocked on the reservation monitor") + } + assert(used.get == blockSize) + assert(spilled.get == 0L) + } finally { + buf.close() + } + } + } + + test("concurrent Arrow and native reservations make progress") { + // Two blocks of budget shared by both consumers, so most requests are short and Spark walks + // its consumer list, calling getUsed on the Arrow listener while holding its own monitor. + withTask(pool = blockSize * 2) { task => + val allocator = CometTaskArrowAllocator.forCurrentTask() + val native = new NativeLikeConsumer(task.taskMemoryManager) + val failure = new AtomicReference[Throwable]() + + val arrowThread = loopingThread("arrow-allocations", failure) { + val buf = allocator.buffer(blockSize) + buf.close() + } + val nativeThread = loopingThread("native-reservations", failure) { + native.release(native.reserve(blockSize)) + } + + Seq(arrowThread, nativeThread).foreach(_.start()) + Seq(arrowThread, nativeThread).foreach { t => + t.join(60000L) + assert(!t.isAlive, s"${t.getName} did not finish; concurrent reservations deadlocked") + } + + Option(failure.get).foreach(e => fail("a worker failed", e)) + assert(allocator.getAllocatedMemory == 0L) + native.release(native.used()) + } + } + + // --------------------------------------------------------------------------------------------- + // Fixtures. + // --------------------------------------------------------------------------------------------- + + /** Holds memory and refuses to give it back, so `trySpillAndAcquire` throws on its behalf. */ + private class FailingSpillConsumer(tmm: TaskMemoryManager, failure: Exception) + extends MemoryConsumer(tmm, 0L, MemoryMode.OFF_HEAP) { + def take(bytes: Long): Long = acquireMemory(bytes) + override def spill(size: Long, trigger: MemoryConsumer): Long = throw failure + } + + /** Any other consumer in the task: takes memory, and cannot give it back. */ + private class PlainConsumer(tmm: TaskMemoryManager) + extends MemoryConsumer(tmm, 0L, MemoryMode.OFF_HEAP) { + def take(bytes: Long): Long = acquireMemory(bytes) + override def spill(size: Long, trigger: MemoryConsumer): Long = 0L + } + + /** + * Runs one action immediately after the usage snapshot the listener takes before asking Spark + * for memory, so a test can drive what happens in the window between that snapshot and the + * acquisition. The real value is read first, which is what makes the window the one under test: + * running the action before the read would fold whatever it does into the snapshot itself. + */ + private class HookedTaskMemoryManager( + memoryManager: MemoryManager, + taskAttemptId: Long, + hook: AtomicReference[Runnable]) + extends TaskMemoryManager(memoryManager, taskAttemptId) { + override def getMemoryConsumptionForThisTask(): Long = { + val held = super.getMemoryConsumptionForThisTask() + val pending = hook.getAndSet(null) + if (pending != null) pending.run() + held + } + } + + /** Stands in for `CometTaskMemoryManager`: reserves from Spark directly and never spills. */ + private class NativeLikeConsumer(tmm: TaskMemoryManager) + extends MemoryConsumer(tmm, 0L, MemoryMode.OFF_HEAP) { + private val reserved = new AtomicLong(0L) + override def spill(size: Long, trigger: MemoryConsumer): Long = 0L + override def getUsed: Long = reserved.get() + def used(): Long = reserved.get() + def reserve(bytes: Long): Long = { + val granted = tmm.acquireExecutionMemory(bytes, this) + reserved.addAndGet(granted) + granted + } + def release(bytes: Long): Unit = { + if (bytes > 0L) { + reserved.addAndGet(-bytes) + tmm.releaseExecutionMemory(bytes, this) + } + } + } + + private case class TaskFixture( + taskAttemptId: Long, + context: TaskContextImpl, + taskMemoryManager: TaskMemoryManager, + snapshotHook: AtomicReference[Runnable]) + + private def reservedFor(task: TaskFixture): Long = + CometTaskArrowAllocator.reservedBytesForTask(task.taskAttemptId) + + private def listenerFor(task: TaskFixture): CometArrowAllocationListener = + CometTaskArrowAllocator.listenerForTask(task.taskAttemptId).get + + private def newTask(offHeap: Boolean = true, pool: Long = poolBytes): TaskFixture = { + val conf = new SparkConf(false) + if (offHeap) { + conf + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", pool.toString) + } + val memoryManager = new TestMemoryManager(conf) + memoryManager.limit(pool) + val taskAttemptId = nextTaskAttemptId.getAndIncrement() + val snapshotHook = new AtomicReference[Runnable]() + val taskMemoryManager = + new HookedTaskMemoryManager(memoryManager, taskAttemptId, snapshotHook) + val context = new TaskContextImpl( + stageId = 0, + stageAttemptNumber = 0, + partitionId = 0, + numPartitions = 1, + taskAttemptId = taskAttemptId, + attemptNumber = 0, + taskMemoryManager = taskMemoryManager, + localProperties = new Properties, + metricsSystem = null, + taskMetrics = TaskMetrics.empty, + cpus = 1, + resources = Map.empty) + TaskFixture(taskAttemptId, context, taskMemoryManager, snapshotHook) + } + + /** Installs the task on this thread, restoring whatever was there before. */ + private def withInstalledTask[T](task: TaskFixture)(f: => T): T = { + val previous = TaskContext.get() + TaskContext.setTaskContext(task.context) + try { + f + } finally { + try { + // Fires the completion listener that drops the reservation; harmless if already run. + task.context.markTaskCompleted(None) + task.taskMemoryManager.cleanUpAllAllocatedMemory() + } finally { + if (previous == null) TaskContext.unset() else TaskContext.setTaskContext(previous) + } + } + } + + private def withTask(offHeap: Boolean = true, pool: Long = poolBytes)( + f: TaskFixture => Unit): Unit = { + val task = newTask(offHeap, pool) + withInstalledTask(task)(f(task)) + } + + /** Runs the body on a fresh thread, which by construction carries no task context. */ + private def onDetachedThread(body: => Unit): Unit = { + val failure = new AtomicReference[Throwable]() + val thread = new Thread(() => { + try body + catch { case t: Throwable => failure.set(t) } + }) + thread.setDaemon(true) + thread.setName("detached-release") + thread.start() + thread.join(30000L) + assert(!thread.isAlive, "the detached release did not finish") + Option(failure.get).foreach(t => throw t) + } + + /** An unstarted daemon thread, so a test can choose the moment it runs. */ + private def daemonThread(name: String)(body: => Unit): Thread = { + val thread = new Thread(() => body) + thread.setDaemon(true) + thread.setName(name) + thread + } + + /** + * Waits until the thread is either blocked on a monitor or finished, whichever happens first, + * so that a test can tell the two orderings apart without depending on timing. + */ + private def awaitBlockedOrFinished(thread: Thread): Unit = { + val deadline = System.currentTimeMillis() + 30000L + var state = thread.getState + while (state != Thread.State.BLOCKED && state != Thread.State.TERMINATED && + System.currentTimeMillis() < deadline) { + Thread.sleep(1L) + state = thread.getState + } + } + + private def loopingThread(name: String, failure: AtomicReference[Throwable])( + body: => Unit): Thread = { + val thread = new Thread(() => { + try { + var i = 0 + while (i < 500) { + body + i += 1 + } + } catch { + case t: Throwable => failure.compareAndSet(null, t) + } + }) + thread.setDaemon(true) + thread.setName(name) + thread + } + + private def rootChildNames(): Set[String] = { + val names = Set.newBuilder[String] + val children = CometArrowAllocator.getChildAllocators.iterator() + while (children.hasNext) { + names += children.next().asInstanceOf[BufferAllocator].getName + } + names.result() + } +}