LocalKey, references, and async functions

[inklesspen]

I discovered task_local and LocalKey today, because axum suggests it as an option for passing state around. However, I quickly ran into an issue when I want to call an async function with the task_local value via LocalKey.wait.

use tokio::time::{Duration, sleep};

async fn wait_a_while_and_greet(delay: Duration, name: &str) -> String {
    println!("...");
    sleep(delay).await;
    format!("Hello, {}.", name)
}

#[tokio::main]
async fn main() {
    let delay = Duration::from_secs(1);
    // This works fine, because the name has a static lifetime.
    let greeting = wait_a_while_and_greet(delay, "Alice").await;
    println!("{}", greeting);

    // This works fine, because although we're putting a reference into the future,
    // we consume the future immediately.
    let greeting = wait_a_while_and_greet(delay, &("Bob".to_string())).await;
    println!("{}", greeting);
}

Ok, this is basic, works as expected. Now let's try using a task_local.

use tokio::time::{Duration, sleep};

async fn wait_a_while_and_greet(delay: Duration, name: &str) -> String {
    println!("...");
    sleep(delay).await;
    format!("Hello, {}.", name)
}

tokio::task_local! {
    static NAME: String
}

// Cloning is no big deal, right?
async fn greet_a_clone(delay: Duration) -> String {
    wait_a_while_and_greet(delay, &NAME.get()).await
}

#[tokio::main]
async fn main() {
    let delay = Duration::from_secs(1);
    // This will work, but what if our clone gets anxiety?
    let greeting = NAME
        .scope(
            "Carol".to_string(),
            async move { greet_a_clone(delay).await },
        )
        .await;
    println!("{}", greeting);
}

This also works, but it requires cloning.

And so we get to my problem: I want to put something into the task_local, and then get a reference to it and call an async function with that reference. It worked fine to call the async function with a short-lived reference ("Bob") earlier, so this should be possible, right?

use tokio::time::{Duration, sleep};

async fn wait_a_while_and_greet(delay: Duration, name: &str) -> String {
    println!("...");
    sleep(delay).await;
    format!("Hello, {}.", name)
}

tokio::task_local! {
    static NAME: String
}

// What if cloning is fraught with ethical problems?
async fn greet_by_reference(delay: Duration) -> String {
    // This doesn't work, because the wrapper reference is used by the future and we get lifetime issues.
    // NAME.with(|name| wait_a_while_and_greet(delay, name)).await

    // What if we have an async closure inside .with?
    // This doesn't work, because of https://github.com/rust-lang/rust/issues/108114
    // NAME.with(async move |name| wait_a_while_and_greet(delay, name).await).await

    todo!("Can this be done?");
}

#[tokio::main]
async fn main() {
    let delay = Duration::from_secs(1);
    let greeting = NAME
        .scope("Dave".to_string(), async move {
            greet_by_reference(delay).await
        })
        .await;
    println!("{}", greeting);
}

I was surprised that LocalKey.async_with doesn't exist; is there a reason that's not feasible?

[Darksonn]

The problem is that the value in the LocalKey can in principle change or be deleted at .await points.

async fn wait_a_while_and_greet(delay: Duration, name: &str) -> String {
    println!("...");
    sleep(delay).await; // <-- the String may be deleted here
    format!("Hello, {}.", name) // <-- the &str may be dangling now
}

However, you can use an Arc to make cloning the value cheap:

use std::sync::Arc;
use tokio::time::{Duration, sleep};

async fn wait_a_while_and_greet(delay: Duration, name: &str) -> String {
    println!("...");
    sleep(delay).await; // <-- value in task local may change
    format!("Hello, {}.", name) // <-- but `name` still valid here (with old value)
}

tokio::task_local! {
    static NAME: Arc<str>
}

// Cloning is no big deal, right?
async fn greet_a_clone(delay: Duration) -> String {
    wait_a_while_and_greet(delay, &NAME.get().clone()).await
}

#[tokio::main]
async fn main() {
    let delay = Duration::from_secs(1);
    // This will work, but what if our clone gets anxiety?
    let greeting = NAME
        .scope(
            "Carol".into(),
            async move { greet_a_clone(delay).await },
        )
        .await;
    println!("{}", greeting);
}

[inklesspen]

Yeah, I know… it's just such a shame that so much async Rust style seems to favor clones and Arcs, when sync Rust is all about relying on the borrow checker and lifetimes.

[For8stt]

Yeah, that tradeoff is annoying, but it’s also the part that keeps this sound.

The problem isn’t just “async Rust likes clones”. It’s that once you hit .await, the current future can be suspended, and the task-local scope may no longer be the same when it resumes. So holding a borrowed &str from LocalKey across .await would let the future keep a reference to something Tokio can’t guarantee is still alive.

For this case I’d probably make the task-local value cheap to clone instead of trying to borrow it across awaits:

use std::sync::Arc;

tokio::task_local! {
    static NAME: Arc<str>;
}

async fn greet(delay: Duration) -> String {
    let name = NAME.get().clone();
    wait_a_while_and_greet(delay, &name).await
}

That clone is just an Arc bump, not cloning the string contents.

If this is request/app state and not really “task-local context”, another option is to pass an Arc<State> explicitly or use axum State. It’s less magical, but usually easier to reason about.

So yeah, clones/Arcs feel less elegant than sync borrowing, but here they’re basically the ownership boundary across .await.