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.
22#[derive(Debug, Clone)]
23pub struct LocalFileFlowTaskQueue {
24    root: PathBuf,
25    lock: Arc<Mutex<()>>,
26}
27
28impl LocalFileFlowTaskQueue {
29    pub fn new(root: impl Into<PathBuf>) -> Self {
30        Self {
31            root: root.into(),
32            lock: Arc::new(Mutex::new(())),
33        }
34    }
35
36    pub fn root(&self) -> &Path {
37        &self.root
38    }
39
40    fn pending_dir(&self) -> PathBuf {
41        self.root.join("pending")
42    }
43
44    fn inflight_dir(&self) -> PathBuf {
45        self.root.join("inflight")
46    }
47
48    fn dead_letter_dir(&self) -> PathBuf {
49        self.root.join("dead")
50    }
51
52    fn temp_path(&self, id: Uuid) -> PathBuf {
53        self.root.join(format!(".{id}.tmp"))
54    }
55
56    fn queue_file_name(now: DateTime<Utc>, id: Uuid) -> String {
57        let timestamp = now
58            .timestamp_nanos_opt()
59            .unwrap_or_else(|| now.timestamp_micros() * 1_000);
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) -> PathBuf {
75        self.inflight_dir().join(lease_id)
76    }
77
78    fn dead_letter_path(&self, name: &str) -> PathBuf {
79        self.dead_letter_dir().join(name)
80    }
81
82    async fn json_files(dir: PathBuf) -> Result<Vec<PathBuf>> {
83        let mut files = Vec::new();
84        let mut dir = match tokio::fs::read_dir(dir).await {
85            Ok(dir) => dir,
86            Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(files),
87            Err(err) => return Err(FlowError::Io(err)),
88        };
89
90        while let Some(entry) = dir.next_entry().await? {
91            let path = entry.path();
92            if path.extension().and_then(|ext| ext.to_str()) == Some("json") {
93                files.push(path);
94            }
95        }
96        files.sort();
97        Ok(files)
98    }
99
100    async fn pending_files(&self) -> Result<Vec<PathBuf>> {
101        Self::json_files(self.pending_dir()).await
102    }
103
104    async fn inflight_files(&self) -> Result<Vec<PathBuf>> {
105        Self::json_files(self.inflight_dir()).await
106    }
107
108    async fn dead_letter_files(&self) -> Result<Vec<PathBuf>> {
109        Self::json_files(self.dead_letter_dir()).await
110    }
111
112    async fn read_task_file(path: &Path) -> Result<FlowTask> {
113        let bytes = tokio::fs::read(path).await?;
114        serde_json::from_slice(&bytes).map_err(|err| {
115            FlowError::Store(format!(
116                "failed to decode queued task from {}: {err}",
117                path.display()
118            ))
119        })
120    }
121
122    async fn write_json_file<T: Serialize>(&self, path: &Path, value: &T) -> Result<()> {
123        let id = Uuid::new_v4();
124        let temp_path = self.temp_path(id);
125
126        let mut file = File::create(&temp_path).await?;
127        file.write_all(serde_json::to_string(value)?.as_bytes())
128            .await?;
129        file.write_all(b"\n").await?;
130        file.flush().await?;
131        file.sync_data().await?;
132        drop(file);
133
134        tokio::fs::rename(temp_path, path).await?;
135        Ok(())
136    }
137
138    async fn requeue_inflight_paths(&self, paths: Vec<PathBuf>) -> Result<usize> {
139        tokio::fs::create_dir_all(self.pending_dir()).await?;
140        let mut count = 0usize;
141        for path in paths {
142            let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else {
143                continue;
144            };
145            tokio::fs::rename(&path, self.pending_path(file_name)).await?;
146            count += 1;
147        }
148        Ok(count)
149    }
150
151    async fn expired_inflight_paths(&self, cutoff: DateTime<Utc>) -> Result<Vec<PathBuf>> {
152        let cutoff = cutoff
153            .timestamp_nanos_opt()
154            .unwrap_or_else(|| cutoff.timestamp_micros() * 1_000);
155        Ok(self
156            .inflight_files()
157            .await?
158            .into_iter()
159            .filter(|path| {
160                Self::file_timestamp_nanos(path).is_some_and(|leased_at| leased_at <= cutoff)
161            })
162            .collect())
163    }
164
165    pub async fn inflight_len(&self) -> Result<usize> {
166        let _guard = self.lock.lock().await;
167        Ok(self.inflight_files().await?.len())
168    }
169
170    pub async fn dead_letter_len(&self) -> Result<usize> {
171        let _guard = self.lock.lock().await;
172        Ok(self.dead_letter_files().await?.len())
173    }
174
175    pub async fn dead_lettered_tasks(&self) -> Result<Vec<LocalFileDeadLetteredTask>> {
176        let _guard = self.lock.lock().await;
177        let mut records = Vec::new();
178        for path in self.dead_letter_files().await? {
179            let bytes = tokio::fs::read(&path).await?;
180            let record = serde_json::from_slice(&bytes).map_err(|err| {
181                FlowError::Store(format!(
182                    "failed to decode dead-lettered task from {}: {err}",
183                    path.display()
184                ))
185            })?;
186            records.push(record);
187        }
188        Ok(records)
189    }
190
191    pub async fn requeue_inflight_older_than(&self, cutoff: DateTime<Utc>) -> Result<usize> {
192        let _guard = self.lock.lock().await;
193        let expired = self.expired_inflight_paths(cutoff).await?;
194        self.requeue_inflight_paths(expired).await
195    }
196
197    pub async fn dead_letter_inflight_older_than(
198        &self,
199        cutoff: DateTime<Utc>,
200        reason: impl Into<String>,
201    ) -> Result<usize> {
202        let _guard = self.lock.lock().await;
203        tokio::fs::create_dir_all(self.dead_letter_dir()).await?;
204        let reason = reason.into();
205        let mut count = 0usize;
206        for path in self.expired_inflight_paths(cutoff).await? {
207            let Some(lease_id) = path
208                .file_name()
209                .and_then(|name| name.to_str())
210                .map(str::to_string)
211            else {
212                continue;
213            };
214            let task = Self::read_task_file(&path).await?;
215            let record = LocalFileDeadLetteredTask {
216                lease_id,
217                task,
218                reason: reason.clone(),
219                dead_lettered_at: Utc::now(),
220            };
221            let dead_path = self.dead_letter_path(&Self::queue_file_name(
222                record.dead_lettered_at,
223                Uuid::new_v4(),
224            ));
225            self.write_json_file(&dead_path, &record).await?;
226            tokio::fs::remove_file(&path).await?;
227            count += 1;
228        }
229        Ok(count)
230    }
231}
232
233#[async_trait]
234impl FlowTaskQueue for LocalFileFlowTaskQueue {
235    async fn enqueue(&self, task: FlowTask) -> Result<()> {
236        let _guard = self.lock.lock().await;
237        tokio::fs::create_dir_all(self.pending_dir()).await?;
238
239        let id = Uuid::new_v4();
240        let file_name = Self::queue_file_name(Utc::now(), id);
241        let temp_path = self.temp_path(id);
242        let pending_path = self.pending_path(&file_name);
243
244        let mut file = File::create(&temp_path).await?;
245        file.write_all(serde_json::to_string(&task)?.as_bytes())
246            .await?;
247        file.write_all(b"\n").await?;
248        file.flush().await?;
249        file.sync_data().await?;
250        drop(file);
251
252        tokio::fs::rename(temp_path, pending_path).await?;
253        Ok(())
254    }
255
256    async fn lease(&self) -> Result<Option<FlowTaskLease>> {
257        let _guard = self.lock.lock().await;
258        tokio::fs::create_dir_all(self.inflight_dir()).await?;
259        let Some(path) = self.pending_files().await?.into_iter().next() else {
260            return Ok(None);
261        };
262        if path.file_name().and_then(|name| name.to_str()).is_none() {
263            return Err(FlowError::Store(format!(
264                "queued task path {} does not have a valid file name",
265                path.display()
266            )));
267        };
268        let lease_id = Self::queue_file_name(Utc::now(), Uuid::new_v4());
269        let inflight_path = self.inflight_path(&lease_id);
270        tokio::fs::rename(&path, &inflight_path).await?;
271
272        let task = Self::read_task_file(&inflight_path).await?;
273        Ok(Some(FlowTaskLease { lease_id, task }))
274    }
275
276    async fn ack(&self, lease_id: &str) -> Result<()> {
277        let _guard = self.lock.lock().await;
278        let path = self.inflight_path(lease_id);
279        match tokio::fs::remove_file(&path).await {
280            Ok(()) => Ok(()),
281            Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
282            Err(err) => Err(FlowError::Io(err)),
283        }
284    }
285
286    async fn requeue_inflight(&self) -> Result<usize> {
287        let _guard = self.lock.lock().await;
288        let paths = self.inflight_files().await?;
289        self.requeue_inflight_paths(paths).await
290    }
291
292    async fn len(&self) -> Result<usize> {
293        let _guard = self.lock.lock().await;
294        Ok(self.pending_files().await?.len())
295    }
296}
a3s_flow/worker/local_file.rs

a3s_flow/worker/
local_file.rs
