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:
2026-08-09 20:05:42 +02:00
parent 7a77d3ffd1
commit 5bb4a8b611
6 changed files with 417 additions and 127 deletions
@@ -8,6 +8,7 @@ mod planning;
mod progress;
mod retry;
mod storage;
mod task_drain;
mod transport;
mod version_ini;
@@ -1,5 +1,6 @@
use std::{collections::HashMap, net::SocketAddr, path::Path, sync::Arc};
use futures::stream::FuturesUnordered;
use tokio::sync::mpsc::UnboundedSender;
use tokio_util::sync::CancellationToken;
@@ -11,6 +12,7 @@ use super::{
progress::{DownloadProgressTracker, sample_download_progress},
retry::{RetryContext, retry_failed_chunks},
storage::{prepare_game_storage, sync_game_storage},
task_drain::collect_or_drain_on_cancel,
transport::download_from_peer,
version_ini::{
VersionIniBuffer,
@@ -198,14 +200,14 @@ async fn download_transfer_chunks(
) -> eyre::Result<()> {
let plans = build_peer_plans(ctx.peers, transfer_descs, ctx.file_peer_map);
let mut tasks = Vec::new();
let tasks = FuturesUnordered::new();
for (peer_addr, plan) in plans {
let game_root = ctx.game_root.clone();
let game_id = ctx.game_id.to_string();
let cancel_token = ctx.cancel_token.clone();
let version_buffer = ctx.version_buffer.clone();
let progress_tracker = ctx.progress_tracker.clone();
tasks.push(tokio::spawn(async move {
tasks.push(async move {
download_from_peer(
peer_addr,
&game_id,
@@ -216,23 +218,19 @@ async fn download_transfer_chunks(
progress_tracker,
)
.await
}));
});
}
let mut failed_chunks: Vec<DownloadChunk> = Vec::new();
let mut last_err: Option<eyre::Report> = None;
for handle in tasks {
for result in collect_or_drain_on_cancel(tasks, ctx.cancel_token, ctx.game_id).await? {
if ctx.cancel_token.is_cancelled() {
eyre::bail!("download cancelled for game {}", ctx.game_id);
}
match handle.await {
Ok(Ok(results)) => {
if ctx.cancel_token.is_cancelled() {
eyre::bail!("download cancelled for game {}", ctx.game_id);
}
match result {
Ok(results) => {
collect_chunk_results(
ctx.game_id,
ctx.tx_notify_ui,
@@ -241,11 +239,10 @@ async fn download_transfer_chunks(
&mut last_err,
);
}
Ok(Err(_)) | Err(_) if ctx.cancel_token.is_cancelled() => {
Err(_) if ctx.cancel_token.is_cancelled() => {
eyre::bail!("download cancelled for game {}", ctx.game_id);
}
Ok(Err(e)) => last_err = Some(e),
Err(e) => last_err = Some(eyre::eyre!("task join error: {e}")),
Err(e) => last_err = Some(e),
}
}
+8 -14
View File
@@ -4,13 +4,14 @@ use std::{
sync::Arc,
};
use futures::{StreamExt, stream::FuturesUnordered};
use futures::stream::FuturesUnordered;
use tokio_util::sync::CancellationToken;
use super::{
confined_fs::ConfinedGameRoot,
planning::{ChunkDownloadResult, DownloadChunk, PeerDownloadPlan, resolve_file_peers},
progress::DownloadProgressTracker,
task_drain::collect_or_drain_on_cancel,
transport::download_from_peer,
version_ini::VersionIniBuffer,
};
@@ -109,9 +110,13 @@ async fn run_retry_batch(
retry_plans: HashMap<SocketAddr, PeerDownloadPlan>,
ctx: &RetryContext<'_>,
) -> eyre::Result<Vec<RetryAttempt>> {
let mut attempts = FuturesUnordered::new();
let attempts = FuturesUnordered::new();
for (peer_addr, plan) in retry_plans {
if ctx.cancel_token.is_cancelled() {
break;
}
let retry_chunks = plan.chunks.clone();
let game_root = ctx.game_root.clone();
let game_id = ctx.game_id.to_string();
@@ -138,18 +143,7 @@ async fn run_retry_batch(
});
}
let mut results = Vec::new();
while !attempts.is_empty() {
let result = tokio::select! {
() = ctx.cancel_token.cancelled() => {
eyre::bail!("download cancelled for game {}", ctx.game_id);
}
result = attempts.next() => result.expect("retry attempt should exist"),
};
results.push(result);
}
Ok(results)
collect_or_drain_on_cancel(attempts, ctx.cancel_token, ctx.game_id).await
}
fn handle_retry_chunk_result(
@@ -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);
}
}
+186 -55
View File
@@ -2,7 +2,7 @@ use std::{collections::VecDeque, net::SocketAddr, sync::Arc};
use futures::{SinkExt, StreamExt, stream::FuturesUnordered};
use s2n_quic::{Connection, stream::ReceiveStream};
use tokio::io::{AsyncSeekExt, AsyncWriteExt};
use tokio::io::{AsyncSeekExt, AsyncWrite, AsyncWriteExt};
use tokio_util::{
codec::{FramedWrite, LengthDelimitedCodec},
sync::CancellationToken,
@@ -72,10 +72,17 @@ async fn open_chunk_stream(
conn: &mut Connection,
game_id: &str,
chunk: &DownloadChunk,
cancel_token: &CancellationToken,
) -> eyre::Result<ReceiveStream> {
use lanspread_proto::{Message, Request};
let stream = conn.open_bidirectional_stream().await?;
let stream = tokio::select! {
biased;
() = cancel_token.cancelled() => {
eyre::bail!("download cancelled for game {game_id}");
}
result = conn.open_bidirectional_stream() => result?,
};
let (rx, tx) = stream.split();
let mut framed_tx = FramedWrite::new(tx, LengthDelimitedCodec::new());
@@ -85,42 +92,91 @@ async fn open_chunk_stream(
offset: chunk.offset,
length: chunk.length,
};
framed_tx.send(request.encode()).await?;
tokio::select! {
biased;
() = cancel_token.cancelled() => {
eyre::bail!("download cancelled for game {game_id}");
}
result = framed_tx.send(request.encode()) => result?,
}
framed_tx.close().await?;
tokio::select! {
biased;
() = cancel_token.cancelled() => {
eyre::bail!("download cancelled for game {game_id}");
}
result = framed_tx.close() => result?,
}
Ok(rx)
}
/// Receives one requested chunk from a peer stream.
async fn receive_chunk(
#[derive(Clone)]
struct ChunkReceiveContext {
peer_addr: SocketAddr,
mut rx: ReceiveStream,
game_root: &ConfinedGameRoot,
chunk: &DownloadChunk,
game_root: ConfinedGameRoot,
game_id: String,
cancel_token: CancellationToken,
version_buffer: Option<Arc<VersionIniBuffer>>,
progress_tracker: Arc<DownloadProgressTracker>,
}
async fn flush_before_propagating<T>(
writer: &mut (impl AsyncWrite + Unpin),
operation_result: eyre::Result<T>,
) -> eyre::Result<T> {
let flush_result = writer.flush().await;
let value = operation_result?;
flush_result?;
Ok(value)
}
async fn receive_chunk(
mut rx: ReceiveStream,
chunk: &DownloadChunk,
ctx: &ChunkReceiveContext,
) -> eyre::Result<()> {
if let Some(buffer) = version_buffer
if let Some(buffer) = &ctx.version_buffer
&& buffer.matches(&chunk.request_path)
{
return download_version_ini_chunk(peer_addr, rx, chunk, &buffer, progress_tracker).await;
return download_version_ini_chunk(rx, chunk, buffer, ctx).await;
}
let mut file = tokio::fs::File::from_std(game_root.open_chunk_file(&chunk.destination).await?);
ensure_download_not_cancelled(&ctx.cancel_token, &ctx.game_id)?;
let mut file =
tokio::fs::File::from_std(ctx.game_root.open_chunk_file(&chunk.destination).await?);
file.seek(std::io::SeekFrom::Start(chunk.offset)).await?;
ensure_download_not_cancelled(&ctx.cancel_token, &ctx.game_id)?;
let mut receive_budget = ReceiveBudget::new(chunk.length);
let mut progress =
progress_tracker.track_chunk(peer_addr, &chunk.request_path, chunk.offset, chunk.length);
let mut progress = ctx.progress_tracker.track_chunk(
ctx.peer_addr,
&chunk.request_path,
chunk.offset,
chunk.length,
);
while let Some(bytes) = rx.receive().await? {
receive_budget.accept(bytes.len())?;
file.write_all(&bytes).await?;
progress.record_bytes(bytes.len());
let receive_result: eyre::Result<()> = async {
loop {
let bytes = tokio::select! {
biased;
() = ctx.cancel_token.cancelled() => {
eyre::bail!("download cancelled for game {}", ctx.game_id);
}
result = rx.receive() => result?,
};
let Some(bytes) = bytes else {
break;
};
receive_budget.accept(bytes.len())?;
file.write_all(&bytes).await?;
progress.record_bytes(bytes.len());
}
receive_budget.finish()
}
receive_budget.finish()?;
.await;
file.flush().await?;
flush_before_propagating(&mut file, receive_result).await?;
// Verify file integrity by checking the file size
verify_chunk_integrity(&file, chunk.offset, chunk.length).await?;
@@ -129,41 +185,43 @@ async fn receive_chunk(
}
async fn receive_chunk_result(
peer_addr: SocketAddr,
game_root: ConfinedGameRoot,
chunk: DownloadChunk,
rx: ReceiveStream,
version_buffer: Option<Arc<VersionIniBuffer>>,
progress_tracker: Arc<DownloadProgressTracker>,
ctx: ChunkReceiveContext,
) -> ChunkDownloadResult {
let result = receive_chunk(
peer_addr,
rx,
&game_root,
&chunk,
version_buffer,
progress_tracker,
)
.await;
let result = receive_chunk(rx, &chunk, &ctx).await;
ChunkDownloadResult {
chunk,
result,
peer_addr,
peer_addr: ctx.peer_addr,
}
}
async fn download_version_ini_chunk(
peer_addr: SocketAddr,
mut rx: ReceiveStream,
chunk: &DownloadChunk,
buffer: &VersionIniBuffer,
progress_tracker: Arc<DownloadProgressTracker>,
ctx: &ChunkReceiveContext,
) -> eyre::Result<()> {
let mut received = Vec::with_capacity(usize::try_from(chunk.length)?);
let mut receive_budget = ReceiveBudget::new(chunk.length);
let mut progress =
progress_tracker.track_chunk(peer_addr, &chunk.request_path, chunk.offset, chunk.length);
while let Some(bytes) = rx.receive().await? {
let mut progress = ctx.progress_tracker.track_chunk(
ctx.peer_addr,
&chunk.request_path,
chunk.offset,
chunk.length,
);
loop {
let bytes = tokio::select! {
biased;
() = ctx.cancel_token.cancelled() => {
eyre::bail!("download cancelled for game {}", ctx.game_id);
}
result = rx.receive() => result?,
};
let Some(bytes) = bytes else {
break;
};
receive_budget.accept(bytes.len())?;
progress.record_bytes(bytes.len());
received.extend_from_slice(&bytes);
@@ -240,14 +298,26 @@ async fn download_chunk_plan(
let mut in_flight = FuturesUnordered::new();
let mut results = Vec::new();
let window = PEER_DOWNLOAD_STREAM_WINDOW.max(1);
let game_root = ctx.game_root.clone();
let receive_ctx = ChunkReceiveContext {
peer_addr: ctx.peer_addr,
game_root: ctx.game_root.clone(),
game_id: ctx.game_id.to_owned(),
cancel_token: ctx.cancel_token.clone(),
version_buffer: ctx.version_buffer.clone(),
progress_tracker: ctx.progress_tracker.clone(),
};
let mut cancelled = false;
while !pending.is_empty() || !in_flight.is_empty() {
while in_flight.len() < window {
if ctx.cancel_token.is_cancelled() {
cancelled = true;
results.clear();
}
while !cancelled && in_flight.len() < window {
let Some(chunk) = pending.pop_front() else {
break;
};
ensure_download_not_cancelled(ctx.cancel_token, ctx.game_id)?;
log::info!(
"Downloading chunk {} (offset {}, length {}) from {}",
@@ -257,16 +327,14 @@ async fn download_chunk_plan(
ctx.peer_addr
);
match open_chunk_stream(conn, ctx.game_id, &chunk).await {
match open_chunk_stream(conn, ctx.game_id, &chunk, ctx.cancel_token).await {
Ok(rx) => {
in_flight.push(receive_chunk_result(
ctx.peer_addr,
game_root.clone(),
chunk,
rx,
ctx.version_buffer.clone(),
ctx.progress_tracker.clone(),
));
in_flight.push(receive_chunk_result(chunk, rx, receive_ctx.clone()));
}
Err(_) if ctx.cancel_token.is_cancelled() => {
cancelled = true;
results.clear();
break;
}
Err(err) => {
let reason = format!("failed to open chunk stream: {err}");
@@ -284,16 +352,34 @@ async fn download_chunk_plan(
}
if in_flight.is_empty() {
if cancelled {
break;
}
continue;
}
let result = tokio::select! {
if cancelled {
let _ = in_flight
.next()
.await
.expect("in-flight chunk stream should exist");
continue;
}
tokio::select! {
biased;
() = ctx.cancel_token.cancelled() => {
eyre::bail!("download cancelled for game {}", ctx.game_id);
cancelled = true;
results.clear();
}
result = in_flight.next() => {
results.push(result.expect("in-flight chunk stream should exist"));
}
result = in_flight.next() => result.expect("in-flight chunk stream should exist"),
};
results.push(result);
}
if cancelled {
eyre::bail!("download cancelled for game {}", ctx.game_id);
}
Ok(results)
@@ -345,7 +431,39 @@ pub(super) async fn download_from_peer(
#[cfg(test)]
mod tests {
use super::ReceiveBudget;
use std::{
io,
pin::Pin,
task::{Context, Poll},
};
use tokio::io::AsyncWrite;
use super::{ReceiveBudget, flush_before_propagating};
#[derive(Default)]
struct FlushProbe {
flushed: bool,
}
impl AsyncWrite for FlushProbe {
fn poll_write(
self: Pin<&mut Self>,
_context: &mut Context<'_>,
bytes: &[u8],
) -> Poll<io::Result<usize>> {
Poll::Ready(Ok(bytes.len()))
}
fn poll_flush(self: Pin<&mut Self>, _context: &mut Context<'_>) -> Poll<io::Result<()>> {
self.get_mut().flushed = true;
Poll::Ready(Ok(()))
}
fn poll_shutdown(self: Pin<&mut Self>, _context: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
}
#[test]
fn receive_budget_accepts_exactly_the_requested_bytes() {
@@ -371,4 +489,17 @@ mod tests {
.finish()
.expect("empty zero-length stream should finish");
}
#[tokio::test]
async fn receive_errors_are_propagated_only_after_flushing() {
let mut probe = FlushProbe::default();
let receive_error: eyre::Result<()> = Err(eyre::eyre!("peer receive failed"));
let error = flush_before_propagating(&mut probe, receive_error)
.await
.expect_err("receive error should be preserved");
assert!(probe.flushed);
assert_eq!(error.to_string(), "peer receive failed");
}
}
+45 -45
View File
@@ -212,27 +212,34 @@ async fn handle_active_downloads_without_peers(
return;
}
let mut changed = false;
for id in active_ids {
if peers_still_have_game(peer_game_db, &id).await {
continue;
}
changed |= active_operations.write().await.remove(&id).is_some();
let Some(cancel_token) = active_downloads.write().await.remove(&id) else {
continue;
let cancelled = {
// An exclusive guard makes the check-and-cancel transition one-shot even when
// concurrent liveness checks remove the last peers at the same time.
let active_downloads = active_downloads.write().await;
let Some(cancel_token) = active_downloads.get(&id) else {
continue;
};
if cancel_token.is_cancelled() {
false
} else {
cancel_token.cancel();
true
}
};
cancel_token.cancel();
if !cancelled {
continue;
}
events::send(
tx_notify_ui,
PeerEvent::DownloadGameFilesAllPeersGone { id: id.clone() },
);
}
if changed {
events::emit_active_operations(active_operations, tx_notify_ui).await;
}
}
async fn peers_still_have_game(peer_game_db: &Arc<RwLock<PeerGameDB>>, game_id: &str) -> bool {
@@ -248,16 +255,10 @@ mod tests {
use tokio_util::sync::CancellationToken;
use super::handle_active_downloads_without_peers;
use crate::{
ActiveOperation,
ActiveOperationKind,
PeerEvent,
context::OperationKind,
peer_db::PeerGameDB,
};
use crate::{PeerEvent, context::OperationKind, peer_db::PeerGameDB};
#[tokio::test]
async fn all_peers_gone_cancels_download_and_emits_peers_gone_then_active_snapshot() {
async fn all_peers_gone_cancels_once_and_leaves_cleanup_to_download_owner() {
let peer_game_db = Arc::new(RwLock::new(PeerGameDB::new()));
let active_operations = Arc::new(RwLock::new(HashMap::from([(
"game".to_string(),
@@ -279,30 +280,38 @@ mod tests {
.await;
assert!(cancel.is_cancelled());
assert!(!active_operations.read().await.contains_key("game"));
assert!(!active_downloads.read().await.contains_key("game"));
assert_eq!(
active_operations.read().await.get("game"),
Some(&OperationKind::Downloading)
);
assert!(active_downloads.read().await.contains_key("game"));
let event = rx.recv().await.expect("peers-gone event should be emitted");
assert!(matches!(
event,
PeerEvent::DownloadGameFilesAllPeersGone { id } if id == "game"
));
let event = rx
.recv()
.await
.expect("active operation snapshot should be emitted");
assert!(matches!(
event,
PeerEvent::ActiveOperationsChanged { active_operations } if active_operations.is_empty()
));
assert!(
rx.try_recv().is_err(),
"peers-gone cancellation must not emit extra events"
"cancellation must not emit a premature active-operation snapshot"
);
handle_active_downloads_without_peers(
&peer_game_db,
&active_operations,
&active_downloads,
&tx,
)
.await;
assert!(
rx.try_recv().is_err(),
"an already-cancelled download must not emit peers-gone twice"
);
}
#[tokio::test]
async fn all_peers_gone_cancels_multiple_downloads_without_stuck_entries() {
async fn all_peers_gone_cancels_multiple_downloads_without_releasing_admission() {
let peer_game_db = Arc::new(RwLock::new(PeerGameDB::new()));
let first_cancel = CancellationToken::new();
let second_cancel = CancellationToken::new();
@@ -328,14 +337,17 @@ mod tests {
assert!(first_cancel.is_cancelled());
assert!(second_cancel.is_cancelled());
let operations = active_operations.read().await;
assert!(!operations.contains_key("first"));
assert!(!operations.contains_key("second"));
assert_eq!(operations.get("first"), Some(&OperationKind::Downloading));
assert_eq!(operations.get("second"), Some(&OperationKind::Downloading));
assert_eq!(
operations.get("installing"),
Some(&OperationKind::Installing)
);
drop(operations);
assert!(active_downloads.read().await.is_empty());
let downloads = active_downloads.read().await;
assert!(downloads.contains_key("first"));
assert!(downloads.contains_key("second"));
drop(downloads);
let mut cancelled_ids = Vec::new();
for _ in 0..2 {
@@ -347,21 +359,9 @@ mod tests {
}
cancelled_ids.sort();
assert_eq!(cancelled_ids, vec!["first", "second"]);
let event = rx
.recv()
.await
.expect("active operation snapshot should be emitted");
assert!(matches!(
event,
PeerEvent::ActiveOperationsChanged { active_operations }
if active_operations == vec![ActiveOperation {
id: "installing".to_string(),
operation: ActiveOperationKind::Installing,
}]
));
assert!(
rx.try_recv().is_err(),
"multiple peers-gone cancellations must not emit extra events"
"multiple cancellations must not emit an active-operation snapshot"
);
}
}