From 0ce934b111381821006c0232ac375d0f56f598cc Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Fri, 14 Aug 2026 21:02:51 +0200 Subject: [PATCH] Refactor sibling annex to use mmap-backed concurrent storage Shift the sibling annex construction pipeline from an in-memory atomic mask to a memory-mapped file backend. This enables lock-free concurrent writes directly into the mapped region, streamlining the two-phase write process to accumulate bits atomically before finalization. Adjusted the cache lookup to return the specific matching layer index rather than a boolean flag, and added tests to verify correct layer tracking and cross-partition resolution in merged indexes. --- src/obikphylo/src/siblings/build.rs | 101 ++++++++++++++------- src/obikphylo/src/siblings/cache.rs | 17 ++-- src/obikphylo/src/siblings/siblingannex.rs | 53 ++++++++++- src/obikphylo/src/siblings/tests.rs | 34 +++++++ 4 files changed, 158 insertions(+), 47 deletions(-) diff --git a/src/obikphylo/src/siblings/build.rs b/src/obikphylo/src/siblings/build.rs index 5a541e5f..c2f7a2f5 100644 --- a/src/obikphylo/src/siblings/build.rs +++ b/src/obikphylo/src/siblings/build.rs @@ -1,5 +1,5 @@ use std::path::Path; -use std::sync::atomic::{AtomicU8, Ordering}; +use std::sync::atomic::Ordering; use std::sync::Arc; use rayon::prelude::*; @@ -102,7 +102,7 @@ impl SiblingAnnexBuildExt for KmerIndex { let mut part_slots: u64 = 0; for l in 0..meta.n_layers { let layer_dir = index_dir.join(format!("layer_{l}")); - part_slots += build_layer_sibling_annex(self, &layer_dir, n_parts, &cache)?; + part_slots += build_layer_sibling_annex(self, &layer_dir, n_parts, l, &cache)?; } total_slots += part_slots; pb.inc(1); @@ -122,6 +122,7 @@ fn build_layer_sibling_annex( index: &KmerIndex, layer_dir: &Path, n_parts: usize, + l: usize, cache: &Arc, ) -> OKIResult { let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); @@ -130,28 +131,51 @@ fn build_layer_sibling_annex( let k = index.kmer_size(); let n = mphf.n(); - // ── Reconciliation state, indexed by this layer's k-mer iteration - // order (the physical layout of `unitigs.bin`), never by MPHF slot - // — this layer's own k-mers are known members by construction, so - // no evidence check, no MPHF slot, and no slot -> k-mer - // reconstruction is needed or legitimate here (see - // `docmd/architecture/siblings.md`: the MPHF is not invertible, and - // evidence answers membership, not identity). The annex is written - // under this same iteration order end to end, so a reader can later - // zip `iter_kmers()` with the annex file directly, with no - // MPHF/slot indirection at read time either. + // Whether this layer's family-member fields can carry a real layer + // number (fast path, exploited by a future scan-time reader) or must + // stay presence-only (today's behaviour) — a single fact about the + // whole index, decided once here, never stored: `n_layers` is + // guaranteed identical across every partition (a merge adds one layer + // to all of them at once), and `FamilyMask`'s 3-bit field only has room + // for layers `0..=6`. + let fast_mode = meta.n_layers <= 7; + if !fast_mode { + tracing::warn!( + "layer {l} ({layer_dir:?}): index has {} layers (>7) — sibling-annex fast layer \ + lookup disabled for this build; consider compacting this index", + meta.n_layers + ); + } + + // ── The annex file itself is the reconciliation state, indexed by + // this layer's k-mer iteration order (the physical layout of + // `unitigs.bin`), never by MPHF slot — this layer's own k-mers are + // known members by construction, so no evidence check, no MPHF + // slot, and no slot -> k-mer reconstruction is needed or legitimate + // here (see `docmd/architecture/siblings.md`: the MPHF is not + // invertible, and evidence answers membership, not identity). No + // separate in-memory accumulator: every concurrent write goes + // straight through `SiblingAnnexBuilder::atomic_slot` into the + // mmap — a layer can hold billions of k-mers, so a `Vec` shadowing + // the whole file in RAM just to copy it out again at the end (the + // previous design) doubles memory for no benefit once the builder + // itself can be written concurrently. // - // `Arc>`, not `Vec` — `Pipe::apply` requires - // its source iterator to be `Send + 'static` (its items are - // dispatched to worker threads that outlive this call), so the - // batch-generating closure below needs an owned handle it can move - // in, not a borrow of a local. `AtomicU8`, not `FamilyMask`, - // because the gather phase below parallelises across destination - // partitions (independent `query_partition_with` calls, safe to run - // concurrently) and their `Found` hits can land on arbitrary, - // possibly-shared source entries — a lock-free `fetch_or` avoids - // needing any synchronisation beyond that. ───────────────────────── - let mask: Arc> = Arc::new((0..n).map(|_| AtomicU8::new(0)).collect()); + // `Arc`, not a bare `SiblingAnnexBuilder` — + // `Pipe::apply` requires its source iterator to be `Send + 'static` + // (its items are dispatched to worker threads that outlive this + // call), so the batch-generating closure below needs an owned + // handle it can move in, not a borrow of a local. `atomic_slot`, + // not `set`, for every concurrent write below: the gather phase + // parallelises across destination partitions (independent + // `cache.find` calls, safe to run concurrently) and their hits can + // land on arbitrary, possibly-shared source entries — a lock-free + // `fetch_or` avoids needing any synchronisation beyond that. + // Reclaimed as a plain owned value (`Arc::try_unwrap`) once every + // concurrent phase below is done, for the sequential finalisation + // pass. ───────────────────────────────────────────────────────── + let annex_path = layer_dir.join(ANNEX_FILE_NAME); + let builder = Arc::new(SiblingAnnexBuilder::new(n, &annex_path)?); // ── obipipeline: a *batch* transform, not a per-k-mer `Flat` one — // the actual cross-partition lookup reuses @@ -196,26 +220,28 @@ fn build_layer_sibling_annex( // Batches stream straight from `unitigs.bin` via // `enumerate_kmers_batch` — never a full-layer `Vec` collect (a // layer can hold billions of k-mers; see the "no full collect" - // rule). Each k-mer's own base is seeded into `mask` in this same - // pass, since this is exactly the iteration order `mask` is keyed - // on — no separate seeding pass needed. - let seed_mask = Arc::clone(&mask); + // rule). Each k-mer's own base is seeded straight into the annex in + // this same pass, since this is exactly the iteration order the + // annex is keyed on — no separate seeding pass needed. + let seed_builder = Arc::clone(&builder); let batches = mphf.enumerate_kmers_batch(BATCH_SIZE).map(move |(start, kmers)| { kmers .into_iter() .enumerate() .map(|(i, kmer)| { - seed_mask[start + i].fetch_or(1 << central_base(kmer, k), Ordering::Relaxed); + let base = central_base(kmer, k); + let bits = if fast_mode { FamilyMask::layer_bits(base, l) } else { FamilyMask::presence_bits(base) }; + seed_builder.atomic_slot(start + i).fetch_or(bits, Ordering::Relaxed); (start + i, kmer) }) .collect::>() }); - // Diagnostic: count still-empty mask slots after seeding + cross-partition + // Diagnostic: count still-empty annex slots after seeding + cross-partition // resolution. A non-zero count here means some k-mer's own base was never // written (should be impossible after the batch_offset fix above). let check_empty = |label: &str| { - let empty = mask.iter().filter(|v| v.load(Ordering::Relaxed) == 0).count(); + let empty = (0..n).filter(|&slot| builder.get(slot).is_none()).count(); if empty > 0 { tracing::warn!("{label}: {empty} mask slots still empty after resolution (layer {layer_dir:?})"); } @@ -294,8 +320,9 @@ fn build_layer_sibling_annex( .collect(); work.par_iter().for_each(|&(dest, chunk)| { for &(variant, source_order, base) in chunk { - if cache.find(dest, variant) { - mask[source_order].fetch_or(1 << base, Ordering::Relaxed); + if let Some(layer) = cache.find(dest, variant) { + let bits = if fast_mode { FamilyMask::layer_bits(base, layer) } else { FamilyMask::presence_bits(base) }; + builder.atomic_slot(source_order).fetch_or(bits, Ordering::Relaxed); } } }); @@ -315,10 +342,14 @@ fn build_layer_sibling_annex( // `unitigs.bin` and re-hashing through the MPHF, the cost // `is_minorant` was cheap to *compute* but expensive to *get the // inputs for* every time). - let annex_path = layer_dir.join(ANNEX_FILE_NAME); - let mut builder = SiblingAnnexBuilder::new(n, &annex_path)?; + // Every concurrent phase above is done and its `Arc` clone (the + // seeding pass's `seed_builder`) has already been dropped along + // with the now-fully-drained `batches` iterator — this is the only + // remaining strong reference, so reclaiming plain ownership for the + // sequential `set` calls below is guaranteed to succeed. + let mut builder = Arc::try_unwrap(builder).unwrap_or_else(|_| panic!("sibling-annex builder still shared after every concurrent phase completed")); for (order, kmer) in mphf.enumerate_kmers() { - let final_mask = FamilyMask::from_bits(mask[order].load(Ordering::Relaxed)); + let final_mask = builder.get(order).expect("seeded by construction"); let minorant = is_minorant(kmer, final_mask, k); builder.set(order, final_mask.with_minorant(minorant)); } diff --git a/src/obikphylo/src/siblings/cache.rs b/src/obikphylo/src/siblings/cache.rs index e3047821..003bd342 100644 --- a/src/obikphylo/src/siblings/cache.rs +++ b/src/obikphylo/src/siblings/cache.rs @@ -151,14 +151,17 @@ impl PartitionCache { Ok(Self { mats: built }) } - /// Existence-only lookup of `variant` in partition `dest_partition`: - /// tries each of the partition's already-open layers in turn, stopping - /// at the first hit. `find_slot`, not `sub_matrix`/`carries` — no data - /// read needed for a plain existence check. - pub(super) fn find(&self, dest_partition: usize, variant: CanonicalKmer) -> bool { + /// Existence lookup of `variant` in partition `dest_partition`: tries + /// each of the partition's already-open layers in turn, stopping at the + /// first hit and reporting which layer it was — `find_slot`, not + /// `sub_matrix`/`carries`, no data read needed for a plain existence + /// check. + pub(super) fn find(&self, dest_partition: usize, variant: CanonicalKmer) -> Option { self.mats - .get(dest_partition) - .is_some_and(|mats| mats.iter().any(|mat| mat.find_slot(variant).is_some())) + .get(dest_partition)? + .iter() + .enumerate() + .find_map(|(li, mat)| mat.find_slot(variant).map(|_| li)) } /// Resolve many `(variant, family_idx, base)` queries against one diff --git a/src/obikphylo/src/siblings/siblingannex.rs b/src/obikphylo/src/siblings/siblingannex.rs index 38f62aee..69637a24 100644 --- a/src/obikphylo/src/siblings/siblingannex.rs +++ b/src/obikphylo/src/siblings/siblingannex.rs @@ -50,6 +50,7 @@ use std::fs::{File, OpenOptions}; use std::io; use std::path::{Path, PathBuf}; +use std::sync::atomic::AtomicU16; use memmap2::{Mmap, MmapMut}; @@ -101,6 +102,24 @@ impl FamilyMask { FamilyMask(cleared | ((value as u16) << shift)) } + /// Raw bits to `fetch_or` into a shared `AtomicU16` accumulator to mark + /// `base` present at `layer`, without allocating a whole `FamilyMask` + /// per write — for `build::build_layer_sibling_annex`'s concurrent + /// construction, where each `(word, base)` field is written by exactly + /// one thread (see that module's docs), so a lock-free OR is safe. + #[inline] + pub(crate) fn layer_bits(base: u8, layer: usize) -> u16 { + debug_assert!(layer < MAX_FIELD_VALUE as usize, "layer out of range: {layer}"); + (layer as u16 + 1) << Self::field_shift(base) + } + + /// Same as [`layer_bits`](Self::layer_bits), presence-only (no layer + /// recorded) — the accumulator-side equivalent of [`with`](Self::with). + #[inline] + pub(crate) fn presence_bits(base: u8) -> u16 { + 1u16 << Self::field_shift(base) + } + /// Mark the member with central base `base` (0=A, 1=C, 2=G, 3=T) as /// present, without recording a layer (compat path: same observable /// effect as the original 1-byte design's `with`). Use @@ -283,12 +302,36 @@ impl SiblingAnnexBuilder { Ok(Self { mmap }) } + /// `None` means the slot has not (yet) been computed — see module + /// docs. Used to read back a slot's concurrently-accumulated value + /// (via [`atomic_slot`](Self::atomic_slot)) before finalising it with + /// [`set`](Self::set). + pub fn get(&self, slot: usize) -> Option { + let off = HEADER_SIZE + slot * 2; + FamilyMask::decode(u16::from_le_bytes(self.mmap[off..off + 2].try_into().unwrap())) + } + + /// Atomic handle to `slot`'s raw storage, for lock-free concurrent + /// construction straight into the mmap — no separate in-memory + /// accumulator needed (see `build::build_layer_sibling_annex`, the + /// only caller). Safe: every concurrent write goes through + /// `fetch_or`/atomic RMW, and by that caller's own invariant, a given + /// 3-bit field within a slot's word is written by exactly one thread — + /// only *different* fields of the same word can race, which `fetch_or` + /// already serialises correctly. Alignment holds: `HEADER_SIZE` is + /// even and every slot is 2 bytes, so `HEADER_SIZE + slot * 2` is + /// always even, and the mmap's own base address is page-aligned. + pub(crate) fn atomic_slot(&self, slot: usize) -> &AtomicU16 { + let off = HEADER_SIZE + slot * 2; + debug_assert!(off + 2 <= self.mmap.len()); + let ptr = unsafe { self.mmap.as_ptr().add(off) as *mut u16 }; + unsafe { AtomicU16::from_ptr(ptr) } + } + + /// Plain (non-atomic) store — for the sequential finalisation pass + /// only, once every concurrent write (via + /// [`atomic_slot`](Self::atomic_slot)) is done. pub fn set(&mut self, slot: usize, mask: FamilyMask) { - // Redundant concurrent writes from independent recomputation paths - // converge to the same encoded value for a given slot, so a plain - // store here is safe even without external synchronisation, as long - // as the write itself doesn't tear — true for a 2-byte-aligned - // `u16` store on every platform this project targets. let off = HEADER_SIZE + slot * 2; self.mmap[off..off + 2].copy_from_slice(&mask.encode().to_le_bytes()); } diff --git a/src/obikphylo/src/siblings/tests.rs b/src/obikphylo/src/siblings/tests.rs index 702d8365..43139e38 100644 --- a/src/obikphylo/src/siblings/tests.rs +++ b/src/obikphylo/src/siblings/tests.rs @@ -297,3 +297,37 @@ fn sibling_histogram_does_not_panic_on_partial_last_batch() { let total: u64 = hist.iter().sum(); assert!(total > 0, "histogram should contain at least one family"); } + +/// `sibling_annex_one_sibling_each`'s fixture happens to put g2's k-mer in +/// layer 0 and g1's in layer 1 (a fresh single-genome index, then one +/// merge that adds exactly one new layer for content absent from the +/// first — confirmed empirically, not assumed) — a real two-layer index, +/// not a contrived one, so it exercises both write paths +/// `build_layer_sibling_annex` now has to get right: the seeding pass (a +/// slot's own base gets *this* layer's index directly) and the +/// cross-partition/cross-layer resolution pass (`PartitionCache::find` +/// reporting which layer the other family member actually lives in, not +/// just that it exists). +#[test] +fn sibling_annex_records_the_real_layer_of_each_family_member() { + let dir = tempdir().unwrap(); + let g1 = build_single_genome_index(dir.path(), "g1", b"AACCGCTTAAG"); + let g2 = build_single_genome_index(dir.path(), "g2", b"AACCGGTTAAG"); + let merged = merge_two(dir.path(), &g1, &g2); + merged.build_sibling_annex().expect("build_sibling_annex"); + + let index_dir = merged.partition().part_dir(0).join(INDEX_SUBDIR); + let meta = PartitionMeta::load(&index_dir).unwrap(); + assert_eq!(meta.n_layers, 2, "fixture assumption: one merge, one new layer"); + + let g1_kmer = canonical(b"AACCGCTTAAG"); // own base C (1), sibling base G (2) — lives in layer 1 + let g2_kmer = canonical(b"AACCGGTTAAG"); // own base G (2), sibling base C (1) — lives in layer 0 + + let a = annex_info_for(&merged, g1_kmer); + assert_eq!(a.layer_value(1), Some(1), "g1's own base: seeded with its own layer (1)"); + assert_eq!(a.layer_value(2), Some(0), "g1's sibling (g2's form): resolved to its real layer (0)"); + + let b = annex_info_for(&merged, g2_kmer); + assert_eq!(b.layer_value(2), Some(0), "g2's own base: seeded with its own layer (0)"); + assert_eq!(b.layer_value(1), Some(1), "g2's sibling (g1's form): resolved to its real layer (1)"); +}