diff --git a/Cargo.lock b/Cargo.lock index 746b5f7..dfb7c64 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -87,6 +87,12 @@ version = "0.28.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4271d37baee1b8c7e4b708028c57d816cf9d2434acb33a549475f78c181f6253" +[[package]] +name = "hermit-abi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" + [[package]] name = "libc" version = "0.2.154" @@ -108,6 +114,16 @@ dependencies = [ "adler", ] +[[package]] +name = "num_cpus" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4161fcb6d602d4d2081af7c3a45852d875a03dd337a6bfdd6e06407b61342a43" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "object" version = "0.32.2" @@ -180,13 +196,14 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1adbebffeca75fcfd058afa480fb6c0b81e165a0323f9c9d39c9697e37c46787" dependencies = [ "backtrace", + "num_cpus", "pin-project-lite", "tokio-macros", ] [[package]] name = "tokio-inherit-task-local" -version = "0.1.0" +version = "0.2.0" dependencies = [ "const-random", "tokio", diff --git a/Cargo.toml b/Cargo.toml index 901bcdc..855a856 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "tokio-inherit-task-local" description = "Task local variables for tokio that can be inherited across a spawn" -version = "0.1.0" +version = "0.2.0" edition = "2021" repository = "https://github.com/Xaeroxe/tokio-inherit-task-local" documentation = "https://docs.rs/tokio-inherit-task-local" @@ -13,4 +13,4 @@ const-random = "0.1.18" tokio = { version = "1.37.0", features = ["rt"] } [dev-dependencies] -tokio = { version = "1.37.0", features = ["rt", "macros"]} +tokio = { version = "1.37.0", features = ["rt", "rt-multi-thread", "macros"]} diff --git a/src/lib.rs b/src/lib.rs index 6c3a1e6..ddac00c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -91,6 +91,32 @@ where } } +/// Returns a closure which has its own copy of the current table for inheritable task locals. +/// Intended for use with [`tokio::task::spawn_blocking`]. +/// +/// # Example +/// ``` +/// use tokio_inherit_task_local::{inherit_task_local, inheritable_task_local}; +/// +/// inheritable_task_local! { +/// static NUMBER: u32; +/// } +/// +/// #[tokio::main] +/// async fn main() { +/// let output = NUMBER.scope(1, async move { +/// tokio::task::spawn_blocking(inherit_task_local(|| NUMBER.get())).await.unwrap() +/// }).await; +/// assert_eq!(output, 1); +/// } +/// ``` +pub fn inherit_task_local(f: F) -> impl FnOnce() -> R + Send + 'static where F: FnOnce() -> R + Send + 'static { + let new_task_locals = INHERITABLE_TASK_LOCALS + .try_with(|task_locals| task_locals.clone()) + .unwrap_or_else(|_| TaskLocalInheritableTable::new(HashMap::new())); + move || INHERITABLE_TASK_LOCALS.sync_scope(new_task_locals, f) +} + tokio::task_local! { static INHERITABLE_TASK_LOCALS: TaskLocalInheritableTable }