kumo_jsonl/
checkpoint.rs

1use camino::{Utf8Path, Utf8PathBuf};
2use serde::{Deserialize, Serialize};
3use std::io::Write;
4use std::time::{Duration, SystemTime};
5
6/// Prefix reserved for checkpoint files whose atomic writes are still in
7/// progress.
8pub const CHECKPOINT_TEMP_PREFIX: &str = ".tmp.checkpoint.";
9
10/// Returns whether `name` would produce a checkpoint file that collides with
11/// the reserved temporary-write namespace. Such a name cannot be used because a
12/// directory scan would treat its checkpoint as an in-progress write and ignore
13/// it.
14pub fn is_reserved_checkpoint_name(name: &str) -> bool {
15    // The 1.. here is because the constant starts with `.` but the
16    // `name` parameter here doesn't have a `.` prefix
17    name.starts_with(&CHECKPOINT_TEMP_PREFIX[1..])
18}
19
20/// Age past which a leftover checkpoint temp file is treated as abandoned. A
21/// committed write renames its temp file into place within microseconds of
22/// creating it. A temp file older than this was never renamed, which means the
23/// process that created it exited before finishing the write.
24pub const CHECKPOINT_TEMP_MAX_AGE: Duration = Duration::from_secs(300);
25
26/// Remove checkpoint temp files in `directory` left behind by a process that
27/// crashed between writing and renaming a checkpoint. Only entries older than
28/// `max_age` are removed, leaving a temp file currently being written by
29/// another consumer process untouched. Best-effort: a missing directory or a
30/// per-entry failure is skipped rather than reported.
31pub fn sweep_orphaned_temp_files(directory: &Utf8Path, max_age: Duration) {
32    let Ok(read_dir) = std::fs::read_dir(directory.as_std_path()) else {
33        return;
34    };
35    let now = SystemTime::now();
36    for entry in read_dir.flatten() {
37        let name = entry.file_name();
38        let Some(name) = name.to_str() else {
39            continue;
40        };
41        if !name.starts_with(CHECKPOINT_TEMP_PREFIX) {
42            continue;
43        }
44        let is_old = entry
45            .metadata()
46            .and_then(|m| m.modified())
47            .ok()
48            .and_then(|modified| now.duration_since(modified).ok())
49            .is_some_and(|age| age >= max_age);
50        if is_old {
51            match std::fs::remove_file(entry.path()) {
52                Ok(()) => {}
53                // We treat not found as success because the goal of the sweep
54                // is an absent file.
55                Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
56                Err(err) => tracing::warn!(
57                    "removing orphaned checkpoint temp file {:?}: {err:#}",
58                    entry.path()
59                ),
60            }
61        }
62    }
63}
64
65/// Persisted checkpoint data recording the current file and line position.
66#[derive(Deserialize, Serialize, Debug, Clone, PartialEq)]
67pub struct CheckpointData {
68    pub file: String,
69    pub line: usize,
70}
71
72impl CheckpointData {
73    /// Load a checkpoint from the given path.
74    /// Returns `Ok(None)` if the file does not exist.
75    pub async fn load(path: &Utf8PathBuf) -> anyhow::Result<Option<Self>> {
76        match tokio::fs::read(path.as_std_path()).await {
77            Ok(bytes) => {
78                let data: Self = serde_json::from_slice(&bytes)?;
79                Ok(Some(data))
80            }
81            Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(None),
82            Err(err) => Err(err.into()),
83        }
84    }
85
86    /// Atomically write a checkpoint file by writing to a temporary
87    /// file in the same directory and then renaming it into place.
88    pub fn save_atomic(
89        checkpoint_path: &Utf8PathBuf,
90        file: &Utf8PathBuf,
91        line: usize,
92    ) -> anyhow::Result<()> {
93        let data = Self {
94            file: file.to_string(),
95            line,
96        };
97        let json = serde_json::to_string(&data)?;
98        let dir = checkpoint_path
99            .parent()
100            .unwrap_or_else(|| Utf8Path::new("."));
101        let mut tmp = tempfile::Builder::new()
102            .prefix(CHECKPOINT_TEMP_PREFIX)
103            .tempfile_in(dir.as_std_path())?;
104        tmp.write_all(json.as_bytes())?;
105        tmp.persist(checkpoint_path.as_std_path())?;
106        Ok(())
107    }
108}
kumo_jsonl/checkpoint.rs

kumo_jsonl/
checkpoint.rs
