comphead commented on code in PR #25763:
URL: https://github.com/apache/datafusion/pull/25763#discussion_r4319197534
##########
datafusion/physical-plan/src/stream.rs:
##########
@@ -278,6 +299,27 @@ impl RecordBatchReceiverStreamBuilder {
self.inner.spawn_on(task, handle)
}
+ /// Like [`Self::spawn`], but dropping the builder or stream also drops the
+ /// task's future right away when no worker is polling it, with the
+ /// destructor contract of
+ ///
[`SpawnedTask::spawn_reclaimable`](datafusion_common_runtime::SpawnedTask::spawn_reclaimable)
+ pub fn spawn_reclaimable<F>(&mut self, task: F)
+ where
+ F: Future<Output = Result<()>>,
+ F: Send + 'static,
+ {
+ self.inner.spawn_reclaimable(task)
+ }
+
+ /// Same as [`Self::spawn_reclaimable`] but it spawns the task on the
provided runtime.
+ pub fn spawn_reclaimable_on<F>(&mut self, task: F, handle: &Handle)
Review Comment:
Neither `spawn_reclaimable_on` on `ReceiverStreamBuilder` nor this one has a
caller in the tree, and the existing `spawn_on` has a single test caller.
Consider dropping both builder variants and keeping
`JoinSet::spawn_reclaimable_on`, which `JoinSet::spawn_reclaimable` needs. That
is two fewer public APIs to keep stable.
##########
datafusion/common-runtime/src/common.rs:
##########
@@ -186,4 +355,126 @@ mod tests {
// The sender was dropped so we receive `None`.
assert!(receiver.recv().await.is_none());
}
+
+ fn wait_until_idle<R>(task: &SpawnedTask<R>) {
+ task.reclaim.as_ref().unwrap().wait_until_idle();
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn drop_releases_idle_task_in_place() {
Review Comment:
This test overlaps with
`reclaimed_future_drops_in_the_runtime_outside_its_task`. Both park a task,
wait until idle and drop it, and the later one already proves the future is
dropped on the dropping thread before `drop` returns. Merging them (add the
`Weak` count assertion there) keeps the coverage with one test fewer.
##########
datafusion/common-runtime/src/join_set.rs:
##########
@@ -56,6 +68,27 @@ impl<T> JoinSet<T> {
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
+
+ fn reclaim_all(&mut self) {
+ for (_, task) in self.reclaim.drain() {
+ task.reclaim();
+ }
+ }
+
+ fn forget<R>(&mut self, joined: Option<&Result<(Id, R), JoinError>>) {
Review Comment:
`forget` now runs on every join, including for sets with no reclaimable
tasks, which is every existing `JoinSet` user. A `HashMap::remove` hashes the
`Id` even when the map is empty (unverified for this hashbrown version). An
early return on `self.reclaim.is_empty()` would keep their path as it was.
Unmeasured.
##########
datafusion/common-runtime/src/join_set.rs:
##########
@@ -142,34 +218,271 @@ impl<T: 'static> JoinSet<T> {
&mut self,
cx: &mut Context<'_>,
) -> Poll<Option<Result<T, JoinError>>> {
- self.inner.poll_join_next(cx)
+ let joined = ready!(self.inner.poll_join_next_with_id(cx));
+ self.forget(joined.as_ref());
+ Poll::Ready(joined.map(|res| res.map(|(_, output)| output)))
}
/// [JoinSet::join_next_with_id](tokio::task::JoinSet::join_next_with_id)
- Await the next completed task with its ID.
pub async fn join_next_with_id(&mut self) -> Option<Result<(Id, T),
JoinError>> {
- self.inner.join_next_with_id().await
+ let joined = self.inner.join_next_with_id().await;
+ self.forget(joined.as_ref());
+ joined
}
///
[JoinSet::try_join_next_with_id](tokio::task::JoinSet::try_join_next_with_id) -
Try to join the next completed task with its ID.
pub fn try_join_next_with_id(&mut self) -> Option<Result<(Id, T),
JoinError>> {
- self.inner.try_join_next_with_id()
+ let joined = self.inner.try_join_next_with_id();
+ self.forget(joined.as_ref());
+ joined
}
///
[JoinSet::poll_join_next_with_id](tokio::task::JoinSet::poll_join_next_with_id)
- Poll for the next completed task with its ID.
pub fn poll_join_next_with_id(
&mut self,
cx: &mut Context<'_>,
) -> Poll<Option<Result<(Id, T), JoinError>>> {
- self.inner.poll_join_next_with_id(cx)
+ let joined = ready!(self.inner.poll_join_next_with_id(cx));
+ self.forget(joined.as_ref());
+ Poll::Ready(joined)
}
/// [JoinSet::shutdown](tokio::task::JoinSet::shutdown) - Abort all tasks
and wait for shutdown.
pub async fn shutdown(&mut self) {
+ self.reclaim_all();
self.inner.shutdown().await
}
/// [JoinSet::join_all](tokio::task::JoinSet::join_all) - Await all tasks.
- pub async fn join_all(self) -> Vec<T> {
- self.inner.join_all().await
+ pub async fn join_all(mut self) -> Vec<T> {
+ // Tokio's loop, run on this set so that dropping the future reclaims
before aborting
+ let mut output = Vec::with_capacity(self.len());
+ while let Some(res) = self.join_next().await {
+ match res {
+ Ok(t) => output.push(t),
+ Err(err) if err.is_panic() =>
std::panic::resume_unwind(err.into_panic()),
+ Err(err) => panic!("{err}"),
+ }
+ }
+ output
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ use std::future::pending;
+
+ use tokio::sync::oneshot;
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn drop_releases_idle_tasks_in_place() {
+ let owned = Arc::new(());
+ let released = Arc::downgrade(&owned);
+ let mut set = JoinSet::new();
+ let mut started = vec![];
+ for _ in 0..4 {
+ let owned = Arc::clone(&owned);
+ let (started_tx, started_rx) = oneshot::channel();
+ started.push(started_rx);
+ set.spawn_reclaimable(async move {
+ started_tx.send(()).unwrap();
+ pending::<()>().await;
+ drop(owned);
+ });
+ }
+ drop(owned);
+ for started_rx in started {
+ started_rx.await.unwrap();
+ }
+ set.reclaim
+ .values()
+ .for_each(ReclaimHandle::wait_until_idle);
+
+ drop(set);
+ assert_eq!(released.strong_count(), 0);
+ }
+
+ #[tokio::test]
+ async fn joined_tasks_are_forgotten() {
+ let mut set = JoinSet::new();
+ for i in 0..3 {
+ set.spawn_reclaimable(async move { i });
+ }
+ let mut outputs = vec![];
+ while let Some(output) = set.join_next().await {
+ outputs.push(output.unwrap());
+ }
+ outputs.sort_unstable();
+ assert_eq!(outputs, vec![0, 1, 2]);
+ assert!(set.reclaim.is_empty());
+ }
+
+ #[tokio::test]
+ async fn detached_tasks_keep_running() {
+ let mut set = JoinSet::new();
+ let (resume_tx, resume_rx) = oneshot::channel::<()>();
+ let (done_tx, done_rx) = oneshot::channel();
+ set.spawn_reclaimable(async move {
+ resume_rx.await.unwrap();
+ done_tx.send(()).unwrap();
+ });
+ set.detach_all();
+ drop(set);
+
+ resume_tx.send(()).unwrap();
+ done_rx.await.unwrap();
+ }
+
+ #[tokio::test]
+ async fn abort_all_reports_a_panic_from_a_reclaimed_destructor() {
+ struct PanicOnDrop;
+ impl Drop for PanicOnDrop {
+ fn drop(&mut self) {
+ panic!("destructor panic");
+ }
+ }
+ let mut set = JoinSet::new();
+ set.spawn_reclaimable(async move {
+ let _panics = PanicOnDrop;
+ pending::<()>().await;
+ });
+ // Lets the task start and park
+ tokio::task::yield_now().await;
+
+ set.abort_all();
+ assert!(set.join_next().await.unwrap().unwrap_err().is_panic());
+ }
+
+ #[test]
+ fn runtime_shutdown_does_not_wait_for_a_reclaiming_destructor() {
+ use std::sync::mpsc::{Receiver, Sender, channel};
+ use std::time::Duration;
+
+ struct BlockInDrop {
+ entered: Sender<()>,
+ resume: Receiver<()>,
+ }
+ impl Drop for BlockInDrop {
+ fn drop(&mut self) {
+ self.entered.send(()).unwrap();
+ let _ = self.resume.recv();
+ }
+ }
+
+ let runtime = tokio::runtime::Builder::new_multi_thread()
+ .worker_threads(1)
+ .build()
+ .unwrap();
+ let (entered_tx, entered_rx) = channel();
+ let (resume_tx, resume_rx) = channel();
+ let (started_tx, started_rx) = channel();
+ let guard = BlockInDrop {
+ entered: entered_tx,
+ resume: resume_rx,
+ };
+ let mut set = JoinSet::new();
+ set.spawn_reclaimable_on(
+ async move {
+ started_tx.send(()).unwrap();
+ pending::<()>().await;
+ drop(guard);
+ },
+ runtime.handle(),
+ );
+ started_rx.recv().unwrap();
+ set.reclaim
+ .values()
+ .for_each(ReclaimHandle::wait_until_idle);
+
+ // The set's drop runs the destructor on this thread until the test
resumes it
+ let dropper = std::thread::spawn(move || drop(set));
+ entered_rx.recv().unwrap();
+ let (shut_down_tx, shut_down_rx) = channel();
+ std::thread::spawn(move || {
+ runtime.shutdown_timeout(Duration::from_secs(10));
+ shut_down_tx.send(()).unwrap();
+ });
+ let shut_down =
shut_down_rx.recv_timeout(Duration::from_secs(5)).is_ok();
+ resume_tx.send(()).unwrap();
+ dropper.join().unwrap();
+ assert!(
+ shut_down,
+ "shutdown waited for a destructor on another thread"
+ );
+ }
+
+ #[test]
+ fn dropping_join_all_reclaims_before_aborting() {
+ use std::sync::mpsc::{Receiver, Sender, channel};
+ use std::time::Duration;
+
+ struct BlockInDrop {
Review Comment:
`BlockInDrop` and the runtime and channel setup are copied from
`runtime_shutdown_does_not_wait_for_a_reclaiming_destructor` above, and
`PanicOnDrop` is defined again in `common.rs`. A small shared `#[cfg(test)]`
helper would remove the duplicates.
##########
datafusion/physical-plan/src/common.rs:
##########
@@ -196,7 +196,7 @@ pub fn spawn_buffered(
let sender = builder.tx();
- builder.spawn(async move {
+ builder.spawn_reclaimable(async move {
Review Comment:
Extending this to the operator opt-ins: the PR keeps the retry loop of
`assert_strong_count_converges_to_zero` in all 11 cancellation tests, so
switching `RepartitionExec`, `CoalescePartitionsExec` or `AnalyzeExec` back to
plain `spawn` would still pass. The 4 tests that needed a retry on main
(`analyze`, `coalesce_partitions` and `repartition` `test_drop_cancel`, and
`record_batch_receiver_stream_drop_cancel`) could assert a zero strong count
right after the drop. That would cover those sites in addition to a
`spawn_buffered` test.
--
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]
