Skip to main content

moonlight_core/storage/
retention.rs

1use crate::ComparisonRun;
2use std::path::Path;
3use tokio::{
4    fs::{self, OpenOptions},
5    io::AsyncWriteExt,
6};
7use uuid::Uuid;
8
9#[derive(Debug, Clone, Copy, Default)]
10pub struct StorageOptions {
11    pub retention_max_runs: Option<usize>,
12    pub retention_max_bytes: Option<u64>,
13}
14
15impl StorageOptions {
16    pub(super) fn is_configured(self) -> bool {
17        self.retention_max_runs.is_some() || self.retention_max_bytes.is_some()
18    }
19}
20
21pub(super) fn retain_runs(
22    mut runs: Vec<ComparisonRun>,
23    options: StorageOptions,
24) -> anyhow::Result<Vec<ComparisonRun>> {
25    if let Some(max_runs) = options.retention_max_runs {
26        runs = retain_by_max_runs(runs, max_runs);
27    }
28    if let Some(max_bytes) = options.retention_max_bytes {
29        runs = retain_by_max_bytes(runs, max_bytes)?;
30    }
31    Ok(runs)
32}
33
34fn retain_by_max_runs(runs: Vec<ComparisonRun>, max_runs: usize) -> Vec<ComparisonRun> {
35    if runs.len() <= max_runs {
36        return runs;
37    }
38
39    let mut retained = runs.into_iter().rev().take(max_runs).collect::<Vec<_>>();
40    retained.reverse();
41    retained
42}
43
44fn retain_by_max_bytes(
45    runs: Vec<ComparisonRun>,
46    max_bytes: u64,
47) -> anyhow::Result<Vec<ComparisonRun>> {
48    let mut retained = Vec::new();
49    let mut total_bytes = 0_u64;
50    for run in runs.into_iter().rev() {
51        let line = serde_json::to_string(&run)?;
52        let line_bytes = line.len() as u64 + 1;
53        if total_bytes + line_bytes <= max_bytes || retained.is_empty() {
54            total_bytes += line_bytes;
55            retained.push(run);
56        } else {
57            break;
58        }
59    }
60    retained.reverse();
61    Ok(retained)
62}
63
64pub(super) fn serialize_runs_jsonl(runs: &[ComparisonRun]) -> anyhow::Result<String> {
65    let mut content = String::new();
66    for run in runs {
67        content.push_str(&serde_json::to_string(run)?);
68        content.push('\n');
69    }
70    Ok(content)
71}
72
73pub(super) async fn atomic_write(path: &Path, content: String) -> anyhow::Result<()> {
74    if let Some(parent) = path.parent() {
75        fs::create_dir_all(parent).await?;
76    }
77    let parent = path.parent().unwrap_or_else(|| Path::new("."));
78    let file_name = path
79        .file_name()
80        .and_then(|value| value.to_str())
81        .unwrap_or("moonlight-runs.jsonl");
82    let temp_path = parent.join(format!(
83        ".{file_name}.{}.{}.tmp",
84        std::process::id(),
85        Uuid::new_v4()
86    ));
87
88    let mut file = OpenOptions::new()
89        .create_new(true)
90        .write(true)
91        .open(&temp_path)
92        .await?;
93    file.write_all(content.as_bytes()).await?;
94    file.flush().await?;
95    file.sync_data().await?;
96    drop(file);
97    fs::rename(&temp_path, path).await?;
98    Ok(())
99}