1use std::sync::Arc;
25
26use crate::run_store::RunStore;
27
28use serde::{Deserialize, Serialize};
29pub use serde_json::Value;
30
31#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
39#[serde(rename_all = "camelCase")]
40pub struct Record {
41 pub run_id: String,
42 pub region: String,
43 #[serde(default)]
44 pub image_identifier: Option<String>,
45 #[serde(default)]
46 pub image_name: Option<String>,
47 #[serde(default)]
48 pub microvm_id: Option<String>,
49 #[serde(default)]
51 pub leaked: Vec<String>,
52 #[serde(default, skip_serializing_if = "Option::is_none")]
56 pub vm_log_group: Option<String>,
57}
58
59#[derive(Clone, Debug)]
61pub struct Ledger {
62 pub record: Record,
63 key: String,
64 store: Arc<dyn RunStore>,
65}
66
67impl Ledger {
68 pub fn from_store(region: &str, run_id: String, store: Arc<dyn RunStore>) -> Self {
74 Self {
75 key: run_id.clone(),
76 record: Record {
77 run_id,
78 region: region.to_string(),
79 ..Record::default()
80 },
81 store,
82 }
83 }
84
85 pub fn open_all_for_vm_with_store(store: Arc<dyn RunStore>, microvm_id: &str) -> Vec<Self> {
98 let mut documents = store.ledgers();
99 documents.sort_by(|a, b| a.key.cmp(&b.key));
100 documents
101 .into_iter()
102 .rev()
103 .filter_map(|document| {
104 let record: Record = serde_json::from_str(document.text.as_deref()?).ok()?;
105 (record.microvm_id.as_deref() == Some(microvm_id)).then(|| Self {
106 record,
107 key: document.key,
108 store: store.clone(),
109 })
110 })
111 .collect()
112 }
113
114 pub fn flush(&self) {
116 if let Ok(text) = serde_json::to_string_pretty(&self.record) {
117 let _ = self.store.write_ledger(&self.key, &text);
118 }
119 }
120
121 pub fn record_image(&mut self, identifier: &str, name: &str) {
123 self.record.image_identifier = Some(identifier.to_string());
124 self.record.image_name = Some(name.to_string());
125 self.flush();
126 }
127
128 pub fn record_vm_log_group(&mut self, group: &str) {
130 self.record.vm_log_group = Some(group.to_string());
131 self.flush();
132 }
133
134 pub fn record_microvm(&mut self, id: &str) {
136 self.record.microvm_id = Some(id.to_string());
137 self.flush();
138 }
139
140 pub fn mark_outstanding(&mut self) {
145 self.record.leaked = [
146 self.record.microvm_id.clone(),
147 self.record.image_identifier.clone(),
148 ]
149 .into_iter()
150 .flatten()
151 .collect();
152 self.flush();
153 }
154
155 pub fn mark_deleted(&mut self, still_undeleted: &[String]) {
161 self.record.leaked = still_undeleted.to_vec();
162 self.flush();
163 }
164
165 pub fn clear(&self) {
167 if !self.record.leaked.is_empty() {
168 return;
169 }
170 let _ = self.store.remove_ledger(&self.key);
171 }
172}
173
174pub fn remove(store: &dyn RunStore, run_id: &str) -> std::io::Result<()> {
182 let legal = !run_id.is_empty()
183 && run_id
184 .bytes()
185 .all(|byte| byte.is_ascii_alphanumeric() || byte == b'-' || byte == b'_');
186 if !legal {
187 return Err(std::io::Error::new(
188 std::io::ErrorKind::InvalidInput,
189 format!("{run_id:?} is not a ledger run id"),
190 ));
191 }
192 store.remove_ledger(run_id)
193}
194
195pub fn read_all(store: &dyn RunStore) -> Vec<serde_json::Value> {
197 let mut documents = store.ledgers();
198 documents.sort_by(|a, b| a.key.cmp(&b.key));
199 documents
200 .into_iter()
201 .map(|document| {
202 document
203 .text
204 .and_then(|text| serde_json::from_str(&text).ok())
205 .unwrap_or_else(
206 || serde_json::json!({"runId": document.key, "unreadable": document.location}),
207 )
208 })
209 .collect()
210}
211
212#[cfg(test)]
213mod tests {
214 use super::*;
215 use crate::run_store::StoredLedger;
216 use crate::testing::MemoryRunStore;
217 use serde_json::json;
218
219 #[test]
220 fn lifecycle_flushes_each_transition_and_keeps_outstanding_runs() {
221 let store = Arc::new(MemoryRunStore::default());
222 store
223 .histories
224 .lock()
225 .unwrap()
226 .insert("vm".into(), "history\n".into());
227 let mut ledger = Ledger::from_store("region", "run-1".into(), store.clone());
228 ledger.record_image("image-arn", "image-name");
229 assert_eq!(
230 read_all(store.as_ref()),
231 vec![json!({
232 "runId": "run-1", "region": "region", "imageIdentifier": "image-arn",
233 "imageName": "image-name", "microvmId": null, "leaked": [],
234 })]
235 );
236 ledger.record_microvm("vm");
237 ledger.record_vm_log_group("logs");
238 ledger.mark_outstanding();
239 assert_eq!(
240 read_all(store.as_ref()),
241 vec![json!({
242 "runId": "run-1", "region": "region", "imageIdentifier": "image-arn",
243 "imageName": "image-name", "microvmId": "vm", "vmLogGroup": "logs",
244 "leaked": ["vm", "image-arn"],
245 })]
246 );
247 ledger.clear();
248 assert!(store.removed.lock().unwrap().is_empty());
249 ledger.mark_deleted(&["image-arn".into()]);
250 assert_eq!(read_all(store.as_ref())[0]["leaked"], json!(["image-arn"]));
251 ledger.clear();
252 assert!(store.removed.lock().unwrap().is_empty());
253 ledger.mark_deleted(&[]);
254 assert_eq!(read_all(store.as_ref())[0]["leaked"], json!([]));
255 ledger.clear();
256 assert_eq!(*store.removed.lock().unwrap(), ["run-1"]);
257 assert!(read_all(store.as_ref()).is_empty());
258 assert_eq!(store.read_history("vm").unwrap(), "history\n");
259 }
260
261 #[test]
262 fn reopening_uses_the_original_key_and_returns_newest_matching_records_first() {
263 let store = Arc::new(MemoryRunStore::default());
264 for (key, vm, run_id) in [
265 ("a", "vm", "old-field"),
266 ("c", "other", "other"),
267 ("b", "vm", "new-field"),
268 ] {
269 store
270 .write_ledger(
271 key,
272 &json!({"runId": run_id, "region": "r", "microvmId": vm}).to_string(),
273 )
274 .unwrap();
275 }
276 store.documents.lock().unwrap().push(StoredLedger {
277 key: "z".into(),
278 location: "memory/z".into(),
279 text: None,
280 });
281 store.write_ledger("y", "broken").unwrap();
282 let mut records = Ledger::open_all_for_vm_with_store(store.clone(), "vm");
283 assert_eq!(
284 records
285 .iter()
286 .map(|ledger| ledger.record.run_id.as_str())
287 .collect::<Vec<_>>(),
288 ["new-field", "old-field"]
289 );
290 records[1].record_microvm("changed");
291 let documents = store.ledgers();
292 assert_eq!(documents.len(), 5);
293 assert!(documents.iter().all(|doc| doc.key != "old-field"));
294 let old: Value = serde_json::from_str(
295 documents
296 .iter()
297 .find(|doc| doc.key == "a")
298 .unwrap()
299 .text
300 .as_deref()
301 .unwrap(),
302 )
303 .unwrap();
304 assert_eq!(old["microvmId"], "changed");
305 }
306
307 #[test]
308 fn listing_orders_documents_and_preserves_unreadable_locations() {
309 let store = MemoryRunStore::default();
310 *store.documents.lock().unwrap() = vec![
311 StoredLedger {
312 key: "c".into(),
313 location: "file-c".into(),
314 text: None,
315 },
316 StoredLedger {
317 key: "b".into(),
318 location: "file-b".into(),
319 text: Some("broken".into()),
320 },
321 StoredLedger {
322 key: "a".into(),
323 location: "file-a".into(),
324 text: Some("{\"kept\":true}".into()),
325 },
326 ];
327 assert_eq!(
328 read_all(&store),
329 vec![
330 json!({"kept": true}),
331 json!({"runId": "b", "unreadable": "file-b"}),
332 json!({"runId": "c", "unreadable": "file-c"})
333 ]
334 );
335 }
336
337 #[test]
338 fn removal_accepts_each_legal_character_and_refuses_every_other_class() {
339 let store = MemoryRunStore::default();
340 for key in ["abc", "ABC", "123", "-", "_", "1700000000-12_3"] {
341 remove(&store, key).unwrap();
342 }
343 assert_eq!(
344 *store.removed.lock().unwrap(),
345 ["abc", "ABC", "123", "-", "_", "1700000000-12_3"]
346 );
347 for key in ["", ".", "../names/x", "a/b", "a\\b", "a:b", "a b", "café"] {
348 assert_eq!(
349 remove(&store, key).unwrap_err().kind(),
350 std::io::ErrorKind::InvalidInput,
351 "{key:?}"
352 );
353 }
354 assert_eq!(store.removed.lock().unwrap().len(), 6);
355 }
356}