fix: park the native scan loop instead of busy-polling while waiting on native I/O

[mixermt]

Which issue does this PR close?

Closes #6091.

Rationale for this change

Java_org_apache_comet_Native_executePlan runs a plan that has any JVM-fed input (a broadcast build side, CometSparkRowToColumnar, a shuffle read) through a loop that polls the stream and, on Pending, pulls the next batches from the JVM. That pull blocks inside JNI while the JVM produces data, so the loop never spun as long as the JVM was the only thing worth waiting for. With native scans reading from S3 or HDFS the stream is also pending on asynchronous I/O, and once every JVM-fed scan holds a batch or has reached EOF the pull is a no-op: the loop re-polls at full speed for the duration of every read, pinning one core per task.

On a production workload (Iceberg on HDFS joined with a broadcast relation) the scan stages used 87 core-hours against 6.6 for Spark alone, ran 3x to 6x longer, and the saturated cores caused HDFS ack timeouts and retries. Details in #6091.

What changes are included in this PR?

• ScanStream and ShuffleScanStream honor the Stream contract. poll_next registers cx.waker() when the buffer is empty, and get_next_batch wakes it after refilling, through an AtomicWaker shared by the exec's clones. EOF stays buffered, so a re-poll returns Ready(None) again and get_next_batch is a no-op once the reader is drained.
• The ScanExec path of executePlan moves its loop into next_batch(stream, on_pending): poll the stream; on Pending, refill the JVM-fed scans and check the metrics interval, then park the block_on task until a waker fires (park_until_woken, a poll_fn that yields once). A refill wakes the task before it parks, so it resumes at once; a stream waiting on native I/O sleeps until that I/O wakes it. There is no timeout, because every Pending now carries a waker.
• Awaiting the stream directly is still not an option: the JVM refill has to run between polls, and the loop is what keeps that step reachable.
• update_metrics_on_interval replaces the per-100-polls gate. It runs on every pending poll and once per returned batch, and the tracing log_memory_usage sample sits behind the same interval, so trace density does not depend on how often the loop turns.
• development.md describes the pull-then-park loop.

How are these changes tested?

• next_batch_parks_while_the_stream_waits_on_native_io: a stream pending on tokio::time::sleep for 50 ms; the pull closure runs once. With the park removed it ran 135,311 times.
• next_batch_resumes_on_a_refill_and_stops_pulling_after_eof: a real ScanExec in test mode. Each park ends only on the refill's wake, under a timeout that turns a lost wake into a failure, and a re-poll after EOF pulls nothing. With EOF cleared on poll it fails.
• refill_wakes_the_pending_poll_and_eof_stays_buffered in shuffle_scan.rs: the empty-buffer poll registers the waker, the refill wakes it, and EOF stays buffered.
• cargo clippy --all-targets -p datafusion-comet -- -D warnings is clean; all 433 core unit tests pass.
• JVM, with the rebuilt library: CometTaskMetricsSuite, CometJoinSuite, CometNativeShuffleInputRDDSuite, CometNativeShuffleSuite and CometIcebergNativeSuite: 235 tests pass and none hang; the one canceled test is the pre-existing Spark 4.1 assume gate (SPARK-55626).
• Not measured here: the CPU reduction on the production workload itself, which needs a run with this build.

This touches the native execution loop, so the Spark SQL suites (run-spark-4.1-tests) should run before merge.

AI Disclosure

Drafted, implemented and tested with AI assistance (Claude Code); reviewed before submission.

🤖 Generated with Claude Code

[mbutrovich]

Thanks @mixermt. The diagnosis in #6091 closes the gap #3553 left open. #3553 parked the executor thread only for plans without JVM-fed inputs, so plans with a ScanExec or ShuffleScanExec still spin while native I/O is pending. You're right that the loop can't .await the stream, because the JVM refill has to run between polls. Parking for one wake-up keeps that refill reachable.

The threading section of development.md (lines 46-48) describes this loop. Can you update it in this PR to say the loop parks until a waker fires when nothing was pulled?

[andygrove]

Thanks for the writeup on this one, the diagnosis in #6091 made it easy to follow. I traced through the alternative you rejected and I agree it isn't safe. If the loop awaits the stream, the wake can land inside the await, an operator drains a JVM buffer and re-polls that scan within the same poll, and the no-waker Pending that comes back has no loop left to rescue it. Parking for exactly one wake-up is the right shape. I also checked that build_runtime already calls .enable_all(), so the timer driver is there for the park, and that spark.comet.metrics.updateInterval defaults to 3000 ms, comfortably above the 100 ms park, so metrics stay timely.

Two things on top of Matt's comments.

[mixermt]

Thanks both. Everything below is in the latest push, and the two threads changed the design rather than patching it.

• Wakers instead of a timeout (@mbutrovich). ScanStream and ShuffleScanStream now register cx.waker() when the buffer is empty, and get_next_batch wakes it after refilling, through an AtomicWaker shared by the exec's clones. Every Pending carries a waker, so the 100 ms timeout, the bool return from get_next_batch and the tokio time feature are gone; park_until_woken is a poll_fn that yields once. I dropped the timeout rather than logging its firings: with wakers in place a firing could not be told apart from a legitimately slow read, so a counter would not have found a lost wake.
• EOF is sticky (@andygrove). poll_next leaves InputBatch::EOF in the buffer, so a re-poll returns Ready(None) again and get_next_batch stays a no-op instead of making another JNI round trip into a drained reader. The resting-state invariant in Native executePlan busy-polls at 100% CPU while a plan with a JVM-fed input waits on native I/O #6091 is now what the code does.
• A loop test that fails without the fix (@mbutrovich). The poll, pull and park steps moved into next_batch(stream, on_pending); update_metrics and prepare_output stay in executePlan. next_batch_parks_while_the_stream_waits_on_native_io drives it with a stream pending on tokio::time::sleep and asserts the pull closure ran a few times; with the park deleted it ran 135,311 times in one 50 ms wait. next_batch_resumes_on_a_refill_and_stops_pulling_after_eof runs a real ScanExec in test mode: each park ends only on the refill's wake, under a timeout that turns a lost wake into a failure, and a re-poll after EOF pulls nothing; with EOF cleared again it fails. refill_wakes_the_pending_poll_and_eof_stays_buffered covers ShuffleScanStream.
• Tracing density (@andygrove). log_memory_usage sits behind the same interval check as update_metrics, in update_metrics_on_interval, which runs on every pending poll and once per returned batch, so in-loop trace density no longer depends on how often the loop turns.
• Docs (@mbutrovich). The JVM data source paragraph in development.md describes the pull-then-park loop. poll_fn is imported next to task::Poll and the Park struct is gone.

Verified with the rebuilt library: clippy with -D warnings is clean, the 433 core unit tests pass, and CometTaskMetricsSuite, CometJoinSuite, CometNativeShuffleInputRDDSuite, CometNativeShuffleSuite and CometIcebergNativeSuite pass (235 tests, no hangs).

This rewrites the native execution loop, so the Spark SQL suites should report here rather than in the merge queue. Could a committer add run-spark-4.1-tests?