robit_agent/
lock.rs

1//! 工作目录独占文件锁
2//!
3//! 确保同一工作目录下只能有一个 robit 程序运行。
4
5use std::fs::{File, OpenOptions, remove_file};
6use std::io::{Write, Read, Seek, SeekFrom};
7use std::path::{Path, PathBuf};
8use std::process;
9
10use fs2::FileExt;
11use serde::{Serialize, Deserialize};
12use thiserror::Error;
13use time::OffsetDateTime;
14
15/// 存储在锁文件中的基本信息
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct LockInfo {
18    /// 程序名称（如 "robit-tui", "robit-qq", "robit-gui"）
19    pub program: String,
20    /// 进程 ID
21    pub pid: u32,
22    /// 启动用户名
23    pub username: String,
24    /// 启动时间（ISO 8601 格式）
25    pub started_at: String,
26}
27
28/// 文件锁错误
29#[derive(Debug, Error)]
30pub enum LockError {
31    /// 目录已被另一个进程锁定
32    #[error("Directory is already locked by {program} (PID {pid}, user {username}, started at {started_at})")]
33    AlreadyLocked {
34        program: String,
35        pid: u32,
36        username: String,
37        started_at: String,
38    },
39    /// IO 错误
40    #[error("IO error: {0}")]
41    Io(#[from] std::io::Error),
42    /// 序列化/反序列化错误
43    #[error("JSON error: {0}")]
44    Json(#[from] serde_json::Error),
45    /// 无法创建 .robit 目录
46    #[error("Failed to create .robit directory: {0}")]
47    CreateDir(std::io::Error),
48}
49
50/// 文件锁守护对象，RAII 模式，析构时自动释放锁
51pub struct DirectoryLock {
52    /// 锁文件路径
53    lock_path: PathBuf,
54    /// 锁文件句柄（保持打开以维持操作系统级锁）
55    file: Option<File>,
56    /// 锁定时的信息
57    info: LockInfo,
58}
59
60impl DirectoryLock {
61    /// 获取工作目录的独占锁
62    ///
63    /// # Arguments
64    /// * `workdir` - 工作目录路径
65    /// * `program_name` - 程序名称（如 "robit-tui"）
66    pub fn acquire(workdir: &Path, program_name: &str) -> Result<Self, LockError> {
67        let robit_dir = workdir.join(".robit");
68        let lock_path = robit_dir.join("LOCK");
69
70        // 确保 .robit 目录存在
71        if !robit_dir.exists() {
72            std::fs::create_dir_all(&robit_dir)
73                .map_err(LockError::CreateDir)?;
74        }
75
76        // 获取当前用户名
77        let username = get_username().unwrap_or_else(|| "unknown".to_string());
78
79        // 创建锁信息
80        let info = LockInfo {
81            program: program_name.to_string(),
82            pid: process::id(),
83            username,
84            started_at: OffsetDateTime::now_utc()
85                .format(&time::format_description::well_known::Rfc3339)
86                .unwrap_or_else(|_| "unknown".to_string()),
87        };
88
89        // 尝试获取锁，最多重试一次（处理孤立锁）
90        match Self::try_acquire(&lock_path, &info) {
91            Ok(lock) => Ok(lock),
92            Err(LockError::AlreadyLocked { program, pid, username, started_at }) => {
93                // 检查进程是否还在运行
94                if !is_process_running(pid) {
95                    // 进程不存在，删除旧锁文件并重试
96                    tracing::warn!(
97                        "Found stale lock file from {} (PID {}, user {}, started at {}), cleaning up",
98                        program,
99                        pid,
100                        username,
101                        started_at
102                    );
103                    let _ = remove_file(&lock_path);
104                    Self::try_acquire(&lock_path, &info)
105                } else {
106                    Err(LockError::AlreadyLocked { program, pid, username, started_at })
107                }
108            }
109            Err(e) => Err(e),
110        }
111    }
112
113    /// 尝试获取锁（单次尝试）
114    fn try_acquire(lock_path: &Path, info: &LockInfo) -> Result<Self, LockError> {
115        // 打开或创建锁文件
116        let mut file = OpenOptions::new()
117            .read(true)
118            .write(true)
119            .create(true)
120            .open(lock_path)?;
121
122        // 尝试获取排他锁（非阻塞）
123        match file.try_lock_exclusive() {
124            Ok(_) => {
125                // 获取锁成功，写入信息
126                file.set_len(0)?;
127                file.seek(SeekFrom::Start(0))?;
128                serde_json::to_writer_pretty(&mut file, info)?;
129                file.flush()?;
130                file.sync_data()?;
131
132                Ok(Self {
133                    lock_path: lock_path.to_path_buf(),
134                    file: Some(file),
135                    info: info.clone(),
136                })
137            }
138            Err(_) => {
139                // 获取锁失败，读取现有锁信息
140                let mut content = String::new();
141                file.seek(SeekFrom::Start(0))?;
142                file.read_to_string(&mut content)?;
143
144                let existing_info: LockInfo = match serde_json::from_str(&content) {
145                    Ok(info) => info,
146                    Err(_) => {
147                        // 锁文件内容损坏，视为可清理
148                        LockInfo {
149                            program: "unknown".to_string(),
150                            pid: 0,
151                            username: "unknown".to_string(),
152                            started_at: "unknown".to_string(),
153                        }
154                    }
155                };
156
157                Err(LockError::AlreadyLocked {
158                    program: existing_info.program,
159                    pid: existing_info.pid,
160                    username: existing_info.username,
161                    started_at: existing_info.started_at,
162                })
163            }
164        }
165    }
166
167    /// 获取当前锁的信息
168    pub fn info(&self) -> &LockInfo {
169        &self.info
170    }
171
172    /// 手动释放锁（通常不需要调用，RAII 会自动处理）
173    pub fn release(&mut self) {
174        if let Some(file) = self.file.take() {
175            let _ = file.unlock();
176        }
177    }
178}
179
180impl Drop for DirectoryLock {
181    fn drop(&mut self) {
182        if let Some(file) = self.file.take() {
183            let _ = file.unlock();
184        }
185        // 删除锁文件
186        let _ = std::fs::remove_file(&self.lock_path);
187    }
188}
189
190/// 检测进程是否还在运行（跨平台）
191fn is_process_running(pid: u32) -> bool {
192    #[cfg(unix)]
193    {
194        use std::ffi::CString;
195        use libc::{kill, c_int, EPERM, ESRCH};
196
197        unsafe {
198            let result = kill(pid as c_int, 0);
199            if result == 0 {
200                return true;
201            }
202            let errno = *libc::__errno_location();
203            errno == EPERM // 没有权限但进程存在
204        }
205    }
206
207    #[cfg(windows)]
208    {
209        use windows_sys::Win32::Foundation::{CloseHandle, ERROR_INVALID_PARAMETER, GetLastError};
210        use windows_sys::Win32::System::Threading::{OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION};
211
212        unsafe {
213            let handle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, 0, pid);
214            if handle != 0 {
215                CloseHandle(handle);
216                return true;
217            }
218            GetLastError() != ERROR_INVALID_PARAMETER
219        }
220    }
221
222    #[cfg(not(any(unix, windows)))]
223    {
224        // 其他平台保守处理：假设进程存在
225        true
226    }
227}
228
229/// 获取当前用户名（跨平台）
230fn get_username() -> Option<String> {
231    // 优先从环境变量获取
232    if let Ok(name) = std::env::var("USER") {
233        if !name.is_empty() {
234            return Some(name);
235        }
236    }
237    if let Ok(name) = std::env::var("USERNAME") {
238        if !name.is_empty() {
239            return Some(name);
240        }
241    }
242
243    // Unix 备选方案
244    #[cfg(unix)]
245    {
246        use std::ffi::CStr;
247        use libc::{getpwuid, getuid};
248
249        unsafe {
250            let uid = getuid();
251            let passwd = getpwuid(uid);
252            if !passwd.is_null() {
253                let name = CStr::from_ptr((*passwd).pw_name);
254                return name.to_str().ok().map(|s| s.to_string());
255            }
256        }
257    }
258
259    None
260}
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265    use tempfile::tempdir;
266
267    #[test]
268    fn test_lock_acquire_and_release() {
269        let dir = tempdir().unwrap();
270        let lock = DirectoryLock::acquire(dir.path(), "test-program").unwrap();
271        assert_eq!(lock.info().program, "test-program");
272        assert!(lock.info().pid > 0);
273    }
274
275    #[test]
276    fn test_username_is_set() {
277        let dir = tempdir().unwrap();
278        let lock = DirectoryLock::acquire(dir.path(), "test-program").unwrap();
279        assert!(!lock.info().username.is_empty());
280    }
281
282    #[test]
283    fn test_started_at_is_set() {
284        let dir = tempdir().unwrap();
285        let lock = DirectoryLock::acquire(dir.path(), "test-program").unwrap();
286        assert!(!lock.info().started_at.is_empty());
287    }
288
289    #[test]
290    fn test_is_process_running_with_our_pid() {
291        // 我们自己的 PID 应该是在运行的
292        assert!(is_process_running(std::process::id()));
293    }
294
295    #[test]
296    fn test_is_process_running_with_invalid_pid() {
297        // 一个很大的 PID 应该不在运行
298        // 注意：这不是 100% 可靠，但对于测试来说足够了
299        assert!(!is_process_running(u32::MAX));
300    }
301}
robit_agent/lock.rs

robit_agent/
lock.rs
