refactor(query): implement throttled per-file streaming pipeline
Replace the flat chunk iterator with a throttled, per-file streaming architecture using `obipipeline::throttle`. The new `GuardedChunkIter` binds file handles to their `ThrottleGuard`, enforcing concurrent open-file limits and tracking active files via an atomic counter. Pipeline stages and the progress spinner are updated to support this resource-aware, parallelized I/O flow.
This commit is contained in:
@@ -10,6 +10,7 @@ use obikindex::KmerIndex;
|
||||
use obikrope::Rope;
|
||||
use obikseq::RoutableSuperKmer;
|
||||
use obilayeredmap::IndexMode;
|
||||
use obipipeline::{Throttled, ThrottleGuard, throttle};
|
||||
use obiread::chunk::read_sequence_chunks_sized;
|
||||
use obiread::record::{SeqRecord, parse_chunk};
|
||||
use obiskbuilder::SuperKmerIter;
|
||||
@@ -19,6 +20,7 @@ use tracing::{debug, info};
|
||||
// ── Pipeline data ─────────────────────────────────────────────────────────────
|
||||
|
||||
enum QueryData {
|
||||
Path(Throttled<PathBuf>),
|
||||
Chunk(Rope),
|
||||
Output(Vec<u8>),
|
||||
}
|
||||
@@ -415,6 +417,30 @@ fn process_chunk(
|
||||
buf
|
||||
}
|
||||
|
||||
// ── GuardedChunkIter — keeps the throttle slot guard alive until the file is exhausted ──
|
||||
|
||||
/// Wraps a per-file `Rope` chunk iterator together with its `ThrottleGuard`,
|
||||
/// so the guard (and the throttle slot it holds) is only released once the
|
||||
/// file has been fully read — never earlier, never held past that point.
|
||||
struct GuardedChunkIter {
|
||||
inner: Box<dyn Iterator<Item = Rope> + Send>,
|
||||
_guard: ThrottleGuard,
|
||||
files_open: Arc<AtomicU32>,
|
||||
}
|
||||
|
||||
impl Iterator for GuardedChunkIter {
|
||||
type Item = Rope;
|
||||
fn next(&mut self) -> Option<Rope> {
|
||||
self.inner.next()
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for GuardedChunkIter {
|
||||
fn drop(&mut self) {
|
||||
self.files_open.fetch_sub(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
// ── Entry point ───────────────────────────────────────────────────────────────
|
||||
|
||||
pub fn run(args: QueryArgs) {
|
||||
@@ -463,36 +489,64 @@ pub fn run(args: QueryArgs) {
|
||||
let force_presence = args.force_presence;
|
||||
let presence_threshold = args.presence_threshold;
|
||||
|
||||
// Flat iterator over all Rope chunks from all input files.
|
||||
// I/O runs in the source thread; chunk processing is parallelised by the pipe.
|
||||
info!("query: chunk_size={}MiB", chunk_bytes / (1024 * 1024));
|
||||
// Throttled iterator over input file paths: at most `effective_max_open()`
|
||||
// files are open at once. Opening + decompressing + chunking each file is
|
||||
// now a Flat pipeline stage, executed across the `n_workers` pool — not
|
||||
// serialised in the pipe's dedicated source thread (see steps::scatter /
|
||||
// cmd::superkmer for the same pattern applied to indexing).
|
||||
info!("query: chunk_size={}MiB, max_open_files={}", chunk_bytes / (1024 * 1024), args.effective_max_open());
|
||||
|
||||
let paths: Vec<PathBuf> = args.inputs.iter().map(PathBuf::from).collect();
|
||||
let all_chunks = paths.into_iter().flat_map(move |path| {
|
||||
let path_str = path.to_str().unwrap_or("").to_owned();
|
||||
match read_sequence_chunks_sized(&path_str, chunk_bytes) {
|
||||
Ok(iter) => Box::new(iter.filter_map(|r| match r {
|
||||
Ok(rope) => Some(rope),
|
||||
Err(e) => {
|
||||
eprintln!("read error: {e}");
|
||||
None
|
||||
}
|
||||
})) as Box<dyn Iterator<Item = Rope> + Send>,
|
||||
Err(e) => {
|
||||
eprintln!("error opening {path_str}: {e}");
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
});
|
||||
let path_source = throttle(paths.into_iter(), args.effective_max_open());
|
||||
|
||||
// Instrumentation: total bytes processed (for the EMA throughput readout)
|
||||
// and the number of chunks currently being processed by a worker — both
|
||||
// read from the spinner loop below, updated from inside the pipe closure.
|
||||
// Instrumentation: total bytes processed (for the EMA throughput readout),
|
||||
// number of files currently open/being chunked, and number of chunks
|
||||
// currently being processed by a worker — all read from the spinner loop
|
||||
// below, updated from inside the pipe closures.
|
||||
let total_bytes = Arc::new(AtomicU64::new(0));
|
||||
let files_open = Arc::new(AtomicU32::new(0));
|
||||
let chunks_active = Arc::new(AtomicU32::new(0));
|
||||
|
||||
let pipe = obipipeline::make_pipe! {
|
||||
QueryData : Rope => Vec<u8>,
|
||||
QueryData : Throttled<PathBuf> => Vec<u8>,
|
||||
|| {
|
||||
let files_open = Arc::clone(&files_open);
|
||||
move |pw: Throttled<PathBuf>| -> GuardedChunkIter {
|
||||
let path = pw.item;
|
||||
let guard = pw.guard;
|
||||
let path_str = path.to_str().unwrap_or("").to_owned();
|
||||
files_open.fetch_add(1, Ordering::Relaxed);
|
||||
let open_start = Instant::now();
|
||||
// Hard-exit on file-open failure (mirrors the previous behaviour):
|
||||
// propagating this as a pipeline Err would hit a known scheduler
|
||||
// hang on early stage errors (obipipeline::scheduler::WorkerPool::run
|
||||
// breaks its main loop without unblocking the still-running source
|
||||
// thread, so the final `h.join()` never returns) — worth fixing in
|
||||
// obipipeline itself, but out of scope here; sidestepping it like the
|
||||
// original code already did is the safe choice for this change.
|
||||
let iter = read_sequence_chunks_sized(&path_str, chunk_bytes).unwrap_or_else(|e| {
|
||||
eprintln!("error opening {path_str}: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
debug!(
|
||||
path = %path_str,
|
||||
open_ms = open_start.elapsed().as_millis() as u64,
|
||||
"opened query input file"
|
||||
);
|
||||
let err_path = path_str.clone();
|
||||
GuardedChunkIter {
|
||||
inner: Box::new(iter.filter_map(move |r| match r {
|
||||
Ok(rope) => Some(rope),
|
||||
Err(e) => {
|
||||
eprintln!("read error: {err_path}: {e}");
|
||||
None
|
||||
}
|
||||
})),
|
||||
_guard: guard,
|
||||
files_open: Arc::clone(&files_open),
|
||||
}
|
||||
}
|
||||
} : Path => Chunk,
|
||||
| {
|
||||
let idx = Arc::clone(&idx);
|
||||
let total_bytes = Arc::clone(&total_bytes);
|
||||
@@ -520,7 +574,7 @@ pub fn run(args: QueryArgs) {
|
||||
const ALPHA: f64 = 0.15;
|
||||
|
||||
let mut out = BufWriter::new(io::stdout());
|
||||
for block in pipe.apply(all_chunks, n_workers, 2) {
|
||||
for block in pipe.apply(path_source, n_workers, 2) {
|
||||
if !block.is_empty() {
|
||||
out.write_all(&block).expect("write error");
|
||||
}
|
||||
@@ -540,7 +594,8 @@ pub fn run(args: QueryArgs) {
|
||||
(format!("{:.0} MB", bp / 1e6), format!("{:.0} MB/s", ema_rate / 1e6))
|
||||
};
|
||||
let active = chunks_active.load(Ordering::Relaxed);
|
||||
pb.set_message(format!("{count_str} {rate_str} [chunks in flight: {active}]"));
|
||||
let open = files_open.load(Ordering::Relaxed);
|
||||
pb.set_message(format!("{count_str} {rate_str} [files open: {open}, chunks in flight: {active}]"));
|
||||
}
|
||||
}
|
||||
out.flush().expect("flush error");
|
||||
|
||||
Reference in New Issue
Block a user