Re: [PR] feat: drop a stream task's future in place when its stream is dropped [datafusion]

["kosiew (via GitHub)" <gi...@apache.org>]

kosiew commented on code in PR #25763:
URL: https://github.com/apache/datafusion/pull/25763#discussion_r4282274219


##########
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:
   Could you also document the panic-unwinding exception here and in the JoinSet overview? `ReclaimHandle::reclaim` skips reclamation when the current thread is already panicking, so even an idle future then relies on Tokio cancellation for cleanup.



##########
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:
   Could we add a focused `spawn_buffered` regression test that verifies dropping the output stream immediately releases resources held by an idle producer? A multi-thread runtime with controlled worker progress and an input stream owning a drop sentinel or memory reservation would help catch an accidental switch back to ordinary `spawn`.



##########
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:
   Could we simplify the three ID-less join methods by delegating to their corresponding `*_with_id` methods and stripping the ID? This would avoid duplicating the `forget` bookkeeping and keep reclaim-map cleanup in one place for each join variant, while preserving `Poll::Pending` in the poll version.



-- 
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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org