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::{timestamp_nanos_saturating, 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 = timestamp_nanos_saturating(now);
60        format!("{timestamp:020}-{id}.json")
61    }
62
63    fn file_timestamp_nanos(path: &Path) -> Option<i64> {
64        path.file_name()
65            .and_then(|name| name.to_str())
66            .and_then(|name| name.split_once('-'))
67            .and_then(|(timestamp, _)| timestamp.parse::<i64>().ok())
68    }
69
70    fn pending_path(&self, name: &str) -> PathBuf {
71        self.pending_dir().join(name)
72    }
73
74    fn inflight_path(&self, lease_id: &str) -> Result<PathBuf> {
75        if !Self::is_canonical_lease_id(lease_id) {
76            return Err(FlowError::LeaseLost(lease_id.to_string()));
77        }
78        Ok(self.inflight_dir().join(lease_id))
79    }
80
81    fn is_canonical_lease_id(lease_id: &str) -> bool {
82        let Some((timestamp, uuid_with_extension)) = lease_id.split_once('-') else {
83            return false;
84        };
85        if timestamp.len() != 20 || !timestamp.bytes().all(|byte| byte.is_ascii_digit()) {
86            return false;
87        }
88        let Ok(timestamp_value) = timestamp.parse::<i64>() else {
89            return false;
90        };
91        if format!("{timestamp_value:020}") != timestamp {
92            return false;
93        }
94
95        let Some(uuid_text) = uuid_with_extension.strip_suffix(".json") else {
96            return false;
97        };
98        let Ok(uuid) = Uuid::parse_str(uuid_text) else {
99            return false;
100        };
101        uuid.get_version_num() == 4
102            && uuid.get_variant() == uuid::Variant::RFC4122
103            && uuid.hyphenated().to_string() == uuid_text
104    }
105
106    fn dead_letter_path(&self, name: &str) -> PathBuf {
107        self.dead_letter_dir().join(name)
108    }
109
110    async fn json_files(dir: PathBuf) -> Result<Vec<PathBuf>> {
111        let mut files = Vec::new();
112        let mut dir = match tokio::fs::read_dir(dir).await {
113            Ok(dir) => dir,
114            Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(files),
115            Err(err) => return Err(FlowError::Io(err)),
116        };
117
118        while let Some(entry) = dir.next_entry().await? {
119            let path = entry.path();
120            if path.extension().and_then(|ext| ext.to_str()) == Some("json") {
121                files.push(path);
122            }
123        }
124        files.sort();
125        Ok(files)
126    }
127
128    async fn pending_files(&self) -> Result<Vec<PathBuf>> {
129        Self::json_files(self.pending_dir()).await
130    }
131
132    async fn inflight_files(&self) -> Result<Vec<PathBuf>> {
133        Self::json_files(self.inflight_dir()).await
134    }
135
136    async fn dead_letter_files(&self) -> Result<Vec<PathBuf>> {
137        Self::json_files(self.dead_letter_dir()).await
138    }
139
140    async fn read_task_file(path: &Path) -> Result<FlowTask> {
141        let bytes = tokio::fs::read(path).await?;
142        serde_json::from_slice(&bytes).map_err(|err| {
143            FlowError::Store(format!(
144                "failed to decode queued task from {}: {err}",
145                path.display()
146            ))
147        })
148    }
149
150    async fn write_json_file<T: Serialize>(&self, path: &Path, value: &T) -> Result<()> {
151        let id = Uuid::new_v4();
152        let temp_path = self.temp_path(id);
153
154        let mut file = File::create(&temp_path).await?;
155        file.write_all(serde_json::to_string(value)?.as_bytes())
156            .await?;
157        file.write_all(b"\n").await?;
158        file.flush().await?;
159        file.sync_data().await?;
160        drop(file);
161
162        tokio::fs::rename(temp_path, path).await?;
163        Ok(())
164    }
165
166    async fn requeue_inflight_paths(&self, paths: Vec<PathBuf>) -> Result<usize> {
167        tokio::fs::create_dir_all(self.pending_dir()).await?;
168        let mut count = 0usize;
169        for path in paths {
170            let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else {
171                continue;
172            };
173            tokio::fs::rename(&path, self.pending_path(file_name)).await?;
174            count += 1;
175        }
176        Ok(count)
177    }
178
179    async fn expired_inflight_paths(&self, cutoff: DateTime<Utc>) -> Result<Vec<PathBuf>> {
180        let cutoff = timestamp_nanos_saturating(cutoff);
181        Ok(self
182            .inflight_files()
183            .await?
184            .into_iter()
185            .filter(|path| {
186                Self::file_timestamp_nanos(path).is_some_and(|leased_at| leased_at <= cutoff)
187            })
188            .collect())
189    }
190
191    pub async fn inflight_len(&self) -> Result<usize> {
192        let _guard = self.lock.lock().await;
193        Ok(self.inflight_files().await?.len())
194    }
195
196    pub async fn dead_letter_len(&self) -> Result<usize> {
197        let _guard = self.lock.lock().await;
198        Ok(self.dead_letter_files().await?.len())
199    }
200
201    pub async fn dead_lettered_tasks(&self) -> Result<Vec<LocalFileDeadLetteredTask>> {
202        let _guard = self.lock.lock().await;
203        let mut records = Vec::new();
204        for path in self.dead_letter_files().await? {
205            let bytes = tokio::fs::read(&path).await?;
206            let record = serde_json::from_slice(&bytes).map_err(|err| {
207                FlowError::Store(format!(
208                    "failed to decode dead-lettered task from {}: {err}",
209                    path.display()
210                ))
211            })?;
212            records.push(record);
213        }
214        Ok(records)
215    }
216
217    pub async fn requeue_inflight_older_than(&self, cutoff: DateTime<Utc>) -> Result<usize> {
218        let _guard = self.lock.lock().await;
219        let expired = self.expired_inflight_paths(cutoff).await?;
220        self.requeue_inflight_paths(expired).await
221    }
222
223    pub async fn dead_letter_inflight_older_than(
224        &self,
225        cutoff: DateTime<Utc>,
226        reason: impl Into<String>,
227    ) -> Result<usize> {
228        let _guard = self.lock.lock().await;
229        tokio::fs::create_dir_all(self.dead_letter_dir()).await?;
230        let reason = reason.into();
231        let mut count = 0usize;
232        for path in self.expired_inflight_paths(cutoff).await? {
233            let Some(lease_id) = path
234                .file_name()
235                .and_then(|name| name.to_str())
236                .map(str::to_string)
237            else {
238                continue;
239            };
240            let task = Self::read_task_file(&path).await?;
241            let record = LocalFileDeadLetteredTask {
242                lease_id,
243                task,
244                reason: reason.clone(),
245                dead_lettered_at: Utc::now(),
246            };
247            let dead_path = self.dead_letter_path(&Self::queue_file_name(
248                record.dead_lettered_at,
249                Uuid::new_v4(),
250            ));
251            self.write_json_file(&dead_path, &record).await?;
252            tokio::fs::remove_file(&path).await?;
253            count += 1;
254        }
255        Ok(count)
256    }
257}
258
259#[async_trait]
260impl FlowTaskQueue for LocalFileFlowTaskQueue {
261    async fn enqueue(&self, task: FlowTask) -> Result<()> {
262        let _guard = self.lock.lock().await;
263        tokio::fs::create_dir_all(self.pending_dir()).await?;
264
265        let id = Uuid::new_v4();
266        let file_name = Self::queue_file_name(Utc::now(), id);
267        let temp_path = self.temp_path(id);
268        let pending_path = self.pending_path(&file_name);
269
270        let mut file = File::create(&temp_path).await?;
271        file.write_all(serde_json::to_string(&task)?.as_bytes())
272            .await?;
273        file.write_all(b"\n").await?;
274        file.flush().await?;
275        file.sync_data().await?;
276        drop(file);
277
278        tokio::fs::rename(temp_path, pending_path).await?;
279        Ok(())
280    }
281
282    async fn lease(&self) -> Result<Option<FlowTaskLease>> {
283        let _guard = self.lock.lock().await;
284        tokio::fs::create_dir_all(self.inflight_dir()).await?;
285        let Some(path) = self.pending_files().await?.into_iter().next() else {
286            return Ok(None);
287        };
288        if path.file_name().and_then(|name| name.to_str()).is_none() {
289            return Err(FlowError::Store(format!(
290                "queued task path {} does not have a valid file name",
291                path.display()
292            )));
293        };
294        let lease_id = Self::queue_file_name(Utc::now(), Uuid::new_v4());
295        let inflight_path = self.inflight_path(&lease_id)?;
296        tokio::fs::rename(&path, &inflight_path).await?;
297
298        let task = Self::read_task_file(&inflight_path).await?;
299        Ok(Some(FlowTaskLease { lease_id, task }))
300    }
301
302    async fn heartbeat(&self, lease_id: &str) -> Result<String> {
303        let _guard = self.lock.lock().await;
304        let current_path = self.inflight_path(lease_id)?;
305        let renewed_lease_id = Self::queue_file_name(Utc::now(), Uuid::new_v4());
306        let renewed_path = self.inflight_path(&renewed_lease_id)?;
307        match tokio::fs::rename(current_path, renewed_path).await {
308            Ok(()) => Ok(renewed_lease_id),
309            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
310                Err(FlowError::LeaseLost(lease_id.to_string()))
311            }
312            Err(err) => Err(FlowError::Io(err)),
313        }
314    }
315
316    async fn ack(&self, lease_id: &str) -> Result<()> {
317        let _guard = self.lock.lock().await;
318        let path = self.inflight_path(lease_id)?;
319        match tokio::fs::remove_file(&path).await {
320            Ok(()) => Ok(()),
321            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
322                Err(FlowError::LeaseLost(lease_id.to_string()))
323            }
324            Err(err) => Err(FlowError::Io(err)),
325        }
326    }
327
328    async fn requeue_inflight(&self) -> Result<usize> {
329        let _guard = self.lock.lock().await;
330        let paths = self.inflight_files().await?;
331        self.requeue_inflight_paths(paths).await
332    }
333
334    async fn len(&self) -> Result<usize> {
335        let _guard = self.lock.lock().await;
336        Ok(self.pending_files().await?.len())
337    }
338}
a3s_flow/worker/local_file.rs

a3s_flow/worker/
local_file.rs
