mbutrovich commented on code in PR #6219:
URL: https://github.com/apache/datafusion-comet/pull/6219#discussion_r4105215159
##########
native/core/src/execution/jni_api.rs:
##########
@@ -1008,17 +1005,62 @@ fn prepare_output(
/// Because the input source could be another native execution stream, which
/// will be executed in another tokio blocking thread. It causes JNI throw
/// Java exception. So we pull input batches here and insert them into scan
-/// operators before polling the stream,
+/// operators before polling the stream. Returns whether any scan made a JNI
call.
#[inline]
-fn pull_input_batches(exec_context: &mut ExecutionContext) -> Result<(),
CometError> {
- exec_context.scans.iter_mut().try_for_each(|scan| {
- scan.get_next_batch()?;
- Ok::<(), CometError>(())
- })?;
- exec_context.shuffle_scans.iter_mut().try_for_each(|scan| {
- scan.get_next_batch()?;
- Ok::<(), CometError>(())
+fn pull_input_batches(exec_context: &mut ExecutionContext) -> Result<bool,
CometError> {
+ let mut pulled = false;
+ for scan in exec_context.scans.iter_mut() {
+ pulled |= scan.get_next_batch()?;
+ }
+ for scan in exec_context.shuffle_scans.iter_mut() {
+ pulled |= scan.get_next_batch()?;
+ }
+ Ok(pulled)
+}
+
+/// Yields once, so the `block_on` thread sleeps until a waker registered by
an earlier poll
+/// fires: a JVM-fed scan refilled by `pull_input_batches`, or native I/O that
completed.
+async fn park_until_woken() {
+ let mut polled = false;
+ poll_fn(|_| {
+ if std::mem::replace(&mut polled, true) {
+ Poll::Ready(())
+ } else {
+ Poll::Pending
+ }
})
+ .await
+}
+
+/// Drives `stream` to its next item. JVM-fed scans return `Pending` until
`on_pending` refills
+/// them, so every pending poll runs it. `on_pending` returns whether it made
a JNI call, and the
+/// loop parks only when it did not, because the stream is then waiting on
native I/O.
+///
+/// After a JNI call the loop polls again without parking. The call can run
another Comet plan on
+/// this thread, as when a native writer's input is itself native, and that
plan's `block_on`
+/// shares this thread's parker. If the native I/O completes while the nested
plan is parked, the
+/// nested park takes the wake-up, and a park here would wait for a wake that
has already fired.
+async fn next_batch<S>(
+ stream: &mut S,
+ mut on_pending: impl FnMut() -> Result<bool, CometError>,
+) -> Result<Option<RecordBatch>, CometError>
+where
+ S: Stream<Item = DataFusionResult<RecordBatch>> + Unpin,
+{
+ loop {
+ match poll!(stream.next()) {
+ Poll::Ready(item) => return Ok(item.transpose()?),
+ Poll::Pending => {
+ // JNI call to pull batches from JVM into ScanExec operators.
+ // block_in_place lets tokio move other tasks off this worker
+ // while we wait for JVM data.
+ let pulled = tokio::task::block_in_place(&mut on_pending)?;
+ if !pulled {
+ park_until_woken().await;
+ }
+ }
+ }
+ }
}
Review Comment:
Skipping the park after a pull is correct for the nesting we know about. It
depends on every call that can run another Comet plan on this thread reporting
itself through `on_pending`'s return value, though. The closure in
`executePlan` also calls `update_metrics_on_interval`, which calls into the JVM
and doesn't report it. That's safe today because a metrics update doesn't run a
plan, but nothing in the code enforces it, and a future JNI call in the closure
would bring the hang back without failing a test. The pull isn't the only way
into the JVM from this thread either. The threading section of
[`development.md`](https://github.com/apache/datafusion-comet/blob/88cd4079fc82ada6e0cda2554cc9201356126f43/docs/source/contributor-guide/development.md#L81-L85)
notes that memory pool operations call `acquireMemory()` over JNI on whatever
thread the operator runs on, and on this path that's inside the stream poll,
where no return value from `on_pending` can report it. Spark can make other
consumers in the task spill to satisfy that request. Every Comet `spill()`
returns 0 today, so this doesn't nest a plan now, but the loop's safety still
depends on that staying true.
The lost wake comes from tokio keeping one wake token per thread.
`CachedParkThread::block_on` polls and then parks on the thread-local
`CURRENT_PARKER` with no per-call state
([park.rs](https://github.com/tokio-rs/tokio/blob/75fef53d0a8590c2d1dbb63672aa7b7d1ef51155/tokio/src/runtime/park.rs#L274-L289)),
while the [`block_in_place`
docs](https://docs.rs/tokio/1.53.1/tokio/task/fn.block_in_place.html) present
`Handle::block_on` inside `block_in_place` as supported. The std [`Wake`
docs](https://doc.rust-lang.org/std/task/trait.Wake.html#examples) call out
this case for their `block_on` example: "production-grade implementations will
also need to handle intermediate calls to `thread::unpark` as well as nested
invocations."
Could `next_batch` track its own wake-ups instead? It can poll the stream
with a waker that sets a flag owned by this call and then forwards to the
`block_on` waker. After the pull, it returns `Pending` only if the flag is
still unset. A nested `block_on` can consume the thread's token but not the
flag, so the loop no longer needs to know which calls can nest, and
`get_next_batch` can keep the `Result<(), CometError>` signature from #6092.
The cost is one `Arc` per `next_batch` call, which is once per output batch
rather than once per poll. Roughly:
```rust
struct WakeFlag {
woken: AtomicBool,
parent: Waker,
}
impl Wake for WakeFlag {
fn wake(self: Arc<Self>) {
self.wake_by_ref()
}
fn wake_by_ref(self: &Arc<Self>) {
self.woken.store(true, Ordering::Release);
self.parent.wake_by_ref();
}
}
async fn next_batch<S>(
stream: &mut S,
mut on_pending: impl FnMut() -> Result<(), CometError>,
) -> Result<Option<RecordBatch>, CometError>
where
S: Stream<Item = DataFusionResult<RecordBatch>> + Unpin,
{
poll_fn(|cx| {
let flag = Arc::new(WakeFlag {
woken: AtomicBool::new(false),
parent: cx.waker().clone(),
});
let waker = Waker::from(Arc::clone(&flag));
let mut stream_cx = Context::from_waker(&waker);
loop {
match stream.poll_next_unpin(&mut stream_cx) {
Poll::Ready(item) => return
Poll::Ready(Ok(item.transpose()?)),
Poll::Pending => {
tokio::task::block_in_place(&mut on_pending)?;
if !flag.woken.swap(false, Ordering::AcqRel) {
return Poll::Pending;
}
}
}
}
})
.await
}
```
This also replaces `park_until_woken`. I tried it in a local checkout of
this branch (keeping the `bool` closures and ignoring the value). The three
`next_batch` tests and `refill_wakes_the_pending_poll_and_eof_stays_buffered`
pass. I also changed the nested `block_on` test's closure to return
`Ok(false)`, which is a nested plan that doesn't get reported. The sketch still
returns in 0.10 s. The head commit fails after 10.0 s with "the park after the
pull lost the stream's wake-up". If you go this way, the nested test's closure
could return `Ok(())`, so that it covers a nested plan without depending on how
the pull reports it. The paragraph this PR adds to `development.md` would then
describe the flag instead of the rule about skipping the park after a JNI call.
--
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]
