use std::fs;
use std::io::{Read, Write};
use std::path::Path;
use std::path::PathBuf;
use std::time::{SystemTime, UNIX_EPOCH};
use serde::{Deserialize, Serialize};
use thiserror::Error;
#[derive(Debug, Error)]
pub enum DreamLockError {
#[error("Lock file exists and process {pid} is still alive")]
AlreadyLocked { pid: u32 },
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("JSON error: {0}")]
Json(#[from] serde_json::Error),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LockContent {
pub pid: u32,
pub timestamp: i64,
#[serde(alias = "session_count")]
pub memory_count: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DreamState {
pub last_run_at: i64,
#[serde(alias = "session_count")]
pub memory_count: u32,
}
pub struct DreamLock {
lock_path: PathBuf,
state_path: PathBuf,
}
impl DreamLock {
pub fn new(data_dir: &Path) -> Self {
Self {
lock_path: data_dir.join("dream.lock"),
state_path: data_dir.join("dream.state"),
}
}
pub fn try_acquire(&self, memory_count: u32) -> Result<bool, DreamLockError> {
if self.lock_path.exists() {
let content = self.read_lock()?;
if self.is_process_alive(content.pid) {
return Ok(false);
}
self.cleanup_stale_lock()?;
}
self.write_lock(memory_count)?;
Ok(true)
}
pub fn release(&self) -> Result<(), DreamLockError> {
if self.lock_path.exists() {
fs::remove_file(&self.lock_path)?;
}
Ok(())
}
fn read_lock(&self) -> Result<LockContent, DreamLockError> {
let mut file = fs::File::open(&self.lock_path)?;
let mut content = String::new();
file.read_to_string(&mut content)?;
let lock: LockContent = serde_json::from_str(&content)?;
Ok(lock)
}
fn write_lock(&self, memory_count: u32) -> Result<(), DreamLockError> {
let content = LockContent {
pid: std::process::id(),
timestamp: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs() as i64,
memory_count,
};
let json = serde_json::to_string(&content)?;
let mut file = fs::File::create(&self.lock_path)?;
file.write_all(json.as_bytes())?;
Ok(())
}
fn is_process_alive(&self, pid: u32) -> bool {
#[cfg(target_os = "linux")]
{
std::path::Path::new(&format!("/proc/{}", pid)).exists()
}
#[cfg(not(target_os = "linux"))]
{
true
}
}
fn cleanup_stale_lock(&self) -> Result<(), DreamLockError> {
fs::remove_file(&self.lock_path)?;
Ok(())
}
pub fn get_lock_content(&self) -> Option<LockContent> {
if self.lock_path.exists() {
self.read_lock().ok()
} else {
None
}
}
pub fn read_state(&self) -> Option<DreamState> {
if !self.state_path.exists() {
return None;
}
let mut file = fs::File::open(&self.state_path).ok()?;
let mut content = String::new();
file.read_to_string(&mut content).ok()?;
serde_json::from_str(&content).ok()
}
pub fn write_state(&self, memory_count: u32) -> Result<(), DreamLockError> {
let state = DreamState {
last_run_at: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs() as i64,
memory_count,
};
let json = serde_json::to_string(&state)?;
let mut file = fs::File::create(&self.state_path)?;
file.write_all(json.as_bytes())?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_lock_content_serde() {
let content = LockContent {
pid: 12345,
timestamp: 1700000000,
memory_count: 5,
};
let json = serde_json::to_string(&content).unwrap();
let parsed: LockContent = serde_json::from_str(&json).unwrap();
assert_eq!(content.pid, parsed.pid);
assert_eq!(content.timestamp, parsed.timestamp);
assert_eq!(content.memory_count, parsed.memory_count);
}
#[test]
fn test_state_serde_roundtrip() {
let state = DreamState {
last_run_at: 1700000000,
memory_count: 5,
};
let json = serde_json::to_string(&state).unwrap();
let parsed: DreamState = serde_json::from_str(&json).unwrap();
assert_eq!(state.last_run_at, parsed.last_run_at);
assert_eq!(state.memory_count, parsed.memory_count);
}
#[test]
fn test_lock_acquire_and_release() {
let dir = tempfile::tempdir().unwrap();
let lock = DreamLock::new(dir.path());
let result = lock.try_acquire(10);
assert!(result.is_ok());
assert!(result.unwrap());
let content = lock.get_lock_content();
assert!(content.is_some());
let content = content.unwrap();
assert_eq!(content.memory_count, 10);
assert!(content.pid > 0);
lock.release().unwrap();
assert!(lock.get_lock_content().is_none());
}
#[test]
fn test_lock_double_acquire_fails() {
let dir = tempfile::tempdir().unwrap();
let lock1 = DreamLock::new(dir.path());
let lock2 = DreamLock::new(dir.path());
lock1.try_acquire(5).unwrap();
let result = lock2.try_acquire(10);
assert!(result.is_ok());
assert!(!result.unwrap());
}
#[test]
fn test_state_persists_after_release() {
let dir = tempfile::tempdir().unwrap();
let lock = DreamLock::new(dir.path());
assert!(lock.read_state().is_none());
lock.try_acquire(7).unwrap();
lock.write_state(7).unwrap();
lock.release().unwrap();
let state = lock.read_state();
assert!(state.is_some());
let state = state.unwrap();
assert_eq!(state.memory_count, 7);
assert!(state.last_run_at > 0);
}
}