loom_rs/
context.rs

1//! Thread-local runtime context.
2//!
3//! This module provides thread-local storage for accessing the current loom runtime
4//! from any thread managed by the runtime. This enables ergonomic access to runtime
5//! methods without passing the runtime reference explicitly.
6//!
7//! # Example
8//!
9//! ```ignore
10//! use loom_rs::LoomBuilder;
11//!
12//! let runtime = LoomBuilder::new().build()?;
13//!
14//! runtime.block_on(async {
15//!     // Works anywhere in the runtime's threads
16//!     let result = loom_rs::spawn_compute(|| {
17//!         expensive_work()
18//!     }).await;
19//! });
20//! ```
21
22use std::cell::RefCell;
23use std::sync::Weak;
24
25use crate::runtime::{LoomRuntime, LoomRuntimeInner};
26
27thread_local! {
28    static CURRENT_RUNTIME: RefCell<Option<Weak<LoomRuntimeInner>>> = const { RefCell::new(None) };
29}
30
31/// Get the current loom runtime from thread-local storage.
32///
33/// Returns `Some` when called from a thread managed by a `LoomRuntime`
34/// (tokio workers, rayon workers, or within `block_on`).
35///
36/// Returns `None` when called from outside a loom runtime context.
37///
38/// # Example
39///
40/// ```ignore
41/// runtime.block_on(async {
42///     // Works anywhere in the runtime's threads
43///     let rt = loom_rs::current_runtime().unwrap();
44///     rt.spawn_compute(|| work()).await;
45/// });
46/// ```
47pub fn current_runtime() -> Option<LoomRuntime> {
48    CURRENT_RUNTIME.with(|rt| {
49        rt.borrow()
50            .as_ref()
51            .and_then(|weak| weak.upgrade())
52            .map(LoomRuntime::from_inner)
53    })
54}
55
56/// Set the current runtime for this thread.
57///
58/// This is called internally by the runtime when starting tokio and rayon threads.
59pub(crate) fn set_current_runtime(runtime: Weak<LoomRuntimeInner>) {
60    CURRENT_RUNTIME.with(|rt| {
61        *rt.borrow_mut() = Some(runtime);
62    });
63}
64
65/// Clear the current runtime for this thread.
66///
67/// This is called internally by the runtime when threads are stopping.
68pub(crate) fn clear_current_runtime() {
69    CURRENT_RUNTIME.with(|rt| {
70        *rt.borrow_mut() = None;
71    });
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    #[test]
79    fn test_current_runtime_none_initially() {
80        // Should be None when not in a runtime context
81        assert!(current_runtime().is_none());
82    }
83
84    #[test]
85    fn test_set_and_clear_runtime() {
86        // Create a dummy inner runtime for testing
87        // We can't easily test this without a real runtime, so just test the API
88        assert!(current_runtime().is_none());
89
90        // After clear, should still be None
91        clear_current_runtime();
92        assert!(current_runtime().is_none());
93    }
94}
loom_rs/context.rs

loom_rs/
context.rs
