1 use super::{Context, CONTEXT};
2
3 use crate::runtime::{scheduler, TryCurrentError};
4 use crate::util::markers::SyncNotSend;
5
6 use std::cell::{Cell, RefCell};
7 use std::marker::PhantomData;
8
9 #[derive(Debug)]
10 #[must_use]
11 pub(crate) struct SetCurrentGuard {
12 // The previous handle
13 prev: Option<scheduler::Handle>,
14
15 // The depth for this guard
16 depth: usize,
17
18 // Don't let the type move across threads.
19 _p: PhantomData<SyncNotSend>,
20 }
21
22 pub(super) struct HandleCell {
23 /// Current handle
24 handle: RefCell<Option<scheduler::Handle>>,
25
26 /// Tracks the number of nested calls to `try_set_current`.
27 depth: Cell<usize>,
28 }
29
30 /// Sets this [`Handle`] as the current active [`Handle`].
31 ///
32 /// [`Handle`]: crate::runtime::scheduler::Handle
try_set_current(handle: &scheduler::Handle) -> Option<SetCurrentGuard>33 pub(crate) fn try_set_current(handle: &scheduler::Handle) -> Option<SetCurrentGuard> {
34 CONTEXT.try_with(|ctx| ctx.set_current(handle)).ok()
35 }
36
with_current<F, R>(f: F) -> Result<R, TryCurrentError> where F: FnOnce(&scheduler::Handle) -> R,37 pub(crate) fn with_current<F, R>(f: F) -> Result<R, TryCurrentError>
38 where
39 F: FnOnce(&scheduler::Handle) -> R,
40 {
41 match CONTEXT.try_with(|ctx| ctx.current.handle.borrow().as_ref().map(f)) {
42 Ok(Some(ret)) => Ok(ret),
43 Ok(None) => Err(TryCurrentError::new_no_context()),
44 Err(_access_error) => Err(TryCurrentError::new_thread_local_destroyed()),
45 }
46 }
47
48 impl Context {
set_current(&self, handle: &scheduler::Handle) -> SetCurrentGuard49 pub(super) fn set_current(&self, handle: &scheduler::Handle) -> SetCurrentGuard {
50 let old_handle = self.current.handle.borrow_mut().replace(handle.clone());
51 let depth = self.current.depth.get();
52
53 if depth == usize::MAX {
54 panic!("reached max `enter` depth");
55 }
56
57 let depth = depth + 1;
58 self.current.depth.set(depth);
59
60 SetCurrentGuard {
61 prev: old_handle,
62 depth,
63 _p: PhantomData,
64 }
65 }
66 }
67
68 impl HandleCell {
new() -> HandleCell69 pub(super) const fn new() -> HandleCell {
70 HandleCell {
71 handle: RefCell::new(None),
72 depth: Cell::new(0),
73 }
74 }
75 }
76
77 impl Drop for SetCurrentGuard {
drop(&mut self)78 fn drop(&mut self) {
79 CONTEXT.with(|ctx| {
80 let depth = ctx.current.depth.get();
81
82 if depth != self.depth {
83 if !std::thread::panicking() {
84 panic!(
85 "`EnterGuard` values dropped out of order. Guards returned by \
86 `tokio::runtime::Handle::enter()` must be dropped in the reverse \
87 order as they were acquired."
88 );
89 } else {
90 // Just return... this will leave handles in a wonky state though...
91 return;
92 }
93 }
94
95 *ctx.current.handle.borrow_mut() = self.prev.take();
96 ctx.current.depth.set(depth - 1);
97 });
98 }
99 }
100