Project: a concurrent file indexer
This page builds one program in four steps. The program walks a directory tree and counts the lines, words and bytes of every file in it. The counting is easy. The work is in running it for many files at once without opening too many, without the walk running far ahead of the readers, and with a way to stop it halfway. Each step’s output below is a real run in the image.
Earlier pages explain the ideas used here: Async Basics and the Tokio guide for tasks and select!, Channels for mpsc, Send and Sync for the Send bound. The tests lean on ideas from More testing.
RememberThree separate limits do three separate jobs. A semaphore limits how much work is in flight. A bounded channel limits how far the producer can run ahead. A cancellation token tells everything to stop.
The first cargo command downloads tokio and tokio-util from crates.io, so it needs network. The versions are pinned by the committed Cargo.lock: tokio 1.54.1 and tokio-util 0.7.20. Cargo.toml alone says 1.54.1, which Cargo reads as that version or a newer 1.x one, so copy Cargo.lock too or Cargo may pick a newer release.
indexer/ ├── sample/ a small tree to index │ ├── notes/ │ │ ├── deep/ │ │ │ └── empty.txt │ │ └── todo.txt │ ├── a.txt │ └── b.txt ├── src/ │ ├── index.rs counting one file │ ├── lib.rs lists the modules │ ├── main.rs the command line │ ├── pipeline.rs bounded concurrency and cancellation │ ├── report.rs the sorted table │ └── walk.rs finds the files ├── Cargo.lock created by Cargo, committed └── Cargo.toml
[package]name = "indexer"version = "0.1.0"edition = "2024"
[dependencies]tokio = { version = "1.54.1", features = ["rt-multi-thread", "macros", "fs", "sync", "time", "signal"] }tokio-util = "0.7.20"
[dev-dependencies]# `test-util` gives tests a clock they can pause: sleeps finish instantly and in a fixed order.tokio = { version = "1.54.1", features = ["test-util"] }The test-util feature sits in [dev-dependencies] because only tests need it. Cargo enables dev-dependency features for builds that include dev-dependencies, such as cargo test (the Cargo book on features).
Step 1: count one file, then all of them in order
Section titled “Step 1: count one file, then all of them in order”Start with no concurrency at all. count is a plain function on a &str, which makes it easy to test, and index_file is the async wrapper that reads the file. The tokio::fs docs say Tokio “will use ordinary blocking file operations behind the scenes”, run with spawn_blocking, so an await on a file read does not hold up the runtime’s own threads.
use std::{io, path::PathBuf};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]pub struct FileStats { pub lines: usize, pub words: usize, pub bytes: usize,}
pub fn count(text: &str) -> FileStats { FileStats { lines: text.lines().count(), words: text.split_whitespace().count(), bytes: text.len(), }}
/// Reads one file and counts it. `tokio::fs` runs the read on a thread meant for blocking/// work, so the async runtime's own threads stay free.pub async fn index_file(path: PathBuf) -> io::Result<FileStats> { let text = tokio::fs::read_to_string(path).await?; Ok(count(&text))}
#[cfg(test)]mod tests { use super::*;
#[test] fn counts_lines_words_and_bytes() { let stats = count("one two\nthree\n"); assert_eq!(stats, FileStats { lines: 2, words: 3, bytes: 14 }); }
#[test] fn empty_text_is_all_zero() { assert_eq!(count(""), FileStats::default()); }
#[test] fn a_last_line_without_a_newline_still_counts() { assert_eq!(count("a b").lines, 1); }
#[tokio::test] async fn a_missing_file_is_a_not_found_error() { let err = index_file("no/such/file".into()).await.unwrap_err(); assert_eq!(err.kind(), io::ErrorKind::NotFound); }}use std::{collections::BTreeMap, io, path::PathBuf};
use crate::index::FileStats;
/// What the indexer found: a `BTreeMap` keeps its keys sorted, so the order never depends/// on which file finished first.pub type Results = BTreeMap<PathBuf, io::Result<FileStats>>;
/// One line per file in path order, then a total. Also returns how many files failed.pub fn render(results: &Results) -> (String, usize) { let mut out = format!("{:<30} {:>5} {:>6} {:>7}\n", "file", "lines", "words", "bytes"); let (mut total, mut files, mut failed) = (FileStats::default(), 0, 0); for (path, result) in results { let name = path.to_string_lossy(); match result { Ok(s) => { out += &format!("{name:<30} {:>5} {:>6} {:>7}\n", s.lines, s.words, s.bytes); total.lines += s.lines; total.words += s.words; total.bytes += s.bytes; files += 1; } Err(e) => { out += &format!("{name:<30} error: {e}\n"); failed += 1; } } } out += &format!("{files} files, {} lines, {} words, {} bytes", total.lines, total.words, total.bytes); if failed > 0 { out += &format!(", {failed} failed"); } out.push('\n'); (out, failed)}
#[cfg(test)]mod tests { use super::*;
#[test] fn rows_come_out_in_path_order_and_failures_are_counted() { let mut results = Results::new(); results.insert("b.txt".into(), Ok(FileStats { lines: 1, words: 2, bytes: 6 })); results.insert("a.txt".into(), Err(io::Error::new(io::ErrorKind::NotFound, "gone"))); let (text, failed) = render(&results); let lines: Vec<&str> = text.lines().collect(); assert_eq!(failed, 1); assert!(lines[1].starts_with("a.txt"), "{text}"); assert!(lines[2].starts_with("b.txt"), "{text}"); assert_eq!(lines[3], "1 files, 1 lines, 2 words, 6 bytes, 1 failed"); }}Two choices here pay off later. Results go into a BTreeMap, which keeps its keys sorted, so the table’s order never depends on which file finished first. And render returns a String instead of printing, so a test can look at it.
use std::{ io, path::{Path, PathBuf},};
/// Every file under `root`, in sorted order. The operating system lists a directory in/// whatever order it likes, so sorting is what makes the result the same on every run.pub async fn list_files(root: &Path) -> io::Result<Vec<PathBuf>> { let mut files = Vec::new(); let mut dirs = vec![root.to_path_buf()]; while let Some(dir) = dirs.pop() { let mut entries = tokio::fs::read_dir(&dir).await?; while let Some(entry) = entries.next_entry().await? { let path = entry.path(); if entry.file_type().await?.is_dir() { dirs.push(path); } else { files.push(path); } } } files.sort(); Ok(files)}
#[cfg(test)]mod tests { use super::*;
#[tokio::test] async fn finds_every_file_in_the_sample_tree() { let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("sample"); let found: Vec<String> = list_files(&root) .await .unwrap() .iter() .map(|p| p.strip_prefix(&root).unwrap().to_string_lossy().into_owned()) .collect(); assert_eq!(found, ["a.txt", "b.txt", "notes/deep/empty.txt", "notes/todo.txt"]); }
#[tokio::test] async fn a_missing_directory_is_an_error() { let err = list_files(Path::new("no/such/dir")).await.unwrap_err(); assert_eq!(err.kind(), io::ErrorKind::NotFound); }}use std::{path::PathBuf, process::ExitCode};
use indexer::{index, report, walk};
#[tokio::main]async fn main() -> ExitCode { let Some(root) = std::env::args().nth(1).map(PathBuf::from) else { eprintln!("usage: indexer <dir>"); return ExitCode::from(2); }; let files = match walk::list_files(&root).await { Ok(files) => files, Err(e) => { eprintln!("cannot read {}: {e}", root.display()); return ExitCode::FAILURE; } };
let mut results = report::Results::new(); for path in files { let stats = index::index_file(path.clone()).await; results.insert(path, stats); }
let (text, failed) = report::render(&results); print!("{text}"); if failed > 0 { ExitCode::FAILURE } else { ExitCode::SUCCESS }}file lines words bytes sample/a.txt 2 9 44 sample/b.txt 1 2 12 sample/notes/deep/empty.txt 0 0 0 sample/notes/todo.txt 3 8 41 4 files, 6 lines, 19 words, 97 bytes$ printf '\377\376 not text' > sample/blob.bin # a file that is not UTF-8 $ cargo run -q -- sample
file lines words bytes sample/a.txt 2 9 44 sample/b.txt 1 2 12 sample/blob.bin error: stream did not contain valid UTF-8 sample/notes/deep/empty.txt 0 0 0 sample/notes/todo.txt 3 8 41 4 files, 6 lines, 19 words, 97 bytes, 1 failed$ cargo run -q -- no-such-dir
cannot read no-such-dir: No such file or directory (os error 2)
The second run adds a file that is not UTF-8. read_to_string fails for it, the table says so on that file’s row, and the other rows are still there (Errors covers io::Result and ?). The program exits with code 1 when any file failed.
The tests for this step, with --test-threads=1 so the lines print in a fixed order:
$ cargo test --lib -- --test-threads=1 # timings replaced by X.XXs
running 7 tests
test index::tests::a_last_line_without_a_newline_still_counts ... ok
test index::tests::a_missing_file_is_a_not_found_error ... ok
test index::tests::counts_lines_words_and_bytes ... ok
test index::tests::empty_text_is_all_zero ... ok
test report::tests::rows_come_out_in_path_order_and_failures_are_counted ... ok
test walk::tests::a_missing_directory_is_an_error ... ok
test walk::tests::finds_every_file_in_the_sample_tree ... ok
test result: ok. 7 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in X.XXs
Step 2: a semaphore bounds the work in flight
Section titled “Step 2: a semaphore bounds the work in flight”Step 1 reads one file at a time. The obvious change, tokio::spawn for every file, has a different problem: a directory with a hundred thousand files would create a hundred thousand tasks, each asking for a file read. A semaphore puts a number on it. The Semaphore docs describe it as a “counting semaphore performing asynchronous permit acquisition”, and unlike a mutex it lets more than one holder in.
The new file is pipeline.rs, and lib.rs gains one line, pub mod pipeline;, so main can import it:
//! A small file indexer: walk a directory and count the lines, words and bytes of every//! file in it, a bounded number of files at a time.pub mod index;pub mod pipeline;pub mod report;pub mod walk;run is generic over the work, so the same code can run the real file reads in main and fake sleeping work in tests. The Send bound on the future is there because tokio::spawn needs it (Send and Sync).
use std::{collections::BTreeMap, future::Future, path::PathBuf, sync::Arc};
use tokio::{sync::Semaphore, task::JoinSet};
/// Runs `work` on every path, with at most `limit` running at the same time.pub async fn run<T, F, Fut>(paths: Vec<PathBuf>, limit: usize, work: F) -> BTreeMap<PathBuf, T>where T: Send + 'static, F: Fn(PathBuf) -> Fut, Fut: Future<Output = T> + Send + 'static,{ let permits = Arc::new(Semaphore::new(limit)); let mut tasks = JoinSet::new(); for path in paths { // Wait here until one of the `limit` permits is free. let permit = permits.clone().acquire_owned().await.expect("the semaphore is never closed"); let job = work(path.clone()); tasks.spawn(async move { let _permit = permit; // released when this task ends (path, job.await) }); }
let mut done = BTreeMap::new(); while let Some(joined) = tasks.join_next().await { let (path, value) = joined.expect("a worker panicked"); done.insert(path, value); } done}
#[cfg(test)]mod tests { use std::{ sync::atomic::{AtomicUsize, Ordering::SeqCst}, time::Duration, };
use super::*;
fn names(n: usize) -> Vec<PathBuf> { (0..n).map(|i| PathBuf::from(format!("f{i:02}"))).collect() }
// `start_paused` stops the clock: a sleep finishes the moment nothing else can run, so // this test takes no real time and gives the same answer every time. #[tokio::test(start_paused = true)] async fn never_more_than_limit_at_once() { let running = Arc::new(AtomicUsize::new(0)); let peak = Arc::new(AtomicUsize::new(0)); let (r, p) = (running.clone(), peak.clone()); let done = run(names(10), 3, move |_| { let (r, p) = (r.clone(), p.clone()); async move { let now = r.fetch_add(1, SeqCst) + 1; p.fetch_max(now, SeqCst); tokio::time::sleep(Duration::from_millis(10)).await; r.fetch_sub(1, SeqCst); } }) .await; assert_eq!(done.len(), 10); assert_eq!(peak.load(SeqCst), 3); }
#[tokio::test(start_paused = true)] async fn results_are_sorted_whatever_order_the_work_finishes_in() { // The earlier the name, the longer the sleep, so the files finish in reverse order. let done = run(names(5), 5, |p: PathBuf| async move { let n = p.to_string_lossy()[1..].parse::<u64>().unwrap(); tokio::time::sleep(Duration::from_millis(100 - n * 10)).await; n }) .await; let order: Vec<u64> = done.values().copied().collect(); assert_eq!(order, [0, 1, 2, 3, 4]); }}use std::{path::PathBuf, process::ExitCode};
use indexer::{index, pipeline, report, walk};
#[tokio::main]async fn main() -> ExitCode { let mut args = std::env::args().skip(1); let Some(root) = args.next().map(PathBuf::from) else { eprintln!("usage: indexer <dir> [jobs]"); return ExitCode::from(2); }; let jobs = match args.next() { None => 4, Some(text) => match text.parse::<usize>() { Ok(n) if n > 0 => n, _ => { eprintln!("jobs must be a positive number, not {text:?}"); return ExitCode::from(2); } }, }; let files = match walk::list_files(&root).await { Ok(files) => files, Err(e) => { eprintln!("cannot read {}: {e}", root.display()); return ExitCode::FAILURE; } };
let results = pipeline::run(files, jobs, index::index_file).await;
let (text, failed) = report::render(&results); print!("{text}"); if failed > 0 { ExitCode::FAILURE } else { ExitCode::SUCCESS }}running 9 tests test index::tests::a_last_line_without_a_newline_still_counts ... ok test index::tests::a_missing_file_is_a_not_found_error ... ok test index::tests::counts_lines_words_and_bytes ... ok test index::tests::empty_text_is_all_zero ... ok test pipeline::tests::never_more_than_limit_at_once ... ok test pipeline::tests::results_are_sorted_whatever_order_the_work_finishes_in ... ok test report::tests::rows_come_out_in_path_order_and_failures_are_counted ... ok test walk::tests::a_missing_directory_is_an_error ... ok test walk::tests::finds_every_file_in_the_sample_tree ... ok test result: ok. 9 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in X.XXs
acquire_owned gives a permit that is not tied to a borrow, so it can move into the spawned task and is released when the task ends. The loop waits for a permit before it spawns, so the loop itself is what slows down. A JoinSet is Tokio’s collection of spawned tasks: spawn adds one, and join_next().await returns the next task to finish, in whatever order they finish, until none are left. That unknown order is why the results go into a BTreeMap.
The first test is the interesting one. It counts how many jobs are running at the same moment and asserts that the peak is exactly the limit, 3 (fetch_max stores the larger of the old and new value, which records the peak). Real sleeps could make a test pass or fail by luck of the scheduler, so #[tokio::test(start_paused = true)] stops the clock: the pause docs say that if the runtime “has no work to do, the clock is auto-advanced to the next pending timer”. The sleeps finish at once and in a fixed order. (Pausing needs the current_thread runtime, which #[tokio::test] uses by default.)
The second test makes the files finish in reverse order and asserts that the map still comes out sorted. Step 3 replaces this test with the producer test below, so nothing after step 2 tests the order. run still collects into a BTreeMap.
Step 3: a bounded channel between the walker and the workers
Section titled “Step 3: a bounded channel between the walker and the workers”Step 2 still builds the whole list of paths before any work starts, which holds every path in memory. The fix is to hand paths over as the walk finds them, through a channel. The bounded mpsc channel is the important part: the mpsc docs say that when it is at capacity the send waits for room, and “the channel provides backpressure”. The walker cannot get further ahead than the channel allows.
Three files change. walk sends into the channel, run receives from it, and main creates the channel and gives the sender to the walker.
use std::{io, path::PathBuf};
use tokio::sync::mpsc;
/// Sends every file under `root` into `tx`, one at a time, as it finds them. When the/// channel is full, `send` waits, so the walk never gets far ahead of the workers.pub async fn walk(root: PathBuf, tx: mpsc::Sender<PathBuf>) -> io::Result<()> { let mut dirs = vec![root]; while let Some(dir) = dirs.pop() { let mut entries = tokio::fs::read_dir(&dir).await?; while let Some(entry) = entries.next_entry().await? { let path = entry.path(); if entry.file_type().await?.is_dir() { dirs.push(path); } else if tx.send(path).await.is_err() { return Ok(()); // the receiver is gone: nobody wants more paths } } } Ok(())}
#[cfg(test)]mod tests { use std::path::Path;
use super::*;
#[tokio::test] async fn finds_every_file_in_the_sample_tree() { let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("sample"); let (tx, mut rx) = mpsc::channel(16); walk(root.clone(), tx).await.unwrap(); let mut found = Vec::new(); while let Some(path) = rx.recv().await { found.push(path.strip_prefix(&root).unwrap().to_string_lossy().into_owned()); } found.sort(); // the walk's order is the operating system's, so sort before comparing assert_eq!(found, ["a.txt", "b.txt", "notes/deep/empty.txt", "notes/todo.txt"]); }
#[tokio::test] async fn a_missing_directory_is_an_error() { let (tx, _rx) = mpsc::channel(1); let err = walk("no/such/dir".into(), tx).await.unwrap_err(); assert_eq!(err.kind(), io::ErrorKind::NotFound); }}use std::{collections::BTreeMap, future::Future, path::PathBuf, sync::Arc};
use tokio::{ sync::{Semaphore, mpsc}, task::JoinSet,};
/// Takes paths from `rx` and runs `work` on them, at most `limit` at a time.pub async fn run<T, F, Fut>(mut rx: mpsc::Receiver<PathBuf>, limit: usize, work: F) -> BTreeMap<PathBuf, T>where T: Send + 'static, F: Fn(PathBuf) -> Fut, Fut: Future<Output = T> + Send + 'static,{ let permits = Arc::new(Semaphore::new(limit)); let mut tasks = JoinSet::new(); loop { // Permit first, path second. While every worker is busy this loop does not call // `recv`, the channel fills up, and the walker waits in `send`. let permit = permits.clone().acquire_owned().await.expect("the semaphore is never closed"); let Some(path) = rx.recv().await else { break }; let job = work(path.clone()); tasks.spawn(async move { let _permit = permit; (path, job.await) }); }
let mut done = BTreeMap::new(); while let Some(joined) = tasks.join_next().await { let (path, value) = joined.expect("a worker panicked"); done.insert(path, value); } done}
#[cfg(test)]mod tests { use std::{ sync::atomic::{AtomicUsize, Ordering::SeqCst}, time::Duration, };
use super::*;
fn names(n: usize) -> Vec<PathBuf> { (0..n).map(|i| PathBuf::from(format!("f{i:02}"))).collect() }
#[tokio::test(start_paused = true)] async fn never_more_than_limit_at_once() { let (tx, rx) = mpsc::channel(2); tokio::spawn(async move { for path in names(10) { tx.send(path).await.unwrap(); } }); let running = Arc::new(AtomicUsize::new(0)); let peak = Arc::new(AtomicUsize::new(0)); let (r, p) = (running.clone(), peak.clone()); let done = run(rx, 3, move |_| { let (r, p) = (r.clone(), p.clone()); async move { let now = r.fetch_add(1, SeqCst) + 1; p.fetch_max(now, SeqCst); tokio::time::sleep(Duration::from_millis(10)).await; r.fetch_sub(1, SeqCst); } }) .await; assert_eq!(done.len(), 10); assert_eq!(peak.load(SeqCst), 3); }
#[tokio::test(start_paused = true)] async fn the_producer_cannot_get_far_ahead_of_the_workers() { let capacity = 2; let (tx, rx) = mpsc::channel(capacity); let sent = Arc::new(AtomicUsize::new(0)); let started = Arc::new(AtomicUsize::new(0)); let worst = Arc::new(AtomicUsize::new(0));
let s = sent.clone(); tokio::spawn(async move { for path in names(10) { tx.send(path).await.unwrap(); s.fetch_add(1, SeqCst); } });
let (s, st, w) = (sent.clone(), started.clone(), worst.clone()); // One worker, so the work is the slow part. At the start of each file, how many // paths has the producer sent that no worker has started yet? let done = run(rx, 1, move |_| { let (s, st, w) = (s.clone(), st.clone(), w.clone()); async move { let started_now = st.fetch_add(1, SeqCst) + 1; w.fetch_max(s.load(SeqCst).saturating_sub(started_now), SeqCst); tokio::time::sleep(Duration::from_millis(10)).await; } }) .await;
assert_eq!(done.len(), 10); assert_eq!(sent.load(SeqCst), 10); assert!(worst.load(SeqCst) <= capacity, "the producer got {} ahead", worst.load(SeqCst)); }}use std::{path::PathBuf, process::ExitCode};
use indexer::{index, pipeline, report, walk};use tokio::sync::mpsc;
#[tokio::main]async fn main() -> ExitCode { let mut args = std::env::args().skip(1); let Some(root) = args.next().map(PathBuf::from) else { eprintln!("usage: indexer <dir> [jobs]"); return ExitCode::from(2); }; let jobs = match args.next() { None => 4, Some(text) => match text.parse::<usize>() { Ok(n) if n > 0 => n, _ => { eprintln!("jobs must be a positive number, not {text:?}"); return ExitCode::from(2); } }, };
let (tx, rx) = mpsc::channel(8); let walker = tokio::spawn(walk::walk(root.clone(), tx)); let results = pipeline::run(rx, jobs, index::index_file).await; if let Err(e) = walker.await.expect("the walker panicked") { eprintln!("cannot read {}: {e}", root.display()); return ExitCode::FAILURE; }
let (text, failed) = report::render(&results); print!("{text}"); if failed > 0 { ExitCode::FAILURE } else { ExitCode::SUCCESS }}running 9 tests test index::tests::a_last_line_without_a_newline_still_counts ... ok test index::tests::a_missing_file_is_a_not_found_error ... ok test index::tests::counts_lines_words_and_bytes ... ok test index::tests::empty_text_is_all_zero ... ok test pipeline::tests::never_more_than_limit_at_once ... ok test pipeline::tests::the_producer_cannot_get_far_ahead_of_the_workers ... ok test report::tests::rows_come_out_in_path_order_and_failures_are_counted ... ok test walk::tests::a_missing_directory_is_an_error ... ok test walk::tests::finds_every_file_in_the_sample_tree ... ok test result: ok. 9 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in X.XXs
Look at the order in run: permit first, path second. If it took the path first and then waited for a permit, it would hold one path in its hands while the channel had room for another, so the producer could be one path further ahead than the capacity. With this order, when every worker is busy the loop does not call recv, the channel fills to its capacity, and the walker waits in send.
The the_producer_cannot_get_far_ahead_of_the_workers test checks it with one slow worker: at the start of every file it measures how many paths were sent but not yet started (saturating_sub stops at 0 instead of going below it), and asserts the worst case is at most the channel’s capacity.
Step 4: cancellation
Section titled “Step 4: cancellation”A tool that reads a whole disk needs a way to stop. This step adds a CancellationToken from tokio-util, whose docs call it “a token which can be used to signal a cancellation request to one or more tasks”. Every clone sees the same cancel. The token goes into three places:
- the walker stops sending,
- the loop in
runstops taking new paths, - each running job races its work against
token.cancelled(), and the loser of aselect!is dropped. Dropping a future cancels it (Async Basics), so a job that was cut off simply returnsNoneand has no entry in the result.
run now returns a Report with the finished work and a cancelled flag, so the caller can tell the run was stopped. In main, a small task waits for Ctrl-C and cancels the token. (That Ctrl-C path was not run: no key was pressed in these sessions, so the tests drive cancellation instead.)
A known limitation. cancelled is read after the workers have stopped, and it only says that the token was cancelled. If the cancel arrives after the last job has finished, the table is complete but cancelled is still true, and main prints “the table above is partial”. The tests cover a cancel in the middle of a run and a run nobody cancels, not this case.
use std::{collections::BTreeMap, future::Future, path::PathBuf, sync::Arc};
use tokio::{ sync::{Semaphore, mpsc}, task::JoinSet,};use tokio_util::sync::CancellationToken;
/// What `run` hands back: the finished work, and whether it was stopped early.pub struct Report<T> { pub done: BTreeMap<PathBuf, T>, pub cancelled: bool,}
/// Takes paths from `rx` and runs `work` on them, at most `limit` at a time, until the/// paths run out or `cancel` is triggered.pub async fn run<T, F, Fut>( mut rx: mpsc::Receiver<PathBuf>, limit: usize, cancel: CancellationToken, work: F,) -> Report<T>where T: Send + 'static, F: Fn(PathBuf) -> Fut, Fut: Future<Output = T> + Send + 'static,{ let permits = Arc::new(Semaphore::new(limit)); let mut tasks = JoinSet::new(); loop { // Permit first, path second. While every worker is busy this loop does not call // `recv`, the channel fills up, and the walker waits in `send`. let permit = tokio::select! { biased; _ = cancel.cancelled() => break, permit = permits.clone().acquire_owned() => permit.expect("the semaphore is never closed"), }; let path = tokio::select! { biased; _ = cancel.cancelled() => break, next = rx.recv() => match next { Some(path) => path, None => break, }, }; let job = work(path.clone()); let token = cancel.clone(); tasks.spawn(async move { let _permit = permit; // Whichever finishes first wins; the loser is dropped, which cancels it. tokio::select! { biased; _ = token.cancelled() => None, value = job => Some((path, value)), } }); }
let mut done = BTreeMap::new(); while let Some(joined) = tasks.join_next().await { if let Some((path, value)) = joined.expect("a worker panicked") { done.insert(path, value); } } Report { done, cancelled: cancel.is_cancelled() }}
#[cfg(test)]mod tests { use std::{ sync::atomic::{AtomicUsize, Ordering::SeqCst}, time::Duration, };
use super::*;
fn names(n: usize) -> Vec<PathBuf> { (0..n).map(|i| PathBuf::from(format!("f{i:02}"))).collect() }
/// A channel of `capacity` that a spawned task fills with `n` names. fn feed(n: usize, capacity: usize) -> mpsc::Receiver<PathBuf> { let (tx, rx) = mpsc::channel(capacity); tokio::spawn(async move { for path in names(n) { if tx.send(path).await.is_err() { break; } } }); rx }
#[tokio::test(start_paused = true)] async fn never_more_than_limit_at_once() { let running = Arc::new(AtomicUsize::new(0)); let peak = Arc::new(AtomicUsize::new(0)); let (r, p) = (running.clone(), peak.clone()); let report = run(feed(10, 2), 3, CancellationToken::new(), move |_| { let (r, p) = (r.clone(), p.clone()); async move { let now = r.fetch_add(1, SeqCst) + 1; p.fetch_max(now, SeqCst); tokio::time::sleep(Duration::from_millis(10)).await; r.fetch_sub(1, SeqCst); } }) .await; assert_eq!(report.done.len(), 10); assert_eq!(peak.load(SeqCst), 3); }
#[tokio::test(start_paused = true)] async fn the_producer_cannot_get_far_ahead_of_the_workers() { let capacity = 2; let (tx, rx) = mpsc::channel(capacity); let sent = Arc::new(AtomicUsize::new(0)); let started = Arc::new(AtomicUsize::new(0)); let worst = Arc::new(AtomicUsize::new(0));
let s = sent.clone(); tokio::spawn(async move { for path in names(10) { tx.send(path).await.unwrap(); s.fetch_add(1, SeqCst); } });
let (s, st, w) = (sent.clone(), started.clone(), worst.clone()); // One worker, so the work is the slow part. At the start of each file, how many // paths has the producer sent that no worker has started yet? let report = run(rx, 1, CancellationToken::new(), move |_| { let (s, st, w) = (s.clone(), st.clone(), w.clone()); async move { let started_now = st.fetch_add(1, SeqCst) + 1; w.fetch_max(s.load(SeqCst).saturating_sub(started_now), SeqCst); tokio::time::sleep(Duration::from_millis(10)).await; } }) .await;
assert_eq!(report.done.len(), 10); assert_eq!(sent.load(SeqCst), 10); assert!(worst.load(SeqCst) <= capacity, "the producer got {} ahead", worst.load(SeqCst)); }
#[tokio::test(start_paused = true)] async fn cancelling_keeps_finished_work_and_drops_the_rest() { let cancel = CancellationToken::new(); let trigger = cancel.clone(); tokio::spawn(async move { tokio::time::sleep(Duration::from_millis(25)).await; trigger.cancel(); }); // Two at a time, 10 ms each: f00 and f01 finish at 10 ms, f02 and f03 at 20 ms, // and f04 and f05 are still running when the cancel arrives at 25 ms. let report = run(feed(6, 2), 2, cancel, |_| async { tokio::time::sleep(Duration::from_millis(10)).await; }) .await; assert!(report.cancelled); let finished: Vec<&str> = report.done.keys().map(|p| p.to_str().unwrap()).collect(); assert_eq!(finished, ["f00", "f01", "f02", "f03"]); }
#[tokio::test(start_paused = true)] async fn a_run_nobody_cancels_reports_not_cancelled() { let report = run(feed(3, 1), 2, CancellationToken::new(), |_| async {}).await; assert!(!report.cancelled); assert_eq!(report.done.len(), 3); }}use std::{io, path::PathBuf};
use tokio::sync::mpsc;use tokio_util::sync::CancellationToken;
/// Sends every file under `root` into `tx`, one at a time, as it finds them. When the/// channel is full, `send` waits, so the walk never gets far ahead of the workers.pub async fn walk(root: PathBuf, tx: mpsc::Sender<PathBuf>, cancel: CancellationToken) -> io::Result<()> { let mut dirs = vec![root]; while let Some(dir) = dirs.pop() { let mut entries = tokio::fs::read_dir(&dir).await?; while let Some(entry) = entries.next_entry().await? { let path = entry.path(); if entry.file_type().await?.is_dir() { dirs.push(path); continue; } tokio::select! { biased; _ = cancel.cancelled() => return Ok(()), sent = tx.send(path) => { if sent.is_err() { return Ok(()); // the receiver is gone: nobody wants more paths } } } } } Ok(())}
#[cfg(test)]mod tests { use std::path::Path;
use super::*;
#[tokio::test] async fn finds_every_file_in_the_sample_tree() { let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("sample"); let (tx, mut rx) = mpsc::channel(16); walk(root.clone(), tx, CancellationToken::new()).await.unwrap(); let mut found = Vec::new(); while let Some(path) = rx.recv().await { found.push(path.strip_prefix(&root).unwrap().to_string_lossy().into_owned()); } found.sort(); // the walk's order is the operating system's, so sort before comparing assert_eq!(found, ["a.txt", "b.txt", "notes/deep/empty.txt", "notes/todo.txt"]); }
#[tokio::test] async fn a_missing_directory_is_an_error() { let (tx, _rx) = mpsc::channel(1); let err = walk("no/such/dir".into(), tx, CancellationToken::new()).await.unwrap_err(); assert_eq!(err.kind(), io::ErrorKind::NotFound); }
#[tokio::test] async fn a_cancelled_walk_sends_nothing() { let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("sample"); let (tx, mut rx) = mpsc::channel(16); let cancel = CancellationToken::new(); cancel.cancel(); walk(root, tx, cancel).await.unwrap(); assert!(rx.recv().await.is_none()); }}use std::{path::PathBuf, process::ExitCode};
use indexer::{index, pipeline, report, walk};use tokio::sync::mpsc;use tokio_util::sync::CancellationToken;
#[tokio::main]async fn main() -> ExitCode { let mut args = std::env::args().skip(1); let Some(root) = args.next().map(PathBuf::from) else { eprintln!("usage: indexer <dir> [jobs]"); return ExitCode::from(2); }; let jobs = match args.next() { None => 4, Some(text) => match text.parse::<usize>() { Ok(n) if n > 0 => n, _ => { eprintln!("jobs must be a positive number, not {text:?}"); return ExitCode::from(2); } }, };
// Ctrl-C cancels the whole pipeline instead of killing the process mid-write. let cancel = CancellationToken::new(); let on_ctrl_c = cancel.clone(); tokio::spawn(async move { if tokio::signal::ctrl_c().await.is_ok() { on_ctrl_c.cancel(); } });
let (tx, rx) = mpsc::channel(8); let walker = tokio::spawn(walk::walk(root.clone(), tx, cancel.clone())); let outcome = pipeline::run(rx, jobs, cancel, index::index_file).await; if let Err(e) = walker.await.expect("the walker panicked") { eprintln!("cannot read {}: {e}", root.display()); return ExitCode::FAILURE; }
let (text, failed) = report::render(&outcome.done); print!("{text}"); if outcome.cancelled { println!("cancelled: the table above is partial"); return ExitCode::FAILURE; } if failed > 0 { ExitCode::FAILURE } else { ExitCode::SUCCESS }}Each select! here starts with biased;. By default select! picks which ready branch to poll first at random (the select! docs, also covered in the Tokio guide). That mattered: a first version of the walk test, which cancels before the walk starts and expects nothing to be sent, failed on some runs, because the send branch was ready too and the random choice sometimes picked it. With biased; the cancel check is first.
The cancellation test fixes the timing with the paused clock. Two jobs at a time, 10 ms each, cancel at 25 ms: the first four files finished at 10 ms and 20 ms, and the last two were still running at 25 ms. So the test asserts that exactly f00 to f03 are in the result. All 12 tests pass:
$ cargo test --lib -- --test-threads=1 # timings replaced by X.XXs
running 12 tests
test index::tests::a_last_line_without_a_newline_still_counts ... ok
test index::tests::a_missing_file_is_a_not_found_error ... ok
test index::tests::counts_lines_words_and_bytes ... ok
test index::tests::empty_text_is_all_zero ... ok
test pipeline::tests::a_run_nobody_cancels_reports_not_cancelled ... ok
test pipeline::tests::cancelling_keeps_finished_work_and_drops_the_rest ... ok
test pipeline::tests::never_more_than_limit_at_once ... ok
test pipeline::tests::the_producer_cannot_get_far_ahead_of_the_workers ... ok
test report::tests::rows_come_out_in_path_order_and_failures_are_counted ... ok
test walk::tests::a_cancelled_walk_sends_nothing ... ok
test walk::tests::a_missing_directory_is_an_error ... ok
test walk::tests::finds_every_file_in_the_sample_tree ... ok
test result: ok. 12 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in X.XXs
Run it
Section titled “Run it”The job count is the second argument. It changes how the work is scheduled, not what comes out: the last step runs the program with 1 job and with 16 and compares the two outputs with diff. Zero jobs is refused, because a semaphore with no permits would wait forever.
$ cargo build -q $ target/debug/indexer sample 4 file lines words bytes sample/a.txt 2 9 44 sample/b.txt 1 2 12 sample/notes/deep/empty.txt 0 0 0 sample/notes/todo.txt 3 8 41 4 files, 6 lines, 19 words, 97 bytes $ indexer sample 1 > one.txt; indexer sample 16 > sixteen.txt; diff one.txt sixteen.txt identical output with 1 job and with 16 $ target/debug/indexer sample 0 jobs must be a positive number, not "0"
Break it
Section titled “Break it”A worker that is not Send. Give run a job that holds an Rc across an .await. The Send bound on Fut rejects it, and the error names the captured value:
use std::{path::PathBuf, rc::Rc};
use indexer::pipeline;use tokio::sync::mpsc;use tokio_util::sync::CancellationToken;
#[tokio::main]async fn main() { let (_tx, rx) = mpsc::channel::<PathBuf>(1); let label = Rc::new("worker"); let outcome = pipeline::run(rx, 2, CancellationToken::new(), move |_path| { let label = label.clone(); async move { tokio::task::yield_now().await; println!("{label}"); } }) .await; println!("{}", outcome.done.len());}error: future cannot be sent between threads safely --> src/main.rs:11:19 | 11 | let outcome = pipeline::run(rx, 2, CancellationToken::new(), move |_path| { | ___________________^ 12 | | let label = label.clone(); 13 | | async move { 14 | | tokio::task::yield_now().await; ... | 17 | | }) | |______^ future created by async block is not `Send` | = help: within `{async block@src/main.rs:13:9: 13:19}`, the trait `Send` is not implemented for `Rc<&str>` note: captured value is not `Send` --> src/main.rs:15:24 | 15 | println!("{label}"); | ^^^^^ has type `Rc<&str>` which is not `Send` note: required by a bound in `run` --> src/pipeline.rs:26:31 | 17 | pub async fn run<T, F, Fut>( | --- required by a bound in this function ... 26 | Fut: Future<Output = T> + Send + 'static, | ^^^^ required by this bound in `run`
A channel that is too big. The backpressure test says the producer may be at most capacity paths ahead. This run changes the capacity to 10 but keeps the assertion at 2. The first command is GNU sed -i, as in the image (on macOS write sed -i '', or edit those two spots by hand). The producer then gets 9 paths ahead before the worker catches up, and the test fails with that number. The second command filters the test output to the failing lines.
$ sed -i 's/let capacity = 2;/let capacity = 10;/; s/assert!(worst.load(SeqCst) <= capacity/assert!(worst.load(SeqCst) <= 2/' src/pipeline.rs $ cargo test the_producer 2>&1 | grep -E '^test pipeline|got [0-9]+ ahead' test pipeline::tests::the_producer_cannot_get_far_ahead_of_the_workers ... FAILED the producer got 9 ahead
What to try next
Section titled “What to try next”- Count something more useful. Replace
FileStatswith a word index: aBTreeMap<String, usize>per file, merged into one map at the end. Keep the merge in the caller so the workers share nothing. - Skip what you do not want. Add a third argument with an extension, such as
txt, and filter inwalk. Add a test that uses thesampletree, which holds only.txtfiles, so a filter forrswould find nothing. - Report progress. Send a message on a second channel after each file, and print a running count from a separate task. Decide whether that channel should be bounded, and what should happen when the printer is slow.
- Stop on a deadline. Wrap the whole
runintokio::time::timeout, or cancel the token after N seconds, and test it with the paused clock as above. - CPU-heavy work. If the per-file job were heavy computation instead of a read, it would hold up the runtime’s own threads. The Rayon and Tokio guide shows how to hand that work to a thread pool.
- Split it up. Move the library into its own crate with a thin binary in a Cargo workspace (Project structure).
Read more
Section titled “Read more”- The Tokio tutorial, including its channels and
selectchapters. tokio::sync::Semaphore,tokio::sync::mpscandtokio_util::sync::CancellationToken.- In this course: Threads, Channels, Send and Sync, Async Basics, Testing and More testing.
The project is the guide project-indexer: tokio 1.54.1 and tokio-util 0.7.20, as pinned in Cargo.lock.
Learn Rust is an independent project, not affiliated with the Rust Project or the Rust Foundation. Code and text are under the MIT licence.Source on GitHub.