sunchao commented on code in PR #6883: URL: https://github.com/apache/datafusion-comet/pull/6883#discussion_r4426169499
########## native/local/src/handle.rs: ########## @@ -0,0 +1,203 @@ +// 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. + +//! Delivery to blocking callers, with cancellation that does not wait for a reader's wait. + +use std::panic::AssertUnwindSafe; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex, TryLockError}; +use std::task::{Context, Poll, Wake, Waker}; +use std::thread::Thread; +use std::time::{Duration, Instant}; + +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::{DataFusionError, Result}; +use futures::StreamExt; +use tokio::runtime::Handle; + +use crate::{lock, LocalQuery, LocalQueryStream}; + +pub enum QueryPoll { + Batch(RecordBatch), + Pending, + Finished, +} + +/// A started query, read by one blocking caller at a time. The reader polls the query's stream +/// on its own thread, so a batch reaches it straight from its partition's task, and waits by +/// parking rather than on a Tokio timer, which needs a free worker. Cancelling or dropping the +/// handle drops the query's stream on the calling thread and unread batches with it. The runtime +/// must remain alive until DataFusion finishes aborting the tasks its operators spawn. +pub struct QueryHandle { + /// The query's stream. A reader holds it for the whole of a poll; `cancel` wakes the reader + /// and so waits only for a poll of the stream it is already in. + stream: Mutex<Option<LocalQueryStream>>, + /// The thread of the reader in a poll, which `cancel` wakes. + reader: Mutex<Option<Thread>>, + cancelled: AtomicBool, + runtime: Handle, +} + +/// Wakes a reader parked in `QueryHandle::poll`. +struct Unpark(Thread); + +impl Wake for Unpark { + fn wake(self: Arc<Self>) { + self.0.unpark(); + } +} + +impl QueryHandle { + pub fn start(query: LocalQuery, runtime: &Handle) -> Result<Self> { + let stream = { + let _guard = runtime.enter(); + query.execute()? + }; + Ok(Self { + stream: Mutex::new(Some(stream)), + reader: Mutex::new(None), + cancelled: AtomicBool::new(false), + runtime: runtime.clone(), + }) + } + + /// A bounded wait lets a JVM caller check Spark task interruption between polls. + /// Only one reader may call this at a time. + pub fn poll(&self, wait: Duration) -> Result<QueryPoll> { + let mut stream = match self.stream.try_lock() { + Ok(stream) => stream, + Err(TryLockError::Poisoned(poisoned)) => poisoned.into_inner(), + Err(TryLockError::WouldBlock) if self.cancelled.load(Ordering::Acquire) => { + return Err(cancelled()) + } + Err(TryLockError::WouldBlock) => { + return Err(DataFusionError::Execution( + "Concurrent local query result readers".into(), + )) + } + }; + *lock(&self.reader) = Some(std::thread::current()); + let result = self.poll_stream(&mut stream, wait); + *lock(&self.reader) = None; + result + } + + fn poll_stream( + &self, + stream: &mut Option<LocalQueryStream>, + wait: Duration, + ) -> Result<QueryPoll> { + let deadline = Instant::now() + wait; + let waker = Waker::from(Arc::new(Unpark(std::thread::current()))); + let mut cx = Context::from_waker(&waker); + // Starting a partition executes its operators, which may spawn tasks. + let _guard = self.runtime.enter(); + loop { + // Checked after the reader is registered, so a cancel either sees the reader and + // wakes it, or is seen here. + if self.cancelled.load(Ordering::Acquire) { + return Err(cancelled()); + } + let Some(query) = stream.as_mut() else { + return Err(cancelled()); + }; + // A panic must not unwind into the JNI caller. + let next = + std::panic::catch_unwind(AssertUnwindSafe(|| query.poll_next_unpin(&mut cx))) Review Comment: **Keep partition startup within the cancellation polling contract.** Advancing `PartitionOrder` here can call `plan.execute` for the next partition on the JNI reader. `SortExec` registers a consumer synchronously, and Comet’s fair pool takes the same lock for registration that it holds during a Spark memory grant. The deadline below cannot interrupt that lock acquisition. In a two-partition Rust probe with a 450 ms simulated grant delay, a 50 ms poll took 462 ms on this head. The previous producer returned `Pending` within 52 ms under the same delay. Both runs used the same cached DataFusion 55.1 build. The pinned 55.2 `SortExec` has the same synchronous registration. This delays Spark’s `checkCancellation()` for as long as the grant takes. Could partition startup remain interruptible from the JNI reader’s perspective? ########## native/core/src/local.rs: ########## @@ -0,0 +1,263 @@ +// 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. + +//! JNI adapter for local execution, which runs a whole admitted Comet plan as one DataFusion +//! graph inside the result task of a local Spark application. The registry holds only live +//! queries, never plans for reuse. EOF, errors and explicit close remove entries; Java task +//! completion also closes. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::{Arc, LazyLock}; +use std::time::{Duration, Instant}; + +use datafusion::physical_plan::union::UnionExec; +use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +use datafusion_comet_jni_bridge::{jni_new_global_ref, JVMClasses}; +use datafusion_comet_local::handle::{QueryHandle, QueryPoll}; +use datafusion_comet_local::LocalQuery; +use jni::objects::{JByteArray, JLongArray, JObject, JObjectArray, JString, ReleaseMode}; +use jni::sys::{jboolean, jint, jlong, JNI_FALSE}; +use jni::{Env, EnvUnowned}; +use parking_lot::Mutex; + +use crate::errors::{try_unwrap_or_throw, CometError, CometResult}; +use crate::execution::jni_api::{ + create_execution_context, create_memory_pool, get_runtime, prepare_output, + prepare_plan_settings, ExecutionContext, PlanInputs, PlanJvmRefs, +}; +use crate::execution::spark_config::{SparkConfig, SPARK_EXECUTOR_CORES}; + +struct Entry { + query: QueryHandle, + columns: usize, + // Holds the plan's memory pool and session until the query is closed. Dropped after the + // query, which cancels the graph that reserves from the pool. Only the creating thread + // uses it; the mutex makes the entry shareable with result pulls. + _context: Mutex<Box<ExecutionContext>>, +} + +// JNI carries numeric IDs instead of raw pointers, so close racing with a result pull +// cannot free an object the pull is using. Last reader releases the removed entry. +static QUERIES: LazyLock<Mutex<HashMap<i64, Arc<Entry>>>> = LazyLock::new(Mutex::default); +static NEXT_ID: AtomicI64 = AtomicI64::new(1); +const MAX_LIVE_QUERIES: usize = 1024; + +fn close(id: i64) { + let entry = QUERIES.lock().remove(&id); + if let Some(entry) = entry { + entry.query.cancel(); + } +} + +/// Creates and starts a local query over `partitions` partitions of a serialized Comet plan +/// without JVM inputs. Its memory pool and DataFusion session are those `Native.createPlan` +/// creates for a plan of the calling Spark task. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_local_NativeLocal_createQuery( + e: EnvUnowned, + _: JObject, + serialized_query: JByteArray, + serialized_spark_configs: JByteArray, + partitions: jint, + columns: jint, + comet_task_memory_manager_obj: JObject, + local_dirs: JObjectArray, + batch_size: jint, + off_heap_mode: jboolean, + memory_pool_type: JString, + memory_limit: jlong, + off_heap_size: jlong, + task_attempt_id: jlong, + task_cpus: jlong, + read_ahead_bytes: jlong, +) -> jlong { + try_unwrap_or_throw(&e, |env| { + if !(1..=MAX_PARTITIONS).contains(&partitions) || columns < 1 { + return Err(CometError::Internal( + "Invalid local query parameters".into(), + )); + } + let start = Instant::now(); + let settings = prepare_plan_settings(&env.convert_byte_array(serialized_spark_configs)?)?; + // Run as many partitions at once as Spark would run tasks, which also bounds how many + // operators share the memory pool. `local[N]` resolves to N here. + let max_running = settings.spark_config.get_usize(SPARK_EXECUTOR_CORES, 1); + JVMClasses::init(env); + let serialized_plan = env.convert_byte_array(serialized_query)?; + let task_memory_manager = + Arc::new(jni_new_global_ref!(env, comet_task_memory_manager_obj)?); + let memory_pool_type = memory_pool_type.try_to_string(env)?; + let mut local_dirs_vec = vec![]; + for i in 0..local_dirs.len(env)? { + let dir = local_dirs.get_element(env, i)?; + let dir = unsafe { JString::from_raw(&*env, dir.into_raw()) }; + local_dirs_vec.push(dir.try_to_string(env)?); + } + let id = next_id()?; + let inputs = PlanInputs { + id, + settings, + serialized_plan, + partition_count: partitions as usize, + batch_size: batch_size as usize, + off_heap_mode: off_heap_mode != JNI_FALSE, + memory_pool_type, + memory_limit, + off_heap_size, + task_attempt_id, + task_cpus: task_cpus as usize, + local_dirs: local_dirs_vec, + metrics_update_interval: 0, + start, + jvm: PlanJvmRefs::default(), + }; + let context = create_execution_context(inputs, |config| { + create_memory_pool(config, task_memory_manager, task_attempt_id) + })?; + let plan = local_plan(&context, partitions)?; + if plan.schema().fields().len() != columns as usize { + return Err(CometError::Internal( + "Local output schema width mismatch".into(), + )); + } + let query = LocalQuery::new( + plan, + context.session_ctx.task_ctx(), + max_running, Review Comment: **Share the partition budget across concurrent local collects.** Each `LocalQuery` receives the executor’s entire `max_running` allowance, while Spark schedules each collect as one task. This also breaks with `spark.task.cpus=1`. On Spark 4.1.3 with `local[16]`, 256 MiB off-heap and `spark.comet.batchSize=17`, I sorted 96,000 rows in 16 partitions by `sha2(cast(id as string), 256)`. One local collect succeeded. When two started together, both failed to reserve memory for `ExternalSorterMerge`, on two runs. The same two concurrent collects each returned all 96,000 rows with local execution disabled. Could local queries share executor-wide partition admission, so that starting a second job cannot double the number of active sorts? ########## docs/source/contributor-guide/local-execution.md: ########## @@ -0,0 +1,146 @@ +<!-- +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. +--> + +# Local Execution + +Local execution is an experimental mode for Spark applications that run with a local master. +Instead of running a query as Spark tasks that each create and run a native plan, a `collect` +runs the whole query as one DataFusion graph inside a single Spark task, for every partition at +once. Other actions run the query as ordinary Comet execution; see result delivery. + +The internal `spark.comet.exec.local.enabled` option, false by default, turns it on. A query +is admitted only as a whole. Every other query uses ordinary Comet execution. + +## What is admitted + +Admission runs in `CometRule`, on the plan that Comet's ordinary conversion produces. A query is +admitted when: + +- the application's `SparkContext` is local (`local[N]`), so the query runs in this process; +- the plan is a batch query and not a subquery being prepared; +- the converted plan is one native block, a tree of Comet native operators whose root carries + the block's serialized plan, with no JVM input, broadcast or shuffle inside it; +- the block's only leaf is a `CometRangeExec`, so the range must be converted to Comet + (`spark.comet.exec.range.enabled`); +- its expressions are deterministic, contain no subquery, and do not call back into the JVM, + which the serialized plan is also checked for; +- it has between 1 and 1024 partitions. + +A root `CollectLimitExec`, which Spark plans for `take`, `head`, `show` and `limit`, is not +admitted: it collects through `executeTake`, which does not run locally, so such a query keeps +ordinary Comet execution. Admitted plans have no exchange, so AQE does not wrap them and they are +admitted whether or not AQE is enabled. + +## Execution + +The admitted plan becomes `CometLocalResultExec` over `CometLocalExec`, a leaf whose RDD has a +single partition. That partition's task creates the native graph: + +- The memory pool, spill directories and DataFusion session are those `Native.createPlan` creates + for any Comet plan of the task: `create_execution_context` builds them, with Comet's function + registrations and session settings, and the pool acquires memory from Spark through the task's + `CometTaskMemoryManager`. +- The block's operators are created once per partition, by the planner for that partition, as a + Spark task would create them. Partition `i` of the graph is the output of the operators created + for partition `i`, so per-partition state such as a range slice is the same as with ordinary + Comet execution. +- `native/local` runs the partitions on Comet's Tokio runtime and delivers them in partition + order. At most `spark.executor.cores` partitions (N for `local[N]`) run at once, and the next + starts when the first is delivered, as Spark runs at most that many tasks. A partition's + operators register their memory consumers when it starts, so this also bounds how many share + the task's memory pool. The result therefore has Spark's row order. A partition's error is + reported when its turn comes. +- The running partitions do not wait for their turn to produce batches, so they all make + progress, as Spark's tasks do. What they hold ahead of delivery is bounded by + `spark.driver.maxResultSize`, the most a collect may bring to the driver; past it, each holds + at most two batches until the earlier partitions are delivered. These batches are not reserved + in the memory pool: a reservation could not be given back when an operator later needs the + memory, so it would fail queries that ordinary execution runs. Spark's own collect also holds + its task results outside the execution pool, on the driver. +- All the partitions' memory consumers share the one task's pool, where ordinary execution gives + each task its own share. Consumers that can spill each get a share of the task's pool, as they + would of a task's. Consumers that cannot spill, such as a sort's merge, are bounded together by + the pool, as the tasks' consumers are together, but no longer each by one task's share: one + partition can take more than a task could, leaving less for the others. + +`CometLocalExec` reports `UnknownPartitioning(1)`: its RDD is one partition, but its rows are +not clustered by anything. Spark's ordering of the block holds within each partition, so +`CometLocalExec` reports it only when it holds across the concatenated partitions: with one +partition, or when the partitions are ranges of a prefix of that ordering, as a range's slices +are. Sorting within partitions, for example, reports no ordering. `CometLocalResultExec` reports +the block's own partitioning and ordering, which hold for the partitions it returns. + +Explain shows the native block under `CometLocalExec`. Its native operators do not report SQL +metrics; `CometLocalExec` counts the rows and batches it delivers. + +## Result delivery + +The JVM's pull polls the query's stream on its own thread, so a batch reaches it straight from +its partition's task; the read-ahead bound above is the only buffering. Exported +Arrow arrays use `prepare_output` and `NativeUtil`; an imported batch stays valid after its +native query is closed, until the JVM closes it. A pull returns "pending" after 50 ms, so that +the task can check whether it was killed. It waits by parking its thread rather than on a Tokio +timer, which only fires once a worker is free. + +`CometLocalResultExec` sits above Spark's columnar-to-row transition. For `collect` it still +runs the single task as a Spark job, which keeps cancellation, job groups and SQL metrics, but +the task copies its rows into a driver-side slot instead of returning an encoded task result. Spark would otherwise encode and compress every row in that one task and decode +them on the driver, all on one thread. While copying, the task counts uncompressed `UnsafeRow` +bytes and stops reading, which closes the native query, once they exceed +`spark.driver.maxResultSize`; the driver then fails the action. Spark counts compressed bytes, +so this check is the stricter one. On a non-local master the node uses Spark's ordinary +collect. + +Every other action reads the result as an RDD of the block's partitions: an iterator fetches +one partition at a time so that it holds at most one, a tail scans partitions from the last one, and a cache, +`df.rdd`, `foreachPartition`, `localCheckpoint` or PySpark's collect and `toLocalIterator` keep Review Comment: **Distinguish classic PySpark `collect()` from `toLocalIterator()`.** Classic `DataFrame.collect()` calls `collectToPython()`, which calls `plan.executeCollect()` and reaches the single-task handoff for an admitted plan. I ran that JVM entry point on Spark 4.1.3 with a four-partition range projection. The `CometLocalExec` output metric advanced from 0 to all 128 rows. `toLocalIterator()` does use the partitioned path. Could the guide and the matching comment on `CometLocalResultExec.scala:43` describe these two actions separately? [Python dispatch](https://github.com/apache/spark/blob/v4.1.3/python/pyspark/sql/classic/dataframe.py#L441-L449) · [JVM dispatch](https://github.com/apache/spark/blob/v4.1.3/sql/core/src/main/scala/org/apache/spark/sql/classic/Dataset.scala#L2079-L2087) ########## native/local/src/handle.rs: ########## @@ -0,0 +1,203 @@ +// 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. + +//! Delivery to blocking callers, with cancellation that does not wait for a reader's wait. + +use std::panic::AssertUnwindSafe; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex, TryLockError}; +use std::task::{Context, Poll, Wake, Waker}; +use std::thread::Thread; +use std::time::{Duration, Instant}; + +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::{DataFusionError, Result}; +use futures::StreamExt; +use tokio::runtime::Handle; + +use crate::{lock, LocalQuery, LocalQueryStream}; + +pub enum QueryPoll { + Batch(RecordBatch), + Pending, + Finished, +} + +/// A started query, read by one blocking caller at a time. The reader polls the query's stream +/// on its own thread, so a batch reaches it straight from its partition's task, and waits by +/// parking rather than on a Tokio timer, which needs a free worker. Cancelling or dropping the +/// handle drops the query's stream on the calling thread and unread batches with it. The runtime +/// must remain alive until DataFusion finishes aborting the tasks its operators spawn. +pub struct QueryHandle { + /// The query's stream. A reader holds it for the whole of a poll; `cancel` wakes the reader + /// and so waits only for a poll of the stream it is already in. + stream: Mutex<Option<LocalQueryStream>>, + /// The thread of the reader in a poll, which `cancel` wakes. + reader: Mutex<Option<Thread>>, + cancelled: AtomicBool, + runtime: Handle, +} + +/// Wakes a reader parked in `QueryHandle::poll`. +struct Unpark(Thread); + +impl Wake for Unpark { + fn wake(self: Arc<Self>) { + self.0.unpark(); + } +} + +impl QueryHandle { + pub fn start(query: LocalQuery, runtime: &Handle) -> Result<Self> { + let stream = { + let _guard = runtime.enter(); + query.execute()? + }; + Ok(Self { + stream: Mutex::new(Some(stream)), + reader: Mutex::new(None), + cancelled: AtomicBool::new(false), + runtime: runtime.clone(), + }) + } + + /// A bounded wait lets a JVM caller check Spark task interruption between polls. + /// Only one reader may call this at a time. + pub fn poll(&self, wait: Duration) -> Result<QueryPoll> { + let mut stream = match self.stream.try_lock() { + Ok(stream) => stream, + Err(TryLockError::Poisoned(poisoned)) => poisoned.into_inner(), + Err(TryLockError::WouldBlock) if self.cancelled.load(Ordering::Acquire) => { + return Err(cancelled()) + } + Err(TryLockError::WouldBlock) => { + return Err(DataFusionError::Execution( + "Concurrent local query result readers".into(), + )) + } + }; + *lock(&self.reader) = Some(std::thread::current()); + let result = self.poll_stream(&mut stream, wait); + *lock(&self.reader) = None; + result + } + + fn poll_stream( + &self, + stream: &mut Option<LocalQueryStream>, + wait: Duration, + ) -> Result<QueryPoll> { + let deadline = Instant::now() + wait; + let waker = Waker::from(Arc::new(Unpark(std::thread::current()))); + let mut cx = Context::from_waker(&waker); + // Starting a partition executes its operators, which may spawn tasks. + let _guard = self.runtime.enter(); + loop { + // Checked after the reader is registered, so a cancel either sees the reader and + // wakes it, or is seen here. + if self.cancelled.load(Ordering::Acquire) { + return Err(cancelled()); + } + let Some(query) = stream.as_mut() else { + return Err(cancelled()); + }; + // A panic must not unwind into the JNI caller. + let next = + std::panic::catch_unwind(AssertUnwindSafe(|| query.poll_next_unpin(&mut cx))) + .map_err(|_| DataFusionError::Execution("Local query panicked".into()))?; + match next { + Poll::Ready(Some(batch)) => return batch.map(QueryPoll::Batch), + Poll::Ready(None) => return Ok(QueryPoll::Finished), + Poll::Pending => (), + } + let now = Instant::now(); + if now >= deadline { + return Ok(QueryPoll::Pending); + } + std::thread::park_timeout(deadline - now); + } + } + + /// Stops the query and drops its stream on the calling thread, rather than leaving that to + /// a Tokio task: a task only runs once a worker is free, and every worker can be waiting, + /// for instance in Spark's `acquireMemory` for memory this stream holds. + pub fn cancel(&self) { + self.cancelled.store(true, Ordering::Release); + if let Some(reader) = lock(&self.reader).as_ref() { + reader.unpark(); + } + let stream = lock(&self.stream).take(); + // A panic while dropping must not unwind into the JNI caller. + let _ = std::panic::catch_unwind(AssertUnwindSafe(|| drop(stream))); + } +} + +impl Drop for QueryHandle { + fn drop(&mut self) { + self.cancel(); + } +} + +fn cancelled() -> DataFusionError { + DataFusionError::Execution("Local query cancelled".into()) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::datatypes::Schema; + use datafusion::execution::TaskContext; + use datafusion::physical_plan::test::exec::BlockingExec; + + #[test] + fn cancel_wakes_a_waiting_reader() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .build() + .unwrap(); + let plan = BlockingExec::new(Arc::new(Schema::empty()), 2); + let refs = plan.refs(); + let query = Arc::new( + QueryHandle::start( + LocalQuery::new( + Arc::new(plan), + Arc::new(TaskContext::default()), + usize::MAX, + usize::MAX, + ), + runtime.handle(), + ) + .unwrap(), + ); + let reader = Arc::clone(&query); + let (sender, receiver) = std::sync::mpsc::channel(); + let thread = std::thread::spawn(move || { + sender.send(reader.poll(Duration::from_secs(60))).unwrap(); + }); + while lock(&query.reader).is_none() { Review Comment: **Make the wakeup test wait until the reader needs a wakeup.** `reader` becomes `Some` before `poll_stream` checks `cancelled` for the first time. If the polling thread is descheduled there, the test can call `cancel()` before the reader ever parks. With `reader.unpark()` removed and a 150 ms pause at that point, this test still passed. A bounded readiness signal after the stream returns `Pending`, immediately before `park_timeout`, would force the wakeup path and also keep setup from spinning indefinitely. -- 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]
