feat: Implement RAII-based file handle throttling #7

Merged
coissac merged 1 commits from push-qtnvlqlooklx into main 2026-05-22 12:03:40 +00:00
3 changed files with 93 additions and 52 deletions
+64 -2
View File
@@ -1,5 +1,6 @@
use std::io;
use std::path::PathBuf;
use std::sync::{Arc, Condvar, Mutex};
use clap::Args;
use obikrope::Rope;
@@ -43,7 +44,7 @@ pub struct CommonArgs {
)]
pub threads: usize,
/// Maximum number of input files open simultaneously during scatter
/// Maximum number of input files open simultaneously
#[arg(long, default_value_t = 20)]
pub max_open_files: usize,
}
@@ -61,10 +62,54 @@ impl CommonArgs {
}
}
// ── Open-file throttling ──────────────────────────────────────────────────────
struct FileSlots {
count: Mutex<usize>,
condvar: Condvar,
max: usize,
}
impl FileSlots {
fn new(max: usize) -> Self {
Self { count: Mutex::new(0), condvar: Condvar::new(), max }
}
fn acquire(&self) {
let mut count = self.count.lock().unwrap();
while *count >= self.max {
count = self.condvar.wait(count).unwrap();
}
*count += 1;
}
fn release(&self) {
let mut count = self.count.lock().unwrap();
*count -= 1;
self.condvar.notify_one();
}
}
struct SlotsGuard(Arc<FileSlots>);
impl Drop for SlotsGuard {
fn drop(&mut self) {
self.0.release();
}
}
// ── Pipeline data carrier ─────────────────────────────────────────────────────
/// A path bundled with an opaque guard token.
/// The guard is acquired in the source thread and dropped by the flat worker
/// once the file is fully read, releasing the open-file slot.
pub struct PathWithSlot {
pub path: PathBuf,
pub _guard: Box<dyn Send + 'static>,
}
pub enum PipelineData {
Path(PathBuf),
Path(PathWithSlot),
RawChunk(Rope),
NormChunk(Rope),
Batch(Vec<RoutableSuperKmer>),
@@ -75,6 +120,23 @@ pub enum PipelineData {
unsafe impl Send for PipelineData {}
unsafe impl Sync for PipelineData {}
/// Wrap a path iterator so that at most `max_open` files are open simultaneously.
/// Acquisition happens in the caller's thread (the pipeline source thread),
/// never inside a worker, preventing deadlocks.
pub fn throttle_paths(
source: impl Iterator<Item = PathBuf> + Send + 'static,
max_open: usize,
) -> impl Iterator<Item = PathWithSlot> + Send + 'static {
let slots = Arc::new(FileSlots::new(max_open));
source.map(move |path| {
slots.acquire();
PathWithSlot {
path,
_guard: Box::new(SlotsGuard(Arc::clone(&slots))),
}
})
}
// ── I/O plumbing ──────────────────────────────────────────────────────────────
pub fn open_chunks(path: PathBuf) -> io::Result<impl Iterator<Item = Rope>> {
+5 -5
View File
@@ -1,11 +1,10 @@
use std::io::{self, BufWriter, Write};
use std::path::PathBuf;
use clap::Args;
use obifastwrite::write_scatter;
use obikseq::{RoutableSuperKmer, set_k, set_m};
use crate::cli::{CommonArgs, PipelineData, open_chunks, partitions_to_bits};
use crate::cli::{CommonArgs, PipelineData, PathWithSlot, open_chunks, partitions_to_bits, throttle_paths};
#[derive(Args)]
pub struct SuperkmerArgs {
@@ -40,15 +39,16 @@ pub fn run(args: SuperkmerArgs) {
let level_max = args.common.level_max;
let partition_bits = partitions_to_bits(args.common.partitions);
let n_workers = args.common.threads.max(1);
let max_open = args.common.max_open_files.max(1);
set_k(k);
set_m(m);
let path_source = args.common.seqfile_paths();
let path_source = throttle_paths(args.common.seqfile_paths(), max_open);
let pipe = obipipeline::make_pipe! {
PipelineData : PathBuf => Vec<RoutableSuperKmer>,
||? { |path| open_chunks(path) } : Path => RawChunk,
PipelineData : PathWithSlot => Vec<RoutableSuperKmer>,
||? { |pw: PathWithSlot| open_chunks(pw.path) } : Path => RawChunk,
|? { move |rope| obiread::normalize_sequence_chunk(rope, k) } : RawChunk => NormChunk,
| { move |rope| obiskbuilder::build_superkmers(rope, k, level_max, theta) } : NormChunk => Batch,
};
+20 -41
View File
@@ -1,51 +1,29 @@
use std::path::PathBuf;
use std::sync::{Arc, Condvar, Mutex};
use std::time::{Duration, Instant};
use indicatif::{ProgressBar, ProgressStyle};
use obikpartitionner::KmerPartition;
use obikrope::Rope;
use obisys::{Reporter, Stage};
use tracing::info;
use crate::cli::{PipelineData, open_chunks};
use crate::cli::{PipelineData, PathWithSlot, open_chunks, throttle_paths};
// ── Counting semaphore ────────────────────────────────────────────────────────
// ── Iterator that keeps the slot guard alive until the file is exhausted ──────
struct Semaphore {
state: Mutex<usize>,
cvar: Condvar,
}
impl Semaphore {
fn new(n: usize) -> Arc<Self> {
Arc::new(Self { state: Mutex::new(n), cvar: Condvar::new() })
}
fn acquire(self: &Arc<Self>) {
let mut n = self.state.lock().unwrap();
while *n == 0 { n = self.cvar.wait(n).unwrap(); }
*n -= 1;
}
fn release(self: &Arc<Self>) {
*self.state.lock().unwrap() += 1;
self.cvar.notify_one();
}
}
struct SemGuard(Arc<Semaphore>);
impl Drop for SemGuard {
fn drop(&mut self) { self.0.release(); }
}
/// Iterator wrapper that holds a semaphore permit for its lifetime.
/// The permit is released when the iterator is exhausted or dropped.
struct GuardedIter<I> {
inner: I,
_guard: SemGuard,
_guard: Box<dyn Send + 'static>,
}
impl<I: Iterator> Iterator for GuardedIter<I> {
type Item = I::Item;
fn next(&mut self) -> Option<Self::Item> { self.inner.next() }
impl<I: Iterator<Item = Rope>> Iterator for GuardedIter<I> {
type Item = Rope;
fn next(&mut self) -> Option<Rope> {
self.inner.next()
}
}
// ── scatter ───────────────────────────────────────────────────────────────────
/// Run scatter: normalise → build superkmers → route to partition → close.
/// Reports the "scatter" stage to `rep`.
@@ -61,15 +39,17 @@ pub fn scatter(
) {
use obikseq::RoutableSuperKmer;
let sem = Semaphore::new(max_open);
// Throttle in the source thread — never in a worker — to prevent deadlock.
let throttled = throttle_paths(path_source, max_open);
let t = Stage::start("scatter");
let pipe = obipipeline::make_pipe! {
PipelineData : PathBuf => Vec<RoutableSuperKmer>,
||? { move |path: PathBuf| {
sem.acquire();
PipelineData : PathWithSlot => Vec<RoutableSuperKmer>,
||? { move |pw: PathWithSlot| {
let PathWithSlot { path, _guard } = pw;
info!("indexing: {}", path.display());
open_chunks(path).map(|iter| GuardedIter { inner: iter, _guard: SemGuard(Arc::clone(&sem)) })
// _guard travels into GuardedIter; released when all chunks are read
open_chunks(path).map(|iter| GuardedIter { inner: iter, _guard })
}} : Path => RawChunk,
|? { move |rope| obiread::normalize_sequence_chunk(rope, k) } : RawChunk => NormChunk,
| { move |rope| obiskbuilder::build_superkmers(rope, k, level_max, theta) } : NormChunk => Batch,
@@ -89,7 +69,7 @@ pub fn scatter(
let mut last_bases: u64 = 0;
const ALPHA: f64 = 0.15;
for batch in pipe.apply(path_source, n_workers, 1) {
for batch in pipe.apply(throttled, n_workers, 1) {
total_bases += batch.iter().map(|sk| sk.seql() as u64).sum::<u64>();
let now = Instant::now();
let dt = now.duration_since(last_t).as_secs_f64();
@@ -115,4 +95,3 @@ pub fn scatter(
kp.close().expect("close error");
rep.push(t.stop());
}