jayanth86 commented on code in PR #5722:
URL: https://github.com/apache/datafusion-comet/pull/5722#discussion_r3942256302
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIterator.scala:
##########
@@ -225,6 +233,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))
Review Comment:
[P1] Keep the decoder alive while a read is in flight
Could creation and decoding be coordinated with this release? On supported
Spark 3.4.3, a Python UDF downstream of a Celeborn exchange consumes this
iterator on PythonRunner's writer thread. Its completion listener [interrupts
and joins that
writer](https://github.com/apache/spark/blob/1eb558c3a6fbdd59e5a305bc3ab12ce748f6511f/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala#L251-L260),
but task cancellation with `interruptThread=true` can interrupt the join.
[TaskContextImpl catches the listener exception and continues
cleanup](https://github.com/apache/spark/blob/1eb558c3a6fbdd59e5a305bc3ab12ce748f6511f/core/src/main/scala/org/apache/spark/TaskContextImpl.scala#L194-L250),
so Comet can reach this release while the writer is still in native
decompression.
`fetchNext()` does not take the `close()` lock. This release therefore frees
the boxed decoder while [the active JNI call still borrows its
`expected_types`](https://github.com/apache/datafusion-comet/blob/cedc3bd99ec4580226ffe77154a18fc7058a1547/native/core/src/execution/jni_api.rs#L1317-L1327),
which are read after decompression in `decode_remote_shuffle_batch`. That
introduces a use-after-free capable of crashing the executor. Previously, the
types were owned locally by the JNI invocation and could not be freed by task
cleanup.
Please protect the handle across creation, active decoding, and release, and
add coverage for interrupted cleanup overlapping a read. A latch-based test
using this exact iterator and mocked JNI reproduces release while decoding is
still in flight. This is source-verified reachability plus a lifecycle
reproduction, not a reproduced native crash.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
