iroh_util/
lock.rs

1use std::{fs::File, io, io::ErrorKind, io::Read, io::Write, path::PathBuf, process, result};
2
3use sysinfo::PidExt;
4use sysinfo::{Pid, ProcessExt, ProcessStatus::*, System, SystemExt};
5use thiserror::Error;
6use tracing::warn;
7
8use crate::exitcodes;
9
10/// Manages a lock file used to track if an iroh program is already running.
11/// Aquired locks write a file to iroh's application data path containing the
12/// process identifier (PID) of the process with the lock.
13/// The lock exclusion test requires both a lockfile AND a running process
14/// listed at the PID in the file
15/// An acquired lock is released either when the object is dropped
16/// or when the program stops, which removes the file
17/// Invalid or corrupt locks are overwritten on acquisition
18#[derive(Debug)]
19pub struct ProgramLock {
20    path: PathBuf,
21    lock: Option<sysinfo::Pid>,
22    system: Option<sysinfo::System>,
23}
24
25impl ProgramLock {
26    /// Create a new lock for the given program. This does not yet acquire the lock.
27    pub fn new(prog_name: &str) -> Result<Self> {
28        let path = crate::iroh_data_path(&format!("{prog_name}.lock"))
29            .map_err(|e| LockError::InvalidPath { source: e })?;
30        Ok(Self {
31            path,
32            lock: None,
33            system: None,
34        })
35    }
36
37    /// Shorthand intended for main functions that need a lock to guard the process
38    pub fn acquire_or_exit(&mut self) -> &mut Self {
39        match self.is_locked() {
40            Ok(false) => {
41                if let Err(e) = self.acquire() {
42                    eprintln!("error locking {}: {}", self.program_name(), e);
43                    process::exit(exitcodes::ERROR);
44                }
45                self
46            }
47            Ok(true) => {
48                eprintln!("{} is already running", self.program_name());
49                process::exit(exitcodes::LOCKED);
50            }
51            Err(err) => {
52                eprintln!("error checking lock {}: {}", self.program_name(), err);
53                process::exit(exitcodes::ERROR);
54            }
55        }
56    }
57
58    pub fn path(&self) -> &PathBuf {
59        &self.path
60    }
61
62    pub fn program_name(&self) -> &str {
63        self.path
64            .file_name()
65            .unwrap_or_else(|| std::ffi::OsStr::new(""))
66            .to_str()
67            .unwrap()
68            .split('.')
69            .next()
70            .unwrap_or("")
71    }
72
73    /// Check if the current program is locked or not.
74    pub fn is_locked(&mut self) -> Result<bool> {
75        if !self.path.exists() {
76            return Ok(false);
77        }
78
79        // path exists, examine lock PID
80        let pid = read_lock(&self.path)?;
81        self.process_is_running(pid)
82    }
83
84    /// returns the PID in the lockfile only if the process is active
85    pub fn active_pid(&mut self) -> Result<Pid> {
86        if !self.path.exists() {
87            return Err(LockError::NoLock(self.path.clone()));
88        }
89
90        // path exists, examine lock PID
91        let pid = read_lock(&self.path)?;
92        let running = self.process_is_running(pid)?;
93        if running {
94            Ok(pid)
95        } else {
96            Err(LockError::NoSuchProcess(pid, self.path.clone()))
97        }
98    }
99
100    /// Try to acquire a lock for this program.
101    pub fn acquire(&mut self) -> Result<()> {
102        match self.is_locked() {
103            Ok(false) => self.write(),
104            Ok(true) => Err(LockError::Locked(self.path.clone())),
105            Err(e) => match e {
106                LockError::CorruptLock(_) => {
107                    // overwrite corrupt locks
108                    self.write()
109                }
110                e => Err(e),
111            },
112        }
113    }
114
115    fn process_is_running(&mut self, pid: Pid) -> Result<bool> {
116        // existentialism is sometimes counterproductive
117        let this_pid = sysinfo::get_current_pid().unwrap();
118        if pid == this_pid {
119            return Ok(true);
120        }
121
122        if self.system.is_none() {
123            self.system = Some(System::new());
124        }
125
126        let system = self.system.as_mut().unwrap();
127        if !system.refresh_process(pid) {
128            return Ok(false);
129        }
130
131        match system.process(pid) {
132            Some(process) => {
133                // see https://docs.rs/sysinfo/0.26.5/sysinfo/enum.ProcessStatus.html
134                // for platform-specific details
135                Ok(matches!(process.status(), Idle | Run | Sleep | Waking))
136            }
137            None => Err(LockError::NoSuchProcess(pid, self.path.clone())),
138        }
139    }
140
141    fn write(&mut self) -> Result<()> {
142        // create lock. ensure path to lock exists
143        std::fs::create_dir_all(crate::iroh_data_root()?)?;
144        let mut file = File::create(&self.path)?;
145        let pid = sysinfo::get_current_pid().unwrap();
146        file.write_all(pid.to_string().as_bytes())?;
147        self.lock = Some(pid);
148        Ok(())
149    }
150
151    pub fn destroy_without_checking(&self) -> Result<()> {
152        std::fs::remove_file(&self.path).map_err(|e| e.into())
153    }
154}
155
156impl Drop for ProgramLock {
157    fn drop(&mut self) {
158        if self.lock.is_some() {
159            if let Err(err) = std::fs::remove_file(&self.path) {
160                warn!("removing lock: {}", err);
161            }
162        }
163    }
164}
165
166/// Report Process ID stored in a lock file
167pub fn read_lock_pid(prog_name: &str) -> Result<Pid> {
168    let path = crate::iroh_data_path(&format!("{prog_name}.lock"))?;
169    read_lock(&path)
170}
171
172fn read_lock(path: &PathBuf) -> Result<Pid> {
173    let mut file = File::open(path).map_err(|e| match e.kind() {
174        ErrorKind::NotFound => LockError::NoLock(path.clone()),
175        _ => e.into(),
176    })?;
177    let mut pid = String::new();
178    file.read_to_string(&mut pid)
179        .map_err(|_| LockError::CorruptLock(path.clone()))?;
180    let pid = pid
181        .parse::<u32>()
182        .map_err(|_| LockError::CorruptLock(path.clone()))?;
183    Ok(Pid::from_u32(pid))
184}
185
186/// Alias for a `Result` with the error type set to `LockError`
187pub type Result<T> = result::Result<T, LockError>;
188
189/// LockError is the set of known program lock errors
190#[derive(Error, Debug)]
191pub enum LockError {
192    /// lock present when one is not expected
193    #[error("Locked")]
194    Locked(PathBuf),
195    #[error("No lock file at {0}")]
196    NoLock(PathBuf),
197    /// Failure to parse contents of lock file
198    #[error("Corrupt lock file contents at {0}")]
199    CorruptLock(PathBuf),
200    #[error("could not find process with id: {0}")]
201    NoSuchProcess(Pid, PathBuf),
202    // location for lock no bueno
203    #[error("invalid path for lock file: {source}")]
204    InvalidPath {
205        #[source]
206        source: anyhow::Error,
207    },
208    #[error("operating system i/o: {source}")]
209    Io {
210        #[from]
211        source: io::Error,
212    },
213    #[error("{source}")]
214    Util {
215        #[from]
216        source: anyhow::Error,
217    },
218}
219
220#[cfg(all(test, unix))]
221mod test {
222    use super::*;
223
224    fn create_test_lock(name: &str) -> ProgramLock {
225        ProgramLock {
226            path: PathBuf::from(name),
227            lock: None,
228            system: None,
229        }
230    }
231
232    #[test]
233    fn test_corrupt_lock() {
234        let path = PathBuf::from("lock.lock");
235        let mut f = File::create(&path).unwrap();
236        write!(f, "oh noes, not a lock file").unwrap();
237        let e = read_lock(&path).err().unwrap();
238        match e {
239            LockError::CorruptLock(_) => (),
240            _e => {
241                panic!("expected CorruptLock")
242            }
243        }
244    }
245
246    #[test]
247    fn test_locks() {
248        use nix::unistd::{fork, ForkResult::*};
249        use std::io::{Read, Write};
250        use std::time::Duration;
251
252        // Start with no lock file.
253        // TODO (b5) - use tempfiles for these tests
254        let _ = std::fs::remove_file("test1.lock");
255
256        let mut lock = create_test_lock("test1.lock");
257        assert!(!lock.is_locked().unwrap());
258        assert!(read_lock(&PathBuf::from("test1.lock")).is_err());
259
260        lock.acquire().unwrap();
261
262        assert!(lock.is_locked().unwrap());
263        // ensure call to is_locked doesn't affect PID reporting
264        assert_eq!(
265            sysinfo::get_current_pid().unwrap(),
266            read_lock(&PathBuf::from("test1.lock")).unwrap()
267        );
268
269        // Spawn a child process to check we can't get the same lock.
270        // assert!() failures in the child are not reported by the test
271        // harness, so we write the result in a file from the child and
272        // read them back in the parent after a reasonable delay :(
273        unsafe {
274            match fork() {
275                Ok(Parent { child: _ }) => {
276                    let _ = std::fs::remove_file("lock_test.result");
277
278                    std::thread::sleep(Duration::from_secs(1));
279
280                    let mut result = std::fs::File::open("lock_test.result").unwrap();
281                    let mut buf = String::new();
282                    let _ = result.read_to_string(&mut buf);
283                    assert_eq!(
284                        buf,
285                        format!(
286                            "locked1=true, locked2=false lock1pid={}",
287                            sysinfo::get_current_pid().unwrap()
288                        )
289                    );
290
291                    let _ = std::fs::remove_file("lock_test.result");
292                }
293                Ok(Child) => {
294                    let mut lock = create_test_lock("test1.lock");
295                    let mut lock2 = create_test_lock("test2.lock");
296                    let pid = read_lock(&PathBuf::from("test1.lock")).unwrap();
297                    {
298                        let mut result = std::fs::File::create("lock_test.result").unwrap();
299                        let _ = result.write_all(
300                            format!(
301                                "locked1={}, locked2={} lock1pid={}",
302                                lock.is_locked().unwrap(),
303                                lock2.is_locked().unwrap(),
304                                pid,
305                            )
306                            .as_bytes(),
307                        );
308                    }
309                }
310                Err(err) => panic!("Failed to fork: {err}"),
311            }
312        }
313    }
314}
iroh_util/lock.rs

iroh_util/
lock.rs
