Implement Future for SpawnedTask. #15653

ashdnazg · 2025-04-09T09:45:16Z

Rationale for this change

Allows polling a SpawnedTask, instead of just joining it.
Required for #15654

What changes are included in this PR?

An impl Future was added to SpawnedTask.

In addition, the underlying implementation is changed from JoinSet to a JoinHandle to simplify
the code, as JoinSet doesn't provide any additional benefits.

Documentation was also updated to make clearer what happens when spawn_blocking is called.

Are these changes tested?

Added two tests that verify the behaviour on spawn.

Are there any user-facing changes?

No, although SpawnedTask is public so it does add to the API.

datafusion/common-runtime/src/common.rs

alamb

Thanks @ashdnazg for this contribution.

Implementing future for spawned task seems fine to me but I am not sure about changing to use a join handle (even though implemeting Drop to call abort might be fine)

Can you please add tests showing that any tasks launched by spawned tasks are actually cancelled (both before and after they have started)?

ashdnazg · 2025-04-09T13:12:38Z

@alamb I added two tests for the two scenarios.

gabotechs · 2025-04-09T15:26:14Z

datafusion/common-runtime/src/common.rs

-use crate::JoinSet;
-use tokio::task::JoinError;
+use tokio::task::{JoinError, JoinHandle};


I think not using crate::JoinSet is going to render the work being done by @geoffreyclaude in #14547 pretty much useless. Is this right @geoffreyclaude?

Maybe we can just add something similar to the spawned task

Good catch @gabotechs! Yes, pretty much. For both to be compatible, the same wrapping I implemented over the tokio::task::JoinSet probably needs to be done over tokio::task::spawn and tokio::task::spawn_blocking.

@ashdnazg I'll try to add unit tests for the instrumentation feature tomorrow which you could use to validate your change doesn't introduce any regression.

@geoffreyclaude I updated the PR with the tracing functions wrapping the spawned tasks.

@ashdnazg Looks good to me! I've opened #15673 to add the required regression tests on my feature.

ashdnazg · 2025-04-09T20:00:15Z

Tests should be more reliable now 🤦.

alamb

This looks really nice to me -- thank you @ashdnazg 🙏

Let's wait a while before merging to give @geoffreyclaude and others a chance to review it again if they want.

I went through this one carefully and I think it is 👨‍🍳 👌 very nice.

alamb · 2025-04-09T21:23:21Z

datafusion/common-runtime/src/common.rs

-        let mut inner = JoinSet::new();
-        inner.spawn(task);
+        #[allow(clippy::disallowed_methods)]
+        let inner = tokio::task::spawn(trace_future(task));


I think some rationale about why it is ok to use the disallowed methods would help. Something like:

Suggested change

let inner = tokio::task::spawn(trace_future(task));

// Ok to use spawn here as SpawnedTask handles aborting/cancelling the task on Drop

let inner = tokio::task::spawn(trace_future(task));

alamb · 2025-04-09T21:23:27Z

datafusion/common-runtime/src/common.rs

