provenant/scanner/process/
spill.rs1use crate::models::FileInfo;
5use anyhow::{Context, Result};
6use std::fs::{self, File};
7use std::io::{Read, Write};
8use tempfile::TempDir;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum MemoryMode {
12 CollectFirst,
13 StreamUnlimited,
14 Limit(usize),
15}
16
17impl std::fmt::Display for MemoryMode {
18 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
19 match self {
20 MemoryMode::CollectFirst => write!(f, "0"),
21 MemoryMode::StreamUnlimited => write!(f, "-1"),
22 MemoryMode::Limit(n) => write!(f, "{n}"),
23 }
24 }
25}
26
27pub(super) fn retain_or_spill_chunk(
28 chunk: Vec<FileInfo>,
29 retained_files: &mut Vec<FileInfo>,
30 spill_store: &mut Option<FileInfoSpillStore>,
31 memory_limit: usize,
32) -> Result<()> {
33 if memory_limit == 0 {
34 spill_store_mut(spill_store)?.spill(chunk)?;
35 return Ok(());
36 }
37
38 let remaining_capacity = memory_limit.saturating_sub(retained_files.len());
39 if remaining_capacity >= chunk.len() && spill_store.is_none() {
40 retained_files.extend(chunk);
41 return Ok(());
42 }
43
44 let mut chunk_iter = chunk.into_iter();
45 retained_files.extend(chunk_iter.by_ref().take(remaining_capacity));
46 let overflow: Vec<FileInfo> = chunk_iter.collect();
47 if !overflow.is_empty() {
48 spill_store_mut(spill_store)?.spill(overflow)?;
49 }
50
51 Ok(())
52}
53
54fn spill_store_mut(
55 spill_store: &mut Option<FileInfoSpillStore>,
56) -> Result<&mut FileInfoSpillStore> {
57 if spill_store.is_none() {
58 *spill_store = Some(FileInfoSpillStore::new()?);
59 }
60
61 Ok(spill_store
62 .as_mut()
63 .expect("spill store is always initialized after creation"))
64}
65
66pub(super) struct FileInfoSpillStore {
67 temp_dir: TempDir,
68 batch_index: usize,
69}
70
71impl FileInfoSpillStore {
72 fn new() -> Result<Self> {
73 Ok(Self {
74 temp_dir: TempDir::new().context("create scanner spill directory")?,
75 batch_index: 0,
76 })
77 }
78
79 fn spill(&mut self, files: Vec<FileInfo>) -> Result<()> {
80 let path = self
81 .temp_dir
82 .path()
83 .join(format!("batch-{:06}.postcard.zst", self.batch_index));
84 self.batch_index += 1;
85
86 let payload = postcard::to_allocvec(&files).context("encode scanner spill batch")?;
87 let file = File::create(&path)
88 .with_context(|| format!("create scanner spill batch file {}", path.display()))?;
89 let mut encoder = zstd::Encoder::new(file, 3)
90 .with_context(|| format!("create scanner spill encoder for {}", path.display()))?;
91 encoder
92 .write_all(&payload)
93 .with_context(|| format!("write scanner spill batch {}", path.display()))?;
94 encoder
95 .finish()
96 .with_context(|| format!("finish scanner spill encoder for {}", path.display()))?;
97
98 Ok(())
99 }
100
101 pub(super) fn load_all(self) -> Result<Vec<FileInfo>> {
102 let spill_dir = self.temp_dir.path();
103 let mut paths = Vec::new();
104 for entry in fs::read_dir(spill_dir)
105 .with_context(|| format!("read scanner spill directory {}", spill_dir.display()))?
106 {
107 let entry = entry.with_context(|| {
108 format!(
109 "read scanner spill directory entry from {}",
110 spill_dir.display()
111 )
112 })?;
113 paths.push(entry.path());
114 }
115 paths.sort();
116
117 let mut files = Vec::new();
118 for path in paths {
119 let file = File::open(&path)
120 .with_context(|| format!("open scanner spill batch {}", path.display()))?;
121 let mut decoder = zstd::Decoder::new(file)
122 .with_context(|| format!("create scanner spill decoder for {}", path.display()))?;
123 let mut payload = Vec::new();
124 decoder
125 .read_to_end(&mut payload)
126 .with_context(|| format!("read scanner spill batch {}", path.display()))?;
127 let mut batch: Vec<FileInfo> = postcard::from_bytes(&payload)
128 .with_context(|| format!("decode scanner spill batch {}", path.display()))?;
129 files.append(&mut batch);
130 }
131 Ok(files)
132 }
133}
134
135#[cfg(test)]
136mod tests {
137 use super::*;
138 use crate::models::{DiagnosticSeverity, FileType, ScanDiagnostic};
139
140 #[test]
141 fn spilled_files_preserve_scan_diagnostic_severity() -> Result<()> {
142 let mut store = FileInfoSpillStore::new()?;
143 let file = FileInfo::new(
144 "custom.txt".to_string(),
145 "custom".to_string(),
146 ".txt".to_string(),
147 "project/custom.txt".to_string(),
148 FileType::File,
149 None,
150 None,
151 10,
152 None,
153 None,
154 None,
155 None,
156 None,
157 Vec::new(),
158 None,
159 Vec::new(),
160 Vec::new(),
161 Vec::new(),
162 Vec::new(),
163 Vec::new(),
164 Vec::new(),
165 Vec::new(),
166 Vec::new(),
167 vec![ScanDiagnostic::warning("custom recoverable warning")],
168 );
169
170 store.spill(vec![file])?;
171 let loaded = store.load_all()?;
172
173 assert_eq!(loaded.len(), 1);
174 assert_eq!(loaded[0].scan_diagnostics.len(), 1);
175 assert_eq!(
176 loaded[0].scan_diagnostics[0].severity,
177 DiagnosticSeverity::Warning
178 );
179 Ok(())
180 }
181
182 #[test]
183 fn corrupt_spill_batch_returns_error() -> Result<()> {
184 let mut store = FileInfoSpillStore::new()?;
185 let path = store.temp_dir.path().join("batch-000000.postcard.zst");
186 fs::write(&path, b"not zstd")?;
187 store.batch_index = 1;
188
189 let error = store
190 .load_all()
191 .expect_err("corrupt spill batch should return an error");
192
193 assert!(error.to_string().contains("scanner spill"));
194 Ok(())
195 }
196}