Skip to main content

async_local_executor/
tls.rs

1use crate::executor::Executor;
2use std::cell::Cell;
3use std::ops::{Deref, DerefMut};
4use std::{mem, ptr};
5
6thread_local! {
7     static EXECUTOR: Cell<Option<&'static mut Executor>> = const { Cell::new(None) };
8}
9
10#[must_use = "deref in order to access the executor"]
11pub struct ExecutorGuard {
12    executor: &'static mut Executor,
13}
14
15impl Drop for ExecutorGuard {
16    fn drop(&mut self) {
17        let executor = unsafe { ptr::read(&raw const self.executor) };
18        EXECUTOR.set(Some(executor));
19    }
20}
21
22impl Deref for ExecutorGuard {
23    type Target = Executor;
24
25    fn deref(&self) -> &Self::Target {
26        self.executor
27    }
28}
29
30impl DerefMut for ExecutorGuard {
31    fn deref_mut(&mut self) -> &mut Self::Target {
32        self.executor
33    }
34}
35
36pub fn try_executor() -> Option<ExecutorGuard> {
37    Some(ExecutorGuard {
38        executor: EXECUTOR.try_with(Cell::take).ok()??,
39    })
40}
41
42#[cold]
43const fn no_executor<T>() -> T {
44    panic!("no executor present");
45}
46
47pub fn executor() -> ExecutorGuard {
48    try_executor().unwrap_or_else(no_executor)
49}
50
51impl Executor {
52    #[inline]
53    pub fn run_in<O, F: FnOnce() -> O>(&mut self, fun: F) -> O {
54        let prev = EXECUTOR.replace(Some(unsafe { mem::transmute(self) }));
55        let ret = fun();
56        EXECUTOR.replace(prev).unwrap_or_else(no_executor);
57        ret
58    }
59}