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