use std::collections::BTreeSet;
use std::fs;
use std::path::Path;
use std::path::PathBuf;
use std::time::SystemTime;
use serde::Deserialize;
use serde::Serialize;
use crate::error::Result;
use crate::error::SnapshotError;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SnapshotRef {
pub turn_id: String,
pub manifest_id: String,
pub at: u64,
pub prev_hash: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ThreadLog {
pub version: u32,
pub entries: Vec<SnapshotRef>,
}
#[derive(Debug, Default)]
pub struct ThreadLogs {
pub logs: Vec<ThreadLog>,
pub incomplete: bool,
}
impl Default for ThreadLog {
fn default() -> Self {
Self {
version: crate::workspace::FORMAT_VERSION,
entries: Vec::new(),
}
}
}
impl SnapshotRef {
pub(crate) fn chained(turn_id: String, manifest_id: String, previous: Option<&Self>) -> Self {
Self {
turn_id,
manifest_id,
at: SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map_or(0, |d| d.as_secs()),
prev_hash: previous.map(Self::digest),
}
}
fn digest(entry: &Self) -> String {
serde_json::to_vec(entry).map_or_else(
|_| String::new(),
|bytes| crate::blob::BlobStore::hash_bytes(&bytes),
)
}
}
fn check_version(kind: &'static str, id: &str, found: u32) -> Result<()> {
if found == crate::workspace::FORMAT_VERSION {
return Ok(());
}
Err(SnapshotError::UnknownRecordVersion {
kind,
id: id.to_string(),
found,
supported: crate::workspace::FORMAT_VERSION,
})
}
pub(crate) const TURN_SUFFIX: &str = ".turn";
pub struct RefStore {
root: PathBuf,
}
impl RefStore {
pub fn open(root: impl Into<PathBuf>) -> Result<Self> {
let root = root.into();
fs::create_dir_all(&root).map_err(|e| SnapshotError::io(&root, e))?;
Ok(Self { root })
}
pub fn append(&self, thread_id: &str, turn_id: String, manifest_id: String) -> Result<()> {
let mut log = self.load(thread_id)?;
let entry = SnapshotRef::chained(turn_id, manifest_id, log.entries.last());
log.entries.push(entry);
let path = self.log_path(thread_id)?;
let tmp = crate::sweep::tmp_name(&path);
let bytes = serde_json::to_vec_pretty(&log)?;
fs::write(&tmp, bytes).map_err(|e| SnapshotError::io(&tmp, e))?;
fs::rename(&tmp, &path).map_err(|e| SnapshotError::io(&path, e))?;
Ok(())
}
pub fn exists(&self, thread_id: &str) -> bool {
self.log_path(thread_id).is_ok_and(|p| p.exists())
}
pub fn ensure(&self, thread_id: &str) -> Result<()> {
let path = self.log_path(thread_id)?;
if path.exists() {
return Ok(());
}
let tmp = crate::sweep::tmp_name(&path);
let bytes = serde_json::to_vec_pretty(&ThreadLog::default())?;
fs::write(&tmp, bytes).map_err(|e| SnapshotError::io(&tmp, e))?;
fs::rename(&tmp, &path).map_err(|e| SnapshotError::io(&path, e))?;
Ok(())
}
pub fn load(&self, thread_id: &str) -> Result<ThreadLog> {
let path = self.log_path(thread_id)?;
match fs::read(&path) {
Ok(bytes) => {
let log: ThreadLog = serde_json::from_slice(&bytes)?;
check_version("thread log", thread_id, log.version)?;
Ok(log)
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(ThreadLog::default()),
Err(e) => Err(SnapshotError::io(&path, e)),
}
}
pub fn remove(&self, thread_id: &str) -> Result<()> {
let path = self.log_path(thread_id)?;
match fs::remove_file(&path) {
Ok(()) => Ok(()),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(e) => Err(SnapshotError::io(&path, e)),
}
}
pub fn thread_logs(&self) -> Result<ThreadLogs> {
let mut out = ThreadLogs::default();
let entries = fs::read_dir(&self.root).map_err(|e| SnapshotError::io(&self.root, e))?;
for entry in entries {
let entry = entry.map_err(|e| SnapshotError::io(&self.root, e))?;
let name = entry.file_name().to_string_lossy().into_owned();
if !name.ends_with(".json") {
continue;
}
match fs::read(entry.path()).map(|b| serde_json::from_slice::<ThreadLog>(&b)) {
Ok(Ok(log)) if log.version == crate::workspace::FORMAT_VERSION => {
out.logs.push(log);
}
_ => out.incomplete = true,
}
}
Ok(out)
}
pub fn thread_ids(&self) -> Result<Vec<String>> {
let entries = fs::read_dir(&self.root).map_err(|e| SnapshotError::io(&self.root, e))?;
let mut out = Vec::new();
for entry in entries {
let entry = entry.map_err(|e| SnapshotError::io(&self.root, e))?;
let name = entry.file_name().to_string_lossy().into_owned();
if let Some(id) = name.strip_suffix(".json") {
out.push(id.to_string());
}
}
out.sort();
Ok(out)
}
fn log_path(&self, thread_id: &str) -> Result<PathBuf> {
crate::id::validate_stored("session id", thread_id)?;
Ok(self.root.join(format!("{thread_id}.json")))
}
}
pub(crate) const MAX_RESTORE_HISTORY: usize = 20;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RestoreRecord {
pub target_manifest_id: String,
pub safety_manifest_id: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RestoreLog {
pub version: u32,
pub entries: Vec<RestoreRecord>,
}
impl Default for RestoreLog {
fn default() -> Self {
Self {
version: crate::workspace::FORMAT_VERSION,
entries: Vec::new(),
}
}
}
#[derive(Debug, Default)]
pub struct RestoreLogs {
pub logs: Vec<RestoreLog>,
pub incomplete: bool,
}
#[derive(Debug, Default)]
pub struct HeldManifests {
pub ids: BTreeSet<String>,
pub incomplete: bool,
}
pub struct TurnIndex {
turns_root: PathBuf,
restores_root: PathBuf,
}
impl TurnIndex {
pub fn open(root: impl Into<PathBuf>) -> Result<Self> {
let root = root.into();
let turns_root = root.join("turns");
let restores_root = root.join("restores");
fs::create_dir_all(&turns_root).map_err(|e| SnapshotError::io(&turns_root, e))?;
fs::create_dir_all(&restores_root).map_err(|e| SnapshotError::io(&restores_root, e))?;
Ok(Self {
turns_root,
restores_root,
})
}
pub fn set_turn(&self, turn_id: &str, manifest_id: &str) -> Result<()> {
let path = self.turn_path(turn_id)?;
write_atomic(&path, manifest_id.as_bytes())
}
pub fn manifest_for_turn(&self, turn_id: &str) -> Result<Option<String>> {
let path = self.turn_path(turn_id)?;
match fs::read_to_string(&path) {
Ok(id) => Ok(Some(id.trim().to_string())),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
Err(e) => Err(SnapshotError::io(&path, e)),
}
}
pub fn all_manifest_ids(&self) -> Result<HeldManifests> {
let mut out = HeldManifests::default();
let entries =
fs::read_dir(&self.turns_root).map_err(|e| SnapshotError::io(&self.turns_root, e))?;
for entry in entries {
let entry = entry.map_err(|e| SnapshotError::io(&self.turns_root, e))?;
if !entry.file_name().to_string_lossy().ends_with(TURN_SUFFIX) {
continue;
}
match fs::read_to_string(entry.path()) {
Ok(id) => {
out.ids.insert(id.trim().to_string());
}
Err(_) => out.incomplete = true,
}
}
let logs = self.all_restore_logs()?;
out.incomplete |= logs.incomplete;
for log in logs.logs {
for record in log.entries {
out.ids.insert(record.target_manifest_id);
out.ids.insert(record.safety_manifest_id);
}
}
Ok(out)
}
pub fn orphan_restore_logs(&self, refs: &RefStore) -> Result<Vec<String>> {
let mut out = Vec::new();
let entries = fs::read_dir(&self.restores_root)
.map_err(|e| SnapshotError::io(&self.restores_root, e))?;
for entry in entries {
let entry = entry.map_err(|e| SnapshotError::io(&self.restores_root, e))?;
let name = entry.file_name().to_string_lossy().into_owned();
let Some(thread_id) = name.strip_suffix(".json") else {
continue;
};
if !refs.exists(thread_id) && crate::sweep::settled(&entry.path()) {
out.push(thread_id.to_string());
}
}
Ok(out)
}
pub fn remove_turn_file(&self, turn_file: &str) -> Result<()> {
let Some(turn_id) = turn_file.strip_suffix(TURN_SUFFIX) else {
return Err(SnapshotError::InvalidId {
kind: "turn record",
id: turn_file.to_string(),
reason: "a turn record's name must end in `.turn`",
});
};
crate::id::validate_stored("turn id", turn_id)?;
let path = self.turns_root.join(turn_file);
match fs::remove_file(&path) {
Ok(()) => Ok(()),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(e) => Err(SnapshotError::io(&path, e)),
}
}
pub fn push_restore(&self, thread_id: &str, record: RestoreRecord) -> Result<()> {
let mut log = self.restore_log(thread_id)?;
log.entries.push(record);
if log.entries.len() > MAX_RESTORE_HISTORY {
let excess = log.entries.len() - MAX_RESTORE_HISTORY;
log.entries.drain(..excess);
}
let path = self.restore_path(thread_id)?;
write_atomic(&path, &serde_json::to_vec_pretty(&log)?)
}
pub fn restore_records(&self, thread_id: &str) -> Result<Vec<RestoreRecord>> {
Ok(self.restore_log(thread_id)?.entries)
}
pub fn last_restore(&self, thread_id: &str) -> Result<Option<RestoreRecord>> {
Ok(self.restore_log(thread_id)?.entries.pop())
}
pub fn pop_restore(&self, thread_id: &str) -> Result<Option<RestoreRecord>> {
let mut log = self.restore_log(thread_id)?;
let popped = log.entries.pop();
if popped.is_some() {
let path = self.restore_path(thread_id)?;
write_atomic(&path, &serde_json::to_vec_pretty(&log)?)?;
}
Ok(popped)
}
pub fn remove_restores(&self, thread_id: &str) -> Result<()> {
let path = self.restore_path(thread_id)?;
match fs::remove_file(&path) {
Ok(()) => Ok(()),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(e) => Err(SnapshotError::io(&path, e)),
}
}
fn restore_log(&self, thread_id: &str) -> Result<RestoreLog> {
let path = self.restore_path(thread_id)?;
match fs::read(&path) {
Ok(bytes) => {
let log: RestoreLog = serde_json::from_slice(&bytes)?;
check_version("restore log", thread_id, log.version)?;
Ok(log)
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(RestoreLog::default()),
Err(e) => Err(SnapshotError::io(&path, e)),
}
}
fn all_restore_logs(&self) -> Result<RestoreLogs> {
let mut out = RestoreLogs::default();
let entries = fs::read_dir(&self.restores_root)
.map_err(|e| SnapshotError::io(&self.restores_root, e))?;
for entry in entries {
let entry = entry.map_err(|e| SnapshotError::io(&self.restores_root, e))?;
if !entry.file_name().to_string_lossy().ends_with(".json") {
continue;
}
match fs::read(entry.path()).map(|b| serde_json::from_slice::<RestoreLog>(&b)) {
Ok(Ok(log)) if log.version == crate::workspace::FORMAT_VERSION => {
out.logs.push(log);
}
_ => out.incomplete = true,
}
}
Ok(out)
}
pub fn retain_turns(&self, live_turn_ids: &BTreeSet<String>) -> Result<()> {
let entries =
fs::read_dir(&self.turns_root).map_err(|e| SnapshotError::io(&self.turns_root, e))?;
for entry in entries {
let entry = entry.map_err(|e| SnapshotError::io(&self.turns_root, e))?;
let name = entry.file_name().to_string_lossy().into_owned();
if !name.ends_with(TURN_SUFFIX)
|| live_turn_ids.contains(&name)
|| !crate::sweep::settled(&entry.path())
{
continue;
}
fs::remove_file(entry.path()).map_err(|e| SnapshotError::io(entry.path(), e))?;
}
Ok(())
}
fn turn_path(&self, turn_id: &str) -> Result<PathBuf> {
crate::id::validate_stored("turn id", turn_id)?;
Ok(self.turns_root.join(format!("{turn_id}{TURN_SUFFIX}")))
}
fn restore_path(&self, thread_id: &str) -> Result<PathBuf> {
crate::id::validate_stored("session id", thread_id)?;
Ok(self.restores_root.join(format!("{thread_id}.json")))
}
}
fn write_atomic(path: &Path, bytes: &[u8]) -> Result<()> {
let tmp = crate::sweep::tmp_name(path);
fs::write(&tmp, bytes).map_err(|e| SnapshotError::io(&tmp, e))?;
fs::rename(&tmp, path).map_err(|e| SnapshotError::io(path, e))
}
pub(crate) fn turn_file_name(turn_id: &str) -> String {
format!("{turn_id}{TURN_SUFFIX}")
}
#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct GcStats {
pub manifests_kept: usize,
pub manifests_removed: usize,
pub blobs_kept: usize,
pub blobs_removed: usize,
}
impl GcStats {
pub(crate) fn plus(self, other: Self) -> Self {
Self {
manifests_kept: self.manifests_kept + other.manifests_kept,
manifests_removed: self.manifests_removed + other.manifests_removed,
blobs_kept: self.blobs_kept + other.blobs_kept,
blobs_removed: self.blobs_removed + other.blobs_removed,
}
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used)]
use super::*;
use pretty_assertions::assert_eq;
#[test]
fn append_and_load_roundtrip() {
let dir = tempfile::tempdir().unwrap();
let refs = RefStore::open(dir.path().join("refs")).unwrap();
assert_eq!(refs.load("t1").unwrap(), ThreadLog::default());
refs.append("t1", "turn-1".into(), "m1".into()).unwrap();
refs.append("t1", "turn-2".into(), "m2".into()).unwrap();
let log = refs.load("t1").unwrap();
assert_eq!(log.entries.len(), 2);
assert_eq!(log.entries[1].manifest_id, "m2");
}
#[test]
fn an_unreadable_log_makes_the_enumeration_incomplete_not_fatal() {
let dir = tempfile::tempdir().unwrap();
let root = dir.path().join("refs");
let refs = RefStore::open(&root).unwrap();
refs.append("good", "turn-1".into(), "m1".into()).unwrap();
fs::write(root.join("bad.json"), b"{ truncated").unwrap();
let logs = refs.thread_logs().unwrap();
assert_eq!(logs.logs.len(), 1, "the readable one still comes back");
assert!(logs.incomplete);
}
#[test]
fn the_chain_starts_at_the_first_entry_and_never_breaks() {
let dir = tempfile::tempdir().unwrap();
let refs = RefStore::open(dir.path()).unwrap();
for i in 0..4 {
refs.append("t1", format!("turn-{i}"), format!("manifest-{i}"))
.unwrap();
}
let log = refs.load("t1").unwrap();
assert_eq!(log.entries.len(), 4);
assert_eq!(
log.entries[0].prev_hash, None,
"the first entry has nothing behind it"
);
for pair in log.entries.windows(2) {
assert_eq!(
pair[1].prev_hash.as_deref(),
Some(SnapshotRef::digest(&pair[0]).as_str()),
"each entry names the one before it"
);
}
}
#[test]
fn entries_carry_the_time_they_were_appended() {
let dir = tempfile::tempdir().unwrap();
let refs = RefStore::open(dir.path()).unwrap();
let before = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap()
.as_secs();
refs.append("t1", "turn-1".into(), "m1".into()).unwrap();
let at = refs.load("t1").unwrap().entries[0].at;
assert!(at >= before, "recorded at least when we started");
}
#[test]
fn inherited_entries_are_rechained() {
let dir = tempfile::tempdir().unwrap();
let refs = RefStore::open(dir.path()).unwrap();
refs.append("src", "turn-1".into(), "m1".into()).unwrap();
refs.append("src", "turn-2".into(), "m2".into()).unwrap();
let source = refs.load("src").unwrap();
for entry in &source.entries {
refs.append("fork", entry.turn_id.clone(), entry.manifest_id.clone())
.unwrap();
}
let fork = refs.load("fork").unwrap();
assert_eq!(fork.entries[0].prev_hash, None);
assert_eq!(
fork.entries[1].prev_hash.as_deref(),
Some(SnapshotRef::digest(&fork.entries[0]).as_str())
);
}
#[test]
fn reading_the_top_of_the_undo_stack_does_not_consume_it() {
let dir = tempfile::tempdir().unwrap();
let turns = TurnIndex::open(dir.path()).unwrap();
let record = |n: &str| RestoreRecord {
target_manifest_id: format!("target-{n}"),
safety_manifest_id: format!("safety-{n}"),
};
turns.push_restore("t1", record("a")).unwrap();
turns.push_restore("t1", record("b")).unwrap();
assert_eq!(turns.last_restore("t1").unwrap(), Some(record("b")));
assert_eq!(
turns.last_restore("t1").unwrap(),
Some(record("b")),
"reading twice reads the same thing"
);
assert_eq!(turns.pop_restore("t1").unwrap(), Some(record("b")));
assert_eq!(
turns.last_restore("t1").unwrap(),
Some(record("a")),
"a second undo walks back another rewind rather than oscillating"
);
assert_eq!(turns.pop_restore("t1").unwrap(), Some(record("a")));
assert_eq!(turns.pop_restore("t1").unwrap(), None);
}
#[test]
fn a_record_from_an_unknown_build_is_refused() {
let dir = tempfile::tempdir().unwrap();
let root = dir.path().join("refs");
let refs = RefStore::open(&root).unwrap();
refs.append("t1", "turn-1".into(), "m1".into()).unwrap();
let raw = fs::read_to_string(root.join("t1.json")).unwrap();
fs::write(
root.join("t1.json"),
raw.replace("\"version\": 1", "\"version\": 99"),
)
.unwrap();
let err = refs.load("t1").unwrap_err();
assert!(
matches!(
&err,
SnapshotError::UnknownRecordVersion { kind, found: 99, .. } if *kind == "thread log"
),
"{err:?}"
);
assert!(refs.thread_logs().unwrap().incomplete);
}
#[test]
fn an_undo_record_from_an_unknown_build_is_refused() {
let dir = tempfile::tempdir().unwrap();
let turns = TurnIndex::open(dir.path()).unwrap();
turns
.push_restore(
"t1",
RestoreRecord {
target_manifest_id: "target".into(),
safety_manifest_id: "safety".into(),
},
)
.unwrap();
let path = dir.path().join("restores").join("t1.json");
let raw = fs::read_to_string(&path).unwrap();
fs::write(&path, raw.replace("\"version\": 1", "\"version\": 99")).unwrap();
let err = turns.last_restore("t1").unwrap_err();
assert!(
matches!(
&err,
SnapshotError::UnknownRecordVersion { kind, found: 99, .. } if *kind == "restore log"
),
"{err:?}"
);
}
#[test]
fn a_forged_turn_record_name_cannot_escape_the_turns_directory() {
let dir = tempfile::tempdir().unwrap();
let turns = TurnIndex::open(dir.path()).unwrap();
let outside = dir.path().join("witness.txt");
fs::write(&outside, b"not ours to remove").unwrap();
for forged in [
"../witness.txt",
"../../etc/passwd.turn",
"..",
"",
"no-suffix",
] {
let err = turns.remove_turn_file(forged).unwrap_err();
assert!(
matches!(err, SnapshotError::InvalidId { .. }),
"{forged:?} was not refused: {err:?}"
);
}
assert!(
outside.exists(),
"a forged name reached outside the partition"
);
turns.set_turn("turn-1", "m1").unwrap();
turns.remove_turn_file(&turn_file_name("turn-1")).unwrap();
assert_eq!(turns.manifest_for_turn("turn-1").unwrap(), None);
}
}