diff --git a/src/obikmer/src/cmd/query.rs b/src/obikmer/src/cmd/query.rs index 4972c13..b2a2af6 100644 --- a/src/obikmer/src/cmd/query.rs +++ b/src/obikmer/src/cmd/query.rs @@ -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), Chunk(Rope), Output(Vec), } @@ -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 + Send>, + _guard: ThrottleGuard, + files_open: Arc, +} + +impl Iterator for GuardedChunkIter { + type Item = Rope; + fn next(&mut self) -> Option { + 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 = 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 + 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, + QueryData : Throttled => Vec, + || { + let files_open = Arc::clone(&files_open); + move |pw: Throttled| -> 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");