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
25pub 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 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 pub fn stop(&self) {
101 let _ = self.0.sender.try_send(ArbiterCommand::Stop);
102 }
103
104 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 pub fn with_name(name: String) -> Arbiter {
112 Arbiter::start(COUNT.fetch_add(1, Ordering::Relaxed), name, None)
113 }
114
115 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 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 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 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 let _ = stop_rx.await;
184 });
185
186 sys2.unregister_arbiter(Id(id));
188 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 pub fn id(&self) -> Id {
243 Id(self.0.id)
244 }
245
246 #[cfg(target_os = "linux")]
247 pub(crate) fn tid(&self) -> i32 {
249 self.0.tid
250 }
251
252 #[cfg(windows)]
253 pub(crate) fn thread(&self) -> Option<&crate::watchdog::ThreadHandle> {
255 self.0.thread.as_ref()
256 }
257
258 pub fn name(&self) -> &str {
260 self.0.name.as_ref()
261 }
262
263 #[inline]
264 pub fn handle(&self) -> &Handle {
266 self.0.hnd.as_ref().unwrap()
267 }
268
269 pub fn latency(&self) -> Option<LatencyMetrics> {
274 self.0.latency.metrics()
275 }
276
277 #[inline]
278 pub fn is_running(&self) -> bool {
280 self.0.running.load(Ordering::Relaxed)
281 }
282
283 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 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 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 pub(crate) fn try_current() -> Option<Arbiter> {
340 ADDR.try_with(|cell| cell.borrow().clone()).ok().flatten()
341 }
342
343 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 let _ = panic::catch_unwind(panic::AssertUnwindSafe(f));
352 } else {
353 f();
354 }
355 }
356 }
357
358 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 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
407pub(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
423pub 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
436pub 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
452pub fn with_item<T: Default + 'static, F, R>(f: F) -> R
458where
459 F: FnOnce(&T) -> R,
460{
461 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)]
480pub unsafe fn remove_all_items() {
486 clear_storage();
487 System::remove_current();
488}
489
490fn 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 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 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 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 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 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 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 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 let err = arb.join().unwrap_err();
704 assert_eq!(err.downcast_ref::<&str>(), Some(&"arbiter boom"));
705 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 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 thread::spawn(|| set_item(UseOnDrop)).join().unwrap();
755
756 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 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}