Repository navigation
feat: drop a stream task's future in place when its stream is dropped - #25763
ranflarion wants to merge 7 commits into
Conversation
github-actions
Bot
added
common
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #25763 +/- ##
========================================
Coverage 82.84% 82.84%
========================================
Files 1151 1151
Lines 455510 455920 +410
Branches 455510 455920 +410
========================================
+ Hits 377362 377715 +353
- Misses 55236 55290 +54
- Partials 22912 22915 +3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
kosiew
left a comment
kosiew
left a comment
There was a problem hiding this comment.
Thanks for working on this. The reclaimable task implementation looks sound, and the tests cover the key lifecycle, cancellation, panic-handling, and tracer behavior. I also checked the synchronous-drop path in RepartitionExec and found no lock re-entrancy issues.
I have a few non-blocking suggestions around documentation, simplifying join bookkeeping, and adding physical-plan regression coverage.
Approved.
| } | ||
| } | ||
|
|
||
| /// Like [`Self::spawn`], but dropping the handle also drops the task's |
There was a problem hiding this 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.
There was a problem hiding this 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.
| /// [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; |
There was a problem hiding this 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.
There was a problem hiding this comment.
done, the three go through their *_with_id versions now, so forget is only called from those
| let sender = builder.tx(); | ||
|
|
||
| builder.spawn(async move { | ||
| builder.spawn_reclaimable(async move { |
There was a problem hiding this 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.
There was a problem hiding this 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.
There was a problem hiding this 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.
There was a problem hiding this 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.
|
This is interesting finding. Let me go through it |
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
Thanks @ranflarion. I read the slot and handle protocol against the callers and found no lock re-entrancy or lost wakeup. The try_lock miss in SlotFuture::poll returns Pending without a waker, which is fine because the handle aborts right after taking the future.
MemoryBufferedStream::newindatafusion/physical-plan/src/buffer.rs(behindBufferExec) still drives its input fromSpawnedTask::spawn, and it shares anArc<MemoryReservation>with that task. It has the same shape as thespawn_bufferedtask, so dropping aBufferExecstream still releases the input and the reservation only when the runtime next polls the task. Please either switch it tospawn_reclaimableor say why it stays out of scope.- The inline comments cover the unused builder variants, the duplicated test helpers and the
forgetcost for sets without reclaimable tasks.
| } | ||
|
|
||
| /// 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) |
There was a problem hiding this 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.
There was a problem hiding this comment.
removed both, JoinSet::spawn_reclaimable_on stays
| } | ||
|
|
||
| #[tokio::test(flavor = "multi_thread", worker_threads = 2)] | ||
| async fn drop_releases_idle_task_in_place() { |
There was a problem hiding this 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.
There was a problem hiding this 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
| use std::sync::mpsc::{Receiver, Sender, channel}; | ||
| use std::time::Duration; | ||
|
|
||
| struct BlockInDrop { |
There was a problem hiding this 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.
There was a problem hiding this comment.
moved PanicOnDrop, BlockInDrop and the one-worker runtime into a #[cfg(test)] test_utils module
| } | ||
| } | ||
|
|
||
| fn forget<R>(&mut self, joined: Option<&Result<(Id, R), JoinError>>) { |
There was a problem hiding this 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.
There was a problem hiding this comment.
It does hash, hashbrown's remove_entry calls make_hash before probing even when the table is empty. Added the early return, so sets without reclaimable tasks skip it.
|
marking as draft as it is waiting on comments to be addressed |
…their _with_id variants, and share test destructors
…mable_on, and assert release right after drop in the cancellation tests
Switched it to
|
|
run benchmark tpch tpcds |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing reclaim-spawned-tasks-on-drop (697144a) to ebf4858 (merge-base) diff Run configurationrun benchmark tpchResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing reclaim-spawned-tasks-on-drop (697144a) to ebf4858 (merge-base) diff Run configurationrun benchmark tpcdsResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing reclaim-spawned-tasks-on-drop (697144a) to ebf4858 (merge-base) diff Run configurationrun benchmark tpchCPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing reclaim-spawned-tasks-on-drop (697144a) to ebf4858 (merge-base) diff Run configurationrun benchmark tpcdsCPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
Which issue does this PR close?
Rationale for this change
Dropping a plan's stream only aborts the tasks that drive its inputs, so what they own (streams, memory reservations and the pool behind them) is released whenever the runtime next gets to them, and a caller that needs those resources back, DataFusion's own cancellation tests included, can only poll for them. The issue has the details.
What changes are included in this PR?
datafusion-common-runtimegets opt-in reclaimable spawns:SpawnedTask::spawn_reclaimable, andJoinSet::spawn_reclaimableandspawn_reclaimable_on. The task's future lives in a slot the worker locks while polling it. Dropping the handle, or callingabort_allorshutdownon the set, takes the future out and drops it on the spot when no worker is polling it, then aborts the task. Otherwise it aborts as today, and the future is dropped when that poll returns.tokio::task::try_id()isNone. That is why it is opt-in, and the constructors document it.spawnandspawn_onare unchanged.AbortHandle, and theJoinSetconstructors return the task'sIdinstead. So apart from runtime shutdown, only the owner cancels them, and it finishes the drop before aborting. The handle holds the slot only while it takes the future out, so no worker ever waits on a destructor running elsewhere.JoinError::Panic. The exception is a runtime shutdown duringabort_all, and it is documented.JoinSetTracerwraps the task rather than the reclaimed future, so tokio still drops the tracer inside the task.JoinSet::join_alljoins through the wrapper, so dropping it reclaims before aborting.In
datafusion-physical-plan,ReceiverStreamBuilderandRecordBatchReceiverStreamBuildergetspawn_reclaimable, and their existingspawnandspawn_onare unchanged. The tasks DataFusion spawns to drive plan streams use the new methods:spawn_buffered, the builder's internal input runner behindCoalescePartitionsExecandAnalyzeExec,RepartitionExec's input tasks,BufferExec'sMemoryBufferedStream,MaterializedCteExec's buffer tasks, andcollect_partitioned.Are these changes tested?
Yes. New unit tests in
datafusion-common-runtimecover:dropreturns, and its destructor runs in the runtime but outside the task, while plainspawnkeeps its task contextabort_allandjoin_nextjoin_allreclaims before abortingAn integration test in its own binary, since the tracer is process-wide, checks that an injected tracer is still dropped inside its task. Every opt-in has a test that fails when its call site goes back to
spawn: thetest_drop_canceltests ofAnalyzeExec,CoalescePartitionsExecandRepartitionExecandrecord_batch_receiver_stream_drop_cancelnow assert that nothing is left right after the drop, where they polled for it before, and new tests coverspawn_buffered,BufferExec,MaterializedCteExecandcollect_partitioned.Benchmarks, measured before
BufferExec,MaterializedCteExecandcollect_partitionedwere switched:dfbenchbuilt with--releasefrom main (6a792c6) and from this branch. TPC-H SF10 parquet on local NVMe, on a 32 vCPU i4i.8xlarge with default target partitions, 5 iterations per run, with runs ordered main, branch, branch, main.compare.pyon each adjacent pair:compare.pytpchSF10, round 1tpchSF10, round 2sort_tpchSF10, round 1sort_tpchSF10, round 2The same binary varied by up to 4.6% per query between its two runs, and the two flagged queries differ between rounds and point in opposite directions, so I read them as noise.
Are there any user-facing changes?
New public APIs:
SpawnedTask::spawn_reclaimable,JoinSet::spawn_reclaimableandspawn_reclaimable_on, andspawn_reclaimableon both stream builders. Existing spawn methods keep their behavior. The tasks DataFusion spawns to drive streams underspawn_buffered,CoalescePartitionsExec,AnalyzeExec,RepartitionExec,BufferExec,MaterializedCteExecandcollect_partitionednow drop their futures on the thread that drops the stream when idle. So destructors of the streams those tasks drive may run outside a tokio task, as they already do whenever a stream is not spawned.