1use crate::{Classification, ComparisonRun, ComparisonRunListItem, LatencyStats, StatsSummary};
2use std::{
3 fs as std_fs,
4 io::{BufRead, BufReader as StdBufReader},
5 path::{Path, PathBuf},
6 sync::Arc,
7};
8use tokio::{
9 fs::{self, File, OpenOptions},
10 io::{AsyncBufReadExt, AsyncWriteExt, BufReader as TokioBufReader, BufWriter},
11 sync::{Mutex, RwLock},
12};
13use uuid::Uuid;
14
15#[derive(Clone)]
16pub struct Storage {
17 write_path: PathBuf,
18 scan_dir: PathBuf,
19 writer: RunWriter,
20 options: StorageOptions,
21 insert_lock: Arc<Mutex<()>>,
22 runs: Arc<RwLock<Vec<ComparisonRun>>>,
23}
24
25#[derive(Clone)]
26pub struct RunWriter {
27 file: Arc<Mutex<BufWriter<File>>>,
28}
29
30#[derive(Clone)]
31pub struct JsonlStorageReader {
32 path: PathBuf,
33}
34
35impl RunWriter {
36 pub async fn open(write_path: PathBuf) -> anyhow::Result<Self> {
37 if let Some(parent) = write_path.parent() {
38 fs::create_dir_all(parent).await?;
39 }
40 let file = OpenOptions::new()
41 .create(true)
42 .append(true)
43 .open(write_path)
44 .await?;
45
46 Ok(Self {
47 file: Arc::new(Mutex::new(BufWriter::new(file))),
48 })
49 }
50
51 pub async fn append(&self, run: &ComparisonRun) -> anyhow::Result<()> {
52 let line = serde_json::to_string(run)?;
53 let mut file = self.file.lock().await;
54 file.write_all(line.as_bytes()).await?;
55 file.write_all(b"\n").await?;
56 Ok(())
57 }
58
59 pub async fn flush(&self) -> anyhow::Result<()> {
60 self.file.lock().await.flush().await?;
61 Ok(())
62 }
63}
64
65impl JsonlStorageReader {
66 pub fn new(path: PathBuf) -> Self {
67 Self { path }
68 }
69
70 pub async fn stats(&self) -> anyhow::Result<StatsSummary> {
71 let mut accumulator = StatsAccumulator::default();
72 self.for_each_run(|run| {
73 accumulator.record(&run);
74 true
75 })
76 .await?;
77 Ok(accumulator.finish())
78 }
79
80 pub async fn list_page(
81 &self,
82 limit: Option<usize>,
83 offset: usize,
84 ) -> anyhow::Result<Vec<ComparisonRunListItem>> {
85 let retained_limit = limit.and_then(|value| value.checked_add(offset));
86 let mut runs = Vec::new();
87
88 self.for_each_run(|run| {
89 runs.push(ComparisonRunListItem::from(&run));
90 if let Some(retained_limit) = retained_limit {
91 if runs.len() > retained_limit {
92 runs.remove(0);
93 }
94 }
95 true
96 })
97 .await?;
98
99 runs.reverse();
100 Ok(match limit {
101 Some(limit) => runs.into_iter().skip(offset).take(limit).collect(),
102 None => runs.into_iter().skip(offset).collect(),
103 })
104 }
105
106 pub async fn get(&self, id: Uuid) -> anyhow::Result<Option<ComparisonRun>> {
107 let mut found = None;
108 self.for_each_run(|run| {
109 if run.id == id {
110 found = Some(run);
111 false
112 } else {
113 true
114 }
115 })
116 .await?;
117 Ok(found)
118 }
119
120 async fn for_each_run(
121 &self,
122 mut visit: impl FnMut(ComparisonRun) -> bool,
123 ) -> anyhow::Result<()> {
124 if !self.path.try_exists()? {
125 return Ok(());
126 }
127
128 let file = std_fs::File::open(&self.path)?;
129 let lines = StdBufReader::new(file).lines();
130 for line in lines {
131 let line = line?;
132 if line.trim().is_empty() {
133 continue;
134 }
135 match serde_json::from_str::<ComparisonRun>(&line) {
136 Ok(run) => {
137 if !visit(run) {
138 break;
139 }
140 }
141 Err(error) => warn_corrupt_line(&self.path, &error),
142 }
143 }
144 Ok(())
145 }
146}
147
148#[derive(Debug, Clone, Copy, Default)]
149pub struct StorageOptions {
150 pub retention_max_runs: Option<usize>,
151 pub retention_max_bytes: Option<u64>,
152}
153
154impl Storage {
155 pub async fn load(write_path: PathBuf) -> anyhow::Result<Self> {
156 Self::load_with_options(write_path, StorageOptions::default()).await
157 }
158
159 pub async fn load_with_options(
160 write_path: PathBuf,
161 options: StorageOptions,
162 ) -> anyhow::Result<Self> {
163 if let Some(parent) = write_path.parent() {
164 fs::create_dir_all(parent).await?;
165 }
166 let scan_dir = write_path
167 .parent()
168 .map(Path::to_path_buf)
169 .unwrap_or_else(|| PathBuf::from("."));
170 let runs = load_runs_from_dir(&scan_dir).await?;
171 let writer = RunWriter::open(write_path.clone()).await?;
172
173 Ok(Self {
174 write_path,
175 scan_dir,
176 writer,
177 options,
178 insert_lock: Arc::new(Mutex::new(())),
179 runs: Arc::new(RwLock::new(runs)),
180 })
181 }
182
183 pub async fn insert(&self, run: ComparisonRun) -> anyhow::Result<()> {
184 let _guard = self.insert_lock.lock().await;
185 self.writer.append(&run).await?;
186 self.writer.flush().await?;
187 self.runs.write().await.push(run);
188 self.apply_retention().await?;
189 Ok(())
190 }
191
192 pub async fn refresh(&self) -> anyhow::Result<()> {
193 let runs = load_runs_from_dir(&self.scan_dir).await?;
194 *self.runs.write().await = runs;
195 Ok(())
196 }
197
198 pub async fn list(&self) -> Vec<ComparisonRunListItem> {
199 self.list_page(usize::MAX, 0).await
200 }
201
202 pub async fn list_page(&self, limit: usize, offset: usize) -> Vec<ComparisonRunListItem> {
203 let runs = self.runs.read().await;
204 runs.iter()
205 .rev()
206 .skip(offset)
207 .take(limit)
208 .map(ComparisonRunListItem::from)
209 .collect()
210 }
211
212 pub async fn get(&self, id: Uuid) -> Option<ComparisonRun> {
213 let runs = self.runs.read().await;
214 runs.iter().find(|run| run.id == id).cloned()
215 }
216
217 pub async fn stats(&self) -> StatsSummary {
218 let runs = self.runs.read().await;
219 let mut accumulator = StatsAccumulator::default();
220
221 for run in runs.iter() {
222 accumulator.record(run);
223 }
224
225 accumulator.finish()
226 }
227
228 async fn apply_retention(&self) -> anyhow::Result<()> {
229 if self.options.retention_max_runs.is_none() && self.options.retention_max_bytes.is_none() {
230 return Ok(());
231 }
232
233 let mut active_runs = Vec::new();
234 load_runs_from_file(&self.write_path, &mut active_runs).await?;
235 active_runs.sort_by_key(|run| run.timestamp);
236
237 if let Some(max_runs) = self.options.retention_max_runs {
238 if active_runs.len() > max_runs {
239 active_runs = active_runs
240 .into_iter()
241 .rev()
242 .take(max_runs)
243 .collect::<Vec<_>>();
244 active_runs.reverse();
245 }
246 }
247
248 if let Some(max_bytes) = self.options.retention_max_bytes {
249 let mut retained = Vec::new();
250 let mut total_bytes = 0_u64;
251 for run in active_runs.into_iter().rev() {
252 let line = serde_json::to_string(&run)?;
253 let line_bytes = line.len() as u64 + 1;
254 if total_bytes + line_bytes <= max_bytes || retained.is_empty() {
255 total_bytes += line_bytes;
256 retained.push(run);
257 } else {
258 break;
259 }
260 }
261 retained.reverse();
262 active_runs = retained;
263 }
264
265 let mut content = String::new();
266 for run in active_runs {
267 content.push_str(&serde_json::to_string(&run)?);
268 content.push('\n');
269 }
270 fs::write(&self.write_path, content).await?;
271 self.refresh().await?;
272 Ok(())
273 }
274}
275
276async fn load_runs_from_dir(scan_dir: &Path) -> anyhow::Result<Vec<ComparisonRun>> {
277 let mut runs = Vec::new();
278 if !fs::try_exists(scan_dir).await? {
279 return Ok(runs);
280 }
281
282 let mut entries = fs::read_dir(scan_dir).await?;
283 while let Some(entry) = entries.next_entry().await? {
284 let path = entry.path();
285 if path.extension().and_then(|value| value.to_str()) != Some("jsonl") {
286 continue;
287 }
288 load_runs_from_file(&path, &mut runs).await?;
289 }
290
291 runs.sort_by_key(|run| run.timestamp);
292 Ok(runs)
293}
294
295async fn load_runs_from_file(path: &Path, runs: &mut Vec<ComparisonRun>) -> anyhow::Result<()> {
296 let file = fs::File::open(path).await?;
297 let mut lines = TokioBufReader::new(file).lines();
298 while let Some(line) = lines.next_line().await? {
299 if line.trim().is_empty() {
300 continue;
301 }
302 match serde_json::from_str::<ComparisonRun>(&line) {
303 Ok(run) => runs.push(run),
304 Err(error) => warn_corrupt_line(path, &error),
305 }
306 }
307 Ok(())
308}
309
310fn warn_corrupt_line(path: &Path, error: &serde_json::Error) {
311 eprintln!(
312 "skipping corrupt moonlight JSONL run in {}: {error}",
313 path.display()
314 );
315}
316
317#[derive(Default)]
318struct StatsAccumulator {
319 total_runs: usize,
320 matches: usize,
321 suspicious_differences: usize,
322 reference_noise: usize,
323 suspicious_with_noise: usize,
324 target_errors: usize,
325 primary_total: u128,
326 candidate_total: u128,
327 secondary_latencies: Vec<u128>,
328 latest_runs: Vec<ComparisonRunListItem>,
329}
330
331impl StatsAccumulator {
332 fn record(&mut self, run: &ComparisonRun) {
333 self.total_runs += 1;
334 match run.comparison.classification {
335 Classification::Match => self.matches += 1,
336 Classification::SuspiciousDifference => self.suspicious_differences += 1,
337 Classification::ReferenceNoise => self.reference_noise += 1,
338 Classification::SuspiciousWithNoise => self.suspicious_with_noise += 1,
339 Classification::TargetError => self.target_errors += 1,
340 }
341 self.primary_total += run.primary.latency_ms;
342 self.candidate_total += run.candidate.latency_ms;
343 if let Some(secondary) = &run.secondary {
344 self.secondary_latencies.push(secondary.latency_ms);
345 }
346
347 self.latest_runs.push(ComparisonRunListItem::from(run));
348 if self.latest_runs.len() > 20 {
349 self.latest_runs.remove(0);
350 }
351 }
352
353 fn finish(mut self) -> StatsSummary {
354 self.latest_runs.reverse();
355 StatsSummary {
356 total_runs: self.total_runs,
357 matches: self.matches,
358 suspicious_differences: self.suspicious_differences,
359 reference_noise: self.reference_noise,
360 suspicious_with_noise: self.suspicious_with_noise,
361 target_errors: self.target_errors,
362 latency: LatencyStats {
363 primary_avg_ms: avg(self.total_runs, self.primary_total),
364 candidate_avg_ms: avg(self.total_runs, self.candidate_total),
365 secondary_avg_ms: avg_opt(&self.secondary_latencies),
366 },
367 latest_runs: self.latest_runs,
368 }
369 }
370}
371
372fn avg(count: usize, total: u128) -> f64 {
373 if count == 0 {
374 0.0
375 } else {
376 total as f64 / count as f64
377 }
378}
379
380fn avg_opt(values: &[u128]) -> Option<f64> {
381 if values.is_empty() {
382 None
383 } else {
384 Some(values.iter().sum::<u128>() as f64 / values.len() as f64)
385 }
386}
387
388#[cfg(test)]
389mod tests;