a3s_flow/worker/
local_file.rs

1use async_trait::async_trait;
2use chrono::{DateTime, Utc};
3use serde::Serialize;
4use std::path::{Path, PathBuf};
5use std::sync::Arc;
6use tokio::fs::File;
7use tokio::io::AsyncWriteExt;
8use tokio::sync::Mutex;
9use uuid::Uuid;
10
11use crate::error::{FlowError, Result};
12
13pub use super::task::LocalFileDeadLetteredTask;
14use super::{FlowTask, FlowTaskLease, FlowTaskQueue};
15
16/// JSON-backed local durable task queue.
17///
18/// Tasks are stored as one JSON file per pending item under `<root>/pending`.
19/// The queue serializes access inside the current process. It is intended for
20/// embedded hosts and local crash/restart durability of pending tasks; it does
21/// not provide cross-process locking. Heartbeats atomically rename inflight
22/// files, making the replacement file name a new fencing token and lease-age
23/// timestamp.
24#[derive(Debug, Clone)]
25pub struct LocalFileFlowTaskQueue {
26    root: PathBuf,
27    lock: Arc<Mutex<()>>,
28}
29
30impl LocalFileFlowTaskQueue {
31    pub fn new(root: impl Into<PathBuf>) -> Self {
32        Self {
33            root: root.into(),
34            lock: Arc::new(Mutex::new(())),
35        }
36    }
37
38    pub fn root(&self) -> &Path {
39        &self.root
40    }
41
42    fn pending_dir(&self) -> PathBuf {
43        self.root.join("pending")
44    }
45
46    fn inflight_dir(&self) -> PathBuf {
47        self.root.join("inflight")
48    }
49
50    fn dead_letter_dir(&self) -> PathBuf {
51        self.root.join("dead")
52    }
53
54    fn temp_path(&self, id: Uuid) -> PathBuf {
55        self.root.join(format!(".{id}.tmp"))
56    }
57
58    fn queue_file_name(now: DateTime<Utc>, id: Uuid) -> String {
59        let timestamp = now
60            .timestamp_nanos_opt()
61            .unwrap_or_else(|| now.timestamp_micros() * 1_000);
62        format!("{timestamp:020}-{id}.json")
63    }
64
65    fn file_timestamp_nanos(path: &Path) -> Option<i64> {
66        path.file_name()
67            .and_then(|name| name.to_str())
68            .and_then(|name| name.split_once('-'))
69            .and_then(|(timestamp, _)| timestamp.parse::<i64>().ok())
70    }
71
72    fn pending_path(&self, name: &str) -> PathBuf {
73        self.pending_dir().join(name)
74    }
75
76    fn inflight_path(&self, lease_id: &str) -> Result<PathBuf> {
77        if !Self::is_canonical_lease_id(lease_id) {
78            return Err(FlowError::LeaseLost(lease_id.to_string()));
79        }
80        Ok(self.inflight_dir().join(lease_id))
81    }
82
83    fn is_canonical_lease_id(lease_id: &str) -> bool {
84        let Some((timestamp, uuid_with_extension)) = lease_id.split_once('-') else {
85            return false;
86        };
87        if timestamp.len() != 20 || !timestamp.bytes().all(|byte| byte.is_ascii_digit()) {
88            return false;
89        }
90        let Ok(timestamp_value) = timestamp.parse::<i64>() else {
91            return false;
92        };
93        if format!("{timestamp_value:020}") != timestamp {
94            return false;
95        }
96
97        let Some(uuid_text) = uuid_with_extension.strip_suffix(".json") else {
98            return false;
99        };
100        let Ok(uuid) = Uuid::parse_str(uuid_text) else {
101            return false;
102        };
103        uuid.get_version_num() == 4
104            && uuid.get_variant() == uuid::Variant::RFC4122
105            && uuid.hyphenated().to_string() == uuid_text
106    }
107
108    fn dead_letter_path(&self, name: &str) -> PathBuf {
109        self.dead_letter_dir().join(name)
110    }
111
112    async fn json_files(dir: PathBuf) -> Result<Vec<PathBuf>> {
113        let mut files = Vec::new();
114        let mut dir = match tokio::fs::read_dir(dir).await {
115            Ok(dir) => dir,
116            Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(files),
117            Err(err) => return Err(FlowError::Io(err)),
118        };
119
120        while let Some(entry) = dir.next_entry().await? {
121            let path = entry.path();
122            if path.extension().and_then(|ext| ext.to_str()) == Some("json") {
123                files.push(path);
124            }
125        }
126        files.sort();
127        Ok(files)
128    }
129
130    async fn pending_files(&self) -> Result<Vec<PathBuf>> {
131        Self::json_files(self.pending_dir()).await
132    }
133
134    async fn inflight_files(&self) -> Result<Vec<PathBuf>> {
135        Self::json_files(self.inflight_dir()).await
136    }
137
138    async fn dead_letter_files(&self) -> Result<Vec<PathBuf>> {
139        Self::json_files(self.dead_letter_dir()).await
140    }
141
142    async fn read_task_file(path: &Path) -> Result<FlowTask> {
143        let bytes = tokio::fs::read(path).await?;
144        serde_json::from_slice(&bytes).map_err(|err| {
145            FlowError::Store(format!(
146                "failed to decode queued task from {}: {err}",
147                path.display()
148            ))
149        })
150    }
151
152    async fn write_json_file<T: Serialize>(&self, path: &Path, value: &T) -> Result<()> {
153        let id = Uuid::new_v4();
154        let temp_path = self.temp_path(id);
155
156        let mut file = File::create(&temp_path).await?;
157        file.write_all(serde_json::to_string(value)?.as_bytes())
158            .await?;
159        file.write_all(b"\n").await?;
160        file.flush().await?;
161        file.sync_data().await?;
162        drop(file);
163
164        tokio::fs::rename(temp_path, path).await?;
165        Ok(())
166    }
167
168    async fn requeue_inflight_paths(&self, paths: Vec<PathBuf>) -> Result<usize> {
169        tokio::fs::create_dir_all(self.pending_dir()).await?;
170        let mut count = 0usize;
171        for path in paths {
172            let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else {
173                continue;
174            };
175            tokio::fs::rename(&path, self.pending_path(file_name)).await?;
176            count += 1;
177        }
178        Ok(count)
179    }
180
181    async fn expired_inflight_paths(&self, cutoff: DateTime<Utc>) -> Result<Vec<PathBuf>> {
182        let cutoff = cutoff
183            .timestamp_nanos_opt()
184            .unwrap_or_else(|| cutoff.timestamp_micros() * 1_000);
185        Ok(self
186            .inflight_files()
187            .await?
188            .into_iter()
189            .filter(|path| {
190                Self::file_timestamp_nanos(path).is_some_and(|leased_at| leased_at <= cutoff)
191            })
192            .collect())
193    }
194
195    pub async fn inflight_len(&self) -> Result<usize> {
196        let _guard = self.lock.lock().await;
197        Ok(self.inflight_files().await?.len())
198    }
199
200    pub async fn dead_letter_len(&self) -> Result<usize> {
201        let _guard = self.lock.lock().await;
202        Ok(self.dead_letter_files().await?.len())
203    }
204
205    pub async fn dead_lettered_tasks(&self) -> Result<Vec<LocalFileDeadLetteredTask>> {
206        let _guard = self.lock.lock().await;
207        let mut records = Vec::new();
208        for path in self.dead_letter_files().await? {
209            let bytes = tokio::fs::read(&path).await?;
210            let record = serde_json::from_slice(&bytes).map_err(|err| {
211                FlowError::Store(format!(
212                    "failed to decode dead-lettered task from {}: {err}",
213                    path.display()
214                ))
215            })?;
216            records.push(record);
217        }
218        Ok(records)
219    }
220
221    pub async fn requeue_inflight_older_than(&self, cutoff: DateTime<Utc>) -> Result<usize> {
222        let _guard = self.lock.lock().await;
223        let expired = self.expired_inflight_paths(cutoff).await?;
224        self.requeue_inflight_paths(expired).await
225    }
226
227    pub async fn dead_letter_inflight_older_than(
228        &self,
229        cutoff: DateTime<Utc>,
230        reason: impl Into<String>,
231    ) -> Result<usize> {
232        let _guard = self.lock.lock().await;
233        tokio::fs::create_dir_all(self.dead_letter_dir()).await?;
234        let reason = reason.into();
235        let mut count = 0usize;
236        for path in self.expired_inflight_paths(cutoff).await? {
237            let Some(lease_id) = path
238                .file_name()
239                .and_then(|name| name.to_str())
240                .map(str::to_string)
241            else {
242                continue;
243            };
244            let task = Self::read_task_file(&path).await?;
245            let record = LocalFileDeadLetteredTask {
246                lease_id,
247                task,
248                reason: reason.clone(),
249                dead_lettered_at: Utc::now(),
250            };
251            let dead_path = self.dead_letter_path(&Self::queue_file_name(
252                record.dead_lettered_at,
253                Uuid::new_v4(),
254            ));
255            self.write_json_file(&dead_path, &record).await?;
256            tokio::fs::remove_file(&path).await?;
257            count += 1;
258        }
259        Ok(count)
260    }
261}
262
263#[async_trait]
264impl FlowTaskQueue for LocalFileFlowTaskQueue {
265    async fn enqueue(&self, task: FlowTask) -> Result<()> {
266        let _guard = self.lock.lock().await;
267        tokio::fs::create_dir_all(self.pending_dir()).await?;
268
269        let id = Uuid::new_v4();
270        let file_name = Self::queue_file_name(Utc::now(), id);
271        let temp_path = self.temp_path(id);
272        let pending_path = self.pending_path(&file_name);
273
274        let mut file = File::create(&temp_path).await?;
275        file.write_all(serde_json::to_string(&task)?.as_bytes())
276            .await?;
277        file.write_all(b"\n").await?;
278        file.flush().await?;
279        file.sync_data().await?;
280        drop(file);
281
282        tokio::fs::rename(temp_path, pending_path).await?;
283        Ok(())
284    }
285
286    async fn lease(&self) -> Result<Option<FlowTaskLease>> {
287        let _guard = self.lock.lock().await;
288        tokio::fs::create_dir_all(self.inflight_dir()).await?;
289        let Some(path) = self.pending_files().await?.into_iter().next() else {
290            return Ok(None);
291        };
292        if path.file_name().and_then(|name| name.to_str()).is_none() {
293            return Err(FlowError::Store(format!(
294                "queued task path {} does not have a valid file name",
295                path.display()
296            )));
297        };
298        let lease_id = Self::queue_file_name(Utc::now(), Uuid::new_v4());
299        let inflight_path = self.inflight_path(&lease_id)?;
300        tokio::fs::rename(&path, &inflight_path).await?;
301
302        let task = Self::read_task_file(&inflight_path).await?;
303        Ok(Some(FlowTaskLease { lease_id, task }))
304    }
305
306    async fn heartbeat(&self, lease_id: &str) -> Result<String> {
307        let _guard = self.lock.lock().await;
308        let current_path = self.inflight_path(lease_id)?;
309        let renewed_lease_id = Self::queue_file_name(Utc::now(), Uuid::new_v4());
310        let renewed_path = self.inflight_path(&renewed_lease_id)?;
311        match tokio::fs::rename(current_path, renewed_path).await {
312            Ok(()) => Ok(renewed_lease_id),
313            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
314                Err(FlowError::LeaseLost(lease_id.to_string()))
315            }
316            Err(err) => Err(FlowError::Io(err)),
317        }
318    }
319
320    async fn ack(&self, lease_id: &str) -> Result<()> {
321        let _guard = self.lock.lock().await;
322        let path = self.inflight_path(lease_id)?;
323        match tokio::fs::remove_file(&path).await {
324            Ok(()) => Ok(()),
325            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
326                Err(FlowError::LeaseLost(lease_id.to_string()))
327            }
328            Err(err) => Err(FlowError::Io(err)),
329        }
330    }
331
332    async fn requeue_inflight(&self) -> Result<usize> {
333        let _guard = self.lock.lock().await;
334        let paths = self.inflight_files().await?;
335        self.requeue_inflight_paths(paths).await
336    }
337
338    async fn len(&self) -> Result<usize> {
339        let _guard = self.lock.lock().await;
340        Ok(self.pending_files().await?.len())
341    }
342}
a3s_flow/worker/local_file.rs

a3s_flow/worker/
local_file.rs