@@ -47,22 +54,20 @@ impl<R: 'static> SpawnedTask<R> {
        T: Send + 'static,
        R: Send,
    {
-        let mut inner = JoinSet::new();
-        inner.spawn_blocking(task);
+        #[allow(clippy::disallowed_methods)]


Suggested change

#[allow(clippy::disallowed_methods)]

// Ok to use spawn here as SpawnedTask handles aborting/cancelling the task on Drop

#[allow(clippy::disallowed_methods)]

alamb · 2025-04-09T21:28:07Z

datafusion/common-runtime/src/common.rs


    #[tokio::test]
    async fn runtime_shutdown() {
        let rt = Runtime::new().unwrap();
+        #[allow(clippy::async_yields_async)]


I think this is ok -- the only reason clippy picks this up now is that SpawnedTask is actually a future where previously one had to call join().await -- so TLDR this change makes sense to me (though the test is perhaps somewhat suspect)

If you want, I can change the rt.spawn to rt.spawn_blocking and then we get rid of both the unnecessary surrounding async and the clippy thing.

Maybe we can add a second test -- I think we should leave the existing test as is as part of ensuring this PR doesn't introduce any regressiosn

A second test wouldn't test anything new IMO so keeping as is.

geoffreyclaude · 2025-04-10T16:59:35Z

This looks really nice to me -- thank you @ashdnazg 🙏

Let's wait a while before merging to give @geoffreyclaude and others a chance to review it again if they want.

I went through this one carefully and I think it is 👨‍🍳 👌 very nice.

Thanks! Looks good to me, especially if the new tests pass. @ashdnazg could you please just rebase on top of main so the CI picks up my merged tests?

It allows polling a SpawnedTask, instead of just joining it. The implementation is changed from `JoinSet` to a `JoinHandle` to simplify the code, as `JoinSet` doesn't provide any additional benefits.

ashdnazg · 2025-04-10T19:04:42Z

Rebased and added docs per @alamb 's suggestions.

alamb · 2025-04-10T19:19:57Z

all the tests are green so let's go go go!

rluvaton · 2025-04-11T09:54:38Z

@ashdnazg thank you for your contribution, hope to see you contribute again!

It allows polling a SpawnedTask, instead of just joining it. The implementation is changed from `JoinSet` to a `JoinHandle` to simplify the code, as `JoinSet` doesn't provide any additional benefits.

ashdnazg mentioned this pull request Apr 9, 2025

Remove waits from blocking threads reading spill files. #15654

Merged

github-actions bot added the common Related to common crate label Apr 9, 2025

rluvaton approved these changes Apr 9, 2025

View reviewed changes

alamb reviewed Apr 9, 2025

View reviewed changes

datafusion/common-runtime/src/common.rs Show resolved Hide resolved

alamb reviewed Apr 9, 2025

View reviewed changes

ashdnazg force-pushed the spawned-future branch from 6242951 to 69041ad Compare April 9, 2025 12:54

gabotechs reviewed Apr 9, 2025

View reviewed changes

ashdnazg force-pushed the spawned-future branch from 69041ad to 496140a Compare April 9, 2025 19:57

alamb approved these changes Apr 9, 2025

View reviewed changes

geoffreyclaude mentioned this pull request Apr 10, 2025

Add tracing regression smoke tests #15672

Closed

Implement Future for SpawnedTask.

b47dbba

It allows polling a SpawnedTask, instead of just joining it. The implementation is changed from `JoinSet` to a `JoinHandle` to simplify the code, as `JoinSet` doesn't provide any additional benefits.

ashdnazg force-pushed the spawned-future branch from 496140a to b47dbba Compare April 10, 2025 18:31

alamb merged commit 18a80f0 into apache:main Apr 10, 2025
27 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement Future for SpawnedTask. #15653

Implement Future for SpawnedTask. #15653

ashdnazg commented Apr 9, 2025 •

edited

Loading

alamb left a comment

ashdnazg commented Apr 9, 2025

gabotechs Apr 9, 2025 •

edited

Loading

alamb Apr 9, 2025

geoffreyclaude Apr 9, 2025

geoffreyclaude Apr 9, 2025

ashdnazg Apr 9, 2025

geoffreyclaude Apr 10, 2025

ashdnazg commented Apr 9, 2025

alamb left a comment

alamb Apr 9, 2025

alamb Apr 9, 2025

alamb Apr 9, 2025

ashdnazg Apr 9, 2025

alamb Apr 10, 2025

ashdnazg Apr 10, 2025

geoffreyclaude commented Apr 10, 2025

ashdnazg commented Apr 10, 2025

alamb commented Apr 10, 2025

rluvaton commented Apr 11, 2025

	let inner = tokio::task::spawn(trace_future(task));
	// Ok to use spawn here as SpawnedTask handles aborting/cancelling the task on Drop
	let inner = tokio::task::spawn(trace_future(task));

	#[allow(clippy::disallowed_methods)]
	// Ok to use spawn here as SpawnedTask handles aborting/cancelling the task on Drop
	#[allow(clippy::disallowed_methods)]

Implement Future for SpawnedTask. #15653

Implement Future for SpawnedTask. #15653

Conversation

ashdnazg commented Apr 9, 2025 • edited Loading

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

alamb left a comment

Choose a reason for hiding this comment

ashdnazg commented Apr 9, 2025

gabotechs Apr 9, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ashdnazg commented Apr 9, 2025

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

geoffreyclaude commented Apr 10, 2025

ashdnazg commented Apr 10, 2025

alamb commented Apr 10, 2025

rluvaton commented Apr 11, 2025

ashdnazg commented Apr 9, 2025 •

edited

Loading

gabotechs Apr 9, 2025 •

edited

Loading