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

a3s_flow/worker/
local_file.rs
