Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add TaskLocalFuture::into_value #5743

Closed
TennyZhuang opened this issue May 31, 2023 · 2 comments · Fixed by #6340
Closed

Add TaskLocalFuture::into_value #5743

TennyZhuang opened this issue May 31, 2023 · 2 comments · Fixed by #6340
Labels
A-tokio Area: The main tokio crate C-feature-request Category: A feature request. M-task Module: tokio/task

Comments

@TennyZhuang
Copy link
Contributor

TennyZhuang commented May 31, 2023

Is your feature request related to a problem? Please describe.

Sometimes we want to use LocalKey with Stream, and it's impossible to yield from an async closure. A workaround is wrap inside the stream iteration:

tokio::task_local! {
    static TASK_LOCAL_ALLOCATED_BYTES: AtomicUsize;
}

let allocated_bytes = AtomicUsize::new(0);
while let Some(message) =
        TASK_LOCAL_ALLOCATED_BYTES.scope(allocated_bytes, input.try_next()).await?
    {
        yield message;
    }

However, the scope will take the ownership of the slot value, and compile failed.

Describe the solution you'd like

I hope we can add a method to return the ownership of slot value instead of always dropping it, so that we can:

    let mut allocated_bytes = Some(AtomicUsize::new(0));

    loop {
        let f = TASK_LOCAL_ALLOCATED_BYTES.scope(allocated_bytes.take().unwrap(), input.try_next());

        let message = (&mut f).await?;
        if let Some(message) = message {
            yield message;
            allocated_bytes = Some(f.into_value());
        }
    }

Describe alternatives you've considered
A clear and concise description of any alternative solutions or features you've considered.

Directly support Stream, such as LocalKey::scope_stream.

Additional context
Add any other context or screenshots about the feature request here.

@TennyZhuang TennyZhuang added A-tokio Area: The main tokio crate C-feature-request Category: A feature request. labels May 31, 2023
@Darksonn Darksonn added the M-task Module: tokio/task label May 31, 2023
@Darksonn
Copy link
Contributor

This seems reasonable to me.

@TennyZhuang TennyZhuang changed the title Add TaskLocalFuture::into_slot Add TaskLocalFuture::into_value Jun 3, 2023
@Darksonn
Copy link
Contributor

Darksonn commented Jun 4, 2023

As pointed out in #5758 (comment), this method would not be callable once the future has been pinned. A take_value method might be a better idea?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
A-tokio Area: The main tokio crate C-feature-request Category: A feature request. M-task Module: tokio/task
Projects
None yet
Development

Successfully merging a pull request may close this issue.

2 participants