Address suspend rail review feedback

Author: Staphylococcus

crates/lg-buddy/src/session/runner.rs

@@ -352,11 +352,13 @@ fn run_lifecycle_monitor_with_bus<W: Write, E: SessionActionExecutor>(
     )?;
 
     let mut sleep_inhibitor = None;
-    acquire_lifecycle_sleep_delay_inhibitor_if_enabled(
+    let mut waiting_for_resume = false;
+    sync_lifecycle_sleep_delay_inhibitor_if_idle(
         writer,
         bus,
         config_path,
         &mut sleep_inhibitor,
+        waiting_for_resume,
     )?;
 
     let started = Instant::now();
@@ -378,6 +380,14 @@ fn run_lifecycle_monitor_with_bus<W: Write, E: SessionActionExecutor>(
                 .min(timeout.saturating_sub(started.elapsed()));
         }
 
+        sync_lifecycle_sleep_delay_inhibitor_if_idle(
+            writer,
+            bus,
+            config_path,
+            &mut sleep_inhibitor,
+            waiting_for_resume,
+        )?;
+
         let Some(signal) =
             bus.process(process_timeout)
                 .map_err(|err| SessionRunnerError::Failed {
@@ -393,6 +403,12 @@ fn run_lifecycle_monitor_with_bus<W: Write, E: SessionActionExecutor>(
         };
 
         let event_kind = LifecycleEvent::from_runtime_event(event);
+        match event_kind {
+            Some(LifecycleEvent::MachinePreparingForSleep) => waiting_for_resume = true,
+            Some(LifecycleEvent::MachineResumed) => waiting_for_resume = false,
+            _ => {}
+        }
+
         if !lifecycle_policy_enabled_from_config(config_path)? {
             release_lifecycle_sleep_delay_inhibitor(writer, &mut sleep_inhibitor)?;
             writeln!(
@@ -409,11 +425,12 @@ fn run_lifecycle_monitor_with_bus<W: Write, E: SessionActionExecutor>(
         match event_kind {
             Some(LifecycleEvent::MachineResumed) => {
                 dispatcher.dispatch_lifecycle_event(writer, event)?;
-                acquire_lifecycle_sleep_delay_inhibitor_if_enabled(
+                sync_lifecycle_sleep_delay_inhibitor_if_idle(
                     writer,
                     bus,
                     config_path,
                     &mut sleep_inhibitor,
+                    waiting_for_resume,
                 )?;
             }
             Some(LifecycleEvent::MachinePreparingForSleep) => {
@@ -436,13 +453,19 @@ fn run_lifecycle_monitor_with_bus<W: Write, E: SessionActionExecutor>(
     }
 }
 
-fn acquire_lifecycle_sleep_delay_inhibitor_if_enabled<W: Write>(
+fn sync_lifecycle_sleep_delay_inhibitor_if_idle<W: Write>(
     writer: &mut W,
     bus: &mut impl SessionBusClient,
     config_path: &Path,
     sleep_inhibitor: &mut Option<OwnedFd>,
+    waiting_for_resume: bool,
 ) -> Result<(), SessionRunnerError> {
-    if sleep_inhibitor.is_some() || !lifecycle_policy_enabled_from_config(config_path)? {
+    if !lifecycle_policy_enabled_from_config(config_path)? {
+        release_lifecycle_sleep_delay_inhibitor(writer, sleep_inhibitor)?;
+        return Ok(());
+    }
+
+    if waiting_for_resume || sleep_inhibitor.is_some() {
         return Ok(());
     }
 
@@ -1572,6 +1595,7 @@ mod tests {
         method_calls: Vec<(String, String, String, String)>,
         inhibitor_write_fds: Vec<i32>,
         disable_config_after_signals: Option<PathBuf>,
+        enable_config_after_idle: Option<PathBuf>,
     }
 
     impl SessionBusClient for FakeSessionBus {
@@ -1636,6 +1660,9 @@ mod tests {
             if let Some(config_path) = self.disable_config_after_signals.take() {
                 write_lifecycle_config(&config_path, "disabled");
             }
+            if let Some(config_path) = self.enable_config_after_idle.take() {
+                write_lifecycle_config(&config_path, "enabled");
+            }
 
             Ok(None)
         }
@@ -1806,6 +1833,7 @@ system_sleep_wake_policy={policy}
             method_calls: Vec::new(),
             inhibitor_write_fds: Vec::new(),
             disable_config_after_signals: Some(config_path.clone()),
+            enable_config_after_idle: None,
         };
         let mut output = Vec::new();
 
@@ -1854,6 +1882,7 @@ system_sleep_wake_policy={policy}
             method_calls: Vec::new(),
             inhibitor_write_fds: Vec::new(),
             disable_config_after_signals: None,
+            enable_config_after_idle: None,
         };
         let mut output = Vec::new();
 
@@ -1871,6 +1900,45 @@ system_sleep_wake_policy={policy}
         std::env::remove_var(super::LIFECYCLE_MONITOR_TEST_EVENT_LIMIT_ENV);
     }
 
+    #[test]
+    fn lifecycle_monitor_reacquires_inhibitor_when_policy_is_reenabled_while_idle() {
+        let _guard = env_lock()
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        std::env::set_var(super::LIFECYCLE_MONITOR_TEST_TIMEOUT_SECS_ENV, "0.2");
+
+        let config_path = unique_config_path("lifecycle-reenabled");
+        write_lifecycle_config(&config_path, "disabled");
+        let executor = FakeActionExecutor::default();
+        let mut bus = FakeLifecycleBus {
+            signals: VecDeque::new(),
+            signal_match_count: 0,
+            method_calls: Vec::new(),
+            inhibitor_write_fds: Vec::new(),
+            disable_config_after_signals: None,
+            enable_config_after_idle: Some(config_path.clone()),
+        };
+        let mut output = Vec::new();
+
+        run_lifecycle_monitor_with_bus(&mut output, executor, &config_path, &mut bus)
+            .expect("lifecycle loop exits cleanly after test timeout");
+
+        let output = String::from_utf8(output).expect("utf8");
+        assert!(output.contains("Using logind system lifecycle source"));
+        assert!(output.contains("Acquired logind sleep delay inhibitor"));
+        assert_eq!(bus.signal_match_count, 1);
+        assert_eq!(
+            bus.method_calls
+                .iter()
+                .map(|(_, _, _, member)| member.as_str())
+                .collect::<Vec<_>>(),
+            vec!["Inhibit"]
+        );
+        assert!(bus.enable_config_after_idle.is_none());
+        fs::remove_file(config_path).expect("remove lifecycle test config");
+        std::env::remove_var(super::LIFECYCLE_MONITOR_TEST_TIMEOUT_SECS_ENV);
+    }
+
     #[test]
     fn unhandled_events_are_logged_without_running_actions() {
         let executor = FakeActionExecutor::default();

crates/lg-buddy/src/state.rs

@@ -9,6 +9,7 @@ use std::os::fd::AsRawFd;
 #[cfg(unix)]
 use std::os::unix::fs::OpenOptionsExt;
 use std::path::{Path, PathBuf};
+use std::process;
 use std::time::{Duration, SystemTime};
 
 const STATE_DIR_NAME: &str = "lg_buddy";
@@ -364,20 +365,54 @@ fn create_marker_file(path: &Path) -> io::Result<()> {
     fs::write(path, [])
 }
 
-#[cfg(unix)]
 fn write_state_file(path: &Path, content: &str) -> io::Result<()> {
-    let mut file = OpenOptions::new()
-        .write(true)
-        .create(true)
-        .truncate(true)
-        .custom_flags(libc::O_NOFOLLOW)
-        .open(path)?;
-    file.write_all(content.as_bytes())
+    let mut last_error = None;
+    for attempt in 0..100 {
+        let temp_path = state_temp_path(path, attempt);
+        let mut options = OpenOptions::new();
+        options.write(true).create_new(true);
+        #[cfg(unix)]
+        options.custom_flags(libc::O_NOFOLLOW);
+
+        let mut file = match options.open(&temp_path) {
+            Ok(file) => file,
+            Err(err) if err.kind() == io::ErrorKind::AlreadyExists => {
+                last_error = Some(err);
+                continue;
+            }
+            Err(err) => return Err(err),
+        };
+
+        let result = (|| {
+            file.write_all(content.as_bytes())?;
+            file.flush()?;
+            file.sync_all()?;
+            drop(file);
+            fs::rename(&temp_path, path)
+        })();
+
+        if let Err(err) = result {
+            let _ = fs::remove_file(&temp_path);
+            return Err(err);
+        }
+
+        return Ok(());
+    }
+
+    Err(last_error.unwrap_or_else(|| {
+        io::Error::new(
+            io::ErrorKind::AlreadyExists,
+            "could not create unique state temporary file",
+        )
+    }))
 }
 
-#[cfg(not(unix))]
-fn write_state_file(path: &Path, content: &str) -> io::Result<()> {
-    fs::write(path, content)
+fn state_temp_path(path: &Path, attempt: u8) -> PathBuf {
+    let file_name = path
+        .file_name()
+        .and_then(|name| name.to_str())
+        .unwrap_or("state");
+    path.with_file_name(format!(".{file_name}.{}.{}.tmp", process::id(), attempt))
 }
 
 #[cfg(unix)]
@@ -397,10 +432,10 @@ fn current_uid() -> Option<u32> {
 #[cfg(test)]
 mod tests {
     use super::{
-        resolve_state_dir, RuntimeDirSources, ScreenOwnershipMarker, StateDirError, StateScope,
-        SystemSleepAttemptState, SystemSleepCycleOutcome, SystemSleepCycleState,
-        SCREEN_OFF_BY_US_MARKER, SYSTEM_SLEEP_ATTEMPTED_MARKER, SYSTEM_SLEEP_ATTEMPT_LOCK,
-        SYSTEM_SLEEP_CYCLE_STATE,
+        resolve_state_dir, state_temp_path, RuntimeDirSources, ScreenOwnershipMarker,
+        StateDirError, StateScope, SystemSleepAttemptState, SystemSleepCycleOutcome,
+        SystemSleepCycleState, SCREEN_OFF_BY_US_MARKER, SYSTEM_SLEEP_ATTEMPTED_MARKER,
+        SYSTEM_SLEEP_ATTEMPT_LOCK, SYSTEM_SLEEP_CYCLE_STATE,
     };
     use std::fs;
     #[cfg(unix)]
@@ -565,6 +600,34 @@ mod tests {
         assert_eq!(state.read_outcome().expect("read cleared outcome"), None);
     }
 
+    #[test]
+    fn system_sleep_cycle_outcome_failed_atomic_write_preserves_existing_file() {
+        let temp_dir = TestDir::new("system-sleep-cycle-atomic-failure");
+        let state = SystemSleepCycleState::new(temp_dir.path().to_path_buf());
+
+        state
+            .write_outcome(SystemSleepCycleOutcome::InProgress)
+            .expect("write initial outcome");
+        for attempt in 0..100 {
+            fs::write(state_temp_path(state.cycle_path(), attempt), "occupied")
+                .expect("occupy atomic temp path");
+        }
+
+        let err = state
+            .write_outcome(SystemSleepCycleOutcome::Completed)
+            .expect_err("colliding temp paths should fail");
+
+        assert_eq!(err.kind(), std::io::ErrorKind::AlreadyExists);
+        assert_eq!(
+            fs::read_to_string(state.cycle_path()).expect("read preserved outcome file"),
+            "outcome=in_progress\n"
+        );
+        assert_eq!(
+            state.read_outcome().expect("read preserved outcome"),
+            Some(SystemSleepCycleOutcome::InProgress)
+        );
+    }
+
     #[test]
     fn system_sleep_cycle_outcome_rejects_unknown_value() {
         let temp_dir = TestDir::new("system-sleep-cycle-bad-outcome");