1use std::sync::atomic::{AtomicU64, Ordering};
24use std::sync::{Arc, Condvar, Mutex, MutexGuard};
25use std::time::{Duration, SystemTime, UNIX_EPOCH};
26
27use tauri::{AppHandle, Emitter, Wry};
28
29use crate::EVENT_STALL;
30
31#[derive(Debug, Clone, Copy)]
32pub struct StallReport {
33 pub stalled_ms: u64,
34 pub threshold_ms: u64,
35}
36
37pub type StallCallback = Arc<dyn Fn(&StallReport) + Send + Sync>;
38
39pub const DEFAULT_POLL_INTERVAL: Duration = Duration::from_millis(250);
42
43#[derive(Clone)]
44pub struct WatchdogConfig {
45 pub threshold: Duration,
47 pub heartbeat_interval: Duration,
49 pub poll_interval: Duration,
54 pub on_stall: Option<StallCallback>,
58 pub emit_event: bool,
67}
68
69impl Default for WatchdogConfig {
70 fn default() -> Self {
71 Self {
72 threshold: Duration::from_secs(5),
73 heartbeat_interval: Duration::from_secs(1),
74 poll_interval: DEFAULT_POLL_INTERVAL,
75 on_stall: None,
76 emit_event: true,
77 }
78 }
79}
80
81#[derive(Debug, Clone)]
84pub struct StallDetector {
85 threshold_ms: u64,
86 fired: bool,
87}
88
89impl StallDetector {
90 pub fn new(threshold: Duration) -> Self {
91 Self {
92 threshold_ms: threshold.as_millis() as u64,
93 fired: false,
94 }
95 }
96
97 pub fn observe(&mut self, last_beat_ms: u64, now_ms: u64) -> Option<StallReport> {
100 if should_fire(last_beat_ms, now_ms, self.fired, self.threshold_ms) {
101 self.fired = true;
102 return Some(StallReport {
103 stalled_ms: now_ms.saturating_sub(last_beat_ms),
104 threshold_ms: self.threshold_ms,
105 });
106 }
107 if self.fired && now_ms.saturating_sub(last_beat_ms) <= self.threshold_ms {
108 self.fired = false;
109 }
110 None
111 }
112}
113
114fn now_ms() -> u64 {
115 SystemTime::now()
116 .duration_since(UNIX_EPOCH)
117 .map(|d| d.as_millis() as u64)
118 .unwrap_or(0)
119}
120
121pub fn should_fire(last_beat_ms: u64, now_ms: u64, already_fired: bool, threshold_ms: u64) -> bool {
123 !already_fired && now_ms.saturating_sub(last_beat_ms) > threshold_ms
124}
125
126static LAST_BEAT_MS: AtomicU64 = AtomicU64::new(0);
128static COUNTERS: PollCounters = PollCounters::new();
130
131pub(crate) struct PollCounters {
133 stalls: AtomicU64,
134 emit_failures: AtomicU64,
135}
136
137impl PollCounters {
138 pub(crate) const fn new() -> Self {
139 Self {
140 stalls: AtomicU64::new(0),
141 emit_failures: AtomicU64::new(0),
142 }
143 }
144}
145
146#[derive(Debug, Clone, Copy, PartialEq, Eq)]
148pub struct WatchdogStatus {
149 pub installed: bool,
151 pub last_beat_age_ms: u64,
153 pub stalls_total: u64,
155 pub emit_failures: u64,
158}
159
160#[derive(Debug, Clone, Copy, PartialEq)]
165pub struct ShellRuntimeStatus {
166 pub watchdog: WatchdogStatus,
167 pub process: Option<rightkit_child_process::ProcessUsage>,
168}
169
170pub fn status() -> ShellRuntimeStatus {
172 let beat = LAST_BEAT_MS.load(Ordering::Relaxed);
173 ShellRuntimeStatus {
174 watchdog: WatchdogStatus {
175 installed: beat != 0,
176 last_beat_age_ms: if beat == 0 {
177 0
178 } else {
179 now_ms().saturating_sub(beat)
180 },
181 stalls_total: COUNTERS.stalls.load(Ordering::Relaxed),
182 emit_failures: COUNTERS.emit_failures.load(Ordering::Relaxed),
183 },
184 process: rightkit_child_process::process_usage().ok(),
185 }
186}
187
188const MIN_PERIOD: Duration = Duration::from_millis(10);
190
191const EPISODE_RECHECK: Duration = Duration::from_secs(1);
194
195pub fn install(app: AppHandle<Wry>, cfg: WatchdogConfig) {
196 let last = Arc::new(AtomicU64::new(now_ms()));
197 LAST_BEAT_MS.store(now_ms(), Ordering::Relaxed);
198 let bell = Arc::new(Bell::new());
199 {
200 let (app, last, bell, interval) = (
201 app.clone(),
202 Arc::clone(&last),
203 Arc::clone(&bell),
204 cfg.heartbeat_interval.max(MIN_PERIOD),
205 );
206 std::thread::Builder::new()
207 .name("rk-shell-heartbeat".into())
208 .spawn(move || {
209 heartbeat_loop(&bell, interval, || {
210 let last = Arc::clone(&last);
211 let bell = Arc::clone(&bell);
212 app.run_on_main_thread(move || {
214 let now = now_ms();
215 last.store(now, Ordering::SeqCst);
216 LAST_BEAT_MS.store(now, Ordering::Relaxed);
217 bell.beat();
219 })
220 .is_ok()
221 })
222 })
223 .ok();
224 }
225 let poll = cfg.poll_interval.max(MIN_PERIOD);
226 let emit = emitter_for(cfg.emit_event, || {
228 let app = app.clone();
229 Box::new(move |report: &StallReport| {
230 app.emit(
231 EVENT_STALL,
232 serde_json::json!({"stalledMs": report.stalled_ms, "thresholdMs": report.threshold_ms}),
233 )
234 .map_err(|e| e.to_string())
235 })
236 });
237 std::thread::Builder::new()
238 .name("rk-shell-watchdog".into())
239 .spawn(move || {
240 poll_loop(
241 &last,
242 StallDetector::new(cfg.threshold),
243 poll,
244 cfg.on_stall.as_ref(),
245 emit.as_deref(),
246 &COUNTERS,
247 &bell,
248 )
249 })
250 .ok();
251}
252
253pub(crate) struct Bell {
258 state: Mutex<BellState>,
259 watch: Condvar,
261 pulse: Condvar,
263 #[cfg(test)]
265 wakes: AtomicU64,
266}
267
268#[derive(Default)]
269struct BellState {
270 stop: bool,
271 beats: u64,
273 watching: bool,
275}
276
277impl Bell {
278 pub(crate) fn new() -> Self {
279 Self {
280 state: Mutex::new(BellState::default()),
281 watch: Condvar::new(),
282 pulse: Condvar::new(),
283 #[cfg(test)]
284 wakes: AtomicU64::new(0),
285 }
286 }
287
288 fn state(&self) -> MutexGuard<'_, BellState> {
289 self.state.lock().unwrap_or_else(|e| e.into_inner())
290 }
291
292 #[cfg(test)]
295 pub(crate) fn stop(&self) {
296 self.state().stop = true;
297 self.watch.notify_all();
298 self.pulse.notify_all();
299 }
300
301 pub(crate) fn is_stopped(&self) -> bool {
302 self.state().stop
303 }
304
305 pub(crate) fn beat(&self) {
308 let mut state = self.state();
309 state.beats += 1;
310 if state.watching {
311 self.watch.notify_all();
312 }
313 }
314
315 pub(crate) fn beats(&self) -> u64 {
317 self.state().beats
318 }
319
320 pub(crate) fn sleep_unless_stopped(&self, period: Duration) -> bool {
322 let state = self.state();
323 let (state, _) = self
324 .pulse
325 .wait_timeout_while(state, period, |s| !s.stop)
326 .unwrap_or_else(|e| e.into_inner());
327 #[cfg(test)]
328 self.wakes.fetch_add(1, Ordering::SeqCst);
329 state.stop
330 }
331
332 pub(crate) fn wait_watch(&self, timeout: Duration, episode: Option<u64>) -> bool {
335 let mut state = self.state();
336 state.watching = episode.is_some();
337 let (mut state, _) = self
338 .watch
339 .wait_timeout_while(state, timeout, |s| {
340 !s.stop && episode.is_none_or(|seen| s.beats == seen)
341 })
342 .unwrap_or_else(|e| e.into_inner());
343 state.watching = false;
344 #[cfg(test)]
345 self.wakes.fetch_add(1, Ordering::SeqCst);
346 state.stop
347 }
348
349 #[cfg(test)]
350 pub(crate) fn wakes(&self) -> u64 {
351 self.wakes.load(Ordering::SeqCst)
352 }
353}
354
355pub(crate) fn heartbeat_loop(bell: &Bell, interval: Duration, mut ping: impl FnMut() -> bool) {
359 while !bell.sleep_unless_stopped(interval) {
360 if !ping() {
361 break;
362 }
363 }
364}
365
366fn until_stall(beat_ms: u64, threshold_ms: u64, now_ms: u64) -> Duration {
371 let wait = beat_ms
372 .saturating_add(threshold_ms)
373 .saturating_add(1)
374 .saturating_sub(now_ms);
375 Duration::from_millis(wait.min(threshold_ms.saturating_add(1)))
376}
377
378pub(crate) type EmitDyn = dyn Fn(&StallReport) -> Result<(), String> + Send;
380pub(crate) type EmitFn = Box<EmitDyn>;
381
382pub(crate) fn emitter_for(emit_event: bool, make: impl FnOnce() -> EmitFn) -> Option<EmitFn> {
385 emit_event.then(make)
386}
387
388pub(crate) fn poll_loop(
397 last: &AtomicU64,
398 mut detector: StallDetector,
399 poll: Duration,
400 on_stall: Option<&StallCallback>,
401 emit: Option<&EmitDyn>,
402 counters: &PollCounters,
403 bell: &Bell,
404) {
405 let threshold_ms = detector.threshold_ms;
406 let episode_wait = Duration::from_millis(threshold_ms).max(EPISODE_RECHECK);
407 let mut reported: Option<u64> = None;
409 while !bell.is_stopped() {
410 let seen = bell.beats();
413 let beat = last.load(Ordering::SeqCst);
414 if detector.fired && reported != Some(beat) {
415 detector.fired = false;
416 }
417 if let Some(report) = detector.observe(beat, now_ms()) {
418 reported = Some(beat);
419 counters.stalls.fetch_add(1, Ordering::Relaxed);
420 if let Some(cb) = on_stall {
421 cb(&report);
422 }
423 if let Some(emit) = emit {
424 if emit(&report).is_err() {
425 counters.emit_failures.fetch_add(1, Ordering::Relaxed);
426 }
427 }
428 }
429 let stopped = if detector.fired {
430 bell.wait_watch(episode_wait, Some(seen))
431 } else {
432 bell.wait_watch(until_stall(beat, threshold_ms, now_ms()).max(poll), None)
433 };
434 if stopped {
435 break;
436 }
437 }
438}
439
440#[cfg(test)]
441mod tests {
442 use super::*;
443 use std::sync::atomic::AtomicUsize;
444 use std::time::Instant;
445
446 const POLL: Duration = Duration::from_millis(10);
447
448 fn wait_for(what: &str, mut ok: impl FnMut() -> bool) {
449 let deadline = Instant::now() + Duration::from_secs(10);
450 while !ok() {
451 assert!(Instant::now() < deadline, "timed out waiting for {what}");
452 std::thread::sleep(Duration::from_millis(5));
453 }
454 }
455
456 #[test]
457 fn a_failing_emit_is_counted_and_the_next_stall_still_fires() {
458 let last = Arc::new(AtomicU64::new(0)); let counters = Arc::new(PollCounters::new());
460 let bell = Arc::new(Bell::new());
461 let fired = Arc::new(AtomicUsize::new(0));
462 let emits = Arc::new(AtomicUsize::new(0));
463 let (l, c, b, f, e) = (
464 last.clone(),
465 counters.clone(),
466 bell.clone(),
467 fired.clone(),
468 emits.clone(),
469 );
470 let thread = std::thread::spawn(move || {
471 let on_stall: StallCallback = Arc::new(move |_| {
472 f.fetch_add(1, Ordering::SeqCst);
473 });
474 let emit: EmitFn = Box::new(move |_| {
475 e.fetch_add(1, Ordering::SeqCst);
476 Err("webview gone".into())
477 });
478 poll_loop(
479 &l,
480 StallDetector::new(Duration::from_millis(50)),
481 POLL,
482 Some(&on_stall),
483 Some(&*emit),
484 &c,
485 &b,
486 );
487 });
488
489 wait_for("the first stall", || fired.load(Ordering::SeqCst) == 1);
490 wait_for("its failed emit to be counted", || {
491 counters.emit_failures.load(Ordering::SeqCst) == 1
492 });
493 last.store(now_ms(), Ordering::SeqCst);
496 bell.beat();
497 wait_for("the second stall", || fired.load(Ordering::SeqCst) == 2);
498 wait_for("its failed emit", || {
499 counters.emit_failures.load(Ordering::SeqCst) == 2
500 });
501 assert!(
502 !thread.is_finished(),
503 "an emit failure must not end the watchdog"
504 );
505 assert_eq!(emits.load(Ordering::SeqCst), 2);
506 assert_eq!(counters.stalls.load(Ordering::SeqCst), 2);
507 bell.stop();
508 thread.join().unwrap();
509 }
510
511 #[test]
512 fn callback_only_mode_never_builds_or_uses_the_handle() {
513 let emit = emitter_for(false, || {
515 panic!("callback-only mode must not touch the app handle")
516 });
517 assert!(emit.is_none());
518
519 let last = Arc::new(AtomicU64::new(0));
520 let counters = Arc::new(PollCounters::new());
521 let bell = Arc::new(Bell::new());
522 let fired = Arc::new(AtomicUsize::new(0));
523 let (l, c, b, f) = (last.clone(), counters.clone(), bell.clone(), fired.clone());
524 let thread = std::thread::spawn(move || {
525 let on_stall: StallCallback = Arc::new(move |_| {
526 f.fetch_add(1, Ordering::SeqCst);
527 });
528 poll_loop(
529 &l,
530 StallDetector::new(Duration::from_millis(50)),
531 POLL,
532 Some(&on_stall),
533 emit.as_deref(),
534 &c,
535 &b,
536 );
537 });
538 wait_for("the callback", || fired.load(Ordering::SeqCst) == 1);
539 assert_eq!(counters.emit_failures.load(Ordering::SeqCst), 0);
540 bell.stop();
541 thread.join().unwrap();
542
543 let built = std::cell::Cell::new(0);
545 let emit = emitter_for(true, || {
546 built.set(built.get() + 1);
547 Box::new(|_| Ok(()))
548 });
549 assert!(emit.is_some());
550 assert_eq!(built.get(), 1);
551 }
552
553 fn finished(rx: &std::sync::mpsc::Receiver<&'static str>, n: usize) -> Vec<&'static str> {
555 let mut done: Vec<&'static str> = (0..n)
556 .map(|_| {
557 rx.recv_timeout(Duration::from_secs(10))
558 .expect("a thread did not end after the stop")
559 })
560 .collect();
561 done.sort_unstable();
562 done
563 }
564
565 #[test]
566 fn stop_wakes_both_threads_out_of_hour_long_waits() {
567 let last = Arc::new(AtomicU64::new(now_ms()));
568 let counters = Arc::new(PollCounters::new());
569 let bell = Arc::new(Bell::new());
570 let (tx, rx) = std::sync::mpsc::channel();
571 let (l, c, b, t) = (last.clone(), counters.clone(), bell.clone(), tx.clone());
572 let watchdog = std::thread::spawn(move || {
573 poll_loop(
574 &l,
575 StallDetector::new(Duration::from_secs(3600)),
576 POLL,
577 None,
578 None,
579 &c,
580 &b,
581 );
582 let _ = t.send("watchdog");
583 });
584 let (b, t) = (bell.clone(), tx);
585 let heartbeat = std::thread::spawn(move || {
586 heartbeat_loop(&b, Duration::from_secs(3600), || true);
587 let _ = t.send("heartbeat");
588 });
589 std::thread::sleep(Duration::from_millis(100));
591 bell.stop();
592 assert_eq!(finished(&rx, 2), ["heartbeat", "watchdog"]);
593 watchdog.join().unwrap();
594 heartbeat.join().unwrap();
595 assert_eq!(counters.stalls.load(Ordering::SeqCst), 0);
596 }
597
598 #[test]
599 fn a_heartbeat_wakes_a_stalled_watchdog_at_once() {
600 let last = Arc::new(AtomicU64::new(0)); let counters = Arc::new(PollCounters::new());
603 let bell = Arc::new(Bell::new());
604 let (l, c, b) = (last.clone(), counters.clone(), bell.clone());
605 let thread = std::thread::spawn(move || {
606 poll_loop(
607 &l,
608 StallDetector::new(Duration::from_secs(600)),
609 POLL,
610 None,
611 None,
612 &c,
613 &b,
614 );
615 });
616 wait_for("the stall", || counters.stalls.load(Ordering::SeqCst) == 1);
617 let woke = bell.wakes();
618 last.store(now_ms(), Ordering::SeqCst);
619 bell.beat();
620 wait_for("the beat to wake the watchdog", || bell.wakes() > woke);
621 assert_eq!(counters.stalls.load(Ordering::SeqCst), 1);
623 bell.stop();
624 thread.join().unwrap();
625 }
626
627 #[test]
628 fn a_healthy_watchdog_is_not_woken_by_beats() {
629 let last = Arc::new(AtomicU64::new(now_ms()));
630 let counters = Arc::new(PollCounters::new());
631 let bell = Arc::new(Bell::new());
632 let (l, c, b) = (last.clone(), counters.clone(), bell.clone());
633 let thread = std::thread::spawn(move || {
634 poll_loop(
635 &l,
636 StallDetector::new(Duration::from_secs(600)),
637 POLL,
638 None,
639 None,
640 &c,
641 &b,
642 );
643 });
644 std::thread::sleep(Duration::from_millis(100));
645 for _ in 0..50 {
646 last.store(now_ms(), Ordering::SeqCst);
647 bell.beat();
648 }
649 std::thread::sleep(Duration::from_millis(100));
650 assert_eq!(
651 bell.wakes(),
652 0,
653 "a beat must not wake a watchdog with no stall open"
654 );
655 bell.stop();
656 thread.join().unwrap();
657 }
658
659 #[test]
660 fn a_quiet_watchdog_wakes_about_once_per_heartbeat() {
661 let last = Arc::new(AtomicU64::new(now_ms()));
662 let counters = Arc::new(PollCounters::new());
663 let bell = Arc::new(Bell::new());
664 let (tx, rx) = std::sync::mpsc::channel();
665 let (l, c, b, t) = (last.clone(), counters.clone(), bell.clone(), tx.clone());
666 let watchdog = std::thread::spawn(move || {
667 poll_loop(
668 &l,
669 StallDetector::new(Duration::from_secs(5)),
670 POLL,
671 None,
672 None,
673 &c,
674 &b,
675 );
676 let _ = t.send("watchdog");
677 });
678 let (l, b, t) = (last.clone(), bell.clone(), tx);
679 let heartbeat = std::thread::spawn(move || {
680 heartbeat_loop(&b, Duration::from_secs(1), || {
682 l.store(now_ms(), Ordering::SeqCst);
683 b.beat();
684 true
685 });
686 let _ = t.send("heartbeat");
687 });
688 let (start, woke_before) = (Instant::now(), bell.wakes());
689 std::thread::sleep(Duration::from_secs(2));
690 let (secs, woke) = (start.elapsed().as_secs_f64(), bell.wakes() - woke_before);
691 let allowed = 3u64.max(secs.ceil() as u64 + 1);
694 assert!(woke <= allowed, "woke {woke} times in {secs:.1} s");
695 assert_eq!(counters.stalls.load(Ordering::SeqCst), 0);
696 bell.stop();
697 finished(&rx, 2);
698 watchdog.join().unwrap();
699 heartbeat.join().unwrap();
700 }
701
702 #[test]
703 fn a_beat_that_goes_silent_is_reported_at_the_threshold() {
704 let last = Arc::new(AtomicU64::new(now_ms()));
707 let counters = Arc::new(PollCounters::new());
708 let bell = Arc::new(Bell::new());
709 let reports = Arc::new(Mutex::new(Vec::<StallReport>::new()));
710 let (l, c, b, r) = (
711 last.clone(),
712 counters.clone(),
713 bell.clone(),
714 reports.clone(),
715 );
716 let thread = std::thread::spawn(move || {
717 let on_stall: StallCallback = Arc::new(move |report: &StallReport| {
718 r.lock().unwrap().push(*report);
719 });
720 poll_loop(
721 &l,
722 StallDetector::new(Duration::from_secs(1)),
723 DEFAULT_POLL_INTERVAL,
724 Some(&on_stall),
725 None,
726 &c,
727 &b,
728 );
729 });
730 wait_for("the stall", || !reports.lock().unwrap().is_empty());
731 let report = reports.lock().unwrap()[0];
732 assert_eq!(report.threshold_ms, 1000);
733 assert!(report.stalled_ms > 1000, "{report:?}");
734 assert!(report.stalled_ms < 2000, "reported late: {report:?}");
736 bell.stop();
737 thread.join().unwrap();
738 }
739
740 #[test]
741 fn the_wait_to_the_next_possible_stall_follows_the_beat_and_never_exceeds_the_threshold() {
742 assert_eq!(
744 until_stall(9_000, 5_000, 10_000),
745 Duration::from_millis(4_001)
746 );
747 assert_eq!(until_stall(1_000, 5_000, 6_001), Duration::ZERO);
749 assert_eq!(until_stall(1_000, 5_000, 99_000), Duration::ZERO);
750 assert_eq!(
752 until_stall(1_000_000, 5_000, 10_000),
753 Duration::from_millis(5_001)
754 );
755 assert_eq!(
757 until_stall(10_000, 5_000, 10_000),
758 Duration::from_millis(5_001)
759 );
760 }
761}