How to set tokio::task_local for server so that every handlers can use the same local static?

[hlf20010508]

Summary

I use

// indenter
pub enum Coroutine {
    Listener,
    AuthServer,
    Message,
    Progress,
    Task,
    TaskWorker(u8),
}

pub struct EventIndenter {
    coroutine: Coroutine,
    indent: Mutex<usize>,
}

task_local! {
    static EVENT_INDENTER: EventIndenter;
}

pub fn set_file_indenter<F>(coroutine: Coroutine, f: F) -> TaskLocalFuture<EventIndenter, F>
where
    F: Future,
{
    EVENT_INDENTER.scope(
        EventIndenter {
            coroutine,
            indent: std::sync::Mutex::new(0),
        },
        f,
    )
}

// server
tokio::spawn(async move {
    indenter::set_file_indenter(indenter::Coroutine::AuthServer, async {
        axum_server::from_tcp(server)
            .handle(shutdown_handle_clone)
            .serve(router.into_make_service())
            .await
            .unwrap();
    })
    .await;
});

to spawn a coroutine for axum server, and use EVENT_INDENTER in each handler.

I expect every handlers share the same EVENT_INDENTER, but different from the one in other coroutines.

But it panic at

thread 'tokio-runtime-worker' panicked at src/trace/indenter.rs:78:24:
cannot access a task-local storage value without setting it first

The indenter works fine in other coroutines except axum.

It seems that handlers are not in the same coroutine with the server.

Do anyone know how to fix this?

axum version

0.7.5

[jplatte]

See this example from the documentation: https://docs.rs/axum/0.7.7/axum/#using-tokios-task_local-macro