andygrove commented on code in PR #5722:
URL: https://github.com/apache/datafusion-comet/pull/5722#discussion_r3944263483
##########
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:
The P1 lifetime concern looks right to me. I reproduced the JVM half of it
locally with a latch-based probe on this iterator and a mocked `Native`:
`releaseRemoteShuffleDecoder` runs to completion while
`decodeShuffleBlockWithValidation` is still inside the JNI call, so nothing in
the iterator serializes the release against an active decode.
`remoteDecoderHandle` is published from `fetchNext()` without the monitor that
`close()` takes, and `decode_remote_shuffle_batch` reads `expected_types` after
`read_ipc_compressed_validated` has already decompressed the block, so the
borrow outlives the free in that ordering. The mirror-image ordering is a leak
rather than a crash, since `close()` reads `0L` and the box is never freed.
What I do not think is true is that this PR introduces the hazard.
`CometExecIterator` has had the same shape for a long time, with a wider window
on a much hotter path. It registers `close()` as a task completion listener at
construction
([`CometExecIterator.scala:189`](https://github.com/apache/datafusion-comet/blob/cedc3bd99ec4580226ffe77154a18fc7058a1547/spark/src/main/scala/org/apache/comet/CometExecIterator.scala#L189)),
`close()` is `synchronized` and calls `nativeLib.releasePlan(plan)`
([`:275`](https://github.com/apache/datafusion-comet/blob/cedc3bd99ec4580226ffe77154a18fc7058a1547/spark/src/main/scala/org/apache/comet/CometExecIterator.scala#L275)
and
[`:305`](https://github.com/apache/datafusion-comet/blob/cedc3bd99ec4580226ffe77154a18fc7058a1547/spark/src/main/scala/org/apache/comet/CometExecIterator.scala#L305)),
and `getNextBatch` calls `nativeLib.executePlan(..., plan, ...)` outside that
monitor ([`:205`](https://github.com/apache/datafusion-comet/blob/cedc3bd99ec45
80226ffe77154a18fc7058a1547/spark/src/main/scala/org/apache/comet/CometExecIterator.scala#L205)).
On the native side `releasePlan` is a `Box::from_raw` drop of the context
([`jni_api.rs:1051`](https://github.com/apache/datafusion-comet/blob/cedc3bd99ec4580226ffe77154a18fc7058a1547/native/core/src/execution/jni_api.rs#L1051))
while `executePlan` holds `get_execution_context(exec_context)`, an unbounded
`&'a mut ExecutionContext` derived from the raw handle, for the duration of an
entire plan execution
([`jni_api.rs:848`](https://github.com/apache/datafusion-comet/blob/cedc3bd99ec4580226ffe77154a18fc7058a1547/native/core/src/execution/jni_api.rs#L848)).
So if the threaded Python consumer plus an interrupted cleanup wait reaches
the shuffle decoder, it already reaches every Comet native plan today.
Given that, I would rather not hold this PR for it. Could we file one issue
covering JVM-owned native handle lifetime for both `NativeBatchDecoderIterator`
and `CometExecIterator`, and fix them the same way? Fixing only the new handle
here would leave the wider window untouched and leave us with two different
ownership conventions for the same problem.
For what it is worth, the minimal fix for this iterator looks like taking
the monitor around just the `nativeUtil.getNextBatch` call and leaving
`readNextBlock()` outside it, so `in.close()` can still unblock a transport
read. Decoding one block is bounded work, so cleanup cannot hang waiting on the
monitor. Whatever we settle on should apply to `releasePlan` as well.
--
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]
