Skip to main content

ntex_rt/
arbiter.rs

1#![allow(clippy::missing_panics_doc)]
2use std::sync::{Arc, atomic::AtomicBool, atomic::AtomicUsize, atomic::Ordering};
3use std::{any::Any, any::TypeId, cell::RefCell, fmt, mem, panic, pin::Pin, rc::Rc, thread};
4
5use async_channel::{Receiver, Sender, unbounded};
6use parking_lot::Mutex;
7
8use crate::metrics::{LatencyMetrics, LatencyStats, start_latency_sampler};
9use crate::{Handle, HashMap, Id, System};
10
11thread_local!(
12    static ADDR: RefCell<Option<Arbiter>> = const { RefCell::new(None) };
13    static STORAGE: RefCell<HashMap<TypeId, Rc<dyn Any>>> = RefCell::new(HashMap::default());
14    static ON_SHUTDOWN: RefCell<Vec<Box<dyn FnOnce()>>> = const { RefCell::new(Vec::new()) };
15);
16
17pub(super) static COUNT: AtomicUsize = AtomicUsize::new(99);
18
19pub(super) enum ArbiterCommand {
20    Stop,
21    #[allow(dead_code)]
22    Execute(Pin<Box<dyn Future<Output = ()> + Send>>),
23}
24
25/// An asynchronous execution environment running on one OS thread.
26///
27/// Creating an arbiter starts a thread with its own local event loop. Futures
28/// spawned on that event loop are not required to implement `Send`.
29pub struct Arbiter(pub(crate) Arc<ArbiterInner>);
30
31type OnCloseStorage = Arc<Mutex<Vec<Box<dyn Fn() + Send + Sync>>>>;
32
33pub(crate) struct ArbiterInner {
34    id: usize,
35    name: Arc<String>,
36    sys_id: usize,
37    hnd: Option<Handle>,
38    pub(crate) sender: Sender<ArbiterCommand>,
39    thread_handle: Mutex<Option<thread::JoinHandle<()>>>,
40    on_stop: OnCloseStorage,
41    running: AtomicBool,
42    pub(crate) latency: Arc<LatencyStats>,
43    #[cfg(target_os = "linux")]
44    tid: i32,
45    #[cfg(windows)]
46    thread: Option<crate::watchdog::ThreadHandle>,
47}
48
49impl fmt::Debug for Arbiter {
50    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
51        write!(f, "Arbiter({:?})", self.0.name.as_ref())
52    }
53}
54
55impl Clone for Arbiter {
56    fn clone(&self) -> Self {
57        Self(self.0.clone())
58    }
59}
60
61impl Default for Arbiter {
62    fn default() -> Self {
63        Self::new()
64    }
65}
66
67impl Arbiter {
68    #[allow(clippy::borrowed_box)]
69    pub(super) fn new_system(id: usize, name: String) -> (Self, ArbiterController) {
70        let (tx, rx) = unbounded();
71
72        let aid = COUNT.fetch_add(1, Ordering::Relaxed);
73        let arb = Arbiter::with_sender(id, aid, Arc::new(name), tx, Arc::default());
74        ADDR.with(|cell| *cell.borrow_mut() = Some(arb.clone()));
75        clear_storage();
76
77        (
78            arb,
79            ArbiterController {
80                rx,
81                sys: None,
82                stop: None,
83            },
84        )
85    }
86
87    /// Returns the arbiter running on the current thread.
88    ///
89    /// # Panics
90    ///
91    /// Panics if no arbiter is running on the current thread.
92    pub fn current() -> Arbiter {
93        ADDR.with(|cell| match *cell.borrow() {
94            Some(ref addr) => addr.clone(),
95            None => panic!("Arbiter is not running"),
96        })
97    }
98
99    /// Requests that the arbiter stop its event loop.
100    pub fn stop(&self) {
101        let _ = self.0.sender.try_send(ArbiterCommand::Stop);
102    }
103
104    /// Starts an arbiter on a new thread with an automatically generated name.
105    pub fn new() -> Arbiter {
106        let id = COUNT.fetch_add(1, Ordering::Relaxed);
107        Arbiter::start(id, format!("{}:arb:{}", System::current().name(), id), None)
108    }
109
110    /// Starts an arbiter on a new thread with the specified name.
111    pub fn with_name(name: String) -> Arbiter {
112        Arbiter::start(COUNT.fetch_add(1, Ordering::Relaxed), name, None)
113    }
114
115    /// Starts an arbiter on a new thread with the specified name.
116    ///
117    /// `init` runs on the new thread before the arbiter's runtime is created,
118    /// for example to set the thread's CPU affinity, so that the runtime's
119    /// memory is allocated on the NUMA node of that CPU.
120    pub fn with_init<F>(name: String, init: F) -> Arbiter
121    where
122        F: FnOnce() + Send + 'static,
123    {
124        Arbiter::start(
125            COUNT.fetch_add(1, Ordering::Relaxed),
126            name,
127            Some(Box::new(init)),
128        )
129    }
130
131    fn start(id: usize, name: String, init: Option<Box<dyn FnOnce() + Send>>) -> Arbiter {
132        let sys = System::current();
133        let name2 = Arc::new(name.clone());
134        let config = sys.config();
135        let (arb_tx, arb_rx) = unbounded();
136
137        let builder = if sys.config().stack_size > 0 {
138            thread::Builder::new()
139                .name(name)
140                .stack_size(sys.config().stack_size)
141        } else {
142            thread::Builder::new().name(name)
143        };
144
145        let name = name2.clone();
146        let sys_id = sys.id();
147        let (arb_hnd_tx, arb_hnd_rx) = oneshot::channel();
148
149        let handle = builder
150            .spawn(move || {
151                if let Some(init) = init {
152                    init();
153                }
154                let name3 = name2.clone();
155                log::info!("Starting {name3:?} arbiter");
156
157                let sys2 = sys.clone();
158                let (stop, stop_rx) = oneshot::channel();
159                clear_storage();
160
161                // `block_on` finalizes the arbiter once its event loop exits
162                let result = crate::driver::block_on(config.runner.as_ref(), async move {
163                    let arb = Arbiter::with_sender(sys_id.0, id, name2, arb_tx, Arc::default());
164                    sys.register_arbiter(arb.clone());
165                    // `Arbiter::latency()` is available once the arbiter is returned
166                    start_latency_sampler(&arb.0.latency, config.latency_interval);
167                    arb_hnd_tx
168                        .send(arb.clone())
169                        .expect("Controller thread has gone");
170
171                    // start arbiter controller
172                    crate::spawn_detached(
173                        ArbiterController {
174                            sys: None,
175                            stop: Some(stop),
176                            rx: arb_rx,
177                        }
178                        .run(sys),
179                    );
180                    ADDR.with(|cell| *cell.borrow_mut() = Some(arb.clone()));
181
182                    // run loop
183                    let _ = stop_rx.await;
184                });
185
186                // unregister arbiter
187                sys2.unregister_arbiter(Id(id));
188                // skipped by `block_on` if the event loop panicked
189                run_shutdown_callbacks();
190                unsafe {
191                    remove_all_items();
192                }
193
194                if let Err(e) = result {
195                    log::error!("Arbiter {name3:?} has panicked.");
196                    panic::resume_unwind(e);
197                }
198                log::info!("Arbiter {name3:?} has stopped");
199            })
200            .unwrap_or_else(|err| panic!("Cannot spawn an arbiter's thread {name:?}: {err:?}"));
201
202        let arb = arb_hnd_rx.recv().expect("Could not start new arbiter");
203        *arb.0.thread_handle.lock() = Some(handle);
204        arb
205    }
206
207    fn with_sender(
208        sys_id: usize,
209        id: usize,
210        name: Arc<String>,
211        sender: Sender<ArbiterCommand>,
212        on_stop: OnCloseStorage,
213    ) -> Self {
214        #[cfg(feature = "tokio")]
215        let hnd = { Handle::new(sender.clone()) };
216
217        #[cfg(feature = "compio")]
218        let hnd = { Handle::new(sender.clone()) };
219
220        #[cfg(all(not(feature = "compio"), not(feature = "tokio")))]
221        let hnd = { Handle::current() };
222
223        Self(Arc::new(ArbiterInner {
224            id,
225            sys_id,
226            name,
227            sender,
228            on_stop,
229            hnd: Some(hnd),
230            thread_handle: Mutex::new(None),
231            running: AtomicBool::new(true),
232            latency: Arc::default(),
233            #[cfg(target_os = "linux")]
234            #[allow(clippy::cast_possible_truncation)]
235            tid: unsafe { libc::syscall(libc::SYS_gettid) } as i32,
236            #[cfg(windows)]
237            thread: crate::watchdog::ThreadHandle::current(),
238        }))
239    }
240
241    /// Returns the arbiter identifier.
242    pub fn id(&self) -> Id {
243        Id(self.0.id)
244    }
245
246    #[cfg(target_os = "linux")]
247    /// TID of the arbiter
248    pub(crate) fn tid(&self) -> i32 {
249        self.0.tid
250    }
251
252    #[cfg(windows)]
253    /// Handle of the arbiter thread
254    pub(crate) fn thread(&self) -> Option<&crate::watchdog::ThreadHandle> {
255        self.0.thread.as_ref()
256    }
257
258    /// Returns the arbiter name.
259    pub fn name(&self) -> &str {
260        self.0.name.as_ref()
261    }
262
263    #[inline]
264    /// Returns a handle to the arbiter's runtime.
265    pub fn handle(&self) -> &Handle {
266        self.0.hnd.as_ref().unwrap()
267    }
268
269    /// Returns the arbiter's scheduling latency.
270    ///
271    /// Returns `None` if latency sampling is disabled, see
272    /// [`Builder::latency_interval`](crate::Builder::latency_interval).
273    pub fn latency(&self) -> Option<LatencyMetrics> {
274        self.0.latency.metrics()
275    }
276
277    #[inline]
278    /// Returns whether the arbiter is running.
279    pub fn is_running(&self) -> bool {
280        self.0.running.load(Ordering::Relaxed)
281    }
282
283    /// Returns a value from thread-local arbiter storage, inserting it if absent.
284    ///
285    /// If the storage has already been destroyed because the thread is
286    /// exiting, the value returned by `f` is not stored.
287    pub fn get_value<T, F>(f: F) -> T
288    where
289        T: Clone + 'static,
290        F: FnOnce() -> T,
291    {
292        let mut f = Some(f);
293        STORAGE
294            .try_with(|cell| {
295                let mut st = cell.borrow_mut();
296                if let Some(boxed) = st.get(&TypeId::of::<T>())
297                    && let Some(val) = (&**boxed as &(dyn Any + 'static)).downcast_ref::<T>()
298                {
299                    return val.clone();
300                }
301                let val = (f.take().unwrap())();
302                st.insert(TypeId::of::<T>(), Rc::new(val.clone()));
303                val
304            })
305            .unwrap_or_else(|_| (f.take().unwrap())())
306    }
307
308    /// Registers a callback to run when the current thread's arbiter shuts down.
309    ///
310    /// Callbacks run once, in registration order, on the arbiter's thread when
311    /// its stop is requested, by [`Arbiter::stop()`] or [`System::stop()`],
312    /// while the event loop is still running. An arbiter that ends without a
313    /// stop request, such as one driven by
314    /// [`SystemRunner::block_on()`](crate::SystemRunner::block_on), runs them
315    /// once its event loop has exited instead. A callback registered by
316    /// another callback runs in the same shutdown.
317    ///
318    /// If the thread is exiting and its thread-local storage has already been
319    /// destroyed, `f` is dropped without running.
320    pub fn on_shutdown<F>(f: F)
321    where
322        F: FnOnce() + 'static,
323    {
324        let f: Box<dyn FnOnce()> = Box::new(f);
325        let _ = ON_SHUTDOWN.try_with(move |cell| cell.borrow_mut().push(f));
326    }
327
328    #[must_use]
329    /// Adds a callback to run after the arbiter stops.
330    pub fn on_stop<F>(self, f: F) -> Self
331    where
332        F: Fn() + Send + Sync + 'static,
333    {
334        self.0.on_stop.lock().push(Box::new(f));
335        self
336    }
337
338    /// Returns the arbiter of the current thread, if any.
339    pub(crate) fn try_current() -> Option<Arbiter> {
340        ADDR.try_with(|cell| cell.borrow().clone()).ok().flatten()
341    }
342
343    /// Marks the arbiter as stopped and runs its `on_stop` callbacks,
344    /// called once its event loop has exited.
345    pub(crate) fn finalize(&self) {
346        self.0.running.store(false, Ordering::Relaxed);
347        let on_stop = mem::take(&mut *self.0.on_stop.lock());
348        for f in on_stop {
349            if thread::panicking() {
350                // a second panic while unwinding would abort the process
351                let _ = panic::catch_unwind(panic::AssertUnwindSafe(f));
352            } else {
353                f();
354            }
355        }
356    }
357
358    /// Waits for the arbiter's thread to stop.
359    ///
360    /// This returns immediately for an arbiter that does not own a thread
361    /// handle, including the system's primary arbiter.
362    pub fn join(&mut self) -> thread::Result<()> {
363        if let Some(thread_handle) = self.0.thread_handle.lock().take() {
364            thread_handle.join()
365        } else {
366            Ok(())
367        }
368    }
369}
370
371impl Eq for Arbiter {}
372
373impl PartialEq for Arbiter {
374    fn eq(&self, other: &Self) -> bool {
375        self.0.id == other.0.id && self.0.sys_id == other.0.sys_id
376    }
377}
378
379pub(crate) struct ArbiterController {
380    sys: Option<System>,
381    rx: Receiver<ArbiterCommand>,
382    stop: Option<oneshot::Sender<i32>>,
383}
384
385impl ArbiterController {
386    pub(super) async fn run(mut self, sys: System) {
387        self.sys = Some(sys);
388        loop {
389            match self.rx.recv().await {
390                Ok(ArbiterCommand::Stop) => {
391                    // the system arbiter has no `stop`, `System::stop()`
392                    // runs its callbacks
393                    if let Some(stop) = self.stop.take() {
394                        run_shutdown_callbacks();
395                        let _ = stop.send(0);
396                    }
397                }
398                Ok(ArbiterCommand::Execute(fut)) => {
399                    crate::spawn_detached(fut);
400                }
401                Err(_) => break,
402            }
403        }
404    }
405}
406
407/// Runs the callbacks registered with [`Arbiter::on_shutdown()`], including
408/// ones registered while they run.
409pub(crate) fn run_shutdown_callbacks() {
410    loop {
411        let callbacks = ON_SHUTDOWN
412            .try_with(|cell| mem::take(&mut *cell.borrow_mut()))
413            .unwrap_or_default();
414        if callbacks.is_empty() {
415            break;
416        }
417        for f in callbacks {
418            f();
419        }
420    }
421}
422
423/// Inserts a value into the current arbiter's thread-local storage.
424///
425/// If the storage has already been destroyed because the thread is
426/// exiting, the value is dropped.
427pub fn set_item<T: 'static>(item: T) {
428    let item: Rc<dyn Any> = Rc::new(item);
429    let old = STORAGE
430        .try_with(move |cell| cell.borrow_mut().insert(TypeId::of::<T>(), item))
431        .ok()
432        .flatten();
433    drop(old);
434}
435
436/// Returns a cloned value from the current arbiter's thread-local storage.
437///
438/// Returns `None` if the storage has already been destroyed because the
439/// thread is exiting.
440pub fn get_item<T: Clone + 'static>() -> Option<T> {
441    STORAGE
442        .try_with(move |cell| {
443            cell.borrow()
444                .get(&TypeId::of::<T>())
445                .and_then(|boxed| boxed.downcast_ref())
446                .cloned()
447        })
448        .ok()
449        .flatten()
450}
451
452/// Provides access to a value in the current arbiter's thread-local storage.
453///
454/// A default value is inserted if the requested type is not already present.
455/// If the storage has already been destroyed because the thread is exiting,
456/// `f` receives a temporary default value that is not stored.
457pub fn with_item<T: Default + 'static, F, R>(f: F) -> R
458where
459    F: FnOnce(&T) -> R,
460{
461    // `f` holds its own reference, so the value stays alive even if `f`
462    // replaces or removes it
463    let val = STORAGE
464        .try_with(|cell| {
465            let existing = cell.borrow().get(&TypeId::of::<T>()).cloned();
466            existing.unwrap_or_else(|| {
467                let val: Rc<dyn Any> = Rc::new(T::default());
468                cell.borrow_mut().insert(TypeId::of::<T>(), val.clone());
469                val
470            })
471        })
472        .ok();
473    match val {
474        Some(val) => f(val.downcast_ref::<T>().unwrap()),
475        None => f(&T::default()),
476    }
477}
478
479#[doc(hidden)]
480/// Remove all items from storage.
481///
482/// # Safety
483///
484/// All outstanding calls to [`with_item`] must have completed.
485pub unsafe fn remove_all_items() {
486    clear_storage();
487    System::remove_current();
488}
489
490/// Removes all items from the storage.
491///
492/// Each item is dropped outside of the storage borrow, so that its destructor
493/// may access the storage. Items it inserts are removed too.
494fn clear_storage() {
495    let _ = STORAGE.try_with(|cell| {
496        loop {
497            let mut items = cell.borrow_mut();
498            let Some(key) = items.keys().next().copied() else {
499                break;
500            };
501            let item = items.remove(&key);
502            drop(items);
503            drop(item);
504        }
505    });
506}
507
508#[cfg(test)]
509mod tests {
510    use std::cell::Cell;
511
512    use super::*;
513    use crate::testing::TestRunner;
514
515    #[derive(Clone, Default)]
516    struct Value(usize);
517
518    #[test]
519    #[should_panic(expected = "Arbiter is not running")]
520    fn current_without_arbiter() {
521        let _ = Arbiter::current();
522    }
523
524    #[test]
525    fn arbiter_with_init() {
526        thread_local! {
527            static INIT: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
528        }
529
530        System::new("arb-init", TestRunner).block_on(async {
531            let (tx, rx) = std::sync::mpsc::channel();
532            let mut arb = Arbiter::with_init("init".to_string(), move || {
533                INIT.with(|init| init.set(true));
534                let _ = tx.send(thread::current().name().map(str::to_string));
535            });
536            assert_eq!(rx.recv().unwrap().as_deref(), Some("init"));
537
538            // `init` ran on the arbiter's thread
539            let (tx, rx) = oneshot::async_channel();
540            arb.handle()
541                .spawn(async move {
542                    let _ = tx.send(INIT.with(std::cell::Cell::get));
543                })
544                .detach();
545            assert!(rx.await.unwrap());
546            assert!(!INIT.with(std::cell::Cell::get));
547
548            arb.stop();
549            arb.join().unwrap();
550        });
551    }
552
553    #[test]
554    fn arbiter_api() {
555        System::new("arb-test", TestRunner).block_on(async {
556            let sys = System::current();
557            let cur = Arbiter::current();
558            assert_eq!(cur.name(), "arb-test");
559            assert_eq!(cur, sys.arbiter());
560            assert!(cur.is_running());
561            assert_eq!(format!("{cur:?}"), "Arbiter(\"arb-test\")");
562            // the primary arbiter does not own a thread
563            assert!(cur.clone().join().is_ok());
564
565            let stopped = Arc::new(AtomicUsize::new(0));
566            let stopped2 = stopped.clone();
567            let mut arb = Arbiter::with_name("named".to_string()).on_stop(move || {
568                stopped2.fetch_add(1, Ordering::Relaxed);
569            });
570            assert_eq!(arb.name(), "named");
571            assert!(arb.is_running());
572            assert_ne!(arb, cur);
573            sys.list_arbiters(|arbs| assert!(arbs.contains(&arb)));
574
575            // shutdown callbacks run on the arbiter's thread
576            let (reg_tx, reg_rx) = oneshot::async_channel();
577            let (tx, rx) = oneshot::channel();
578            arb.handle()
579                .spawn(async move {
580                    let name = Arbiter::current().name().to_string();
581                    Arbiter::on_shutdown(move || {
582                        let _ = tx.send(name);
583                    });
584                    let _ = reg_tx.send(());
585                })
586                .detach();
587            reg_rx.await.unwrap();
588
589            arb.stop();
590            arb.join().unwrap();
591            assert!(!arb.is_running());
592            assert_eq!(stopped.load(Ordering::Relaxed), 1);
593            assert_eq!(rx.recv().unwrap(), "named");
594            sys.list_arbiters(|arbs| assert!(!arbs.contains(&arb)));
595
596            let mut arb = Arbiter::default();
597            assert!(arb.name().starts_with("arb-test:arb:"));
598            arb.stop();
599            arb.join().unwrap();
600        });
601    }
602
603    #[test]
604    fn handle_spawn_cancel_before_run() {
605        System::new("cancel-before", TestRunner).block_on(async {
606            let executed = Arc::new(AtomicUsize::new(0));
607            let executed2 = executed.clone();
608            let hnd = Handle::current();
609            hnd.spawn(async move {
610                executed2.fetch_add(1, Ordering::Relaxed);
611            })
612            .cancel();
613            // commands run in order, the cancelled one was processed first
614            hnd.spawn(async {}).await.unwrap();
615            hnd.spawn(async {}).await.unwrap();
616            assert_eq!(executed.load(Ordering::Relaxed), 0);
617        });
618    }
619
620    #[test]
621    fn handle_spawn_cancel_while_running() {
622        struct Token(Option<oneshot::Sender<()>>);
623        impl Drop for Token {
624            fn drop(&mut self) {
625                let _ = self.0.take().unwrap().send(());
626            }
627        }
628
629        System::new("cancel-running", TestRunner).block_on(async {
630            let mut arb = Arbiter::new();
631            let resumed = Arc::new(AtomicBool::new(false));
632            let resumed2 = resumed.clone();
633            let (started_tx, started_rx) = oneshot::async_channel();
634            let (drop_tx, drop_rx) = oneshot::async_channel();
635            let hnd = arb.handle().spawn(async move {
636                let _token = Token(Some(drop_tx));
637                let _ = started_tx.send(());
638                std::future::pending::<()>().await;
639                resumed2.store(true, Ordering::Relaxed);
640            });
641            started_rx.await.unwrap();
642            hnd.cancel();
643            // the future is dropped on the arbiter thread
644            let mut drop_rx = std::pin::pin!(drop_rx);
645            let mut timeout =
646                std::pin::pin!(futures_timer::Delay::new(std::time::Duration::from_secs(5)));
647            std::future::poll_fn(|cx| {
648                if drop_rx.as_mut().poll(cx).is_ready() {
649                    std::task::Poll::Ready(())
650                } else if timeout.as_mut().poll(cx).is_ready() {
651                    panic!("cancelled task is not dropped");
652                } else {
653                    std::task::Poll::Pending
654                }
655            })
656            .await;
657            assert!(!resumed.load(Ordering::Relaxed));
658
659            // detached tasks keep running
660            let (tx, rx) = oneshot::async_channel();
661            arb.handle()
662                .spawn(async move {
663                    let _ = tx.send(());
664                })
665                .detach();
666            rx.await.unwrap();
667
668            arb.stop();
669            arb.join().unwrap();
670        });
671    }
672
673    #[cfg(all(not(feature = "tokio"), not(feature = "compio")))]
674    #[test]
675    fn panicking_arbiter() {
676        /// Returns a panic of the event loop as an error, like the ntex-net
677        /// runners.
678        struct CatchRunner;
679
680        impl crate::Runner for CatchRunner {
681            fn block_on(
682                &self,
683                fut: crate::BlockFuture,
684            ) -> Result<(), Box<dyn std::any::Any + Send>> {
685                panic::catch_unwind(panic::AssertUnwindSafe(|| TestRunner.block_on(fut)))?
686            }
687        }
688
689        System::new("arb-panic", CatchRunner).block_on(async {
690            let sys = System::current();
691            let mut arb = Arbiter::with_name("panics".to_string());
692            let (tx, rx) = oneshot::channel();
693            arb.handle()
694                .spawn(async move {
695                    Arbiter::on_shutdown(move || {
696                        let _ = tx.send(());
697                    });
698                    panic!("arbiter boom");
699                })
700                .detach();
701
702            // the panic of the event loop is returned by `join`
703            let err = arb.join().unwrap_err();
704            assert_eq!(err.downcast_ref::<&str>(), Some(&"arbiter boom"));
705            // the arbiter is unregistered and its shutdown callbacks ran
706            assert!(rx.recv().is_ok());
707            assert!(!arb.is_running());
708            sys.list_arbiters(|arbs| assert!(!arbs.contains(&arb)));
709        });
710    }
711
712    #[test]
713    fn arbiter_values() {
714        System::new("test", TestRunner).block_on(async {
715            assert_eq!(Arbiter::get_value(|| Value(1)).0, 1);
716            assert_eq!(Arbiter::get_value(|| Value(2)).0, 1);
717        });
718    }
719
720    #[test]
721    fn on_shutdown_runs_after_block_on() {
722        let called = Rc::new(Cell::new(0));
723        let called2 = called.clone();
724        System::new("test", TestRunner).block_on(async move {
725            Arbiter::on_shutdown(move || {
726                called2.set(called2.get() + 1);
727                // registered by a callback, runs in the same shutdown
728                Arbiter::on_shutdown(move || called2.set(called2.get() + 1));
729            });
730        });
731        assert_eq!(called.get(), 2);
732    }
733
734    fn use_storage() {
735        set_item(Value(1));
736        let _ = get_item::<Value>();
737        with_item::<Value, _, _>(|v| v.0);
738        Arbiter::get_value(|| Value(2));
739    }
740
741    struct UseOnDrop;
742
743    impl Drop for UseOnDrop {
744        fn drop(&mut self) {
745            use_storage();
746        }
747    }
748
749    thread_local!(static HOLD: RefCell<Option<UseOnDrop>> = const { RefCell::new(None) });
750
751    #[test]
752    fn storage_access_during_thread_exit() {
753        // item stored in STORAGE accesses STORAGE while it is destroyed
754        thread::spawn(|| set_item(UseOnDrop)).join().unwrap();
755
756        // other thread local accesses STORAGE, in both destruction orders
757        thread::spawn(|| {
758            HOLD.with(|h| *h.borrow_mut() = Some(UseOnDrop));
759            use_storage();
760        })
761        .join()
762        .unwrap();
763        thread::spawn(|| {
764            use_storage();
765            HOLD.with(|h| *h.borrow_mut() = Some(UseOnDrop));
766        })
767        .join()
768        .unwrap();
769    }
770
771    #[test]
772    fn with_item_value_outlives_replacement() {
773        #[derive(Clone, Default)]
774        struct Item(std::rc::Rc<Vec<u8>>);
775
776        thread::spawn(|| {
777            set_item(Item(std::rc::Rc::new(vec![1; 64])));
778            let len = with_item::<Item, _, _>(|item| {
779                // replaces and frees the stored value while `f` holds its clone
780                set_item(Item::default());
781                unsafe { remove_all_items() };
782                item.0.len()
783            });
784            assert_eq!(len, 64);
785            assert!(get_item::<Item>().is_none());
786            assert_eq!(with_item::<Item, _, _>(|item| item.0.len()), 0);
787        })
788        .join()
789        .unwrap();
790    }
791
792    #[test]
793    fn remove_all_items_drops_outside_borrow() {
794        struct Item;
795
796        impl Drop for Item {
797            fn drop(&mut self) {
798                let _ = get_item::<u32>();
799                set_item(2u64);
800            }
801        }
802
803        thread::spawn(|| {
804            set_item(Item);
805            set_item(1u32);
806            unsafe { remove_all_items() };
807            assert!(get_item::<u32>().is_none());
808            assert!(get_item::<u64>().is_none(), "item inserted by a destructor");
809        })
810        .join()
811        .unwrap();
812    }
813
814    #[test]
815    fn storage_fallback_values() {
816        struct Check;
817
818        impl Drop for Check {
819            fn drop(&mut self) {
820                set_item(Value(5));
821                assert!(get_item::<Value>().is_none());
822                assert_eq!(with_item::<Value, _, _>(|v| v.0), 0);
823                assert_eq!(Arbiter::get_value(|| Value(3)).0, 3);
824                assert_eq!(Arbiter::get_value(|| Value(4)).0, 4);
825            }
826        }
827        thread::spawn(|| set_item(Check)).join().unwrap();
828    }
829}