virtual_desktop_manager\tray_plugins/
panic_notifier.rs

1use crate::{
2    dynamic_gui::DynamicUiHooks,
3    tray::{SystemTray, TrayPlugin, TrayRoot},
4};
5use nwd::NwgPartial;
6use std::{
7    any::TypeId,
8    cell::{OnceCell, RefCell},
9    rc::{Rc, Weak},
10    sync::{mpsc, Arc, Mutex},
11};
12
13struct ThreadLocalPanicHandler {
14    tray: RefCell<Weak<SystemTray>>,
15    notice: OnceCell<Arc<Mutex<nwg::NoticeSender>>>,
16    panic_messages: OnceCell<mpsc::Receiver<String>>,
17}
18impl ThreadLocalPanicHandler {
19    pub const fn new() -> Self {
20        Self {
21            tray: RefCell::new(Weak::new()),
22            notice: OnceCell::new(),
23            panic_messages: OnceCell::new(),
24        }
25    }
26    thread_local! {
27        static LOCAL: ThreadLocalPanicHandler = const { ThreadLocalPanicHandler::new() };
28    }
29}
30
31/// Sets a panic hook that displays any panic as a notification.
32#[derive(Default, NwgPartial)]
33pub struct PanicNotifier {
34    #[nwg_control]
35    #[nwg_events(OnNotice: [Self::on_panic_notice])]
36    panic_notice: nwg::Notice,
37}
38impl DynamicUiHooks<SystemTray> for PanicNotifier {
39    fn before_partial_build(
40        &mut self,
41        tray_ui: &Rc<SystemTray>,
42        _should_build: &mut bool,
43    ) -> Option<(nwg::ControlHandle, TypeId)> {
44        Some((tray_ui.root().window.handle, TypeId::of::<TrayRoot>()))
45    }
46    fn after_partial_build(&mut self, tray_ui: &Rc<SystemTray>) {
47        let shared_sender =
48            ThreadLocalPanicHandler::LOCAL.with(|state: &ThreadLocalPanicHandler| {
49                state.tray.replace(Rc::downgrade(tray_ui));
50
51                let sender = self.panic_notice.sender();
52                if let Some(shared) = state.notice.get() {
53                    // Updated existing panic hook:
54                    *shared.lock().unwrap() = sender;
55                    None
56                } else {
57                    let shared = Arc::new(Mutex::new(sender));
58                    state.notice.set(shared.clone()).ok().unwrap();
59                    Some(shared)
60                }
61            });
62        let Some(shared_sender) = shared_sender else {
63            return; // Already has hook
64        };
65
66        let prev = std::panic::take_hook();
67        let (tx, rx) = mpsc::channel();
68
69        ThreadLocalPanicHandler::LOCAL.with(|state: &ThreadLocalPanicHandler| {
70            state
71                .panic_messages
72                .set(rx)
73                .expect("Failed to initialize panic notifier");
74        });
75        std::panic::set_hook(Box::new(move |info| {
76            prev(info);
77
78            ThreadLocalPanicHandler::LOCAL.with(|shared: &ThreadLocalPanicHandler| {
79                if let Some(this) = { shared.tray.borrow().upgrade() } {
80                    // Panic on main thread so can display notification immediately:
81                    Self::display_panic_notification(&this, &info);
82                } else {
83                    // Send error to main thread and notify the user:
84                    if tx.send(info.to_string()).is_ok() {
85                        shared_sender.lock().unwrap().notice();
86                    }
87                }
88            });
89        }));
90    }
91}
92impl TrayPlugin for PanicNotifier {}
93impl PanicNotifier {
94    fn display_panic_notification(tray_ui: &SystemTray, info: &dyn std::fmt::Display) {
95        tray_ui.show_notification(
96            "Virtual Desktop Manager Panicked!",
97            "Virtual Desktop Manager encountered a panic and might no longer work correctly, it is recommended to restart the program.",
98        );
99        // Show panic message in separate notification since we
100        // can't fit more text in the first one:
101        tray_ui.show_notification("Virtual Desktop Manager Panic Info:", &format!("{info}"));
102    }
103    fn on_panic_notice(&self) {
104        ThreadLocalPanicHandler::LOCAL.with(|shared: &ThreadLocalPanicHandler| {
105            while let Some(msg) = shared
106                .panic_messages
107                .get()
108                .and_then(|rx| rx.try_recv().ok())
109            {
110                if let Some(tray_ui) = { shared.tray.borrow().upgrade() } {
111                    Self::display_panic_notification(&tray_ui, &msg);
112                }
113            }
114        });
115    }
116}
virtual_desktop_manager\tray_plugins/panic_notifier.rs

virtual_desktop_manager\tray_plugins/
panic_notifier.rs
