OBIKmers/obikmer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 34 Wiki Activity

Introduce session persistence and artifact caching

Browse Source 
This change implements a robust session management system allowing users to persist distance calculation samples and parameters across invocations. It includes concurrency-safe directory locking, deterministic parameter encoding, artifact integrity checks, and caching mechanisms to skip resampling when valid data is available.
This commit is contained in:
Eric Coissac
2026-09-12 07:46:55 +02:00
parent a8bcf5ffac
commit 7861e886e7
15 changed files with 746 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/Cargo.lock
Generated
+146
View File
@@ -213,6 +213,29 @@ version = "3.20.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb"
[[package]]
name = "bytecheck"
version = "0.8.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "26333eeac754f0ad8a6bcd0eb0ac012156302e4e16b852b72ee399aea4f12c29"
dependencies = [
"bytecheck_derive",
"ptr_meta",
"rancor",
"simdutf8",
]
[[package]]
name = "bytecheck_derive"
version = "0.8.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "46d07918caa9eeaaf06b7873925c53a61daac173539b4f7715090745e44e4e69"
dependencies = [
"proc-macro2",
"quote",
"syn 3.0.3",
]
[[package]]
name = "byteorder"
version = "1.5.0"
@@ -1236,6 +1259,26 @@ dependencies = [
"windows 0.48.0",
]
[[package]]
name = "munge"
version = "0.4.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5e17401f259eba956ca16491461b6e8f72913a0a114e39736ce404410f915a0c"
dependencies = [
"munge_macro",
]
[[package]]
name = "munge_macro"
version = "0.4.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4568f25ccbd45ab5d5603dc34318c1ec56b117531781260002151b8530a9f931"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.117",
]
[[package]]
name = "nanorand"
version = "0.6.1"
@@ -1499,6 +1542,7 @@ dependencies = [
"obikrope",
"obikselect",
"obikseq",
"obiksession",
"obikstats",
"obipipeline",
"obiread",
@@ -1541,6 +1585,7 @@ dependencies = [
"obikindex",
"obikindexer",
"obikseq",
"obiksession",
"obipipeline",
"obiread",
"obiskbuilder",
@@ -1549,6 +1594,7 @@ dependencies = [
"petgraph",
"rand 0.10.2",
"rayon",
"rkyv",
"speedytree",
"tempfile",
"tracing",
@@ -1615,6 +1661,17 @@ dependencies = [
"xxhash-rust",
]
[[package]]
name = "obiksession"
version = "0.1.0"
dependencies = [
"memmap2",
"obisys",
"tempfile",
"tracing",
"xxhash-rust",
]
[[package]]
name = "obikstats"
version = "0.1.0"
@@ -1963,6 +2020,26 @@ dependencies = [
"xxhash-rust",
]
[[package]]
name = "ptr_meta"
version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "743da816b98c921cdbe8628ef7381b76f25ecf4da599fc80aca90eae7ef70cc0"
dependencies = [
"ptr_meta_derive",
]
[[package]]
name = "ptr_meta_derive"
version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1c8d9ca532f185d5d4db7a7c9d51420b452168ea1c2b913953281bd6fe1fcbd0"
dependencies = [
"proc-macro2",
"quote",
"syn 3.0.3",
]
[[package]]
name = "quote"
version = "1.0.45"
@@ -1990,6 +2067,15 @@ version = "0.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09"
[[package]]
name = "rancor"
version = "0.1.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9b534442d0fcdb55d66f373d9cac6d33b6293a2335bc2136dbd06ce0e87d2572"
dependencies = [
"ptr_meta",
]
[[package]]
name = "rand"
version = "0.8.6"
@@ -2151,6 +2237,15 @@ version = "0.8.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4"
[[package]]
name = "rend"
version = "0.5.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "663ba70707f96e871406fe10d68128412e619b06d1d47cb91c3a4c6501176240"
dependencies = [
"bytecheck",
]
[[package]]
name = "ring"
version = "0.17.14"
@@ -2165,6 +2260,36 @@ dependencies = [
"windows-sys 0.52.0",
]
[[package]]
name = "rkyv"
version = "0.8.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d9776093b7ca170454ab1406954f7b7d97a57c51dc6c0642957fb2ef25c2d399"
dependencies = [
"bytecheck",
"bytes",
"hashbrown",
"indexmap",
"munge",
"ptr_meta",
"rancor",
"rend",
"rkyv_derive",
"tinyvec",
"uuid",
]
[[package]]
name = "rkyv_derive"
version = "0.8.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1c25ef604ac7dd839d44d64648952ea23c97866f124ff671b0ed2cf3ad9bb06e"
dependencies = [
"proc-macro2",
"quote",
"syn 3.0.3",
]
[[package]]
name = "rustc-hash"
version = "2.1.2"
@@ -2359,6 +2484,12 @@ version = "0.3.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214"
[[package]]
name = "simdutf8"
version = "0.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e"
[[package]]
name = "smallvec"
version = "1.15.1"
@@ -2587,6 +2718,21 @@ dependencies = [
"serde_json",
]
[[package]]
name = "tinyvec"
version = "1.13.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4cf0ded5c4e56918d8f8a339e1bb67d038d3bc6d144ac407904015ba2e4cde9b"
dependencies = [
"tinyvec_macros",
]
[[package]]
name = "tinyvec_macros"
version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
[[package]]
name = "tracing"
version = "0.1.44"
src/Cargo.toml
+1 -1
View File
@@ -1,5 +1,5 @@
[workspace]
resolver = "3"
members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obiskio","obidebruinj", "obicompactvec", "obisys", "obikindex", "obikindexer", "obikquery", "obikdump", "obikfilter", "obikselect", "obikrebuild", "obikmerge", "obikstats", "obikidxcache", "obitaxonomy", "obikentropy", "obikphylo", "obikalgorithm"]
members = ["obikseq", "obiread", "obiskbuilder", "obifastwrite", "obikmer","obikrope","obipipeline", "obiskio","obidebruinj", "obicompactvec", "obisys", "obikindex", "obikindexer", "obikquery", "obikdump", "obikfilter", "obikselect", "obikrebuild", "obikmerge", "obikstats", "obikidxcache", "obitaxonomy", "obikentropy", "obikphylo", "obikalgorithm", "obiksession"]
[profile.release]
debug = 1
src/obikmer/Cargo.toml
+1
View File
@@ -24,6 +24,7 @@ obikstats = { path = "../obikstats" }
obikquery = { path = "../obikquery" }
obikidxcache = { path = "../obikidxcache" }
obikphylo = { path = "../obikphylo" }
obiksession = { path = "../obiksession" }
obikrope = { path = "../obikrope" }
obifastwrite = { path = "../obifastwrite" }
obiskbuilder = { path = "../obiskbuilder" }
src/obikmer/src/cmd/phylo/args.rs
+21
View File
@@ -224,6 +224,27 @@ pub struct PhyloArgs {
#[arg(long)]
pub entropy_sd: Option<f64>,
/// Persist the `snp-*` `--distance` sample (the pairwise substitution
/// tally) in `DIR` across separate `obikmer phylo` invocations: a later
/// run against the same index with the exact same
/// `--subsample`/`--free-loss`/`--no-ambiguity`/`--exclude-genome`/
/// `--min-shared-family`/`--entropy`/`--entropy-sd` restores it instead
/// of resampling. `DIR` is created if it doesn't exist. If it already
/// holds a sample built under *different* selection parameters, this
/// run is rejected with an error listing what changed, unless
/// `--session-force` is also given (which discards the old sample and
/// starts fresh under this run's parameters). Only covers a `snp-*`
/// `--distance` today — `--sankoff`/`--tnt`/`--phyg`/`--iqtree` don't
/// read or write a `--session` yet.
#[arg(long, value_name = "DIR")]
pub session: Option<PathBuf>,
/// With `--session DIR`: overwrite its saved selection parameters (and
/// discard its cached sample) instead of erroring out when this run's
/// parameters don't match. No effect without `--session`.
#[arg(long, requires = "session")]
pub session_force: bool,
/// Calibrate a 16-state Sankoff cost matrix (and its matching
/// pseudo-alignment) from an already-built sibling annex — requires
/// `--subsample <N>`, and shares `--free-loss`/`--no-ambiguity`/
src/obikmer/src/cmd/phylo/mod.rs
+79
View File
@@ -24,6 +24,43 @@ use phylip::write_phylip_relaxed;
use sankoff::{write_sankoff_alignment_fasta, write_sankoff_matrix_csv, write_sankoff_params};
use tnt::write_sankoff_tnt;
/// Deterministic byte encoding of the `snp-*` `--distance` selection
/// parameters, for `obiksession::Session::open`'s params-conflict check.
/// `obiksession` treats this as an opaque blob (see its own crate docs) —
/// only this module needs to know the layout, and only well enough to
/// produce/compare it, never to pretty-print a per-field diff (v1 keeps
/// the conflict message generic rather than decoding both sides field by
/// field — see its use in `run` below).
fn session_params_blob(
n: Option<usize>,
free_loss: bool,
no_ambiguity: bool,
excluded: &[bool],
entropy_bias: Option<EntropyBias>,
) -> Vec<u8> {
let mut buf = Vec::new();
match n {
Some(v) => {
buf.push(1);
buf.extend_from_slice(&(v as u64).to_le_bytes());
}
None => buf.push(0),
}
buf.push(free_loss as u8);
buf.push(no_ambiguity as u8);
buf.extend_from_slice(&(excluded.len() as u64).to_le_bytes());
buf.extend(excluded.iter().map(|&b| b as u8));
match entropy_bias {
Some(EntropyBias { mu, sigma }) => {
buf.push(1);
buf.extend_from_slice(&mu.to_le_bytes());
buf.extend_from_slice(&sigma.to_le_bytes());
}
None => buf.push(0),
}
buf
}
pub use args::PhyloArgs;
pub fn run(args: PhyloArgs) {
@@ -236,6 +273,47 @@ pub fn run(args: PhyloArgs) {
None
};
// ── `--session`: persist the snp-* sample across separate invocations ──
// Only covers `snp-*` `--distance` today (see `--session`'s own docs);
// `None` when `--session` isn't given, matching today's always-resample
// default exactly.
let session = args.session.as_ref().map(|dir| {
let params = session_params_blob(
args.subsample,
args.free_loss,
args.no_ambiguity,
&snp_exclude_mask,
entropy_bias,
);
match obiksession::Session::open(dir, &params, args.session_force) {
Ok(Ok(session)) => {
match session.outcome() {
obiksession::SessionOutcome::Created => {
info!("--session {}: starting fresh", dir.display());
}
obiksession::SessionOutcome::Reused => {
info!("--session {}: found, will reuse its cached sample if present", dir.display());
}
}
session
}
Ok(Err(_conflict)) => {
eprintln!(
"error: --session {} was built with different selection parameters \
(--subsample/--free-loss/--no-ambiguity/--exclude-genome/--min-shared-family/\
--entropy/--entropy-sd) than this run — pass --session-force to discard its \
cached sample and start fresh, or point --session at a different directory",
dir.display()
);
std::process::exit(1);
}
Err(e) => {
eprintln!("error opening --session {}: {e}", dir.display());
std::process::exit(1);
}
}
});
// ── SNP pseudo-alignment (`--pseudo-alignment`) ─────────────────────────────
if args.pseudo_alignment {
let Some(subsample_n) = args.subsample else {
@@ -389,6 +467,7 @@ pub fn run(args: PhyloArgs) {
&snp_exclude_mask,
entropy_bias,
gamma_shape,
session.as_ref(),
)
}
}
src/obikphylo/Cargo.toml
+2
View File
@@ -12,12 +12,14 @@ obicompactvec = { path = "../obicompactvec" }
obikidxcache = { path = "../obikidxcache" }
obiskbuilder = { path = "../obiskbuilder" }
obipipeline = { path = "../obipipeline" }
obiksession = { path = "../obiksession" }
kodama = "0.3.0"
memmap2 = "0.9"
ndarray = "0.17"
petgraph = "0.6.4"
rand = "0.10"
rayon = "1"
rkyv = "0.8.18"
speedytree = "0.1"
tracing = "0.1.44"
src/obikphylo/src/siblings/algorithms/pairwise.rs
+15 -2
View File
@@ -18,6 +18,7 @@
//! machines), not something to optimise away.
use ndarray::Array2;
use rkyv::{Archive, Deserialize, Serialize};
use super::subsample::SurvivingFamily;
@@ -30,7 +31,12 @@ use super::subsample::SurvivingFamily;
/// [`BasePairTally`] the latter (they are *not* the same population — see
/// `BasePairTally::same`'s own docs). `cardinality` is restricted to
/// variable families only, matching [`CardinalityTally`]'s own scope.
#[derive(Default, Clone, Copy)]
///
/// `Archive`/`Serialize`/`Deserialize`: part of [`PairwiseTally`]'s own
/// `--session` persistence (see that struct's docs) — `rkyv` needs every
/// field type to derive these too, hence deriving here on an otherwise
/// purely internal struct.
#[derive(Default, Clone, Copy, Archive, Serialize, Deserialize)]
struct PairStats {
subst: [[u64; 4]; 4],
same_all: [u64; 4],
@@ -40,6 +46,13 @@ struct PairStats {
/// See the module docs. `pub(crate)`: consumed by
/// `algorithms::sankoff`'s orchestration, not yet exposed past this crate.
///
/// `Archive`/`Serialize`/`Deserialize` (`rkyv`): lets `--session` dump/
/// restore a tally across separate CLI invocations as opaque bytes handed
/// to `obiksession::Session::store`/`restore` — `obiksession` itself knows
/// nothing about this type, it only stores/checksums/mmaps the bytes this
/// crate produces (see `obiksession`'s own crate docs on that split).
#[derive(Archive, Serialize, Deserialize)]
pub(crate) struct PairwiseTally {
n_genomes: usize,
/// Upper triangle only (`i < j`), flat-indexed via [`Self::flat_index`].
@@ -300,7 +313,7 @@ pub struct CardinalityTally {
/// applied here either, matching `reduce_pairwise`'s own raw per-pair fold
/// (`PairStats` accumulates over every genome pair unconditionally;
/// exclusion is a filter applied only in derived, post-hoc views).
#[derive(Default, Clone)]
#[derive(Default, Clone, Archive, Serialize, Deserialize)]
pub(crate) struct PartitionDispersion {
/// One entry per partition id seen so far (grows lazily — partitions
/// are visited in increasing order in practice, but nothing here
src/obikphylo/src/siblings/algorithms/snp_distance.rs
+86 -24
View File
@@ -23,12 +23,51 @@
use ndarray::Array2;
use obikidxcache::index_cache::IndexCache;
use obikindex::OKIResult;
use obikindex::{OKIError, OKIResult};
use obiksession::Session;
use super::pairwise::{PairwiseTally, PartitionDispersion};
use super::sibling_family_size_histogram;
use super::subsample::{EntropyBias, sample_index};
const TALLY_ARTIFACT: &str = "pairwise_tally";
const DISPERSION_ARTIFACT: &str = "partition_dispersion";
/// Restore a `PairwiseTally`/`PartitionDispersion` pair from `session`, if
/// both artifacts are cached there — `None` on any miss (not cached, I/O
/// error, or a deserialization failure), each treated identically:
/// something's missing or unusable, fall back to recomputing rather than
/// surfacing a hard error over what's meant to be a transparent cache.
///
/// Full owned deserialization (`rkyv::from_bytes`), not zero-copy
/// (`rkyv::access`) — simpler to integrate with `PairwiseTally`'s existing
/// query methods, which all take `&self` by value semantics, not an
/// `Archived<PairwiseTally>`. Still skips the expensive part (`sample_index`
/// re-scanning the sibling annex); the `O(n²)` deserialize of an already-
/// in-memory-sized byte buffer is comparatively cheap. Making this
/// genuinely zero-copy later would mean every `PairwiseTally` accessor
/// (`categories`, `base_freq`, ...) working generically over
/// `Archived<PairwiseTally>` too — a bigger, separate change.
fn restore_tally(session: &Session) -> Option<(PairwiseTally, PartitionDispersion)> {
let tally_bytes = session.restore(TALLY_ARTIFACT).ok().flatten()?;
let dispersion_bytes = session.restore(DISPERSION_ARTIFACT).ok().flatten()?;
let tally = rkyv::from_bytes::<PairwiseTally, rkyv::rancor::Error>(&tally_bytes[..]).ok()?;
let dispersion =
rkyv::from_bytes::<PartitionDispersion, rkyv::rancor::Error>(&dispersion_bytes[..]).ok()?;
tracing::info!("--session: restored cached SNP tally — skipping resampling");
Some((tally, dispersion))
}
fn store_tally(session: &Session, tally: &PairwiseTally, dispersion: &PartitionDispersion) -> OKIResult<()> {
let tally_bytes = rkyv::to_bytes::<rkyv::rancor::Error>(tally)
.map_err(|e| OKIError::Io(std::io::Error::other(e.to_string())))?;
session.store(TALLY_ARTIFACT, &tally_bytes).map_err(OKIError::Io)?;
let dispersion_bytes = rkyv::to_bytes::<rkyv::rancor::Error>(dispersion)
.map_err(|e| OKIError::Io(std::io::Error::other(e.to_string())))?;
session.store(DISPERSION_ARTIFACT, &dispersion_bytes).map_err(OKIError::Io)?;
Ok(())
}
/// One `snp-*` `--distance` value. `pub`: part of
/// [`crate::siblings::extensions::SiblingExt::snp_distance`]'s public
/// signature.
@@ -231,6 +270,17 @@ fn formula(kind: SnpDistanceKind) -> fn(&PairwiseTally, usize, usize, Option<f64
/// estimation finds no measurable signal (see its own docs for when), the
/// correction is silently left off rather than applied with a fabricated
/// value — logged either way.
///
/// `session`: `--session`, `None` means no persistence (today's default
/// behavior — always resample). `Some(session)` is expected to already be
/// open under the *exact* selection parameters this call is about to use
/// (`n`/`free_loss`/`no_ambiguity`/`excluded`/`entropy_bias`) — checking
/// that is the caller's job (`obiksession::Session::open`'s own
/// params-conflict mechanism), not re-validated here. If a cached tally is
/// found it's restored and `sample_index` is skipped entirely; otherwise
/// the fresh tally is stored back into the session before this returns, so
/// a later call under the same session/params restores instead of
/// resampling.
pub(crate) fn snp_distance(
cache: &IndexCache,
kind: SnpDistanceKind,
@@ -240,6 +290,7 @@ pub(crate) fn snp_distance(
excluded: &[bool],
entropy_bias: Option<EntropyBias>,
gamma_shape: GammaShape,
session: Option<&Session>,
) -> OKIResult<Array2<f64>> {
// Fail fast, before paying for `sample_index`, not just inside
// `distance_matrix` (which runs after sampling either way — the right
@@ -251,33 +302,44 @@ pub(crate) fn snp_distance(
));
}
let n_genomes = cache.meta().genomes().len();
let mut tally = PairwiseTally::new(n_genomes);
let mut dispersion = PartitionDispersion::default();
let (target, entropy_bias) = match n {
Some(target) => (target, entropy_bias),
let cached = session.and_then(restore_tally);
let (tally, dispersion) = match cached {
Some(pair) => pair,
None => {
let counts = sibling_family_size_histogram(cache)?;
let total_eligible = (counts[1] + counts[2] + counts[3]) as usize;
(total_eligible, None)
let n_genomes = cache.meta().genomes().len();
let mut tally = PairwiseTally::new(n_genomes);
let mut dispersion = PartitionDispersion::default();
let (target, entropy_bias) = match n {
Some(target) => (target, entropy_bias),
None => {
let counts = sibling_family_size_histogram(cache)?;
let total_eligible = (counts[1] + counts[2] + counts[3]) as usize;
(total_eligible, None)
}
};
if target > 0 {
sample_index(
cache,
target,
free_loss,
no_ambiguity,
excluded,
entropy_bias,
|partition, _layer, survivors| {
super::pairwise::reduce_pairwise(&survivors, partition, &mut tally, &mut dispersion);
},
)?;
}
if let Some(session) = session {
store_tally(session, &tally, &dispersion)?;
}
(tally, dispersion)
}
};
if target > 0 {
sample_index(
cache,
target,
free_loss,
no_ambiguity,
excluded,
entropy_bias,
|partition, _layer, survivors| {
super::pairwise::reduce_pairwise(&survivors, partition, &mut tally, &mut dispersion);
},
)?;
}
distance_matrix(&tally, &dispersion, kind, gamma_shape)
}
src/obikphylo/src/siblings/extensions/sibling_ext.rs
+10 -1
View File
@@ -147,6 +147,13 @@ pub trait SiblingExt {
/// ([`SnpDistanceKind::supports_gamma`]). `Auto` estimates `alpha` from
/// the data itself (method-of-moments over per-partition substitution
/// rate dispersion) rather than requiring a user-supplied value.
///
/// `session`: `--session`, `None` disables persistence (always
/// resample). `Some` must already be open under these exact selection
/// parameters (the caller's responsibility, via
/// `obiksession::Session::open`'s params-conflict check) — restores a
/// cached tally when present instead of resampling, and stores a
/// freshly computed one back for next time otherwise.
fn snp_distance(
&self,
kind: SnpDistanceKind,
@@ -156,6 +163,7 @@ pub trait SiblingExt {
excluded: &[bool],
entropy_bias: Option<EntropyBias>,
gamma_shape: GammaShape,
session: Option<&obiksession::Session>,
) -> OKIResult<Array2<f64>>;
/// The Family Overlap annex — number of shared *variable* families per
@@ -335,8 +343,9 @@ impl SiblingExt for IndexCache {
excluded: &[bool],
entropy_bias: Option<EntropyBias>,
gamma_shape: GammaShape,
session: Option<&obiksession::Session>,
) -> OKIResult<Array2<f64>> {
snp_distance(self, kind, n, free_loss, no_ambiguity, excluded, entropy_bias, gamma_shape)
snp_distance(self, kind, n, free_loss, no_ambiguity, excluded, entropy_bias, gamma_shape, session)
}
fn family_overlap(&self) -> OKIResult<FamilyOverlap> {
src/obiksession/Cargo.toml
+13
View File
@@ -0,0 +1,13 @@
[package]
name = "obiksession"
version = "0.1.0"
edition = "2024"
[dependencies]
obisys = { path = "../obisys", default-features = false }
memmap2 = "0.9.11"
xxhash-rust = { version = "0.8.18", features = ["xxh3"] }
tracing = "0.1"
[dev-dependencies]
tempfile = "3"
src/obiksession/src/lib.rs
+19
View File
@@ -0,0 +1,19 @@
//! Domain-agnostic session persistence: a named directory that survives
//! across separate CLI invocations, holding an opaque parameter blob (used
//! to detect stale/mismatched reuse) plus a checksummed cache of named
//! binary artifacts.
//!
//! Deliberately knows nothing about what the parameters or artifacts
//! *mean* — no k-mer/genome/tally concept anywhere in this crate. A
//! caller (e.g. `obikphylo`, for its `--session` support around
//! `PairwiseTally`/`PartitionDispersion`/`SnpAlignment`) serializes its own
//! parameter struct and artifacts to bytes however it likes (`rkyv`,
//! `bincode`, ...) and hands raw `&[u8]` to this crate; [`Session`] only
//! deals in byte blobs, directory lifecycle, and locking. Same split
//! already used elsewhere in this project between a generic mechanism
//! crate and the domain crate that plugs into it (`obisys`/`obikindex`,
//! `obicompactvec`/`obikindexer`).
mod session;
pub use session::{ParamsConflict, Session, SessionOutcome};
src/obiksession/src/session.rs
+249
View File
@@ -0,0 +1,249 @@
use std::fs;
use std::io;
use std::path::{Path, PathBuf};
use memmap2::Mmap;
use obisys::DirLock;
const PARAMS_FILE: &str = "params.bin";
const ARTIFACTS_DIR: &str = "artifacts";
const CHECKSUM_EXT: &str = "xxh3";
/// Whether [`Session::open`] found a fresh (just-created, or just-reset by
/// `force`) directory or reused one whose saved parameters matched the
/// requested ones exactly.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SessionOutcome {
/// New directory, or an existing one reset because its saved
/// parameters didn't match and `force: true` was passed. No cached
/// artifact in it can be trusted — there either isn't one yet, or the
/// ones that were there have just been deleted.
Created,
/// Existing directory whose saved parameters matched byte-for-byte.
/// Every artifact already cached under it was produced under these
/// same parameters and is safe to [`Session::restore`].
Reused,
}
/// Returned by [`Session::open`] when the directory already holds a
/// parameter blob that doesn't match the one just requested, and `force`
/// wasn't set. This crate has no idea what the bytes mean, so it can't
/// produce a useful diagnostic on its own — the caller, who serialized
/// both, is expected to deserialize `saved`/`requested` itself and report
/// exactly which field(s) differ before deciding whether to error out,
/// point at a different directory, or retry [`Session::open`] with
/// `force: true`.
#[derive(Debug)]
pub struct ParamsConflict {
pub saved: Vec<u8>,
pub requested: Vec<u8>,
}
/// An open session directory: exclusively locked (via [`obisys::DirLock`])
/// for as long as this value lives, so no other `obiksession`-using
/// process can read a half-written artifact or reset the directory out
/// from under this one.
#[derive(Debug)]
pub struct Session {
dir: PathBuf,
outcome: SessionOutcome,
_lock: DirLock,
}
impl Session {
/// Open (or create, or reset) a session directory — see
/// [`SessionOutcome`]/[`ParamsConflict`] for the three possible
/// outcomes. Blocks until the directory's lock is free (see
/// [`obisys::DirLock::acquire`]) — a second process pointed at the
/// same session directory waits rather than racing it.
pub fn open(dir: &Path, params: &[u8], force: bool) -> io::Result<Result<Session, ParamsConflict>> {
let lock = DirLock::acquire(dir)?;
let params_path = dir.join(PARAMS_FILE);
let saved = read_checked(&params_path)?;
match saved {
Some(saved) if saved == params => Ok(Ok(Session {
dir: dir.to_path_buf(),
outcome: SessionOutcome::Reused,
_lock: lock,
})),
Some(saved) if !force => Ok(Err(ParamsConflict { saved, requested: params.to_vec() })),
_ => {
// No saved params yet, or a mismatch with `force: true`:
// start clean. Any artifact cached under different params
// is stale by construction — deleting the whole
// `artifacts/` directory is simpler and safer than trying
// to figure out which entries are still valid.
let artifacts_dir = dir.join(ARTIFACTS_DIR);
if artifacts_dir.exists() {
fs::remove_dir_all(&artifacts_dir)?;
}
fs::create_dir_all(&artifacts_dir)?;
write_checked(&params_path, params)?;
Ok(Ok(Session {
dir: dir.to_path_buf(),
outcome: SessionOutcome::Created,
_lock: lock,
}))
}
}
}
pub fn dir(&self) -> &Path {
&self.dir
}
pub fn outcome(&self) -> SessionOutcome {
self.outcome
}
fn artifact_path(&self, name: &str) -> PathBuf {
self.dir.join(ARTIFACTS_DIR).join(name)
}
/// Restore a previously [`store`](Self::store)d artifact, `mmap`ed
/// read-only for zero-copy access by the caller (e.g. `rkyv::access`
/// straight over the returned bytes).
///
/// `Ok(None)` if `name` isn't cached yet, *or* if its checksum doesn't
/// match what was recorded at [`store`](Self::store) time — treated
/// the same as "not cached" rather than an error, so a partial dump
/// left behind by a killed process just triggers a recompute on the
/// next run instead of a hard failure.
pub fn restore(&self, name: &str) -> io::Result<Option<Mmap>> {
verify_mmap(&self.artifact_path(name))
}
/// Store `bytes` under `name`, plus a checksum sidecar, atomically:
/// both are written to a temporary file and renamed into place, so a
/// crash mid-write can never leave a corrupt file that a later
/// [`restore`](Self::restore) mistakes for valid (the checksum file is
/// only ever renamed into place *after* the data file's rename
/// succeeds, so the two can't disagree about which write completed).
pub fn store(&self, name: &str, bytes: &[u8]) -> io::Result<()> {
write_checked(&self.artifact_path(name), bytes)
}
}
fn checksum_path_for(path: &Path) -> PathBuf {
let mut os_str = path.as_os_str().to_owned();
os_str.push(".");
os_str.push(CHECKSUM_EXT);
PathBuf::from(os_str)
}
/// `mmap` `path` and verify it against its checksum sidecar. `Ok(None)`
/// when `path` doesn't exist, its checksum sidecar doesn't exist/parse, or
/// the checksum doesn't match — every case here means "nothing usable is
/// cached," never an error the caller has to handle specially.
fn verify_mmap(path: &Path) -> io::Result<Option<Mmap>> {
if !path.exists() {
return Ok(None);
}
let Ok(expected_str) = fs::read_to_string(checksum_path_for(path)) else {
return Ok(None);
};
let Ok(expected) = expected_str.trim().parse::<u64>() else {
return Ok(None);
};
let file = fs::File::open(path)?;
// SAFETY: the file is only ever written by `write_checked` (atomic
// temp-then-rename) and this whole directory is held under this
// process's exclusive `DirLock` for the `Session`'s entire lifetime —
// no other process can be concurrently modifying it underneath this
// mapping.
let mmap = unsafe { Mmap::map(&file)? };
let actual = xxhash_rust::xxh3::xxh3_64(&mmap[..]);
if actual != expected {
tracing::warn!(
path = %path.display(),
"session artifact failed checksum verification — treating as absent"
);
return Ok(None);
}
Ok(Some(mmap))
}
fn read_checked(path: &Path) -> io::Result<Option<Vec<u8>>> {
Ok(verify_mmap(path)?.map(|mmap| mmap[..].to_vec()))
}
fn write_checked(path: &Path, bytes: &[u8]) -> io::Result<()> {
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)?;
}
let tmp = path.with_extension("tmp");
fs::write(&tmp, bytes)?;
fs::rename(&tmp, path)?;
let checksum = xxhash_rust::xxh3::xxh3_64(bytes);
let checksum_path = checksum_path_for(path);
let checksum_tmp = checksum_path.with_extension("tmp");
fs::write(&checksum_tmp, checksum.to_string())?;
fs::rename(&checksum_tmp, checksum_path)?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn fresh_directory_is_created() {
let tmp = tempfile::tempdir().unwrap();
let dir = tmp.path().join("session");
let session = Session::open(&dir, b"params-v1", false).unwrap().unwrap();
assert_eq!(session.outcome(), SessionOutcome::Created);
assert!(session.restore("missing").unwrap().is_none());
}
#[test]
fn matching_params_reuse_and_restore_artifact() {
let tmp = tempfile::tempdir().unwrap();
let dir = tmp.path().join("session");
{
let session = Session::open(&dir, b"params-v1", false).unwrap().unwrap();
session.store("tally", b"some serialized bytes").unwrap();
}
let session = Session::open(&dir, b"params-v1", false).unwrap().unwrap();
assert_eq!(session.outcome(), SessionOutcome::Reused);
let restored = session.restore("tally").unwrap().unwrap();
assert_eq!(&restored[..], b"some serialized bytes");
}
#[test]
fn mismatched_params_without_force_conflicts() {
let tmp = tempfile::tempdir().unwrap();
let dir = tmp.path().join("session");
Session::open(&dir, b"params-v1", false).unwrap().unwrap();
let conflict = Session::open(&dir, b"params-v2", false).unwrap().unwrap_err();
assert_eq!(conflict.saved, b"params-v1");
assert_eq!(conflict.requested, b"params-v2");
}
#[test]
fn mismatched_params_with_force_resets_artifacts() {
let tmp = tempfile::tempdir().unwrap();
let dir = tmp.path().join("session");
{
let session = Session::open(&dir, b"params-v1", false).unwrap().unwrap();
session.store("tally", b"old data").unwrap();
}
let session = Session::open(&dir, b"params-v2", true).unwrap().unwrap();
assert_eq!(session.outcome(), SessionOutcome::Created);
assert!(session.restore("tally").unwrap().is_none());
}
#[test]
fn corrupted_artifact_is_treated_as_absent() {
let tmp = tempfile::tempdir().unwrap();
let dir = tmp.path().join("session");
let session = Session::open(&dir, b"params-v1", false).unwrap().unwrap();
session.store("tally", b"good data").unwrap();
// Corrupt the data file after the fact without touching the
// checksum sidecar — simulates a partial/garbled write.
std::fs::write(dir.join("artifacts").join("tally"), b"corrupted!").unwrap();
assert!(session.restore("tally").unwrap().is_none());
}
}
src/obisys/src/lock.rs
+1
View File
@@ -14,6 +14,7 @@ use tracing::info;
/// Windows) via `std::fs::File::lock`/`try_lock`, not a hand-rolled PID
/// file: the OS releases it automatically on process exit, including a
/// crash — no stale-lock cleanup logic needed.
#[derive(Debug)]
pub struct DirLock {
_file: std::fs::File,
}