Remove in-place flag, optimize presence counting, and handle NUMA panics
The `obikmer select` CLI no longer supports in-place index rewriting; the `--output` flag is now required, with benchmarks updated to use temporary directories for atomic replacement. Added `--dense` and `--force-copy` flags. Introduced `batch_presence_counts` to compute presence counts across multiple column groups in a single pass, eliminating redundant I/O. Refactored the aggregation pipeline to branch on layer content, applying the optimized batched counting for `Presence` layers. Enhanced the NUMA runner to catch worker panics, track them, and re-raise after thread join to prevent indefinite blocking.
This commit is contained in:
@@ -1,9 +1,9 @@
|
||||
# select
|
||||
|
||||
Project and/or aggregate the genome columns of an index into a new (or in-place) index. Where [`filter`](filter.md) selects rows (kmers), `select` operates on columns (genomes): grouping several genomes into one aggregated column, reordering columns, or dropping some.
|
||||
Project and/or aggregate the genome columns of an index into a new index. Where [`filter`](filter.md) selects rows (kmers), `select` operates on columns (genomes): grouping several genomes into one aggregated column, reordering columns, or dropping some.
|
||||
|
||||
```bash
|
||||
obikmer select SOURCE (--output OUTPUT | --in-place) [OPTIONS]
|
||||
obikmer select SOURCE --output OUTPUT [OPTIONS]
|
||||
```
|
||||
|
||||
## Arguments
|
||||
@@ -16,8 +16,7 @@ obikmer select SOURCE (--output OUTPUT | --in-place) [OPTIONS]
|
||||
|
||||
| Option | Default | Description |
|
||||
|---|---|---|
|
||||
| `--output` | — | Output index directory (mutually exclusive with `--in-place`) |
|
||||
| `--in-place` | off | Rewrite the source index in place (mutually exclusive with `--output`) |
|
||||
| `-o, --output` | — | Output index directory (required) |
|
||||
| `-f, --force` | off | Overwrite an existing output directory |
|
||||
| `--group NAME:PRED` | none | Define a named group of genomes by predicate (repeatable; mutually exclusive with `--aggregate-by`) |
|
||||
| `--group-op NAME:OP` | none | Aggregation operator for a named group |
|
||||
@@ -25,6 +24,8 @@ obikmer select SOURCE (--output OUTPUT | --in-place) [OPTIONS]
|
||||
| `--aggregate-op OP` | none | Aggregation operator applied to every auto-generated group |
|
||||
| `--select COL,...` | all columns | Output columns, in order (group names or genome labels) |
|
||||
| `--presence-threshold` | `0` | Minimum count for a genome to be considered a carrier (logical operators only) |
|
||||
| `--dense` | off | Pack the output's presence matrices in the dense format instead of the default sparse one |
|
||||
| `--force-copy` | off | Copy each layer's unchanged kmer-identity files (mphf/unitigs/evidence/fingerprint) instead of hard-linking them |
|
||||
|
||||
## Aggregation operators
|
||||
|
||||
@@ -32,4 +33,16 @@ obikmer select SOURCE (--output OUTPUT | --in-place) [OPTIONS]
|
||||
|
||||
A `select` never changes the underlying kmer set — only the per-genome data (counts or presence) is rewritten, so an unaggregated pass-through column (a plain genome label in `--select`) is a cheap copy.
|
||||
|
||||
At least one of `--output`/`--in-place` is required, and at least one output column must be defined; every name listed in `--select` must resolve to either a defined group or an existing genome label. See [Genome predicates and taxonomy paths](predicates.md) for the predicate syntax used by `--group`.
|
||||
At least one output column must be defined; every name listed in `--select` must resolve to either a defined group or an existing genome label. See [Genome predicates and taxonomy paths](predicates.md) for the predicate syntax used by `--group`.
|
||||
|
||||
## Disk usage
|
||||
|
||||
`select` always writes to a new output directory — there is no in-place mode. Each layer's kmer-identity files (MPHF, unitigs, evidence, fingerprint) never change under a column projection/aggregation, so they are hard-linked into the output rather than copied: no extra disk is used for them, even on a very large index. Linking falls back to a real copy automatically if it fails (e.g. `SOURCE`/`OUTPUT` on different filesystems). Use `--force-copy` to always copy instead — needed when the output must be able to survive independently of the source on disk (a hard link shares the same underlying data, so overwriting one path outside `select` itself would affect the other).
|
||||
|
||||
To replace an index with a selected version of itself, select to a temporary directory and swap it in:
|
||||
|
||||
```bash
|
||||
obikmer select INDEX --output INDEX.tmp --group ... --group-op ... --select ...
|
||||
rm -rf INDEX
|
||||
mv INDEX.tmp INDEX
|
||||
```
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
#!/usr/bin/env bash
|
||||
# Usage: filter_one_count.sh SPECIES
|
||||
# Filters global_index_count to keep only kmers specific to SPECIES,
|
||||
# then selects the SPECIES column in-place.
|
||||
# then selects the SPECIES column (to a temp dir, swapped over the
|
||||
# original — obikmer select has no --in-place mode, but its output hard-
|
||||
# links each layer's unchanged kmer-identity files rather than copying
|
||||
# them, so this costs no extra disk for those).
|
||||
# Outputs:
|
||||
# specific_index_count/SPECIES/index.done (written by obikmer select)
|
||||
# stats/specific_kmer_count/SPECIES.stats (one CSV data row, no header)
|
||||
@@ -37,8 +40,11 @@ trap 'rm -f "${LOG_FILTER}" "${LOG_SELECT}"' EXIT
|
||||
|
||||
cat "${LOG_FILTER}" >&2
|
||||
|
||||
SELECT_TMP="${OUTPUT}.select_tmp"
|
||||
rm -rf "${SELECT_TMP}"
|
||||
|
||||
"${BINARY}" select \
|
||||
--in-place \
|
||||
--output "${SELECT_TMP}" \
|
||||
--group "${SPECIES}:species=${SPECIES}" \
|
||||
--group-op "${SPECIES}:any" \
|
||||
--select "${SPECIES}" \
|
||||
@@ -47,6 +53,9 @@ cat "${LOG_FILTER}" >&2
|
||||
|
||||
cat "${LOG_SELECT}" >&2
|
||||
|
||||
rm -rf "${OUTPUT}"
|
||||
mv "${SELECT_TMP}" "${OUTPUT}"
|
||||
|
||||
python3 - "${SPECIES}" "${LOG_FILTER}" "${LOG_SELECT}" <<'PYEOF' >"${STATS_FILE}"
|
||||
import sys, re
|
||||
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
#!/usr/bin/env bash
|
||||
# Usage: filter_one_presence.sh SPECIES
|
||||
# Filters global_index_presence to keep only kmers specific to SPECIES,
|
||||
# then selects the SPECIES column in-place.
|
||||
# then selects the SPECIES column (to a temp dir, swapped over the
|
||||
# original — obikmer select has no --in-place mode, but its output hard-
|
||||
# links each layer's unchanged kmer-identity files rather than copying
|
||||
# them, so this costs no extra disk for those).
|
||||
# Outputs:
|
||||
# specific_index_presence/SPECIES/index.done (written by obikmer select)
|
||||
# stats/specific_kmer_presence/SPECIES.stats (one CSV data row, no header)
|
||||
@@ -37,8 +40,11 @@ trap 'rm -f "${LOG_FILTER}" "${LOG_SELECT}"' EXIT
|
||||
|
||||
cat "${LOG_FILTER}" >&2
|
||||
|
||||
SELECT_TMP="${OUTPUT}.select_tmp"
|
||||
rm -rf "${SELECT_TMP}"
|
||||
|
||||
"${BINARY}" select \
|
||||
--in-place \
|
||||
--output "${SELECT_TMP}" \
|
||||
--group "${SPECIES}:species=${SPECIES}" \
|
||||
--group-op "${SPECIES}:any" \
|
||||
--select "${SPECIES}" \
|
||||
@@ -47,6 +53,9 @@ cat "${LOG_FILTER}" >&2
|
||||
|
||||
cat "${LOG_SELECT}" >&2
|
||||
|
||||
rm -rf "${OUTPUT}"
|
||||
mv "${SELECT_TMP}" "${OUTPUT}"
|
||||
|
||||
python3 - "${SPECIES}" "${LOG_FILTER}" "${LOG_SELECT}" <<'PYEOF' >"${STATS_FILE}"
|
||||
import sys, re
|
||||
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use std::collections::HashMap;
|
||||
use std::io;
|
||||
|
||||
use crate::colgroup::{chunked_presence_count, ColGroup, MatrixGroupOps};
|
||||
@@ -6,6 +7,110 @@ use crate::tempintvec::{TempCompactIntVec, TempCompactIntVecBuilder};
|
||||
|
||||
use super::persistent::PersistentBitMatrix;
|
||||
|
||||
// ── batch_presence_counts ───────────────────────────────────────────────────
|
||||
//
|
||||
// One shared pass computing every group's presence count at once, instead
|
||||
// of `select`'s old per-spec loop over `MatrixGroupOps::partial_group_*`
|
||||
// (one independent `col_view`-driven pass per output column). Two reasons
|
||||
// this exists, not just one:
|
||||
//
|
||||
// - `Sparse` has no on-disk column representation at all (`col_view` is a
|
||||
// deliberate `panic!` there — see `PersistentBitMatrix::col_view`'s own
|
||||
// docs) — a real column, however cheap, could only ever be synthesised by
|
||||
// scanning every row, which is exactly the row-major work this function
|
||||
// already does directly, natively, via
|
||||
// `PersistentSparseBitMatrix::for_each_genome_in_row` (touches only the
|
||||
// columns actually present per row — 1 for a singleton row, k̄ for a
|
||||
// multi-genome row).
|
||||
// - Even for `Columnar`/`Packed`, doing this once for *every* group
|
||||
// together is strictly less work than the old per-group loop whenever two
|
||||
// groups share a column (`--aggregate-by`'s auto-groups are disjoint, but
|
||||
// nothing stops `--group` from overlapping) — a column referenced by `g`
|
||||
// groups is now read from disk once, not `g` times.
|
||||
//
|
||||
// For a bit matrix, presence count at threshold 1 is the *only* primitive
|
||||
// needed: `sum` (bit matrix: sum of 0/1 = presence count), `any`
|
||||
// (count ≥ 1), `max` (= any), `min`/`all` (count == group size) are all
|
||||
// cheap derivations from this one vector — see
|
||||
// `obikselect::select_layer::agg_result_from_count`, the caller this
|
||||
// exists for.
|
||||
pub fn batch_presence_counts(
|
||||
mat: &PersistentBitMatrix,
|
||||
groups: &[ColGroup],
|
||||
) -> io::Result<Vec<TempCompactIntVec>> {
|
||||
let n = mat.n();
|
||||
|
||||
// Column -> every group index that references it — built once, so a
|
||||
// shared column is attributed to every one of its groups from a single
|
||||
// read, whichever storage format this turns out to be.
|
||||
let mut col_to_groups: HashMap<usize, Vec<usize>> = HashMap::new();
|
||||
for (gi, g) in groups.iter().enumerate() {
|
||||
for &c in &g.indices {
|
||||
col_to_groups.entry(c).or_default().push(gi);
|
||||
}
|
||||
}
|
||||
|
||||
let mut builders: Vec<TempCompactIntVecBuilder> = groups
|
||||
.iter()
|
||||
.map(|_| TempCompactIntVecBuilder::new(n))
|
||||
.collect::<io::Result<_>>()?;
|
||||
|
||||
let bump = |builders: &mut [TempCompactIntVecBuilder], gi: usize, slot: usize| {
|
||||
let b = &mut builders[gi];
|
||||
let v = b.get(slot);
|
||||
b.set(slot, v + 1);
|
||||
};
|
||||
|
||||
match mat {
|
||||
PersistentBitMatrix::Sparse(m) => {
|
||||
// Native row-major decode — the whole reason this exists.
|
||||
for slot in 0..n {
|
||||
m.for_each_genome_in_row(slot, |col| {
|
||||
if let Some(gs) = col_to_groups.get(&col) {
|
||||
for &gi in gs {
|
||||
bump(&mut builders, gi, slot);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
PersistentBitMatrix::Implicit { n_rows, .. } => {
|
||||
// Single column (index 0), present at every slot.
|
||||
if let Some(gs) = col_to_groups.get(&0) {
|
||||
for &gi in gs {
|
||||
for slot in 0..*n_rows {
|
||||
bump(&mut builders, gi, slot);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
PersistentBitMatrix::Columnar(_) | PersistentBitMatrix::Packed(_) => {
|
||||
// One `col_view` per *distinct referenced column*, not per
|
||||
// group — the common case (a column belongs to exactly one
|
||||
// group) reuses the existing fast bulk primitive directly;
|
||||
// only a column shared by several groups falls back to a
|
||||
// per-bit scan.
|
||||
for (&col, gs) in &col_to_groups {
|
||||
let view = mat.col_view(col);
|
||||
match gs.as_slice() {
|
||||
[gi] => builders[*gi].inc_present_fast(view),
|
||||
_ => {
|
||||
for (slot, present) in view.iter().enumerate() {
|
||||
if present {
|
||||
for &gi in gs {
|
||||
bump(&mut builders, gi, slot);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
builders.into_iter().map(|b| b.freeze()).collect()
|
||||
}
|
||||
|
||||
// ── MatrixGroupOps ────────────────────────────────────────────────────────────
|
||||
|
||||
impl MatrixGroupOps for PersistentBitMatrix {
|
||||
|
||||
@@ -20,6 +20,7 @@ mod persistent;
|
||||
mod sparse;
|
||||
|
||||
pub use builder::PersistentBitMatrixBuilder;
|
||||
pub use group_ops::batch_presence_counts;
|
||||
pub use packed::pack_bit_matrix;
|
||||
pub use persistent::PersistentBitMatrix;
|
||||
pub use sparse::{PersistentSparseBitMatrix, PersistentSparseBitMatrixBuilder, pack_sparse_bit_matrix};
|
||||
|
||||
@@ -21,7 +21,7 @@ mod views;
|
||||
|
||||
pub use bitmatrix::{
|
||||
PersistentBitMatrix, PersistentBitMatrixBuilder, PersistentSparseBitMatrix,
|
||||
PersistentSparseBitMatrixBuilder, pack_bit_matrix, pack_sparse_bit_matrix,
|
||||
PersistentSparseBitMatrixBuilder, batch_presence_counts, pack_bit_matrix, pack_sparse_bit_matrix,
|
||||
};
|
||||
pub use bitvec::{BitIter, PersistentBitVec, PersistentBitVecBuilder};
|
||||
pub use builder::PersistentCompactIntVecBuilder;
|
||||
|
||||
@@ -3,6 +3,7 @@ use tempfile::tempdir;
|
||||
use crate::{
|
||||
ColGroup, MatrixGroupOps, PersistentBitMatrix, PersistentBitMatrixBuilder,
|
||||
PersistentIntMatrix, PersistentCompactIntMatrixBuilder,
|
||||
PersistentSparseBitMatrixBuilder, batch_presence_counts,
|
||||
};
|
||||
use crate::{PersistentBitVecBuilder, PersistentCompactIntVec, PersistentCompactIntVecBuilder};
|
||||
|
||||
@@ -234,3 +235,66 @@ fn int_presence_count_additive_across_split() {
|
||||
assert_eq!(pb.get(1), 2);
|
||||
assert_eq!(pb.get(2), 1);
|
||||
}
|
||||
|
||||
// ── batch_presence_counts: Columnar vs Sparse, disjoint and overlapping groups ──
|
||||
|
||||
/// `batch_presence_counts` must agree with the naive "count group members
|
||||
/// present at this row" definition, for both a `Columnar` (col_view-driven)
|
||||
/// and a `Sparse` (native row-major) source — and in particular must still
|
||||
/// be correct when two groups share a column (`g1`/`g3` share column 1,
|
||||
/// `g2`/`g3` share column 2), the case that forces the per-bit fallback
|
||||
/// instead of the single-target `inc_present_fast` fast path.
|
||||
#[test]
|
||||
fn batch_presence_counts_matches_naive_for_columnar_and_sparse() {
|
||||
// Column-major: cols[c][row]
|
||||
let cols: [&[bool]; 4] = [
|
||||
&[true, false, true, false, true], // col0
|
||||
&[true, true, false, false, true], // col1
|
||||
&[false, true, true, true, false], // col2
|
||||
&[false, false, false, true, true], // col3
|
||||
];
|
||||
let (_dcol, columnar) = make_bit_matrix(&cols);
|
||||
|
||||
// Same data, row-major, for the sparse builder.
|
||||
let n_rows = 5;
|
||||
let n_cols = 4;
|
||||
let dir = tempdir().unwrap();
|
||||
let sparse_dir = dir.path().join("sparse");
|
||||
let mut b = PersistentSparseBitMatrixBuilder::new(n_rows, n_cols, &sparse_dir).unwrap();
|
||||
let mut genomes = Vec::new();
|
||||
for row in 0..n_rows {
|
||||
genomes.clear();
|
||||
genomes.extend((0..n_cols).filter(|&c| cols[c][row]).map(|c| c as u32));
|
||||
b.push_row(&genomes);
|
||||
}
|
||||
let sparse = PersistentBitMatrix::Sparse(b.finish().unwrap());
|
||||
|
||||
let groups = [
|
||||
ColGroup::new("g1", vec![0, 1]),
|
||||
ColGroup::new("g2", vec![2, 3]),
|
||||
ColGroup::new("g3", vec![1, 2]), // overlaps g1 (col1) and g2 (col2)
|
||||
];
|
||||
|
||||
let expected: [[u32; 5]; 3] = [
|
||||
[2, 1, 1, 0, 2], // g1: col0+col1
|
||||
[0, 1, 1, 2, 1], // g2: col2+col3
|
||||
[1, 2, 1, 1, 1], // g3: col1+col2
|
||||
];
|
||||
|
||||
for (label, mat) in [("columnar", &columnar), ("sparse", &sparse)] {
|
||||
let counts = batch_presence_counts(mat, &groups).unwrap();
|
||||
assert_eq!(counts.len(), groups.len(), "{label}: wrong number of results");
|
||||
for (gi, exp) in expected.iter().enumerate() {
|
||||
for (row, &e) in exp.iter().enumerate() {
|
||||
assert_eq!(counts[gi].get(row), e, "{label}: group {gi} row {row}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_presence_counts_empty_groups_is_empty() {
|
||||
let (_d, m) = make_bit_matrix(&[&[true, false], &[false, true]]);
|
||||
let counts = batch_presence_counts(&m, &[]).unwrap();
|
||||
assert!(counts.is_empty());
|
||||
}
|
||||
|
||||
@@ -12,7 +12,7 @@ use std::path::Path;
|
||||
|
||||
use obicompactvec::{
|
||||
ColGroup, MatrixBuilder, MatrixGroupOps, PersistentBitMatrix, PersistentIntMatrix, TempBitVec,
|
||||
TempCompactIntVec,
|
||||
TempBitVecBuilder, TempCompactIntVec, TempCompactIntVecBuilder, batch_presence_counts,
|
||||
};
|
||||
use obikindex::layer::{KmerLayer, LayerContent};
|
||||
use obikindex::{KmerIndex, OKIError, OKIResult};
|
||||
@@ -93,6 +93,44 @@ fn compute_group(
|
||||
})
|
||||
}
|
||||
|
||||
/// A bit-matrix source's every `AggOp` is a cheap derivation of one shared
|
||||
/// presence count (see [`obicompactvec::batch_presence_counts`]'s own
|
||||
/// docs for why: `sum` = the count itself, `any`/`max` = `count ≥ 1`,
|
||||
/// `all`/`min` = `count == group size`, `none` = `count == 0`) — no further
|
||||
/// matrix access needed once `count` is in hand.
|
||||
fn agg_result_from_count(op: AggOp, group_len: usize, count: TempCompactIntVec) -> io::Result<AggResult> {
|
||||
let n = count.len();
|
||||
let group_len = group_len as u32;
|
||||
Ok(match op {
|
||||
AggOp::Sum => AggResult::Int(count),
|
||||
AggOp::Any => {
|
||||
let mut b = TempBitVecBuilder::new(n)?;
|
||||
b.or_where(count.view(), |v| v >= 1);
|
||||
AggResult::Bit(b.freeze()?)
|
||||
}
|
||||
AggOp::All => {
|
||||
let mut b = TempBitVecBuilder::new(n)?;
|
||||
b.or_where(count.view(), |v| v == group_len);
|
||||
AggResult::Bit(b.freeze()?)
|
||||
}
|
||||
AggOp::None => {
|
||||
let mut b = TempBitVecBuilder::new(n)?;
|
||||
b.or_where(count.view(), |v| v == 0);
|
||||
AggResult::Bit(b.freeze()?)
|
||||
}
|
||||
AggOp::Min => {
|
||||
let mut b = TempCompactIntVecBuilder::new(n)?;
|
||||
b.inc_predicate_fast(count.view(), |v| v == group_len);
|
||||
AggResult::Int(b.freeze()?)
|
||||
}
|
||||
AggOp::Max => {
|
||||
let mut b = TempCompactIntVecBuilder::new(n)?;
|
||||
b.inc_predicate_fast(count.view(), |v| v >= 1);
|
||||
AggResult::Int(b.freeze()?)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ── AggResult → MatrixBuilder ─────────────────────────────────────────────────
|
||||
|
||||
/// Add one already-aggregated column to `mb` — the only piece `MatrixBuilder`
|
||||
@@ -161,15 +199,6 @@ pub(crate) fn select_partition(
|
||||
.to_path_buf();
|
||||
copy_layer_files(src_layer.dir(), &dst_layer_dir, force_copy).map_err(OKIError::Io)?;
|
||||
|
||||
let group_mat: Box<dyn MatrixGroupOps> = match src_layer.content() {
|
||||
LayerContent::Count => {
|
||||
Box::new(PersistentIntMatrix::open(src_layer.dir()).map_err(OKIError::Io)?)
|
||||
}
|
||||
LayerContent::Presence => {
|
||||
Box::new(PersistentBitMatrix::open(src_layer.dir()).map_err(OKIError::Io)?)
|
||||
}
|
||||
};
|
||||
|
||||
let data_subdir = if output_presence {
|
||||
"presence"
|
||||
} else {
|
||||
@@ -180,10 +209,37 @@ pub(crate) fn select_partition(
|
||||
|
||||
let mut builder =
|
||||
MatrixBuilder::new(output_presence, n, &data_dir).map_err(OKIError::Io)?;
|
||||
for spec in specs {
|
||||
let r = compute_group(group_mat.as_ref(), spec, threshold).map_err(OKIError::Io)?;
|
||||
add_result(&mut builder, r).map_err(OKIError::Io)?;
|
||||
|
||||
match src_layer.content() {
|
||||
// One shared row-major (Sparse) / deduplicated column-major
|
||||
// (Columnar/Packed/Implicit) pass covers every spec at once —
|
||||
// see `batch_presence_counts`'s own docs. Every `AggOp` for a
|
||||
// bit-matrix source is a cheap derivation of that one count.
|
||||
LayerContent::Presence => {
|
||||
let mat = PersistentBitMatrix::open(src_layer.dir()).map_err(OKIError::Io)?;
|
||||
let groups: Vec<ColGroup> = specs
|
||||
.iter()
|
||||
.map(|s| ColGroup::new(s.label.clone(), s.indices.clone()))
|
||||
.collect();
|
||||
let counts = batch_presence_counts(&mat, &groups).map_err(OKIError::Io)?;
|
||||
for (spec, count) in specs.iter().zip(counts) {
|
||||
let r = agg_result_from_count(spec.op, spec.indices.len(), count)
|
||||
.map_err(OKIError::Io)?;
|
||||
add_result(&mut builder, r).map_err(OKIError::Io)?;
|
||||
}
|
||||
}
|
||||
// Count-matrix `sum`/`min`/`max` are genuine per-value
|
||||
// reductions, not derivable from a single presence count —
|
||||
// unchanged, one `col_view`-driven pass per spec.
|
||||
LayerContent::Count => {
|
||||
let mat = PersistentIntMatrix::open(src_layer.dir()).map_err(OKIError::Io)?;
|
||||
for spec in specs {
|
||||
let r = compute_group(&mat, spec, threshold).map_err(OKIError::Io)?;
|
||||
add_result(&mut builder, r).map_err(OKIError::Io)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
builder.close().map_err(OKIError::Io)?;
|
||||
}
|
||||
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
use std::any::Any;
|
||||
use std::panic::{self, AssertUnwindSafe};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
@@ -164,6 +166,7 @@ impl PartitionRunner {
|
||||
let f = &f;
|
||||
|
||||
let mut first_err: Option<E> = None;
|
||||
let mut first_panic: Option<Box<dyn Any + Send + 'static>> = None;
|
||||
|
||||
std::thread::scope(|s| {
|
||||
// ── Timer thread ──────────────────────────────────────────────────
|
||||
@@ -208,12 +211,27 @@ impl PartitionRunner {
|
||||
for i in &prx {
|
||||
debug!(?tid, partition = i, "PartitionRunner worker: picked partition");
|
||||
let t = Instant::now();
|
||||
let r = match &pool {
|
||||
// Caught, not left to unwind straight through this
|
||||
// spawned thread: a panicking `f(i)` would
|
||||
// otherwise never reach `etx.send(...)` below, so
|
||||
// the controller's `completed < n_total` loop
|
||||
// waits forever for an event this partition can
|
||||
// no longer produce (see `run`'s own docs on the
|
||||
// termination protocol) — silently hanging
|
||||
// instead of surfacing the panic. Caught here and
|
||||
// re-raised on the caller's thread once `run`
|
||||
// returns, so the original message/backtrace
|
||||
// still surfaces, just from the right place.
|
||||
let outcome = panic::catch_unwind(AssertUnwindSafe(|| match &pool {
|
||||
Some(p) => p.install(|| f(i)),
|
||||
None => f(i),
|
||||
};
|
||||
}));
|
||||
debug!(?tid, partition = i, "PartitionRunner worker: partition done");
|
||||
etx.send(WorkerEvent::Completed(i, r, t.elapsed())).ok();
|
||||
let event = match outcome {
|
||||
Ok(r) => WorkerEvent::Completed(i, r, t.elapsed()),
|
||||
Err(payload) => WorkerEvent::Panicked(i, payload),
|
||||
};
|
||||
etx.send(event).ok();
|
||||
}
|
||||
debug!(?tid, "PartitionRunner worker: no more partitions, exiting");
|
||||
});
|
||||
@@ -260,6 +278,21 @@ impl PartitionRunner {
|
||||
n_total,
|
||||
);
|
||||
}
|
||||
WorkerEvent::Panicked(_i, payload) => {
|
||||
// Counts toward `completed` like any other outcome —
|
||||
// this partition will never produce a `Completed`
|
||||
// event, so not counting it here is exactly what
|
||||
// used to hang the controller forever. The payload
|
||||
// is re-raised on the caller's thread once `run`
|
||||
// returns (see below), not here: unwinding out of
|
||||
// this `recv` loop would leak the still-running
|
||||
// worker/timer threads this `thread::scope` owns.
|
||||
if first_panic.is_none() {
|
||||
first_panic = Some(payload);
|
||||
}
|
||||
completed += 1;
|
||||
reset_tx.send(()).ok();
|
||||
}
|
||||
WorkerEvent::TimerTick => {
|
||||
maybe_activate(
|
||||
&mut activation,
|
||||
@@ -279,6 +312,15 @@ impl PartitionRunner {
|
||||
drop(reset_tx);
|
||||
});
|
||||
|
||||
// A panic takes priority over a plain `Err`: it means `f` itself hit
|
||||
// a bug (an unhandled case, an assertion) rather than a normal,
|
||||
// typed failure — worth surfacing with its original message/
|
||||
// backtrace via unwinding, not silently downgraded to whatever `Err`
|
||||
// another, unrelated partition happened to return first.
|
||||
if let Some(payload) = first_panic {
|
||||
panic::resume_unwind(payload);
|
||||
}
|
||||
|
||||
match first_err {
|
||||
Some(e) => Err(e),
|
||||
None => Ok(()),
|
||||
@@ -290,6 +332,11 @@ impl PartitionRunner {
|
||||
|
||||
enum WorkerEvent<R, E> {
|
||||
Completed(usize, Result<R, E>, Duration),
|
||||
/// `f(i)` panicked instead of returning — see `run`'s own docs on why
|
||||
/// this is caught at all (never letting the panic unwind straight
|
||||
/// through the spawned worker thread) rather than a bug fix that could
|
||||
/// be skipped.
|
||||
Panicked(usize, Box<dyn Any + Send + 'static>),
|
||||
TimerTick,
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user