tls_api/
thread_local_context.rs

1use std::cell::Cell;
2use std::ptr;
3use std::task::Context;
4
5thread_local! {
6    pub static CONTEXT: Cell<*mut ()> = const { Cell::new(ptr::null_mut()) };
7}
8
9struct RestoreOnDrop(*mut ());
10
11impl Drop for RestoreOnDrop {
12    fn drop(&mut self) {
13        CONTEXT.with(|cell| {
14            cell.set(self.0);
15        })
16    }
17}
18
19/// Store future context in the thread local.
20pub fn save_context<R>(context: &mut Context<'_>, callback: impl FnOnce() -> R) -> R {
21    CONTEXT.with(|cell| {
22        let prev = cell.replace(context as *mut Context<'_> as *mut ());
23        let _restore_on_drop = RestoreOnDrop(prev);
24        callback()
25    })
26}
27
28/// Fetch future context from the thread local.
29pub fn restore_context<R>(callback: impl FnOnce(&mut Context<'_>) -> R) -> R {
30    CONTEXT.with(|cell| {
31        let context = cell.replace(ptr::null_mut());
32        let _restore_on_drop = RestoreOnDrop(context);
33        assert!(!context.is_null());
34        cell.set(ptr::null_mut());
35        callback(unsafe { &mut *(context as *mut Context) })
36    })
37}
tls_api/thread_local_context.rs

tls_api/
thread_local_context.rs
