use std::sync::Arc;
use crate::run_store::RunStore;
use serde::{Deserialize, Serialize};
pub use serde_json::Value;
#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct Record {
pub run_id: String,
pub region: String,
#[serde(default)]
pub image_identifier: Option<String>,
#[serde(default)]
pub image_name: Option<String>,
#[serde(default)]
pub microvm_id: Option<String>,
#[serde(default)]
pub leaked: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub vm_log_group: Option<String>,
}
#[derive(Clone, Debug)]
pub struct Ledger {
pub record: Record,
key: String,
store: Arc<dyn RunStore>,
}
impl Ledger {
pub fn from_store(region: &str, run_id: String, store: Arc<dyn RunStore>) -> Self {
Self {
key: run_id.clone(),
record: Record {
run_id,
region: region.to_string(),
..Record::default()
},
store,
}
}
pub fn open_all_for_vm_with_store(store: Arc<dyn RunStore>, microvm_id: &str) -> Vec<Self> {
let mut documents = store.ledgers();
documents.sort_by(|a, b| a.key.cmp(&b.key));
documents
.into_iter()
.rev()
.filter_map(|document| {
let record: Record = serde_json::from_str(document.text.as_deref()?).ok()?;
(record.microvm_id.as_deref() == Some(microvm_id)).then(|| Self {
record,
key: document.key,
store: store.clone(),
})
})
.collect()
}
pub fn flush(&self) {
if let Ok(text) = serde_json::to_string_pretty(&self.record) {
let _ = self.store.write_ledger(&self.key, &text);
}
}
pub fn record_image(&mut self, identifier: &str, name: &str) {
self.record.image_identifier = Some(identifier.to_string());
self.record.image_name = Some(name.to_string());
self.flush();
}
pub fn record_vm_log_group(&mut self, group: &str) {
self.record.vm_log_group = Some(group.to_string());
self.flush();
}
pub fn record_microvm(&mut self, id: &str) {
self.record.microvm_id = Some(id.to_string());
self.flush();
}
pub fn mark_outstanding(&mut self) {
self.record.leaked = [
self.record.microvm_id.clone(),
self.record.image_identifier.clone(),
]
.into_iter()
.flatten()
.collect();
self.flush();
}
pub fn mark_deleted(&mut self, still_undeleted: &[String]) {
self.record.leaked = still_undeleted.to_vec();
self.flush();
}
pub fn clear(&self) {
if !self.record.leaked.is_empty() {
return;
}
let _ = self.store.remove_ledger(&self.key);
}
}
pub fn remove(store: &dyn RunStore, run_id: &str) -> std::io::Result<()> {
let legal = !run_id.is_empty()
&& run_id
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || byte == b'-' || byte == b'_');
if !legal {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("{run_id:?} is not a ledger run id"),
));
}
store.remove_ledger(run_id)
}
pub fn read_all(store: &dyn RunStore) -> Vec<serde_json::Value> {
let mut documents = store.ledgers();
documents.sort_by(|a, b| a.key.cmp(&b.key));
documents
.into_iter()
.map(|document| {
document
.text
.and_then(|text| serde_json::from_str(&text).ok())
.unwrap_or_else(
|| serde_json::json!({"runId": document.key, "unreadable": document.location}),
)
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::run_store::StoredLedger;
use crate::testing::MemoryRunStore;
use serde_json::json;
#[test]
fn lifecycle_flushes_each_transition_and_keeps_outstanding_runs() {
let store = Arc::new(MemoryRunStore::default());
store
.histories
.lock()
.unwrap()
.insert("vm".into(), "history\n".into());
let mut ledger = Ledger::from_store("region", "run-1".into(), store.clone());
ledger.record_image("image-arn", "image-name");
assert_eq!(
read_all(store.as_ref()),
vec![json!({
"runId": "run-1", "region": "region", "imageIdentifier": "image-arn",
"imageName": "image-name", "microvmId": null, "leaked": [],
})]
);
ledger.record_microvm("vm");
ledger.record_vm_log_group("logs");
ledger.mark_outstanding();
assert_eq!(
read_all(store.as_ref()),
vec![json!({
"runId": "run-1", "region": "region", "imageIdentifier": "image-arn",
"imageName": "image-name", "microvmId": "vm", "vmLogGroup": "logs",
"leaked": ["vm", "image-arn"],
})]
);
ledger.clear();
assert!(store.removed.lock().unwrap().is_empty());
ledger.mark_deleted(&["image-arn".into()]);
assert_eq!(read_all(store.as_ref())[0]["leaked"], json!(["image-arn"]));
ledger.clear();
assert!(store.removed.lock().unwrap().is_empty());
ledger.mark_deleted(&[]);
assert_eq!(read_all(store.as_ref())[0]["leaked"], json!([]));
ledger.clear();
assert_eq!(*store.removed.lock().unwrap(), ["run-1"]);
assert!(read_all(store.as_ref()).is_empty());
assert_eq!(store.read_history("vm").unwrap(), "history\n");
}
#[test]
fn reopening_uses_the_original_key_and_returns_newest_matching_records_first() {
let store = Arc::new(MemoryRunStore::default());
for (key, vm, run_id) in [
("a", "vm", "old-field"),
("c", "other", "other"),
("b", "vm", "new-field"),
] {
store
.write_ledger(
key,
&json!({"runId": run_id, "region": "r", "microvmId": vm}).to_string(),
)
.unwrap();
}
store.documents.lock().unwrap().push(StoredLedger {
key: "z".into(),
location: "memory/z".into(),
text: None,
});
store.write_ledger("y", "broken").unwrap();
let mut records = Ledger::open_all_for_vm_with_store(store.clone(), "vm");
assert_eq!(
records
.iter()
.map(|ledger| ledger.record.run_id.as_str())
.collect::<Vec<_>>(),
["new-field", "old-field"]
);
records[1].record_microvm("changed");
let documents = store.ledgers();
assert_eq!(documents.len(), 5);
assert!(documents.iter().all(|doc| doc.key != "old-field"));
let old: Value = serde_json::from_str(
documents
.iter()
.find(|doc| doc.key == "a")
.unwrap()
.text
.as_deref()
.unwrap(),
)
.unwrap();
assert_eq!(old["microvmId"], "changed");
}
#[test]
fn listing_orders_documents_and_preserves_unreadable_locations() {
let store = MemoryRunStore::default();
*store.documents.lock().unwrap() = vec![
StoredLedger {
key: "c".into(),
location: "file-c".into(),
text: None,
},
StoredLedger {
key: "b".into(),
location: "file-b".into(),
text: Some("broken".into()),
},
StoredLedger {
key: "a".into(),
location: "file-a".into(),
text: Some("{\"kept\":true}".into()),
},
];
assert_eq!(
read_all(&store),
vec![
json!({"kept": true}),
json!({"runId": "b", "unreadable": "file-b"}),
json!({"runId": "c", "unreadable": "file-c"})
]
);
}
#[test]
fn removal_accepts_each_legal_character_and_refuses_every_other_class() {
let store = MemoryRunStore::default();
for key in ["abc", "ABC", "123", "-", "_", "1700000000-12_3"] {
remove(&store, key).unwrap();
}
assert_eq!(
*store.removed.lock().unwrap(),
["abc", "ABC", "123", "-", "_", "1700000000-12_3"]
);
for key in ["", ".", "../names/x", "a/b", "a\\b", "a:b", "a b", "café"] {
assert_eq!(
remove(&store, key).unwrap_err().kind(),
std::io::ErrorKind::InvalidInput,
"{key:?}"
);
}
assert_eq!(store.removed.lock().unwrap().len(), 6);
}
}