nexus_async_rt/
context.rs

1//! Thread-local runtime context.
2//!
3//! All runtime state is accessible via free functions that read from
4//! thread-local storage. The TLS slots are set by [`Runtime::block_on`]
5//! and cleared on exit. All const-initialized for zero first-access cost.
6//!
7//! ```ignore
8//! use nexus_async_rt::{spawn, with_world, sleep, io, shutdown_signal};
9//!
10//! rt.block_on(async {
11//!     spawn(async {
12//!         with_world(|world| { /* ... */ });
13//!         sleep(Duration::from_secs(1)).await;
14//!         let listener = TcpListener::bind(addr, io());
15//!     });
16//!     shutdown_signal().await;
17//! });
18//! ```
19
20use std::cell::Cell;
21use std::sync::Arc;
22use std::sync::atomic::AtomicBool;
23use std::time::{Duration, Instant};
24
25use crate::io::{IoDriver, IoHandle};
26use crate::timer::{TimerDriver, TimerHandle};
27
28// =============================================================================
29// TLS slots — const-initialized, zero first-access cost
30// =============================================================================
31
32thread_local! {
33    static CTX_WORLD: Cell<*mut nexus_rt::World> =
34        const { Cell::new(std::ptr::null_mut()) };
35    static CTX_IO: Cell<*mut IoDriver> =
36        const { Cell::new(std::ptr::null_mut()) };
37    static CTX_TIMER: Cell<*mut TimerDriver> =
38        const { Cell::new(std::ptr::null_mut()) };
39    static CTX_EVENT_TIME: Cell<*const Cell<Instant>> =
40        const { Cell::new(std::ptr::null()) };
41    static CTX_SHUTDOWN: Cell<*const AtomicBool> =
42        const { Cell::new(std::ptr::null()) };
43    static CTX_SHUTDOWN_WAKER: Cell<*const Arc<std::sync::Mutex<Option<std::task::Waker>>>> =
44        const { Cell::new(std::ptr::null()) };
45}
46
47// =============================================================================
48// Install / clear (called by Runtime::block_on)
49// =============================================================================
50
51/// Install runtime context into TLS. Called by RuntimeBuilder::build().
52/// The context stays installed until the guard is dropped (Runtime::drop).
53pub(crate) fn install(
54    world: *mut nexus_rt::World,
55    io: *mut IoDriver,
56    timer: *mut TimerDriver,
57    event_time: *const Cell<Instant>,
58    shutdown_flag: *const AtomicBool,
59    shutdown_waker: *const Arc<std::sync::Mutex<Option<std::task::Waker>>>,
60) -> ContextGuard {
61    let prev = PrevContext {
62        world: CTX_WORLD.with(|c| c.replace(world)),
63        io: CTX_IO.with(|c| c.replace(io)),
64        timer: CTX_TIMER.with(|c| c.replace(timer)),
65        event_time: CTX_EVENT_TIME.with(|c| c.replace(event_time)),
66        shutdown: CTX_SHUTDOWN.with(|c| c.replace(shutdown_flag)),
67        shutdown_waker: CTX_SHUTDOWN_WAKER.with(|c| c.replace(shutdown_waker)),
68    };
69    ContextGuard { prev }
70}
71
72struct PrevContext {
73    world: *mut nexus_rt::World,
74    io: *mut IoDriver,
75    timer: *mut TimerDriver,
76    event_time: *const Cell<Instant>,
77    shutdown: *const AtomicBool,
78    shutdown_waker: *const Arc<std::sync::Mutex<Option<std::task::Waker>>>,
79}
80
81pub(crate) struct ContextGuard {
82    prev: PrevContext,
83}
84
85impl Drop for ContextGuard {
86    fn drop(&mut self) {
87        CTX_WORLD.with(|c| c.set(self.prev.world));
88        CTX_IO.with(|c| c.set(self.prev.io));
89        CTX_TIMER.with(|c| c.set(self.prev.timer));
90        CTX_EVENT_TIME.with(|c| c.set(self.prev.event_time));
91        CTX_SHUTDOWN.with(|c| c.set(self.prev.shutdown));
92        CTX_SHUTDOWN_WAKER.with(|c| c.set(self.prev.shutdown_waker));
93    }
94}
95
96/// Assert that we're inside a runtime context. Panics with `msg` if not.
97pub(crate) fn assert_in_runtime(msg: &str) {
98    let ptr = CTX_WORLD.with(Cell::get);
99    assert!(!ptr.is_null(), "{msg}");
100}
101
102// =============================================================================
103// Public free functions — the user-facing API
104// =============================================================================
105
106/// Access the [`World`](nexus_rt::World) with exclusive access.
107///
108/// Runs the closure synchronously inline. Must be called from within
109/// [`crate::Runtime::block_on`].
110///
111/// # Panics
112///
113/// Panics if called outside a runtime context.
114pub fn with_world<R>(f: impl FnOnce(&mut nexus_rt::World) -> R) -> R {
115    let ptr = CTX_WORLD.with(Cell::get);
116    assert!(
117        !ptr.is_null(),
118        "with_world() called outside Runtime::block_on"
119    );
120    // SAFETY: ptr set by install(), valid for Runtime lifetime.
121    // Single-threaded — exclusive access.
122    let world = unsafe { &mut *ptr };
123    f(world)
124}
125
126/// Access the [`World`](nexus_rt::World) with shared access.
127///
128/// # Panics
129///
130/// Panics if called outside a runtime context.
131pub fn with_world_ref<R>(f: impl FnOnce(&nexus_rt::World) -> R) -> R {
132    let ptr = CTX_WORLD.with(Cell::get);
133    assert!(
134        !ptr.is_null(),
135        "with_world_ref() called outside Runtime::block_on"
136    );
137    let world = unsafe { &*ptr };
138    f(world)
139}
140
141/// Returns the IO handle for registering mio sources.
142///
143/// # Panics
144///
145/// Panics if called outside a runtime context.
146pub fn io() -> IoHandle {
147    let ptr = CTX_IO.with(Cell::get);
148    assert!(!ptr.is_null(), "io() called outside Runtime::block_on");
149    // SAFETY: ptr valid for Runtime lifetime.
150    IoHandle::new(unsafe { &mut *ptr })
151}
152
153/// Create a [`Sleep`](crate::Sleep) future that completes after `duration`.
154///
155/// # Panics
156///
157/// Panics if called outside a runtime context.
158pub fn sleep(duration: Duration) -> crate::Sleep {
159    let ptr = CTX_TIMER.with(Cell::get);
160    assert!(!ptr.is_null(), "sleep() called outside Runtime::block_on");
161    // SAFETY: ptr valid for Runtime lifetime.
162    let handle = TimerHandle::new(unsafe { &mut *ptr });
163    handle.sleep(duration)
164}
165
166/// Create a [`Sleep`](crate::Sleep) future that completes at `deadline`.
167pub fn sleep_until(deadline: Instant) -> crate::Sleep {
168    let ptr = CTX_TIMER.with(Cell::get);
169    assert!(
170        !ptr.is_null(),
171        "sleep_until() called outside Runtime::block_on"
172    );
173    let handle = TimerHandle::new(unsafe { &mut *ptr });
174    handle.sleep_until(deadline)
175}
176
177/// Timestamp taken after the most recent IO poll cycle.
178///
179/// All events dispatched within the same cycle share this timestamp.
180/// One clock read per cycle, not per event.
181pub fn event_time() -> Instant {
182    let ptr = CTX_EVENT_TIME.with(Cell::get);
183    assert!(
184        !ptr.is_null(),
185        "event_time() called outside Runtime::block_on"
186    );
187    // SAFETY: ptr valid for Runtime lifetime.
188    unsafe { &*ptr }.get()
189}
190
191/// Wrap a future with a deadline. Returns `Err(Elapsed)` if the
192/// deadline expires before the future completes.
193///
194/// # Panics
195///
196/// Panics if called outside a runtime context.
197pub fn timeout<F: std::future::Future>(duration: Duration, future: F) -> crate::timer::Timeout<F> {
198    crate::timer::Timeout::new(future, sleep(duration))
199}
200
201/// Create an interval that ticks at a fixed period.
202///
203/// The first tick completes after `period`. Subsequent ticks are
204/// spaced `period` apart. If processing takes longer than `period`,
205/// behavior is controlled by [`MissedTickBehavior`](crate::MissedTickBehavior).
206///
207/// # Panics
208///
209/// Panics if `period` is zero. Polling the interval (via `tick().await`)
210/// requires an active runtime context and will panic otherwise.
211pub fn interval(period: Duration) -> crate::timer::Interval {
212    crate::timer::Interval::new(period)
213}
214
215/// Run a future no earlier than `deadline`.
216///
217/// Waits until `deadline`, then polls the future. Useful for
218/// scheduling deferred work at a specific time.
219///
220/// Polling requires an active runtime context.
221pub async fn after<F: std::future::Future>(deadline: Instant, future: F) -> F::Output {
222    sleep_until(deadline).await;
223    future.await
224}
225
226/// Run a future after `duration` elapses.
227///
228/// Waits for `duration`, then polls the future.
229///
230/// Polling requires an active runtime context.
231pub async fn after_delay<F: std::future::Future>(duration: Duration, future: F) -> F::Output {
232    sleep(duration).await;
233    future.await
234}
235
236/// Wrap a future with an absolute deadline. Returns `Err(Elapsed)` if
237/// the deadline passes before the future completes.
238///
239/// Like [`timeout`] but takes an [`Instant`] instead of a [`Duration`].
240///
241/// # Panics
242///
243/// Panics if called outside a runtime context.
244pub fn timeout_at<F: std::future::Future>(
245    deadline: Instant,
246    future: F,
247) -> crate::timer::Timeout<F> {
248    crate::timer::Timeout::new(future, sleep_until(deadline))
249}
250
251/// Create an interval that starts ticking at `start`, then every `period`.
252///
253/// If `start` is in the past, the first tick fires immediately.
254///
255/// # Panics
256///
257/// Panics if `period` is zero. Polling the interval requires an active
258/// runtime context.
259pub fn interval_at(start: Instant, period: Duration) -> crate::timer::Interval {
260    crate::timer::Interval::new_at(start, period)
261}
262
263/// Cooperatively yield the current task.
264///
265/// Returns `Pending` once, wakes itself, then completes on the next
266/// poll. Other ready tasks get a turn before this task resumes.
267pub fn yield_now() -> crate::timer::YieldNow {
268    crate::timer::YieldNow(false)
269}
270
271/// Returns a future that resolves when shutdown is triggered.
272pub fn shutdown_signal() -> crate::ShutdownSignal {
273    let ptr = CTX_SHUTDOWN.with(Cell::get);
274    assert!(
275        !ptr.is_null(),
276        "shutdown_signal() called outside Runtime::block_on"
277    );
278    let waker_ptr = CTX_SHUTDOWN_WAKER.with(Cell::get);
279    // SAFETY: waker_ptr was set by install() and is valid for Runtime lifetime.
280    let task_waker = unsafe { (*waker_ptr).clone() };
281    crate::ShutdownSignal {
282        flag: ptr,
283        task_waker,
284    }
285}
nexus_async_rt/context.rs

nexus_async_rt/
context.rs
