diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 84abd98f2b..a1069acfcd 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -1252,10 +1252,56 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlock( }) } +#[no_mangle] +/// Parse the expected schema once for a remote shuffle iterator. +/// +/// The iterator owns the returned decoder and releases it when the input is closed. +pub extern "system" fn Java_org_apache_comet_Native_createRemoteShuffleDecoder( + e: EnvUnowned, + _class: JClass, + expected_schema: JByteArray, +) -> jlong { + try_unwrap_or_throw(&e, |env| { + let bytes = env.convert_byte_array(expected_schema)?; + let schema = ShuffleScan::decode(bytes.as_slice()).map_err(|error| { + CometError::Internal(format!("Invalid expected remote shuffle schema: {error}")) + })?; + let decoder = RemoteShuffleDecoder { + expected_types: schema.fields.iter().map(to_arrow_datatype).collect(), + }; + Ok(Box::into_raw(Box::new(decoder)) as jlong) + }) +} + +/// Immutable decoding state owned by one JVM remote shuffle iterator, not shared across tasks. +struct RemoteShuffleDecoder { + expected_types: Vec, +} + +#[no_mangle] +/// Release a remote shuffle iterator's decoder. +/// +/// # Safety +/// A nonzero handle must have been returned by `createRemoteShuffleDecoder`, must not have +/// been released, and must not be in use by a concurrent decode call. +pub unsafe extern "system" fn Java_org_apache_comet_Native_releaseRemoteShuffleDecoder( + e: EnvUnowned, + _class: JClass, + decoder_handle: jlong, +) { + try_unwrap_or_throw(&e, |_| { + if decoder_handle != 0 { + drop(unsafe { Box::from_raw(decoder_handle as *mut RemoteShuffleDecoder) }); + } + Ok(()) + }) +} + #[no_mangle] /// Decode a remote native shuffle block with Arrow array and logical type validation enabled. /// # Safety -/// This function is inherently unsafe since it deals with raw pointers passed from JNI. +/// Buffer and output pointers must be valid. The decoder handle must have been returned by +/// `createRemoteShuffleDecoder` and must remain alive for the duration of this call. pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlockWithValidation( e: EnvUnowned, _class: JClass, @@ -1264,22 +1310,21 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlockWit array_addrs: JLongArray, schema_addrs: JLongArray, tracing_enabled: jboolean, - expected_schema: JByteArray, + decoder_handle: jlong, ) -> jlong { try_unwrap_or_throw(&e, |env| { with_trace("decodeShuffleBlock", tracing_enabled != JNI_FALSE, || { - let bytes = env.convert_byte_array(expected_schema)?; - let schema = ShuffleScan::decode(bytes.as_slice()).map_err(|error| { - CometError::Internal(format!("Invalid expected remote shuffle schema: {error}")) - })?; - let expected_types: Vec<_> = schema.fields.iter().map(to_arrow_datatype).collect(); + let decoder = unsafe { (decoder_handle as *const RemoteShuffleDecoder).as_ref() } + .ok_or_else(|| { + CometError::Internal("Remote shuffle decoder is not initialized".to_owned()) + })?; decode_shuffle_block( env, byte_buffer, length, array_addrs, schema_addrs, - Some(&expected_types), + Some(&decoder.expected_types), ) }) }) diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index e609001618..af406632f1 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -201,9 +201,17 @@ class Native extends NativeBase { tracingEnabled: Boolean): Long /** - * Decode a remote shuffle block with Arrow buffer/offset and logical type validation. The - * expected schema is serialized as a ShuffleScan protobuf. Keep the existing local decoder - * entry point unchanged so trusted local shuffle reads retain their fast path. + * Create a remote shuffle decoder that retains the expected Spark types for one iterator. The + * expected schema is serialized as a ShuffleScan protobuf. + */ + @native def createRemoteShuffleDecoder(expectedSchema: Array[Byte]): Long + + /** Release a remote shuffle decoder after its iterator finishes reading or closes early. */ + @native def releaseRemoteShuffleDecoder(decoderHandle: Long): Unit + + /** + * Decode a remote shuffle block with Arrow buffer/offset and logical type validation, using the + * expected Spark types retained by the decoder. */ @native def decodeShuffleBlockWithValidation( shuffleBlock: ByteBuffer, @@ -211,7 +219,7 @@ class Native extends NativeBase { arrayAddrs: Array[Long], schemaAddrs: Array[Long], tracingEnabled: Boolean, - expectedSchema: Array[Byte]): Long + decoderHandle: Long): Long /** * Log the beginning of an event. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala index 514a3f1095..6227da4bf4 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala @@ -45,11 +45,14 @@ case class NativeBatchDecoderIterator( expectedSchema: Option[Array[Byte]] = None) extends Iterator[ColumnarBatch] { + // One consumer reads this iterator, while task completion may close it from another thread. + // The monitor protects decoder and batch ownership; transport reads stay outside it. private var isClosed = false private val longBuf = ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN) private var currentBatch: ColumnarBatch = null private var batch: Option[ColumnarBatch] = None private val validateRemoteFrames = in.isInstanceOf[CometShuffleReadFailureHandler] + private var remoteDecoderHandle = 0L require( !validateRemoteFrames || expectedSchema.exists(_ != null), @@ -64,25 +67,23 @@ case class NativeBatchDecoderIterator( } override def hasNext: Boolean = { - if (channel == null || isClosed) { - return false - } - if (batch.isDefined) { - return true - } + synchronized { + if (channel == null || isClosed) { + return false + } + if (batch.isDefined) { + return true + } - // Release the previous batch. - if (currentBatch != null) { - currentBatch.close() - currentBatch = null + // Clear ownership before cleanup so a failed close cannot release this batch twice. + if (currentBatch != null) { + val previous = currentBatch + currentBatch = null + previous.close() + } } - batch = fetchNext() - if (batch.isEmpty) { - close() - return false - } - true + fetchNext() } def next(): ColumnarBatch = { @@ -90,52 +91,76 @@ case class NativeBatchDecoderIterator( throw new NoSuchElementException } - val nextBatch = batch.get - - currentBatch = nextBatch - batch = None - currentBatch + synchronized { + // Completion may have closed the iterator after hasNext() returned. + if (isClosed) { + throw new NoSuchElementException + } + currentBatch = batch.get + batch = None + currentBatch + } } - private def fetchNext(): Option[ColumnarBatch] = { + private def fetchNext(): Boolean = { // The remote input owns metadata, transport, and frame-boundary failure classification. Do // not turn deliberately unreported metadata timeouts or stream-close failures into corruption. + // Read outside the monitor so close() can unblock the underlying stream. val block = readNextBlock() - block.flatMap { case (fieldCount, dataBuf, bytesToRead) => - // Allocation and Arrow import failures are not evidence of a corrupt remote shuffle. Only - // forward failures from native codec/IPC decoding and logical type validation. - val startTime = System.nanoTime() - val decoded = nativeUtil.getNextBatch( - fieldCount, - (arrayAddrs, schemaAddrs) => { - handleReadFailure { - if (validateRemoteFrames) { - nativeLib.decodeShuffleBlockWithValidation( - dataBuf, - bytesToRead, - arrayAddrs, - schemaAddrs, - tracingEnabled, - expectedSchema.get) - } else { - nativeLib.decodeShuffleBlock( - dataBuf, - bytesToRead, - arrayAddrs, - schemaAddrs, - tracingEnabled) - } + var nativeFailure: Throwable = null + try { + synchronized { + // Cleanup may have finished during the read. Do not create a decoder or import a batch + // after close(), and publish each decoded batch before cleanup can inspect ownership. + if (isClosed) { + return false + } + batch = block.flatMap { case (fieldCount, dataBuf, bytesToRead) => + val startTime = System.nanoTime() + // Invalid expected schemas are setup failures, not evidence of corrupt persisted data. + if (validateRemoteFrames && remoteDecoderHandle == 0L) { + remoteDecoderHandle = nativeLib.createRemoteShuffleDecoder(expectedSchema.get) } - }) - decodeTime.add(System.nanoTime() - startTime) - decoded - } - } - - private def handleReadFailure[T](read: => T): T = { - try read - catch { - case NonFatal(failure) => + val decoded = nativeUtil.getNextBatch( + fieldCount, + (arrayAddrs, schemaAddrs) => { + try { + if (validateRemoteFrames) { + nativeLib.decodeShuffleBlockWithValidation( + dataBuf, + bytesToRead, + arrayAddrs, + schemaAddrs, + tracingEnabled, + remoteDecoderHandle) + } else { + nativeLib.decodeShuffleBlock( + dataBuf, + bytesToRead, + arrayAddrs, + schemaAddrs, + tracingEnabled) + } + } catch { + case NonFatal(failure) => + // Record only native decode failures; creation, allocation and import failures + // must not be reported as corrupt remote data. Let NativeUtil clean up first. + nativeFailure = failure + throw failure + } + }) + decodeTime.add(System.nanoTime() - startTime) + decoded + } + if (batch.isEmpty) { + close() + } + batch.isDefined + } + } catch { + case NonFatal(failure) if failure eq nativeFailure => + // Reporting can make an RPC. It must run after Arrow cleanup and outside the monitor so + // task completion can release the decoder while that reporting is in progress. in match { case handler: CometShuffleReadFailureHandler => handler.onShuffleReadFailure(failure) case _ => @@ -145,10 +170,6 @@ case class NativeBatchDecoderIterator( } private def readNextBlock(): Option[(Int, ByteBuffer, Int)] = { - if (channel == null || isClosed) { - return None - } - // read compressed batch size from header longBuf.clear() while (longBuf.hasRemaining && channel.read(longBuf) >= 0) {} @@ -212,6 +233,8 @@ case class NativeBatchDecoderIterator( currentBatch = null val prefetched = batch batch = None + val decoderHandle = remoteDecoderHandle + remoteDecoderHandle = 0L var failure: Throwable = null def release(resource: => Unit): Unit = { @@ -225,6 +248,7 @@ case class NativeBatchDecoderIterator( if (previous != null) release(previous.close()) prefetched.filterNot(_ eq previous).foreach(pending => release(pending.close())) + if (decoderHandle != 0L) release(nativeLib.releaseRemoteShuffleDecoder(decoderHandle)) if (in != null) release(in.close()) release(resetDataBuf()) if (failure != null) throw failure diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala index d764137635..54448ec203 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShuffleReaderSuite.scala @@ -1442,6 +1442,25 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { } } + test("JVM reader decodes multiple remote frames with the same expected Spark types") { + val frame = intFrame() + val attributes = Seq(AttributeReference("value", IntegerType)()) + val (rows, failure, reports) = + readRemoteFrame(frame ++ frame ++ frame, attributes, raw = false) + assert(failure.isEmpty) + assert(rows.map(_.getInt(0)) == Seq.fill(3)(Seq(-1, 0, 1)).flatten) + assert(reports == 0) + } + + test("JVM reader validates later frames against the retained expected Spark types") { + val attributes = Seq(AttributeReference("value", IntegerType)()) + val (rows, failure, reports) = + readRemoteFrame(intFrame() ++ dictionaryFrame(), attributes, raw = false) + assert(rows.map(_.getInt(0)) == Seq(-1, 0, 1)) + assert(failure.exists(_.toLowerCase(java.util.Locale.ROOT).contains("type mismatch"))) + assert(reports == 1) + } + test("both remote consumption paths validate Arrow offsets before exposing decoded arrays") { val allocator = new RootAllocator(Long.MaxValue) val strings = new VarCharVector("value", allocator) @@ -1552,8 +1571,16 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { classOf[Native] .getDeclaredMethod( "decodeShuffleBlockWithValidation", - (parameters :+ classOf[Array[Byte]]): _*) + (parameters :+ java.lang.Long.TYPE): _*) + .getReturnType == java.lang.Long.TYPE) + assert( + classOf[Native] + .getDeclaredMethod("createRemoteShuffleDecoder", classOf[Array[Byte]]) .getReturnType == java.lang.Long.TYPE) + assert( + classOf[Native] + .getDeclaredMethod("releaseRemoteShuffleDecoder", java.lang.Long.TYPE) + .getReturnType == java.lang.Void.TYPE) val local = new CometShuffleBlockIterator(new ByteArrayInputStream(Array.empty[Byte])) assert(!local.requiresValidation()) @@ -1608,6 +1635,26 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { NativeBatchDecoderIteratorLifecycleChecks.closesDeliveredBatch() } + test("decoder cleanup waits for native decoder creation") { + NativeBatchDecoderIteratorConcurrencyChecks.closeWaitsForDecoderCreation() + } + + test("decoder cleanup waits for active native decoding even when interrupted") { + NativeBatchDecoderIteratorConcurrencyChecks.closeWaitsForNativeDecoding() + } + + test("decoder cleanup waits for Arrow import and releases the resulting batch") { + NativeBatchDecoderIteratorConcurrencyChecks.closeWaitsForBatchImportAndPublication() + } + + test("decoder cleanup can close a blocked transport without decoding its returned block") { + NativeBatchDecoderIteratorConcurrencyChecks.closeUnblocksTransportAndPreventsDecoding() + } + + test("decoder unwinds Arrow import before reporting failures without blocking cleanup") { + NativeBatchDecoderIteratorConcurrencyChecks.closeProceedsDuringFailureReporting() + } + test("decoder cleanup releases remaining resources and preserves suppressed failures") { NativeBatchDecoderIteratorLifecycleChecks.preservesCleanupFailures() } @@ -1629,6 +1676,18 @@ class CometCelebornShuffleReaderSuite extends CometTestBase { NativeBatchDecoderIteratorLifecycleChecks.selectsValidationOnlyForRemoteStreams() } + test("decoder does not allocate native state for empty or unused remote streams") { + NativeBatchDecoderIteratorLifecycleChecks.doesNotAllocateUnusedRemoteDecoder() + } + + test("decoder releases remote native state once and preserves cleanup failures") { + NativeBatchDecoderIteratorLifecycleChecks.preservesRemoteDecoderCleanupFailures() + } + + test("decoder does not report native decoder creation failures as remote data corruption") { + NativeBatchDecoderIteratorLifecycleChecks.doesNotReportRemoteDecoderCreationFailures() + } + test("decoder rejects a missing remote schema before consuming or reporting persisted data") { NativeBatchDecoderIteratorLifecycleChecks.rejectsRemoteStreamsWithoutExpectedSchema() } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorConcurrencyChecks.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorConcurrencyChecks.scala new file mode 100644 index 0000000000..13607f82c5 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorConcurrencyChecks.scala @@ -0,0 +1,384 @@ +/* + * 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.sql.comet.execution.shuffle + +import java.io.{ByteArrayInputStream, IOException} +import java.nio.{ByteBuffer, ByteOrder} +import java.util.concurrent.{CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger, AtomicReference} + +import org.apache.spark.sql.execution.metric.SQLMetric +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} + +import org.apache.comet.{CometShuffleReadFailureHandler, Native} +import org.apache.comet.vector.NativeUtil + +/** Exercises one consuming thread racing with task cleanup, without calling native code. */ +private[shuffle] object NativeBatchDecoderIteratorConcurrencyChecks { + + private val timeoutSeconds = 10L + + private def await(latch: CountDownLatch): Unit = { + assert(latch.await(timeoutSeconds, TimeUnit.SECONDS), "Timed out waiting for worker") + } + + private final class Worker(name: String)(run: => Unit) { + val finished = new CountDownLatch(1) + val failure = new AtomicReference[Throwable]() + val thread = new Thread( + () => { + try run + catch { + case caught: Throwable => failure.set(caught) + } finally { + finished.countDown() + } + }, + name) + thread.setDaemon(true) + thread.start() + } + + private def join(workers: Seq[Worker]): Unit = { + // Join every worker before propagating errors so one failure cannot skip another join. + workers.foreach(_.thread.join(TimeUnit.SECONDS.toMillis(timeoutSeconds))) + workers.foreach(worker => assert(!worker.thread.isAlive, s"${worker.thread.getName} hung")) + workers.foreach { worker => + Option(worker.failure.get()).foreach(throw _) + } + } + + private def awaitBlocked(worker: Worker): Unit = { + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(timeoutSeconds) + while (worker.thread.getState != Thread.State.BLOCKED && + worker.finished.getCount != 0 && System.nanoTime() < deadline) { + Thread.`yield`() + } + assert( + worker.thread.getState == Thread.State.BLOCKED, + "Cleanup must wait for the consuming thread to leave the decoder monitor") + } + + private final class TrackingBatch extends ColumnarBatch(Array.empty[ColumnVector], 1) { + val closeCalls = new AtomicInteger() + + override def close(): Unit = { + closeCalls.incrementAndGet() + super.close() + } + } + + private class TrackingNative extends Native { + val handle = 42L + val createCalls = new AtomicInteger() + val decodeCalls = new AtomicInteger() + val releaseCalls = new AtomicInteger() + + override def createRemoteShuffleDecoder(expectedSchema: Array[Byte]): Long = { + createCalls.incrementAndGet() + handle + } + + override def decodeShuffleBlockWithValidation( + block: ByteBuffer, + length: Int, + arrays: Array[Long], + schemas: Array[Long], + tracing: Boolean, + decoderHandle: Long): Long = { + assert(decoderHandle == handle) + assert(releaseCalls.get() == 0) + decodeCalls.incrementAndGet() + 1L + } + + override def releaseRemoteShuffleDecoder(decoderHandle: Long): Unit = { + assert(decoderHandle == handle) + releaseCalls.incrementAndGet() + } + } + + private val frame = ByteBuffer + .allocate(20) + .order(ByteOrder.LITTLE_ENDIAN) + .putLong(12L) + .putLong(0L) + .putInt(0) + .array() + + private sealed trait DecodeStage + private case object CreateDecoder extends DecodeStage + private case object DecodeBlock extends DecodeStage + private case object ImportBatch extends DecodeStage + + private def closeWaitsFor(stage: DecodeStage, interruptClose: Boolean = false): Unit = { + val entered = new CountDownLatch(1) + val proceed = new CountDownLatch(1) + def pauseAt(current: DecodeStage): Unit = { + if (stage == current) { + entered.countDown() + await(proceed) + } + } + + val batch = new TrackingBatch + val inputCloseCalls = new AtomicInteger() + val input = new ByteArrayInputStream(frame) with CometShuffleReadFailureHandler { + override def close(): Unit = { + inputCloseCalls.incrementAndGet() + super.close() + } + + override def onShuffleReadFailure(failure: Throwable): Unit = throw failure + } + val nativeLib = new TrackingNative { + override def createRemoteShuffleDecoder(expectedSchema: Array[Byte]): Long = { + val created = super.createRemoteShuffleDecoder(expectedSchema) + pauseAt(CreateDecoder) + assert(releaseCalls.get() == 0) + created + } + + override def decodeShuffleBlockWithValidation( + block: ByteBuffer, + length: Int, + arrays: Array[Long], + schemas: Array[Long], + tracing: Boolean, + decoderHandle: Long): Long = { + val rows = super.decodeShuffleBlockWithValidation( + block, + length, + arrays, + schemas, + tracing, + decoderHandle) + pauseAt(DecodeBlock) + assert(releaseCalls.get() == 0, "Cleanup released a decoder that is still in use") + rows + } + } + val util = new NativeUtil { + override def getNextBatch( + numOutputCols: Int, + decode: (Array[Long], Array[Long]) => Long): Option[ColumnarBatch] = { + assert(decode(Array.empty[Long], Array.empty[Long]) == 1L) + pauseAt(ImportBatch) + assert(batch.closeCalls.get() == 0) + Some(batch) + } + } + val decoder = NativeBatchDecoderIterator( + input, + new SQLMetric("nsTiming", 0L), + nativeLib, + util, + tracingEnabled = false, + expectedSchema = Some(Array.empty[Byte])) + val consumer = new Worker("shuffle decoder")({ + decoder.hasNext + () + }) + var closer: Option[Worker] = None + try { + await(entered) + val cleanup = new Worker("shuffle cleanup")({ + decoder.close() + if (interruptClose) assert(Thread.currentThread().isInterrupted) + }) + closer = Some(cleanup) + awaitBlocked(cleanup) + if (interruptClose) { + cleanup.thread.interrupt() + awaitBlocked(cleanup) + } + assert(nativeLib.releaseCalls.get() == 0) + assert(batch.closeCalls.get() == 0) + assert(inputCloseCalls.get() == 0) + } finally { + proceed.countDown() + try join(Seq(consumer) ++ closer) + finally { + decoder.close() + util.close() + } + } + assert(nativeLib.createCalls.get() == 1) + assert(nativeLib.decodeCalls.get() == 1) + assert(nativeLib.releaseCalls.get() == 1) + assert(batch.closeCalls.get() == 1, "Cleanup missed the batch produced by the active read") + assert(inputCloseCalls.get() == 1) + assert(!decoder.hasNext) + } + + def closeWaitsForDecoderCreation(): Unit = closeWaitsFor(CreateDecoder) + + def closeWaitsForNativeDecoding(): Unit = { + closeWaitsFor(DecodeBlock) + closeWaitsFor(DecodeBlock, interruptClose = true) + } + + def closeWaitsForBatchImportAndPublication(): Unit = closeWaitsFor(ImportBatch) + + def closeProceedsDuringFailureReporting(): Unit = { + val reporting = new CountDownLatch(1) + val finishReporting = new CountDownLatch(1) + val importCleanedUp = new AtomicBoolean() + val cleanupBeforeReporting = new AtomicBoolean() + val decodeFailure = new IOException("native decode failed") + val reportedFailure = new IOException("remote fetch failed", decodeFailure) + val inputCloseCalls = new AtomicInteger() + val input = new ByteArrayInputStream(frame) with CometShuffleReadFailureHandler { + override def onShuffleReadFailure(failure: Throwable): Unit = { + assert(failure eq decodeFailure) + cleanupBeforeReporting.set(importCleanedUp.get()) + reporting.countDown() + await(finishReporting) + throw reportedFailure + } + + override def close(): Unit = { + inputCloseCalls.incrementAndGet() + super.close() + } + } + val nativeLib = new TrackingNative { + override def decodeShuffleBlockWithValidation( + block: ByteBuffer, + length: Int, + arrays: Array[Long], + schemas: Array[Long], + tracing: Boolean, + decoderHandle: Long): Long = throw decodeFailure + } + val util = new NativeUtil { + override def getNextBatch( + numOutputCols: Int, + decode: (Array[Long], Array[Long]) => Long): Option[ColumnarBatch] = { + try { + decode(Array.empty[Long], Array.empty[Long]) + None + } finally { + importCleanedUp.set(true) + } + } + } + val decoder = NativeBatchDecoderIterator( + input, + new SQLMetric("nsTiming", 0L), + nativeLib, + util, + tracingEnabled = false, + expectedSchema = Some(Array.empty[Byte])) + val consumer = new Worker("shuffle failure reporting")({ + try { + decoder.hasNext + throw new AssertionError("Expected the remote fetch failure") + } catch { + case caught: IOException => assert(caught eq reportedFailure) + } + }) + var closer: Option[Worker] = None + try { + await(reporting) + assert( + cleanupBeforeReporting.get(), + "Arrow import must unwind before reporting data errors") + val cleanup = new Worker("shuffle failure cleanup")({ decoder.close() }) + closer = Some(cleanup) + await(cleanup.finished) + assert(nativeLib.releaseCalls.get() == 1) + assert(inputCloseCalls.get() == 1) + } finally { + finishReporting.countDown() + try join(Seq(consumer) ++ closer) + finally { + decoder.close() + util.close() + } + } + assert(nativeLib.releaseCalls.get() == 1) + assert(inputCloseCalls.get() == 1) + assert(!decoder.hasNext) + } + + def closeUnblocksTransportAndPreventsDecoding(): Unit = { + val readingBody = new CountDownLatch(1) + val streamClosed = new CountDownLatch(1) + val inputCloseCalls = new AtomicInteger() + val input = new ByteArrayInputStream(frame) with CometShuffleReadFailureHandler { + override def read(bytes: Array[Byte], offset: Int, length: Int): Int = { + if (available() == 4) { + readingBody.countDown() + await(streamClosed) + } + // A transport read can finish successfully even after cleanup has closed the stream. + super.read(bytes, offset, length) + } + + override def close(): Unit = { + inputCloseCalls.incrementAndGet() + streamClosed.countDown() + super.close() + } + + override def onShuffleReadFailure(failure: Throwable): Unit = throw failure + } + val nativeLib = new TrackingNative + val importCalls = new AtomicInteger() + val util = new NativeUtil { + override def getNextBatch( + numOutputCols: Int, + decode: (Array[Long], Array[Long]) => Long): Option[ColumnarBatch] = { + importCalls.incrementAndGet() + None + } + } + val decoder = NativeBatchDecoderIterator( + input, + new SQLMetric("nsTiming", 0L), + nativeLib, + util, + tracingEnabled = false, + expectedSchema = Some(Array.empty[Byte])) + val consumer = new Worker("shuffle transport read")({ assert(!decoder.hasNext) }) + var closer: Option[Worker] = None + try { + await(readingBody) + val cleanup = new Worker("shuffle transport cleanup")({ decoder.close() }) + closer = Some(cleanup) + await(cleanup.finished) + await(consumer.finished) + } finally { + streamClosed.countDown() + try join(Seq(consumer) ++ closer) + finally { + decoder.close() + util.close() + } + } + assert(inputCloseCalls.get() == 1) + assert(nativeLib.createCalls.get() == 0) + assert(nativeLib.decodeCalls.get() == 0) + assert(nativeLib.releaseCalls.get() == 0) + assert(importCalls.get() == 0) + assert(!decoder.hasNext) + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorLifecycleChecks.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorLifecycleChecks.scala index e9362b58e0..52d11d94e3 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorLifecycleChecks.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorLifecycleChecks.scala @@ -45,6 +45,42 @@ private[shuffle] object NativeBatchDecoderIteratorLifecycleChecks { } } + private class TrackingNative extends Native { + val handle = 42L + var createCalls = 0 + var releaseCalls = 0 + + override def createRemoteShuffleDecoder(expectedSchema: Array[Byte]): Long = { + createCalls += 1 + handle + } + + override def releaseRemoteShuffleDecoder(decoderHandle: Long): Unit = { + assert(decoderHandle == handle) + releaseCalls += 1 + } + } + + private final class TrackingRemoteInput( + bytes: Array[Byte], + closeFailure: Option[Throwable] = None) + extends ByteArrayInputStream(bytes) + with CometShuffleReadFailureHandler { + var closeCalls = 0 + var reports = 0 + + override def close(): Unit = { + closeCalls += 1 + super.close() + closeFailure.foreach(throw _) + } + + override def onShuffleReadFailure(failure: Throwable): Unit = { + reports += 1 + throw failure + } + } + private val frame = ByteBuffer .allocate(20) .order(ByteOrder.LITTLE_ENDIAN) @@ -152,14 +188,17 @@ private[shuffle] object NativeBatchDecoderIteratorLifecycleChecks { throw new AssertionError("Expected native decoding to fail") } } - val nativeLib = new Native { + val nativeLib = new TrackingNative { override def decodeShuffleBlockWithValidation( block: ByteBuffer, length: Int, arrays: Array[Long], schemas: Array[Long], tracing: Boolean, - expectedSchema: Array[Byte]): Long = throw decodeFailure + decoderHandle: Long): Long = { + assert(decoderHandle == handle) + throw decodeFailure + } } val decoder = NativeBatchDecoderIterator( input, @@ -186,6 +225,9 @@ private[shuffle] object NativeBatchDecoderIteratorLifecycleChecks { } assert(closeError eq closeFailure) assert(reports == 1) + decoder.close() + assert(nativeLib.createCalls == 1) + assert(nativeLib.releaseCalls == 1) } finally { util.close() } @@ -197,28 +239,37 @@ private[shuffle] object NativeBatchDecoderIteratorLifecycleChecks { .addFields(QueryPlanSerde.serializeDataType(IntegerType).get) .build() .toByteArray - Seq(None, Some(Array.empty[Byte]), Some(intSchema)).foreach { expectedSchema => - val remote = expectedSchema.isDefined + Seq( + false -> None, + false -> Some(intSchema), + true -> Some(Array.empty[Byte]), + true -> Some(intSchema)).foreach { case (remote, expectedSchema) => var localCalls = 0 var validatedCalls = 0 val input = if (remote) { - new ByteArrayInputStream(frame) with CometShuffleReadFailureHandler { + new ByteArrayInputStream(frame ++ frame) with CometShuffleReadFailureHandler { override def onShuffleReadFailure(failure: Throwable): Unit = throw failure } } else { - new ByteArrayInputStream(frame) + new ByteArrayInputStream(frame ++ frame) } - val batch = new TrackingBatch() + val batches = Seq.fill(2)(new TrackingBatch()) + val pending = batches.iterator val util = new NativeUtil { override def getNextBatch( numOutputCols: Int, decode: (Array[Long], Array[Long]) => Long): Option[ColumnarBatch] = { assert(decode(Array.empty[Long], Array.empty[Long]) == 1L) - Some(batch) + Some(pending.next()) } } - val nativeLib = new Native { + val nativeLib = new TrackingNative { + override def createRemoteShuffleDecoder(schema: Array[Byte]): Long = { + assert(expectedSchema.exists(_.sameElements(schema))) + super.createRemoteShuffleDecoder(schema) + } + override def decodeShuffleBlock( block: ByteBuffer, length: Int, @@ -235,8 +286,10 @@ private[shuffle] object NativeBatchDecoderIteratorLifecycleChecks { arrays: Array[Long], schemas: Array[Long], tracing: Boolean, - schema: Array[Byte]): Long = { - assert(expectedSchema.exists(_.sameElements(schema))) + decoderHandle: Long): Long = { + assert(decoderHandle == handle) + assert(createCalls == 1) + assert(releaseCalls == 0) validatedCalls += 1 1L } @@ -249,14 +302,22 @@ private[shuffle] object NativeBatchDecoderIteratorLifecycleChecks { tracingEnabled = false, expectedSchema = expectedSchema) try { - assert(decoder.hasNext) - assert(localCalls == (if (remote) 0 else 1)) - assert(validatedCalls == (if (remote) 1 else 0)) + assert(nativeLib.createCalls == 0) + batches.foreach { batch => + assert(decoder.hasNext) + assert(decoder.hasNext) + assert(decoder.next() eq batch) + } + assert(!decoder.hasNext) + assert(localCalls == (if (remote) 0 else 2)) + assert(validatedCalls == (if (remote) 2 else 0)) } finally { decoder.close() util.close() } - assert(batch.closeCalls == 1) + assert(batches.forall(_.closeCalls == 1)) + assert(nativeLib.createCalls == (if (remote) 1 else 0)) + assert(nativeLib.releaseCalls == (if (remote) 1 else 0)) } } @@ -295,6 +356,114 @@ private[shuffle] object NativeBatchDecoderIteratorLifecycleChecks { } } + def doesNotAllocateUnusedRemoteDecoder(): Unit = { + Seq(false, true).foreach { closeBeforeReading => + val input = new TrackingRemoteInput(if (closeBeforeReading) frame else Array.empty[Byte]) + val nativeLib = new TrackingNative + val decoder = NativeBatchDecoderIterator( + input, + new SQLMetric("nsTiming", 0L), + nativeLib, + null, + tracingEnabled = false, + expectedSchema = Some(Array.empty[Byte])) + if (!closeBeforeReading) assert(!decoder.hasNext) + decoder.close() + decoder.close() + assert(!decoder.hasNext) + assert(nativeLib.createCalls == 0) + assert(nativeLib.releaseCalls == 0) + assert(input.closeCalls == 1) + assert(input.reports == 0) + } + } + + def preservesRemoteDecoderCleanupFailures(): Unit = { + for (delivered <- Seq(false, true); failBatch <- Seq(false, true)) { + val batchFailure = new IOException("batch release failed") + val nativeFailure = new IOException("native decoder release failed") + val streamFailure = new IOException("stream release failed") + val batch = new TrackingBatch(if (failBatch) Some(batchFailure) else None) + val input = new TrackingRemoteInput(frame, Some(streamFailure)) + val nativeLib = new TrackingNative { + override def releaseRemoteShuffleDecoder(decoderHandle: Long): Unit = { + super.releaseRemoteShuffleDecoder(decoderHandle) + throw nativeFailure + } + } + val util = new NativeUtil { + override def getNextBatch( + numOutputCols: Int, + decode: (Array[Long], Array[Long]) => Long): Option[ColumnarBatch] = Some(batch) + } + val decoder = NativeBatchDecoderIterator( + input, + new SQLMetric("nsTiming", 0L), + nativeLib, + util, + tracingEnabled = false, + expectedSchema = Some(Array.empty[Byte])) + try { + assert(decoder.hasNext) + if (delivered) assert(decoder.next() eq batch) + val caught = + try { + decoder.close() + throw new AssertionError("Expected a resource release failure") + } catch { + case failure: IOException => failure + } + assert(caught eq (if (failBatch) batchFailure else nativeFailure)) + val suppressed = + if (failBatch) Seq(nativeFailure, streamFailure) else Seq(streamFailure) + assert(caught.getSuppressed.toSeq == suppressed) + } finally { + decoder.close() + util.close() + } + assert(batch.closeCalls == 1) + assert(nativeLib.createCalls == 1) + assert(nativeLib.releaseCalls == 1) + assert(input.closeCalls == 1) + assert(input.reports == 0) + assert(!decoder.hasNext) + } + } + + def doesNotReportRemoteDecoderCreationFailures(): Unit = { + val expected = new IllegalArgumentException("invalid expected Spark schema") + val input = new TrackingRemoteInput(frame) + val nativeLib = new TrackingNative { + override def createRemoteShuffleDecoder(schema: Array[Byte]): Long = { + super.createRemoteShuffleDecoder(schema) + throw expected + } + } + val decoder = NativeBatchDecoderIterator( + input, + new SQLMetric("nsTiming", 0L), + nativeLib, + null, + tracingEnabled = false, + expectedSchema = Some(Array.empty[Byte])) + try { + val caught = + try { + decoder.hasNext + throw new AssertionError("Expected the native decoder creation failure") + } catch { + case failure: IllegalArgumentException => failure + } + assert(caught eq expected) + } finally { + decoder.close() + } + assert(nativeLib.createCalls == 1) + assert(nativeLib.releaseCalls == 0) + assert(input.closeCalls == 1) + assert(input.reports == 0) + } + def doesNotReportAllocationOrImportFailures(): Unit = { Seq( new OutOfMemoryException("Arrow allocator exhausted"), @@ -308,10 +477,11 @@ private[shuffle] object NativeBatchDecoderIteratorLifecycleChecks { numOutputCols: Int, decode: (Array[Long], Array[Long]) => Long): Option[ColumnarBatch] = throw expected } + val nativeLib = new TrackingNative val decoder = NativeBatchDecoderIterator( input, new SQLMetric("nsTiming", 0L), - null, + nativeLib, util, tracingEnabled = false, expectedSchema = Some(Array.empty[Byte])) @@ -329,6 +499,8 @@ private[shuffle] object NativeBatchDecoderIteratorLifecycleChecks { decoder.close() util.close() } + assert(nativeLib.createCalls == 1) + assert(nativeLib.releaseCalls == 1) } }