ranflarion commented on code in PR #25763:
URL: https://github.com/apache/datafusion/pull/25763#discussion_r4363492014
##########
datafusion/common-runtime/src/common.rs:
##########
@@ -46,7 +55,36 @@ impl<R: 'static> SpawnedTask<R> {
// Ok to use spawn here as SpawnedTask handles aborting/cancelling the
task on Drop
#[expect(clippy::disallowed_methods)]
let inner = tokio::task::spawn(trace_future(task));
- Self { inner }
+ Self {
+ inner,
+ reclaim: None,
+ }
+ }
+
+ /// Like [`Self::spawn`], but dropping the handle also drops the task's
Review Comment:
Removed the exception instead of documenting it. The guard assumed a
destructor panic would abort the process while the thread unwinds, but a panic
caught by `catch_unwind` inside a destructor that runs during unwinding doesn't
abort (checked on 1.95, the MSRV, and 1.99), so `reclaim` now drops the future
in place during unwinding as well. `drop_while_unwinding_releases_in_place`
drops a handle from a panicking closure with a destructor that panics too, and
fails with the guard back. The `!panicking()` check in `SlotFuture::drop`
stays, since re-raising out of a drop during unwinding does abort.
##########
datafusion/common-runtime/src/join_set.rs:
##########
@@ -119,21 +189,27 @@ impl<T: 'static> JoinSet<T> {
/// [JoinSet::join_next](tokio::task::JoinSet::join_next) - Await the next
completed task.
pub async fn join_next(&mut self) -> Option<Result<T, JoinError>> {
- self.inner.join_next().await
+ let joined = self.inner.join_next_with_id().await;
Review Comment:
done, the three go through their `*_with_id` versions now, so `forget` is
only called from those
##########
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:
Added `test_spawn_buffered_drop_releases_idle_input`. It runs on a
one-worker runtime and parks that worker in a blocking task across the drop, so
the worker can neither be polling the buffered task nor drop its future, and
the input's strong count has to be 0 right after the drop. With `builder.spawn`
it fails, and it can't pass by timing, since the only worker is held.
##########
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:
done, the four assert a zero strong count right after the drop. They run on
a current-thread runtime, so no worker can be polling the task at that point.
Each fails when its call site goes back to `spawn`: `run_input` breaks the
analyze, coalesce and receiver stream tests, and either of `RepartitionExec`'s
two spawns breaks its test.
##########
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:
removed both, `JoinSet::spawn_reclaimable_on` stays
##########
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:
merged, `reclaimed_future_drops_in_the_runtime_outside_its_task` now also
checks the strong count on the dropping thread right after `drop` returns
--
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]
