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::{Arc, Weak};
24
25use crate::runtime::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<Arc<LoomRuntimeInner>> {
48    CURRENT_RUNTIME.with(|rt| rt.borrow().as_ref().and_then(|weak| weak.upgrade()))
49}
50
51/// Set the current runtime for this thread.
52///
53/// This is called internally by the runtime when starting tokio and rayon threads.
54pub(crate) fn set_current_runtime(runtime: Weak<LoomRuntimeInner>) {
55    CURRENT_RUNTIME.with(|rt| {
56        *rt.borrow_mut() = Some(runtime);
57    });
58}
59
60/// Clear the current runtime for this thread.
61///
62/// This is called internally by the runtime when threads are stopping.
63pub(crate) fn clear_current_runtime() {
64    CURRENT_RUNTIME.with(|rt| {
65        *rt.borrow_mut() = None;
66    });
67}
68
69#[cfg(test)]
70mod tests {
71    use super::*;
72
73    #[test]
74    fn test_current_runtime_none_initially() {
75        // Should be None when not in a runtime context
76        assert!(current_runtime().is_none());
77    }
78
79    #[test]
80    fn test_set_and_clear_runtime() {
81        // Create a dummy inner runtime for testing
82        // We can't easily test this without a real runtime, so just test the API
83        assert!(current_runtime().is_none());
84
85        // After clear, should still be None
86        clear_current_runtime();
87        assert!(current_runtime().is_none());
88    }
89}
loom_rs/context.rs

loom_rs/
context.rs
