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<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");