1use super::effector::Applied;
17use super::policy::is_scarce;
18use super::types::{Action, Level, Sample};
19use common::GuardConfig;
20use std::collections::{HashMap, HashSet};
21use std::process::Command;
22use std::sync::mpsc;
23use std::time::{Duration, Instant};
24
25pub const APP_ID: &str = "io.github.rlm.gtk";
27pub const APP_NAME: &str = "rlm";
29
30pub const PRESSURE_TITLE: &str = "Memory is running low";
32pub const PRESSURE_BODY: &str = "rlm will step in if an app keeps growing.";
33
34pub const PRESSURE_INTERVAL_MS: u64 = 60_000;
36
37pub const PRESSURE_KEY: &str = "/pressure";
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub enum Memory {
44 Unknown,
46 Fine,
48 Low,
51}
52
53pub fn memory_state(level: Level, sample: &Sample, trigger: &common::GuardTrigger) -> Memory {
55 if matches!(level, Level::High | Level::Critical) && is_scarce(sample, trigger) {
56 Memory::Low
57 } else {
58 Memory::Fine
59 }
60}
61
62pub trait NotifySink {
65 fn show(&mut self, key: &str, title: &str, body: &str);
66 fn close(&mut self, key: &str);
67}
68
69pub fn paused_text(app: &str, hold_secs: u64) -> (String, String) {
71 let unit = if hold_secs == 1 { "second" } else { "seconds" };
72 (
73 format!("{app} paused"),
74 format!("Paused for {hold_secs} {unit} while memory is low."),
77 )
78}
79
80pub fn slowed_text(app: &str, cap_bytes: u64) -> (String, String) {
82 (
83 format!("{app} slowed down"),
84 format!(
85 "Held to about {} until memory frees up.",
86 format_size(cap_bytes)
87 ),
88 )
89}
90
91pub fn format_size(bytes: u64) -> String {
93 const GB: f64 = 1e9;
94 const MB: f64 = 1e6;
95 let b = bytes as f64;
96 if b >= GB {
97 format!("{:.1} GB", b / GB)
98 } else {
99 format!("{:.0} MB", (b / MB).max(1.0))
100 }
101}
102
103pub fn display_name(key: &str, desktop: &HashMap<String, String>, comm: Option<&str>) -> String {
110 let base = key.split('@').next().unwrap_or(key);
111 let has_letters = |s: &str| s.chars().any(char::is_alphabetic);
112 let program = match comm {
113 Some(c) if !has_letters(base) && has_letters(c) => c.trim(),
114 _ => base,
115 };
116 if let Some(name) = desktop.get(program).or_else(|| desktop.get(base)) {
117 return name.clone();
118 }
119 if program.is_empty() {
120 return "An app".to_string();
121 }
122 let mut chars = program.chars();
123 match chars.next() {
124 Some(first) => first.to_uppercase().chain(chars).collect(),
125 None => String::new(),
126 }
127}
128
129pub struct AppNames {
133 desktop: HashMap<String, String>,
134 loading: Option<mpsc::Receiver<HashMap<String, String>>>,
136}
137
138impl AppNames {
139 pub fn new() -> Self {
141 Self {
142 desktop: crate::desktop::names_by_program(),
143 loading: None,
144 }
145 }
146
147 pub fn in_background() -> Self {
151 let (tx, rx) = mpsc::channel();
152 let spawned = std::thread::Builder::new()
153 .name("rlm-app-names".into())
154 .spawn(move || {
155 let _ = tx.send(crate::desktop::names_by_program());
156 });
157 Self {
158 desktop: HashMap::new(),
159 loading: spawned.ok().map(|_| rx),
160 }
161 }
162
163 pub fn name(&mut self, key: &str, cgroup: &str) -> String {
165 if let Some(rx) = &self.loading {
166 match rx.try_recv() {
167 Ok(desktop) => {
168 self.desktop = desktop;
169 self.loading = None;
170 }
171 Err(mpsc::TryRecvError::Disconnected) => self.loading = None,
172 Err(mpsc::TryRecvError::Empty) => {}
173 }
174 }
175 let base = key.split('@').next().unwrap_or(key);
176 let comm = if base.chars().any(char::is_alphabetic) {
177 None
178 } else {
179 comm_in(cgroup, base)
180 };
181 display_name(key, &self.desktop, comm.as_deref())
182 }
183}
184
185impl Default for AppNames {
186 fn default() -> Self {
187 Self::new()
188 }
189}
190
191fn comm_in(cgroup: &str, exe: &str) -> Option<String> {
193 super::cgfs::pids_under(cgroup)
194 .into_iter()
195 .find(|&p| super::cgfs::exe_basename(p).as_deref() == Some(exe))
196 .and_then(|p| std::fs::read_to_string(format!("/proc/{p}/comm")).ok())
197 .map(|c| c.trim().to_string())
198 .filter(|c| !c.is_empty())
199}
200
201#[derive(Debug, Clone, Copy, PartialEq, Eq)]
203pub struct NotifyFlags {
204 pub notify: bool,
206 pub pressure: bool,
208}
209
210impl NotifyFlags {
211 pub fn from_config(g: &GuardConfig) -> Self {
212 Self {
213 notify: g.notify,
214 pressure: g.notify_pressure,
215 }
216 }
217}
218
219pub fn flags_after_reload(
223 current: NotifyFlags,
224 reloaded: std::result::Result<&GuardConfig, &common::Error>,
225) -> NotifyFlags {
226 match reloaded {
227 Ok(g) => NotifyFlags::from_config(g),
228 Err(_) => current,
229 }
230}
231
232#[derive(Debug, Clone)]
234struct Held {
235 app: String,
236 cap_bytes: Option<u64>,
238}
239
240pub struct Notifier<S: NotifySink> {
249 sink: S,
250 flags: NotifyFlags,
251 freeze_hold_secs: u64,
252 held: HashMap<String, Held>,
254 names: HashMap<String, String>,
256 shown: HashMap<String, (String, String)>,
258 thawed: HashSet<String>,
260 pressure_shown: bool,
261 last_pressure_ms: Option<u64>,
262}
263
264impl<S: NotifySink> Notifier<S> {
265 pub fn new(sink: S, cfg: &GuardConfig) -> Self {
266 Self {
267 sink,
268 flags: NotifyFlags::from_config(cfg),
269 freeze_hold_secs: cfg.timing.freeze_hold_secs,
270 held: HashMap::new(),
271 names: HashMap::new(),
272 shown: HashMap::new(),
273 thawed: HashSet::new(),
274 pressure_shown: false,
275 last_pressure_ms: None,
276 }
277 }
278
279 pub fn sink(&self) -> &S {
280 &self.sink
281 }
282
283 pub fn flags(&self) -> NotifyFlags {
284 self.flags
285 }
286
287 pub fn set_flags(&mut self, flags: NotifyFlags) {
290 if !flags.notify {
291 self.close_all();
292 } else if !flags.pressure {
293 self.close_pressure();
294 }
295 self.flags = flags;
296 }
297
298 pub fn record(&mut self, action: &Action, applied: Option<Applied>) {
302 match action {
303 Action::Freeze { res, name: app } | Action::Cap { res, name: app } => {
304 let Some(applied) = applied else {
305 return;
306 };
307 let cap_bytes = match action {
308 Action::Cap { .. } => Some(applied.cap_bytes.unwrap_or(0)),
309 _ => None,
310 };
311 self.held.insert(
312 res.cgroup.clone(),
313 Held {
314 app: app.clone(),
315 cap_bytes,
316 },
317 );
318 }
319 Action::Thaw { res } => {
322 if let Some(h) = self.held.remove(&res.cgroup) {
323 self.thawed.insert(h.app);
324 }
325 }
326 Action::LiftCap { res } => {
327 self.held.remove(&res.cgroup);
328 }
329 }
330 }
331
332 pub fn end_tick(
337 &mut self,
338 now_ms: u64,
339 memory: Memory,
340 name: &mut dyn FnMut(&str, &str) -> String,
341 ) {
342 let thawed = std::mem::take(&mut self.thawed);
343 if !self.flags.notify {
344 return;
345 }
346
347 let mut held: Vec<(&String, &Held)> = self.held.iter().collect();
348 held.sort_by(|a, b| a.0.cmp(b.0));
349 for (cgroup, h) in &held {
350 if !self.names.contains_key(&h.app) {
351 self.names.insert(h.app.clone(), name(&h.app, cgroup));
352 }
353 }
354 let mut want: HashMap<String, (String, String)> = HashMap::new();
355 for (_, h) in held {
356 let display = self
357 .names
358 .get(&h.app)
359 .map_or(h.app.as_str(), |n| n.as_str());
360 let capped: u64 = self
361 .held
362 .values()
363 .filter(|o| o.app == h.app)
364 .filter_map(|o| o.cap_bytes)
365 .sum();
366 let any_capped = self
367 .held
368 .values()
369 .any(|o| o.app == h.app && o.cap_bytes.is_some());
370 let text = if any_capped {
371 slowed_text(display, capped)
372 } else {
373 paused_text(display, self.freeze_hold_secs)
374 };
375 want.entry(h.app.clone()).or_insert(text);
376 }
377
378 let mut gone: Vec<String> = self
379 .shown
380 .keys()
381 .filter(|app| !want.contains_key(*app) && !thawed.contains(*app))
382 .cloned()
383 .collect();
384 gone.sort();
385 for app in gone {
386 self.sink.close(&app);
387 self.shown.remove(&app);
388 }
389 let (held, shown) = (&self.held, &self.shown);
390 self.names
391 .retain(|app, _| shown.contains_key(app) || held.values().any(|h| h.app == *app));
392 let mut changed: Vec<(String, (String, String))> = want
393 .into_iter()
394 .filter(|(app, text)| self.shown.get(app) != Some(text))
395 .collect();
396 changed.sort();
397 for (app, (title, body)) in changed {
398 self.sink.show(&app, &title, &body);
399 self.shown.insert(app, (title, body));
400 }
401
402 let pressing = memory == Memory::Low;
403 if self.held.is_empty() && self.shown.is_empty() && pressing {
404 let due = self
405 .last_pressure_ms
406 .is_none_or(|last| now_ms.saturating_sub(last) >= PRESSURE_INTERVAL_MS);
407 if self.flags.pressure && due {
408 self.sink.show(PRESSURE_KEY, PRESSURE_TITLE, PRESSURE_BODY);
409 self.pressure_shown = true;
410 self.last_pressure_ms = Some(now_ms);
411 }
412 } else if memory != Memory::Unknown || !self.held.is_empty() {
413 self.close_pressure();
414 }
415 }
416
417 pub fn close_all(&mut self) {
419 let mut apps: Vec<String> = self.shown.drain().map(|(app, _)| app).collect();
420 apps.sort();
421 for app in apps {
422 self.sink.close(&app);
423 }
424 self.close_pressure();
425 }
426
427 fn close_pressure(&mut self) {
428 if self.pressure_shown {
429 self.sink.close(PRESSURE_KEY);
430 self.pressure_shown = false;
431 }
432 }
433}
434
435const CALL_TIMEOUT: Duration = Duration::from_secs(2);
437
438const EXPIRE_TIMEOUT: i32 = -1;
443const QUEUE: usize = 64;
445
446const NOTIFY_DEST: &str = "org.freedesktop.Notifications";
447const NOTIFY_PATH: &str = "/org/freedesktop/Notifications";
448
449enum Cmd {
450 Show {
451 key: String,
452 title: String,
453 body: String,
454 },
455 Close {
456 key: String,
457 },
458 Flush(mpsc::Sender<()>),
459}
460
461pub struct DesktopSink {
470 tx: Option<mpsc::SyncSender<Cmd>>,
471}
472
473impl DesktopSink {
474 pub fn spawn() -> Self {
477 let (tx, rx) = mpsc::sync_channel(QUEUE);
478 let started = std::thread::Builder::new()
479 .name("rlm-notify".into())
480 .spawn(move || sender(rx));
481 match started {
482 Ok(_) => Self { tx: Some(tx) },
483 Err(e) => {
484 tracing::debug!(error = %e, "cannot start the notification thread");
485 Self { tx: None }
486 }
487 }
488 }
489
490 pub fn flush(&self, timeout: Duration) -> bool {
493 let (done, wait) = mpsc::channel();
494 if !self.send(Cmd::Flush(done)) {
495 return false;
496 }
497 wait.recv_timeout(timeout).is_ok()
498 }
499
500 fn send(&self, cmd: Cmd) -> bool {
501 match &self.tx {
502 Some(tx) => match tx.try_send(cmd) {
503 Ok(()) => true,
504 Err(e) => {
505 tracing::debug!(error = %e, "notification dropped");
506 false
507 }
508 },
509 None => false,
510 }
511 }
512}
513
514impl NotifySink for DesktopSink {
515 fn show(&mut self, key: &str, title: &str, body: &str) {
516 self.send(Cmd::Show {
517 key: key.to_string(),
518 title: title.to_string(),
519 body: body.to_string(),
520 });
521 }
522
523 fn close(&mut self, key: &str) {
524 self.send(Cmd::Close {
525 key: key.to_string(),
526 });
527 }
528}
529
530fn sender(rx: mpsc::Receiver<Cmd>) {
533 let mut sender = Sender::new(SessionBus);
534 for cmd in rx {
535 sender.handle(cmd, Instant::now());
536 }
537}
538
539const RECONNECT_AFTER: Duration = Duration::from_secs(30);
542
543#[derive(Debug)]
545enum BusError {
546 Gone(String),
548 Call(String),
550}
551
552trait Bus {
554 fn notify(&self, replaces_id: u32, title: &str, body: &str) -> Result<u32, BusError>;
555 fn close(&self, id: u32) -> Result<(), BusError>;
556}
557
558trait Connector {
560 type Bus: Bus;
561 fn connect(&mut self) -> Option<Self::Bus>;
562 fn fallback(&mut self, title: &str, body: &str);
564}
565
566struct Sender<C: Connector> {
569 connector: C,
570 bus: Option<C::Bus>,
571 retry_at: Option<Instant>,
573 ids: HashMap<String, u32>,
574}
575
576impl<C: Connector> Sender<C> {
577 fn new(connector: C) -> Self {
578 Self {
579 connector,
580 bus: None,
581 retry_at: None,
582 ids: HashMap::new(),
583 }
584 }
585
586 fn ensure_bus(&mut self, now: Instant) {
588 if self.bus.is_some() || self.retry_at.is_some_and(|t| now < t) {
589 return;
590 }
591 self.bus = self.connector.connect();
592 self.retry_at = if self.bus.is_none() {
593 Some(now + RECONNECT_AFTER)
594 } else {
595 None
596 };
597 }
598
599 fn reset(&mut self, why: &str) {
602 tracing::debug!(error = why, "notification connection lost; reconnecting");
603 self.bus = None;
604 self.retry_at = None;
605 self.ids.clear();
606 }
607
608 fn handle(&mut self, cmd: Cmd, now: Instant) {
609 if let Cmd::Flush(done) = cmd {
610 let _ = done.send(());
611 return;
612 }
613 self.ensure_bus(now);
614 match cmd {
615 Cmd::Show { key, title, body } => {
616 let Some(bus) = &self.bus else {
617 self.connector.fallback(&title, &body);
618 return;
619 };
620 let replaces = self.ids.get(&key).copied().unwrap_or(0);
621 match bus.notify(replaces, &title, &body) {
622 Ok(id) => {
623 self.ids.insert(key, id);
624 }
625 Err(BusError::Gone(e)) => {
626 self.reset(&e);
627 self.ensure_bus(now);
629 match &self.bus {
630 Some(bus) => match bus.notify(0, &title, &body) {
631 Ok(id) => {
632 self.ids.insert(key, id);
633 }
634 Err(e) => tracing::debug!(error = ?e, "Notify failed"),
635 },
636 None => self.connector.fallback(&title, &body),
637 }
638 }
639 Err(BusError::Call(e)) => tracing::debug!(error = e, "Notify failed"),
640 }
641 }
642 Cmd::Close { key } => {
643 let (Some(bus), Some(id)) = (&self.bus, self.ids.remove(&key)) else {
644 return;
645 };
646 match bus.close(id) {
647 Ok(()) => {}
648 Err(BusError::Gone(e)) => self.reset(&e),
649 Err(BusError::Call(e)) => {
650 tracing::debug!(error = e, "CloseNotification failed")
651 }
652 }
653 }
654 Cmd::Flush(_) => {}
655 }
656 }
657}
658
659struct SessionBus;
661
662impl Connector for SessionBus {
663 type Bus = zbus::blocking::Connection;
664
665 fn connect(&mut self) -> Option<Self::Bus> {
666 connect()
667 }
668
669 fn fallback(&mut self, title: &str, body: &str) {
670 notify_send(title, body);
671 }
672}
673
674impl Bus for zbus::blocking::Connection {
675 fn notify(&self, replaces_id: u32, title: &str, body: &str) -> Result<u32, BusError> {
676 dbus_notify(self, replaces_id, title, body, EXPIRE_TIMEOUT).map_err(classify)
677 }
678
679 fn close(&self, id: u32) -> Result<(), BusError> {
680 self.call_method(
681 Some(NOTIFY_DEST),
682 NOTIFY_PATH,
683 Some(NOTIFY_DEST),
684 "CloseNotification",
685 &(id,),
686 )
687 .map(|_| ())
688 .map_err(classify)
689 }
690}
691
692fn classify(e: zbus::Error) -> BusError {
697 let gone = match &e {
698 zbus::Error::InputOutput(io) => io.kind() != std::io::ErrorKind::TimedOut,
699 zbus::Error::MethodError(name, _, _) => matches!(
700 name.as_str(),
701 "org.freedesktop.DBus.Error.ServiceUnknown"
702 | "org.freedesktop.DBus.Error.Disconnected"
703 | "org.freedesktop.DBus.Error.NoServer"
704 ),
705 _ => false,
706 };
707 if gone {
708 BusError::Gone(e.to_string())
709 } else {
710 BusError::Call(e.to_string())
711 }
712}
713
714fn connect() -> Option<zbus::blocking::Connection> {
717 let conn = within(CALL_TIMEOUT, || {
718 zbus::blocking::connection::Builder::session()
719 .map(|b| b.method_timeout(CALL_TIMEOUT))
720 .and_then(|b| b.build())
721 .map_err(|e| tracing::debug!(error = %e, "cannot connect to the session bus"))
722 .ok()
723 });
724 if conn.is_none() {
725 tracing::debug!("no session bus; notifications use notify-send");
726 }
727 conn
728}
729
730fn within<T: Send + 'static>(
733 timeout: Duration,
734 f: impl FnOnce() -> Option<T> + Send + 'static,
735) -> Option<T> {
736 let (tx, rx) = mpsc::channel();
737 std::thread::Builder::new()
738 .name("rlm-notify-connect".into())
739 .spawn(move || {
740 let _ = tx.send(f());
741 })
742 .ok()?;
743 rx.recv_timeout(timeout).ok().flatten()
744}
745
746fn dbus_notify(
747 conn: &zbus::blocking::Connection,
748 replaces_id: u32,
749 title: &str,
750 body: &str,
751 expire_timeout: i32,
752) -> zbus::Result<u32> {
753 use zbus::zvariant::Value;
754 let mut hints: HashMap<&str, Value<'_>> = HashMap::new();
755 hints.insert("desktop-entry", Value::from(APP_ID));
756 hints.insert("urgency", Value::U8(1));
758 let actions: Vec<&str> = Vec::new();
759 let reply = conn.call_method(
760 Some(NOTIFY_DEST),
761 NOTIFY_PATH,
762 Some(NOTIFY_DEST),
763 "Notify",
764 &(
765 APP_NAME,
766 replaces_id,
767 APP_ID,
768 title,
769 body,
770 actions,
771 hints,
772 expire_timeout,
773 ),
774 )?;
775 reply.body().deserialize::<u32>()
776}
777
778fn notify_send(title: &str, body: &str) {
781 match Command::new("notify-send")
782 .args(["-a", APP_NAME, "-i", APP_ID, title, body])
783 .spawn()
784 {
785 Ok(mut child) => {
786 std::thread::spawn(move || {
787 let _ = child.wait();
788 });
789 }
790 Err(e) => tracing::debug!(error = %e, "notify-send unavailable; skipping notification"),
791 }
792}
793
794#[cfg(test)]
795mod tests {
796 use super::super::resolve::{Coverage, Mechanism, Resolution, Verdict};
797 use super::*;
798
799 #[derive(Debug, Clone, PartialEq, Eq)]
800 enum Call {
801 Show(String, String, String),
802 Close(String),
803 }
804
805 #[derive(Default)]
806 struct Fake(Vec<Call>);
807
808 impl NotifySink for Fake {
809 fn show(&mut self, key: &str, title: &str, body: &str) {
810 self.0
811 .push(Call::Show(key.into(), title.into(), body.into()));
812 }
813 fn close(&mut self, key: &str) {
814 self.0.push(Call::Close(key.into()));
815 }
816 }
817
818 impl Notifier<Fake> {
819 fn take(&mut self) -> Vec<Call> {
820 std::mem::take(&mut self.sink.0)
821 }
822 }
823
824 const GB: u64 = 1_000_000_000;
825
826 fn res(cg: &str) -> Resolution {
827 Resolution {
828 cgroup: cg.into(),
829 unit: None,
830 verdict: Verdict::Freeze,
831 coverage: Coverage::Full,
832 mechanism: Mechanism::Raw,
833 }
834 }
835
836 fn freeze(cg: &str) -> Action {
837 Action::Freeze {
838 res: res(cg),
839 name: "firefox".into(),
840 }
841 }
842
843 fn cap(cg: &str) -> Action {
844 Action::Cap {
845 res: res(cg),
846 name: "firefox".into(),
847 }
848 }
849
850 fn ok() -> Option<Applied> {
851 Some(Applied { cap_bytes: None })
852 }
853
854 fn capped(bytes: u64) -> Option<Applied> {
855 Some(Applied {
856 cap_bytes: Some(bytes),
857 })
858 }
859
860 fn notifier(notify: bool, pressure: bool) -> Notifier<Fake> {
861 let cfg = GuardConfig {
862 notify,
863 notify_pressure: pressure,
864 ..GuardConfig::default()
865 };
866 Notifier::new(Fake::default(), &cfg)
867 }
868
869 fn names(key: &str, _cg: &str) -> String {
870 display_name(key, &HashMap::new(), None)
871 }
872
873 fn paused() -> Call {
874 let (t, b) = paused_text("Firefox", 5);
875 Call::Show("firefox".into(), t, b)
876 }
877
878 fn slowed(bytes: u64) -> Call {
879 let (t, b) = slowed_text("Firefox", bytes);
880 Call::Show("firefox".into(), t, b)
881 }
882
883 #[test]
884 fn texts_read_as_specified() {
885 assert_eq!(
886 paused_text("Firefox", 5),
887 (
888 "Firefox paused".to_string(),
889 "Paused for 5 seconds while memory is low.".to_string()
890 )
891 );
892 assert_eq!(
893 slowed_text("Firefox", 3_200_000_000),
894 (
895 "Firefox slowed down".to_string(),
896 "Held to about 3.2 GB until memory frees up.".to_string()
897 )
898 );
899 assert!(paused_text("X", 1).1.contains("for 1 second while"));
900 assert_eq!(format_size(512_000_000), "512 MB");
901 assert_eq!(format_size(268_435_456), "268 MB");
902 }
903
904 #[test]
905 fn freeze_shows_paused_once_for_all_of_an_apps_cgroups() {
906 let mut n = notifier(true, false);
907 n.record(&freeze("/a"), ok());
908 n.record(&freeze("/b"), ok());
909 n.end_tick(0, Memory::Low, &mut names);
910 assert_eq!(n.take(), vec![paused()]);
911 n.end_tick(1_000, Memory::Low, &mut names);
912 assert!(n.take().is_empty(), "nothing changed, nothing sent");
913 }
914
915 #[test]
916 fn thaw_then_cap_in_one_tick_replaces_without_close() {
917 let mut n = notifier(true, false);
918 n.record(&freeze("/a"), ok());
919 n.end_tick(0, Memory::Low, &mut names);
920 n.take();
921 n.record(&Action::Thaw { res: res("/a") }, ok());
922 n.record(&cap("/a"), capped(3 * GB));
923 n.end_tick(5_000, Memory::Low, &mut names);
924 assert_eq!(n.take(), vec![slowed(3 * GB)]);
925 }
926
927 #[test]
928 fn cap_on_the_tick_after_a_thaw_still_replaces() {
929 let mut n = notifier(true, false);
930 n.record(&freeze("/a"), ok());
931 n.end_tick(0, Memory::Low, &mut names);
932 n.take();
933 n.record(&Action::Thaw { res: res("/a") }, ok());
934 n.end_tick(5_000, Memory::Low, &mut names);
935 assert!(n.take().is_empty(), "the close waits a tick");
936 n.record(&cap("/a"), capped(2 * GB));
937 n.end_tick(6_000, Memory::Low, &mut names);
938 assert_eq!(n.take(), vec![slowed(2 * GB)]);
939 }
940
941 #[test]
942 fn thaw_alone_closes() {
943 let mut n = notifier(true, false);
944 n.record(&freeze("/a"), ok());
945 n.end_tick(0, Memory::Low, &mut names);
946 n.take();
947 n.record(&Action::Thaw { res: res("/a") }, ok());
948 n.end_tick(5_000, Memory::Fine, &mut names);
949 n.end_tick(6_000, Memory::Fine, &mut names);
950 assert_eq!(n.take(), vec![Call::Close("firefox".into())]);
951 }
952
953 #[test]
954 fn lift_closes_in_the_same_tick() {
955 let mut n = notifier(true, false);
956 n.record(&cap("/a"), capped(GB));
957 n.end_tick(0, Memory::Low, &mut names);
958 assert_eq!(n.take(), vec![slowed(GB)]);
959 n.record(&Action::LiftCap { res: res("/a") }, ok());
960 n.end_tick(40_000, Memory::Fine, &mut names);
961 assert_eq!(n.take(), vec![Call::Close("firefox".into())]);
962 }
963
964 #[test]
965 fn a_release_that_reported_an_error_still_closes() {
966 let mut n = notifier(true, false);
967 n.record(&cap("/a"), capped(GB));
968 n.end_tick(0, Memory::Low, &mut names);
969 n.take();
970 n.record(&Action::LiftCap { res: res("/a") }, None);
971 n.end_tick(1_000, Memory::Fine, &mut names);
972 assert_eq!(n.take(), vec![Call::Close("firefox".into())]);
973 }
974
975 #[test]
976 fn failed_interventions_send_nothing() {
977 let mut n = notifier(true, false);
978 n.record(&freeze("/a"), None);
979 n.record(&cap("/b"), None);
980 n.end_tick(0, Memory::Low, &mut names);
981 assert!(n.take().is_empty());
982 }
983
984 #[test]
985 fn caps_of_one_app_add_up() {
986 let mut n = notifier(true, false);
987 n.record(&cap("/a"), capped(GB));
988 n.record(&cap("/b"), capped(2 * GB));
989 n.end_tick(0, Memory::Low, &mut names);
990 assert_eq!(n.take(), vec![slowed(3 * GB)]);
991 n.record(&Action::LiftCap { res: res("/a") }, ok());
992 n.end_tick(1_000, Memory::Low, &mut names);
993 assert_eq!(n.take(), vec![slowed(2 * GB)], "updated, not closed");
994 }
995
996 #[test]
997 fn disabled_notify_sends_nothing() {
998 let mut n = notifier(false, true);
999 n.record(&freeze("/a"), ok());
1000 n.end_tick(0, Memory::Low, &mut names);
1001 n.record(&Action::Thaw { res: res("/a") }, ok());
1002 n.end_tick(5_000, Memory::Low, &mut names);
1003 n.end_tick(6_000, Memory::Low, &mut names);
1004 n.close_all();
1005 assert!(n.take().is_empty());
1006 }
1007
1008 #[test]
1009 fn turning_notify_off_clears_what_is_shown() {
1010 let mut n = notifier(true, false);
1011 n.record(&freeze("/a"), ok());
1012 n.end_tick(0, Memory::Low, &mut names);
1013 n.take();
1014 n.set_flags(NotifyFlags {
1015 notify: false,
1016 pressure: false,
1017 });
1018 assert_eq!(n.take(), vec![Call::Close("firefox".into())]);
1019 n.end_tick(1_000, Memory::Low, &mut names);
1020 assert!(n.take().is_empty());
1021 }
1022
1023 #[test]
1024 fn close_all_closes_every_app() {
1025 let mut n = notifier(true, false);
1026 n.record(&freeze("/a"), ok());
1027 n.end_tick(0, Memory::Low, &mut names);
1028 n.take();
1029 n.close_all();
1030 assert_eq!(n.take(), vec![Call::Close("firefox".into())]);
1031 }
1032
1033 fn warning() -> Call {
1034 Call::Show(
1035 PRESSURE_KEY.into(),
1036 PRESSURE_TITLE.into(),
1037 PRESSURE_BODY.into(),
1038 )
1039 }
1040
1041 #[test]
1042 fn pressure_warning_is_off_by_default() {
1043 let mut n = Notifier::new(Fake::default(), &GuardConfig::default());
1044 n.end_tick(0, Memory::Low, &mut names);
1045 assert!(n.take().is_empty());
1046 }
1047
1048 #[test]
1049 fn pressure_warning_is_rate_limited_and_only_at_high() {
1050 let mut n = notifier(true, true);
1051 n.end_tick(0, Memory::Fine, &mut names);
1052 assert!(n.take().is_empty(), "Warn is not enough");
1053 n.end_tick(1_000, Memory::Low, &mut names);
1054 assert_eq!(n.take(), vec![warning()]);
1055 n.end_tick(30_000, Memory::Low, &mut names);
1056 assert!(n.take().is_empty(), "at most once a minute");
1057 n.end_tick(61_000, Memory::Low, &mut names);
1058 assert_eq!(n.take(), vec![warning()]);
1059 }
1060
1061 #[test]
1062 fn memory_is_low_only_when_pressure_is_high_and_memory_short() {
1063 let t = common::GuardTrigger::default();
1064 let sample = |avail| Sample {
1065 some_avg10: 40.0,
1066 full_avg10: 0.0,
1067 mem_available_mb: avail,
1068 mem_total_mb: 16_000,
1069 source: super::super::types::PsiSource::AppSlice,
1070 };
1071 assert_eq!(memory_state(Level::High, &sample(1_000), &t), Memory::Low);
1072 assert_eq!(
1073 memory_state(Level::High, &sample(8_000), &t),
1074 Memory::Fine,
1075 "a stall with half the RAM free is not low memory"
1076 );
1077 assert_eq!(memory_state(Level::Warn, &sample(1_000), &t), Memory::Fine);
1078 assert_eq!(memory_state(Level::Critical, &sample(300), &t), Memory::Low);
1079 }
1080
1081 #[test]
1082 fn pressure_warning_gives_way_to_an_intervention() {
1083 let mut n = notifier(true, true);
1084 n.end_tick(0, Memory::Low, &mut names);
1085 n.take();
1086 n.record(&freeze("/a"), ok());
1087 n.end_tick(1_000, Memory::Low, &mut names);
1088 assert_eq!(n.take(), vec![paused(), Call::Close(PRESSURE_KEY.into())]);
1089 n.end_tick(70_000, Memory::Low, &mut names);
1090 assert!(n.take().is_empty(), "no warning while an app is held");
1091 }
1092
1093 #[test]
1094 fn pressure_warning_needs_notify() {
1095 let mut n = notifier(false, true);
1096 n.end_tick(0, Memory::Low, &mut names);
1097 assert!(n.take().is_empty());
1098 }
1099
1100 #[test]
1101 fn reload_takes_only_valid_notification_flags() {
1102 let current = NotifyFlags {
1103 notify: true,
1104 pressure: false,
1105 };
1106 let mut changed = GuardConfig {
1107 notify_pressure: true,
1108 enabled: false,
1109 ..GuardConfig::default()
1110 };
1111 changed.timing.freeze_hold_secs = 30;
1112 assert_eq!(
1113 flags_after_reload(current, Ok(&changed)),
1114 NotifyFlags {
1115 notify: true,
1116 pressure: true
1117 }
1118 );
1119 let err = common::Error::Config("bad".into());
1120 assert_eq!(flags_after_reload(current, Err(&err)), current);
1121 }
1122
1123 use std::cell::{Cell, RefCell};
1124 use std::rc::Rc;
1125
1126 #[derive(Default)]
1128 struct Log {
1129 calls: RefCell<Vec<String>>,
1130 connects: RefCell<Vec<bool>>,
1133 fail_next: RefCell<Option<BusError>>,
1135 next_id: Cell<u32>,
1136 }
1137
1138 struct FakeBus(Rc<Log>, u32);
1139
1140 impl Bus for FakeBus {
1141 fn notify(&self, replaces_id: u32, title: &str, _body: &str) -> Result<u32, BusError> {
1142 if let Some(e) = self.0.fail_next.borrow_mut().take() {
1143 return Err(e);
1144 }
1145 let id = if replaces_id == 0 {
1146 self.0.next_id.set(self.0.next_id.get() + 1);
1147 self.0.next_id.get()
1148 } else {
1149 replaces_id
1150 };
1151 self.0
1152 .calls
1153 .borrow_mut()
1154 .push(format!("conn{} notify {replaces_id}->{id} {title}", self.1));
1155 Ok(id)
1156 }
1157
1158 fn close(&self, id: u32) -> Result<(), BusError> {
1159 if let Some(e) = self.0.fail_next.borrow_mut().take() {
1160 return Err(e);
1161 }
1162 self.0
1163 .calls
1164 .borrow_mut()
1165 .push(format!("conn{} close {id}", self.1));
1166 Ok(())
1167 }
1168 }
1169
1170 struct FakeConnector(Rc<Log>, u32);
1171
1172 impl Connector for FakeConnector {
1173 type Bus = FakeBus;
1174
1175 fn connect(&mut self) -> Option<FakeBus> {
1176 let ok = {
1177 let mut answers = self.0.connects.borrow_mut();
1178 if answers.is_empty() {
1179 true
1180 } else {
1181 answers.remove(0)
1182 }
1183 };
1184 self.0
1185 .calls
1186 .borrow_mut()
1187 .push(format!("connect {}", if ok { "ok" } else { "failed" }));
1188 ok.then(|| {
1189 self.1 += 1;
1190 FakeBus(Rc::clone(&self.0), self.1)
1191 })
1192 }
1193
1194 fn fallback(&mut self, title: &str, _body: &str) {
1195 self.0
1196 .calls
1197 .borrow_mut()
1198 .push(format!("notify-send {title}"));
1199 }
1200 }
1201
1202 fn show(key: &str, title: &str) -> Cmd {
1203 Cmd::Show {
1204 key: key.into(),
1205 title: title.into(),
1206 body: String::new(),
1207 }
1208 }
1209
1210 fn sender_with(log: &Rc<Log>) -> Sender<FakeConnector> {
1211 Sender::new(FakeConnector(Rc::clone(log), 0))
1212 }
1213
1214 fn take(log: &Log) -> Vec<String> {
1215 std::mem::take(&mut *log.calls.borrow_mut())
1216 }
1217
1218 #[test]
1219 fn a_failed_connect_is_retried_after_the_backoff() {
1220 let log = Rc::new(Log::default());
1221 log.connects.borrow_mut().push(false);
1222 let mut s = sender_with(&log);
1223 let t0 = Instant::now();
1224 s.handle(show("a", "A paused"), t0);
1225 s.handle(show("a", "A slowed"), t0 + Duration::from_secs(10));
1226 assert_eq!(
1227 take(&log),
1228 [
1229 "connect failed",
1230 "notify-send A paused",
1231 "notify-send A slowed"
1232 ],
1233 "no new attempt during the backoff"
1234 );
1235 s.handle(show("a", "A slowed"), t0 + RECONNECT_AFTER);
1236 s.handle(show("a", "A slowed again"), t0 + RECONNECT_AFTER);
1237 s.handle(
1238 Cmd::Close { key: "a".into() },
1239 t0 + RECONNECT_AFTER + Duration::from_secs(1),
1240 );
1241 assert_eq!(
1242 take(&log),
1243 [
1244 "connect ok",
1245 "conn1 notify 0->1 A slowed",
1246 "conn1 notify 1->1 A slowed again",
1247 "conn1 close 1"
1248 ]
1249 );
1250 }
1251
1252 #[test]
1253 fn a_broken_connection_is_dropped_and_rebuilt() {
1254 let log = Rc::new(Log::default());
1255 let mut s = sender_with(&log);
1256 let t0 = Instant::now();
1257 s.handle(show("a", "A paused"), t0);
1258 take(&log);
1259 *log.fail_next.borrow_mut() = Some(BusError::Gone("broken pipe".into()));
1260 s.handle(show("a", "A slowed"), t0);
1261 assert_eq!(
1262 take(&log),
1263 ["connect ok", "conn2 notify 0->2 A slowed"],
1264 "reconnects and shows it as a new notification"
1265 );
1266 *log.fail_next.borrow_mut() = Some(BusError::Gone("gone".into()));
1267 s.handle(Cmd::Close { key: "a".into() }, t0);
1268 s.handle(show("b", "B paused"), t0);
1269 assert_eq!(
1270 take(&log),
1271 ["connect ok", "conn3 notify 0->3 B paused"],
1272 "a failed close also resets"
1273 );
1274 }
1275
1276 #[test]
1277 fn a_timed_out_call_keeps_the_connection() {
1278 let log = Rc::new(Log::default());
1279 let mut s = sender_with(&log);
1280 let t0 = Instant::now();
1281 s.handle(show("a", "A paused"), t0);
1282 *log.fail_next.borrow_mut() = Some(BusError::Call("timed out".into()));
1283 s.handle(show("a", "A slowed"), t0);
1284 s.handle(show("a", "A slowed"), t0);
1285 assert_eq!(
1286 take(&log),
1287 [
1288 "connect ok",
1289 "conn1 notify 0->1 A paused",
1290 "conn1 notify 1->1 A slowed"
1291 ]
1292 );
1293 }
1294
1295 #[test]
1296 fn io_errors_other_than_timeouts_mean_the_connection_is_gone() {
1297 let io = |kind| zbus::Error::InputOutput(std::sync::Arc::new(std::io::Error::from(kind)));
1298 assert!(matches!(
1299 classify(io(std::io::ErrorKind::BrokenPipe)),
1300 BusError::Gone(_)
1301 ));
1302 assert!(matches!(
1303 classify(io(std::io::ErrorKind::TimedOut)),
1304 BusError::Call(_)
1305 ));
1306 assert!(matches!(
1307 classify(zbus::Error::InvalidReply),
1308 BusError::Call(_)
1309 ));
1310 }
1311
1312 #[test]
1313 fn a_slow_connect_gives_up() {
1314 let start = std::time::Instant::now();
1315 let got = within(Duration::from_millis(50), || {
1316 std::thread::sleep(Duration::from_secs(5));
1317 Some(1)
1318 });
1319 assert_eq!(got, None);
1320 assert!(start.elapsed() < Duration::from_secs(2));
1321 assert_eq!(within(Duration::from_secs(2), || Some(7)), Some(7));
1322 }
1323
1324 #[test]
1325 fn notifications_use_the_server_default_lifetime() {
1326 assert_eq!(EXPIRE_TIMEOUT, -1);
1327 }
1328
1329 #[test]
1330 fn display_names() {
1331 let mut desktop = HashMap::new();
1332 desktop.insert("code".to_string(), "Visual Studio Code".to_string());
1333 desktop.insert("claude".to_string(), "Claude".to_string());
1334 assert_eq!(display_name("code", &desktop, None), "Visual Studio Code");
1335 assert_eq!(display_name("firefox", &desktop, None), "Firefox");
1336 assert_eq!(display_name("node@app-x.scope", &desktop, None), "Node");
1337 assert_eq!(
1338 display_name("python3@run-u12.service", &desktop, None),
1339 "Python3"
1340 );
1341 assert_eq!(
1342 display_name("2.1.283", &HashMap::new(), Some("claude")),
1343 "Claude"
1344 );
1345 assert_eq!(display_name("2.1.283", &desktop, Some("claude")), "Claude");
1346 assert_eq!(display_name("2.1.283", &desktop, None), "2.1.283");
1347 assert_eq!(
1348 display_name("firefox", &desktop, Some("Isolated Web Co")),
1349 "Firefox",
1350 "comm is only used when the basename has no letters"
1351 );
1352 assert_eq!(display_name("élan", &desktop, None), "Élan");
1353 assert_eq!(display_name("", &desktop, None), "An app");
1354 }
1355}