Skip to main content

foxtive_supervisor/persistence/
fs.rs

1use super::{PersistedTaskState, TaskStateStore};
2use crate::error::SupervisorResult;
3use std::path::{Path, PathBuf};
4use tokio::fs;
5use tokio::io::{AsyncReadExt, AsyncWriteExt};
6
7/// A file system-based implementation of TaskStateStore.
8/// Each task's state is stored as a JSON file in a specified directory.
9pub struct FsStateStore {
10    base_path: PathBuf,
11}
12
13impl FsStateStore {
14    pub async fn new(base_path: impl AsRef<Path>) -> SupervisorResult<Self> {
15        let path = base_path.as_ref().to_path_buf();
16        fs::create_dir_all(&path).await?;
17        Ok(Self { base_path: path })
18    }
19
20    fn state_file_path(&self, task_id: &str) -> PathBuf {
21        self.base_path.join(format!("{}.json", task_id))
22    }
23}
24
25#[async_trait::async_trait]
26impl TaskStateStore for FsStateStore {
27    async fn save_state(&self, state: PersistedTaskState) -> SupervisorResult<()> {
28        let file_path = self.state_file_path(&state.task_id);
29        let json = serde_json::to_string(&state)?;  // Compact JSON for production
30        let mut file = fs::File::create(&file_path).await?;
31        file.write_all(json.as_bytes()).await?;
32        file.sync_data().await?; // Sync data only (faster than sync_all)
33        Ok(())
34    }
35
36    async fn load_state(&self, task_id: &str) -> SupervisorResult<Option<PersistedTaskState>> {
37        let file_path = self.state_file_path(task_id);
38        
39        // Open directly and handle NotFound to avoid TOCTOU race
40        let mut file = match fs::File::open(&file_path).await {
41            Ok(file) => file,
42            Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
43            Err(e) => return Err(e.into()),
44        };
45
46        let mut json = String::new();
47        file.read_to_string(&mut json).await?;
48
49        if json.is_empty() {
50            return Ok(None);
51        }
52
53        let state: PersistedTaskState = serde_json::from_str(&json)?;
54        Ok(Some(state))
55    }
56
57    async fn load_all_states(&self) -> SupervisorResult<Vec<PersistedTaskState>> {
58        let mut states = Vec::new();
59        let mut entries = fs::read_dir(&self.base_path).await?;
60
61        while let Some(entry) = entries.next_entry().await? {
62            let path = entry.path();
63            if path.is_file() && path.extension().is_some_and(|ext| ext == "json") {
64                let mut file = fs::File::open(&path).await?;
65                let mut json = String::new();
66                file.read_to_string(&mut json).await?;
67
68                if json.is_empty() {
69                    continue;
70                }
71
72                if let Ok(state) = serde_json::from_str::<PersistedTaskState>(&json) {
73                    states.push(state);
74                } else {
75                    // Log error for malformed state files, but don't fail the whole operation
76                    eprintln!("Warning: Could not parse state file: {:?}", path);
77                }
78            }
79        }
80        Ok(states)
81    }
82
83    async fn delete_state(&self, task_id: &str) -> SupervisorResult<()> {
84        let file_path = self.state_file_path(task_id);
85        
86        // Remove directly and handle NotFound to avoid TOCTOU race
87        match fs::remove_file(&file_path).await {
88            Ok(_) => Ok(()),
89            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
90            Err(e) => Err(e.into()),
91        }
92    }
93}