use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex as StdMutex, OnceLock, Weak};
use crate::{StorageError, StorageResult};
pub const GENERATION_SEP: &str = "__g";
pub fn logical_name(name: &str) -> &str {
if let Some(pos) = name.rfind(GENERATION_SEP) {
let suffix = &name[pos + GENERATION_SEP.len()..];
if !suffix.is_empty() && suffix.bytes().all(|b| b.is_ascii_digit()) {
return &name[..pos];
}
}
name
}
pub fn generation_name(logical: &str, generation: u64) -> String {
format!("{logical}{GENERATION_SEP}{generation}")
}
pub fn parse_generation(name: &str) -> Option<u64> {
let pos = name.rfind(GENERATION_SEP)?;
let suffix = &name[pos + GENERATION_SEP.len()..];
if suffix.is_empty() || !suffix.bytes().all(|b| b.is_ascii_digit()) {
return None;
}
suffix.parse().ok()
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GenerationInfo {
pub generation: u64,
pub metadata_path: PathBuf,
}
pub async fn list_generations(base: &Path, logical: &str) -> StorageResult<Vec<GenerationInfo>> {
let mut out = Vec::new();
for (generation, metadata_path) in scan_generation_paths(base, logical).await? {
if metadata_path.is_file() {
out.push(GenerationInfo {
generation,
metadata_path,
});
}
}
out.sort_by_key(|g| g.generation);
Ok(out)
}
pub async fn list_artifact_generations(base: &Path, logical: &str) -> StorageResult<Vec<u64>> {
let mut gens = std::collections::BTreeSet::new();
let mut rd = tokio::fs::read_dir(base)
.await
.map_err(|e| StorageError::Io(e.to_string()))?;
let artifact_prefix = format!("{logical}{GENERATION_SEP}");
while let Some(entry) = rd
.next_entry()
.await
.map_err(|e| StorageError::Io(e.to_string()))?
{
let name = entry.file_name();
let name = name.to_string_lossy();
let Some(rest) = name.strip_prefix(&artifact_prefix) else {
continue;
};
let Some(num) = rest.split('_').next() else {
continue;
};
if !num.is_empty() && num.bytes().all(|b| b.is_ascii_digit()) {
gens.insert(num.parse::<u64>().unwrap_or(u64::MAX));
}
}
Ok(gens.into_iter().collect())
}
pub async fn delete_generation(base: &Path, logical: &str, generation: u64) -> StorageResult<()> {
let prefix = format!("{logical}{GENERATION_SEP}{generation}_");
let metadata_path = base.join(format!("{prefix}metadata.json"));
let state = crate::commit::weak_lookup(
&GENERATION_STATES,
metadata_path.to_string_lossy().to_string(),
Arc::new(StdMutex::new(GenerationState::default())),
);
sweep_gate(SWEEP_GATE_PRE_LOCK);
let base = base.to_path_buf();
let logical = logical.to_string();
tokio::task::spawn_blocking(move || -> StorageResult<()> {
let state = state
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
if state.pins > 0 {
return Err(StorageError::InvalidState(format!(
"generation {generation} of '{logical}' is pinned by an in-flight reader; \
sweep refused until all readers release"
)));
}
sweep_gate(SWEEP_GATE_POST_CHECK);
let rd = std::fs::read_dir(&base)
.map_err(|e| StorageError::Io(format!("read {base:?}: {e}")))?;
let mut removed = Vec::new();
for entry in rd {
let entry = entry.map_err(|e| StorageError::Io(e.to_string()))?;
if !entry.file_name().to_string_lossy().starts_with(&prefix) {
continue;
}
removed.push(entry.path());
}
for path in removed {
if path.is_dir() {
std::fs::remove_dir_all(&path)
.map_err(|e| StorageError::Io(format!("remove {path:?}: {e}")))?;
} else {
std::fs::remove_file(&path)
.map_err(|e| StorageError::Io(format!("remove {path:?}: {e}")))?;
}
}
Ok(())
})
.await
.map_err(|e| StorageError::Io(format!("generation sweep task failed: {e}")))?
}
#[derive(Debug, Default)]
struct GenerationState {
pins: usize,
}
static GENERATION_STATES: OnceLock<StdMutex<HashMap<String, Weak<StdMutex<GenerationState>>>>> =
OnceLock::new();
#[derive(Debug)]
pub struct GenerationGuard {
state: Arc<StdMutex<GenerationState>>,
}
impl Drop for GenerationGuard {
fn drop(&mut self) {
let mut state = self
.state
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
state.pins -= 1;
}
}
pub fn pin_generation(info: &GenerationInfo) -> StorageResult<GenerationGuard> {
let state = crate::commit::weak_lookup(
&GENERATION_STATES,
info.metadata_path.to_string_lossy().to_string(),
Arc::new(StdMutex::new(GenerationState::default())),
);
{
let mut guard = state
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
if !info.metadata_path.is_file() {
return Err(StorageError::Invalid(format!(
"generation {} is not committed (no metadata at {:?}); \
only committed generations can be pinned",
info.generation, info.metadata_path
)));
}
guard.pins += 1;
}
Ok(GenerationGuard { state })
}
pub fn write_json_atomic(path: &Path, contents: &str) -> StorageResult<()> {
use std::io::Write;
let tmp = path.with_extension(format!("json.tmp.{}", uuid::Uuid::new_v4().simple()));
let parent = path.parent().unwrap_or_else(|| Path::new("."));
let parent = if parent.as_os_str().is_empty() {
Path::new(".")
} else {
parent
};
if !parent.exists() {
std::fs::create_dir_all(parent).map_err(|e| StorageError::Io(e.to_string()))?;
if let Some(grandparent) = parent.parent()
&& !grandparent.as_os_str().is_empty()
{
fsync_dir(grandparent)?;
}
fsync_dir(parent)?;
}
{
let mut f = std::fs::File::create(&tmp).map_err(|e| StorageError::Io(e.to_string()))?;
f.write_all(contents.as_bytes())
.map_err(|e| StorageError::Io(e.to_string()))?;
f.sync_all().map_err(|e| StorageError::Io(e.to_string()))?;
}
std::fs::rename(&tmp, path).map_err(|e| StorageError::Io(e.to_string()))?;
fsync_dir(parent)
}
pub(crate) fn fsync_dir(dir: &Path) -> StorageResult<()> {
#[cfg(unix)]
{
let f = std::fs::File::open(dir)
.map_err(|e| StorageError::Io(format!("open dir {dir:?} for fsync: {e}")))?;
f.sync_all()
.map_err(|e| StorageError::Io(format!("fsync dir {dir:?}: {e}")))?;
}
#[cfg(not(unix))]
let _ = dir;
Ok(())
}
async fn scan_generation_paths(base: &Path, logical: &str) -> StorageResult<Vec<(u64, PathBuf)>> {
let mut out = Vec::new();
let mut rd = tokio::fs::read_dir(base)
.await
.map_err(|e| StorageError::Io(e.to_string()))?;
let prefix = format!("{logical}{GENERATION_SEP}");
while let Some(entry) = rd
.next_entry()
.await
.map_err(|e| StorageError::Io(e.to_string()))?
{
let name = entry.file_name();
let name = name.to_string_lossy().into_owned();
let Some(rest) = name.strip_prefix(&prefix) else {
continue;
};
let Some(meta) = rest.strip_suffix("_metadata.json") else {
continue;
};
if !meta.is_empty() && meta.bytes().all(|b| b.is_ascii_digit()) {
out.push((meta.parse().unwrap_or(u64::MAX), entry.path()));
}
}
Ok(out)
}
pub(crate) const SWEEP_GATE_PRE_LOCK: usize = 0;
pub(crate) const SWEEP_GATE_POST_CHECK: usize = 1;
#[cfg(not(test))]
fn sweep_gate(_stage: usize) {}
#[cfg(test)]
fn sweep_gates() -> &'static StdMutex<[Option<SweepGate>; 2]> {
static GATES: OnceLock<StdMutex<[Option<SweepGate>; 2]>> = OnceLock::new();
GATES.get_or_init(|| StdMutex::new([None, None]))
}
#[cfg(test)]
struct SweepGate {
arrived: std::sync::mpsc::Sender<()>,
release: std::sync::mpsc::Receiver<()>,
}
#[cfg(test)]
pub(crate) fn arm_sweep_gate(
stage: usize,
) -> (std::sync::mpsc::Receiver<()>, std::sync::mpsc::Sender<()>) {
let (arrived_tx, arrived_rx) = std::sync::mpsc::channel();
let (release_tx, release_rx) = std::sync::mpsc::channel();
let mut gates = sweep_gates()
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
gates[stage] = Some(SweepGate {
arrived: arrived_tx,
release: release_rx,
});
(arrived_rx, release_tx)
}
#[cfg(test)]
fn sweep_gate(stage: usize) {
let gate = sweep_gates()
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.get_mut(stage)
.unwrap()
.take();
if let Some(gate) = gate {
let _ = gate.arrived.send(());
let _ = gate.release.recv();
}
}