CODEWITHSUNDEEP
rustrust-tokio
Tauta
Tauta

Reputation: 936

How to initialize thread local variables using async function

I want to initialize thread local variables for all 4 threads at the beginning of the program.

thread_local! {
  static local: i32
}

#[tokio::main(worker_threads = 4)]
async fn main() {
   // local = get_local().await;
}

Upvotes: 0

Views: 1772

Answers (1)

sebpuetz
sebpuetz

Reputation: 2618

Your tokio runtime is configured to have 4 worker threads, your thread local is provided to the main thread but not to the worker threads.

If you intend to initialize the gRPC client just once, a OnceCell could be appropriate:

use once_cell::sync::OnceCell;

pub static CLIENT: OnceCell<hello_world::greeter_client::GreeterClient<tonic::transport::Channel>> =
    OnceCell::new();

pub fn client() -> hello_world::greeter_client::GreeterClient<tonic::transport::Channel> {
    CLIENT.get().unwrap().clone()
}

#[tokio::main]
async fn main() {
    let channel = tonic::transport::Endpoint::new("http://helloworld")
        .unwrap()
        .connect_lazy();
    let client = hello_world::greeter_client::GreeterClient::new(channel);
    CLIENT.set(client).unwrap();
    main_().await;
}

async fn main_() {
    let _ = client()
        .say_hello(hello_world::HelloRequest { name: "foo".into() })
        .await;
}

pub mod hello_world {
    tonic::include_proto!("helloworld");
}

If you want to stick to something more similar to a thread local or you need more control over the client values, then you can use tokio's task local. It allows you to provide context to tasks, but keep in mind that tokio::spawn introduces new tasks, so this context is lost when you use tokio::spawn.

The following snippet makes a tonic client available through a client() helper function that internally calls .with() on the task local. This function panics if the task local is not set, there is also try_with() which returns a Result if the value is not provided.

use tokio::task_local;

task_local! {
    pub static CLIENT: hello_world::greeter_client::GreeterClient<tonic::transport::Channel>
}

pub fn client() -> hello_world::greeter_client::GreeterClient<tonic::transport::Channel> {
    CLIENT.with(|c| c.clone())
}

#[tokio::main]
async fn main() {
    let channel = tonic::transport::Endpoint::new("http://helloworld")
        .unwrap()
        .connect_lazy();
    let client = hello_world::greeter_client::GreeterClient::new(channel);
    CLIENT.scope(client, main_()).await;
}

async fn main_() {
    let _ = client()
        .say_hello(hello_world::HelloRequest { name: "foo".into() })
        .await;
}

pub mod hello_world {
    tonic::include_proto!("helloworld");
}

Upvotes: 1

Related Questions