fix(peer): drain cancelled download workers

Keep initial peer transfers, retry attempts, and chunk receivers structurally owned until they quiesce. Cancellation now stops opening new streams, flushes accepted file writes, and drains active work before ownership rollback can begin.

Leave operation admission owned by the running download task when liveness detects that every source disappeared, and emit the peers-gone notification only once.

Test Plan:
- just fmt (Rust and configured formatters completed; 39 pre-existing rumdl findings remain)
- just clippy
- just test
- git diff --cached --check
This commit is contained in:
ddidderr committed 2026-08-09 20:05:42 +02:00
1 parent 7a77d3ffd1
commit 5bb4a8b611
6 files changed
+417 -127

No files matched your search

@@ -0,0 +1,167 @@
use std::future::Future;
use futures::{StreamExt, stream::FuturesUnordered};
use tokio_util::sync::CancellationToken;
/// Collects structured child futures, draining them before reporting cancellation.
///
/// Children share `cancel_token` and are responsible for cooperatively settling
/// their own resources. Keeping them unspawned also guarantees that dropping the
/// collector synchronously drops every child instead of detaching background work.
pub(super) async fn collect_or_drain_on_cancel<F>(
mut in_flight: FuturesUnordered<F>,
cancel_token: &CancellationToken,
game_id: &str,
) -> eyre::Result<Vec<F::Output>>
where
F: Future,
{
let mut results = Vec::with_capacity(in_flight.len());
let mut cancelled = cancel_token.is_cancelled();
while !in_flight.is_empty() {
if cancelled {
let _ = in_flight
.next()
.await
.expect("in-flight future should exist");
continue;
}
tokio::select! {
biased;
() = cancel_token.cancelled() => {
cancelled = true;
results.clear();
}
result = in_flight.next() => {
results.push(result.expect("in-flight future should exist"));
}
}
}
if cancelled {
eyre::bail!("download cancelled for game {game_id}");
}
Ok(results)
}
#[cfg(test)]
mod tests {
use std::{
future::Future,
pin::Pin,
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
time::Duration,
};
use futures::stream::FuturesUnordered;
use tokio::sync::{Notify, Semaphore, mpsc};
use tokio_util::sync::CancellationToken;
use super::collect_or_drain_on_cancel;
type Worker = Pin<Box<dyn Future<Output = ()> + Send>>;
struct DropProbe(Arc<AtomicUsize>);
impl Drop for DropProbe {
fn drop(&mut self) {
self.0.fetch_add(1, Ordering::SeqCst);
}
}
#[tokio::test]
async fn cancellation_waits_for_every_child_to_quiesce() {
let cancel_token = CancellationToken::new();
let started = Arc::new(AtomicUsize::new(0));
let completed = Arc::new(AtomicUsize::new(0));
let started_notify = Arc::new(Notify::new());
let release = Arc::new(Semaphore::new(0));
let workers = FuturesUnordered::<Worker>::new();
for _ in 0..2 {
let started = Arc::clone(&started);
let completed = Arc::clone(&completed);
let started_notify = Arc::clone(&started_notify);
let release = Arc::clone(&release);
workers.push(Box::pin(async move {
started.fetch_add(1, Ordering::SeqCst);
started_notify.notify_one();
release
.acquire_owned()
.await
.expect("test semaphore should remain open")
.forget();
completed.fetch_add(1, Ordering::SeqCst);
}));
}
let collector_token = cancel_token.clone();
let mut collector = tokio::spawn(async move {
collect_or_drain_on_cancel(workers, &collector_token, "game").await
});
while started.load(Ordering::SeqCst) != 2 {
started_notify.notified().await;
}
cancel_token.cancel();
assert!(
tokio::time::timeout(Duration::from_millis(50), &mut collector)
.await
.is_err(),
"cancellation must wait for child cleanup"
);
release.add_permits(2);
let error = collector
.await
.expect("collector task should join")
.expect_err("cancelled collection should fail after draining");
assert!(
error
.to_string()
.contains("download cancelled for game game")
);
assert_eq!(completed.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn dropping_collector_drops_every_child() {
let cancel_token = CancellationToken::new();
let dropped = Arc::new(AtomicUsize::new(0));
let (started_tx, mut started_rx) = mpsc::unbounded_channel();
let workers = FuturesUnordered::<Worker>::new();
for worker_id in 0..2 {
let probe = DropProbe(dropped.clone());
let worker_started = started_tx.clone();
workers.push(Box::pin(async move {
let _probe = probe;
worker_started
.send(worker_id)
.expect("collector should wait for workers");
std::future::pending::<()>().await;
}));
}
drop(started_tx);
let collector_token = cancel_token.clone();
let collector = tokio::spawn(async move {
collect_or_drain_on_cancel(workers, &collector_token, "game").await
});
started_rx.recv().await.expect("first worker should start");
started_rx.recv().await.expect("second worker should start");
collector.abort();
let join_error = collector
.await
.expect_err("collector should have been cancelled");
assert!(join_error.is_cancelled());
assert_eq!(dropped.load(Ordering::SeqCst), 2);
}
}