From 0448137d849baad30ee44ef1d23cc8b0790d1e74 Mon Sep 17 00:00:00 2001 From: Hammed Abass Date: Fri, 24 Jul 2026 00:11:58 +0200 Subject: [PATCH] fix(core): abort send when provider stream closes --- crates/vnidrop/src/runtime/mod.rs | 2 + crates/vnidrop/src/runtime/provider.rs | 107 +++++++++++++++++-------- crates/vnidrop/src/tests/runtime.rs | 42 +++++++++- crates/vnidrop/tests/transfer.rs | 29 ++++++- 4 files changed, 144 insertions(+), 36 deletions(-) diff --git a/crates/vnidrop/src/runtime/mod.rs b/crates/vnidrop/src/runtime/mod.rs index aee54ce..cb13930 100644 --- a/crates/vnidrop/src/runtime/mod.rs +++ b/crates/vnidrop/src/runtime/mod.rs @@ -16,6 +16,8 @@ mod share; mod storage; pub use facade::VnidropCore; +#[cfg(test)] +pub(crate) use provider::{consume_request_updates, RequestStreamOutcome}; use std::{ collections::{HashMap, HashSet}, diff --git a/crates/vnidrop/src/runtime/provider.rs b/crates/vnidrop/src/runtime/provider.rs index d6263fa..9447626 100644 --- a/crates/vnidrop/src/runtime/provider.rs +++ b/crates/vnidrop/src/runtime/provider.rs @@ -10,6 +10,31 @@ use tokio::sync::mpsc; use super::CoreInner; use crate::access_policy::AccessDecision; +#[derive(Debug, PartialEq, Eq)] +pub(crate) enum RequestStreamOutcome { + TerminalUpdateReceived, + Aborted, +} + +pub(crate) async fn consume_request_updates( + mut rx: irpc::channel::mpsc::Receiver, + mut handle_update: impl FnMut(RequestUpdate), +) -> RequestStreamOutcome { + let mut terminal_update_received = false; + while let Ok(Some(update)) = rx.recv().await { + terminal_update_received |= matches!( + update, + RequestUpdate::Completed(_) | RequestUpdate::Aborted(_) + ); + handle_update(update); + } + if terminal_update_received { + RequestStreamOutcome::TerminalUpdateReceived + } else { + RequestStreamOutcome::Aborted + } +} + impl CoreInner { pub(super) async fn spawn_provider_event_task( self: &Arc, @@ -302,7 +327,7 @@ impl CoreInner { transfer_id: u64, connection_id: u64, request_id: u64, - mut rx: irpc::channel::mpsc::Receiver, + rx: irpc::channel::mpsc::Receiver, ) { // Request update tasks are tied to individual provider streams. Router // shutdown closes those streams; only the long-lived provider receiver @@ -318,35 +343,35 @@ impl CoreInner { .cloned(); let core = self.clone(); tokio::spawn(async move { - while let Ok(Some(update)) = rx.recv().await { - match update { - RequestUpdate::Started(started) => core.emit_transfer( - transfer_id, - "send", - "transfer", - "started", - json!({ - "connection_id": connection_id, - "request_id": request_id, - "endpoint_id": endpoint_id, - "hash": started.hash.to_string(), - "size": started.size, - "index": started.index, - }), - ), - RequestUpdate::Progress(progress) => core.emit_transfer( - transfer_id, - "send", - "transfer", - "progress", - json!({ - "connection_id": connection_id, - "request_id": request_id, - "endpoint_id": endpoint_id, - "end_offset": progress.end_offset, - }), - ), - RequestUpdate::Completed(_) => core.emit_transfer( + let outcome = consume_request_updates(rx, |update| match update { + RequestUpdate::Started(started) => core.emit_transfer( + transfer_id, + "send", + "transfer", + "started", + json!({ + "connection_id": connection_id, + "request_id": request_id, + "endpoint_id": endpoint_id, + "hash": started.hash.to_string(), + "size": started.size, + "index": started.index, + }), + ), + RequestUpdate::Progress(progress) => core.emit_transfer( + transfer_id, + "send", + "transfer", + "progress", + json!({ + "connection_id": connection_id, + "request_id": request_id, + "endpoint_id": endpoint_id, + "end_offset": progress.end_offset, + }), + ), + RequestUpdate::Completed(_) => { + core.emit_transfer( transfer_id, "send", "transfer", @@ -356,8 +381,10 @@ impl CoreInner { "request_id": request_id, "endpoint_id": endpoint_id, }), - ), - RequestUpdate::Aborted(_) => core.emit_transfer( + ); + } + RequestUpdate::Aborted(_) => { + core.emit_transfer( transfer_id, "send", "transfer", @@ -367,8 +394,22 @@ impl CoreInner { "request_id": request_id, "endpoint_id": endpoint_id, }), - ), + ); } + }) + .await; + if outcome == RequestStreamOutcome::Aborted { + core.emit_transfer( + transfer_id, + "send", + "transfer", + "aborted", + json!({ + "connection_id": connection_id, + "request_id": request_id, + "endpoint_id": endpoint_id, + }), + ); } }); } diff --git a/crates/vnidrop/src/tests/runtime.rs b/crates/vnidrop/src/tests/runtime.rs index 8b3ca76..9c3bee8 100644 --- a/crates/vnidrop/src/tests/runtime.rs +++ b/crates/vnidrop/src/tests/runtime.rs @@ -1,9 +1,16 @@ -use std::sync::Arc; +use std::{sync::Arc, time::Duration}; -use iroh_blobs::Hash; +use iroh_blobs::{ + provider::{ + events::{RequestUpdate, TransferCompleted}, + TransferStats, + }, + Hash, +}; use crate::{ repository::{PendingDeliveryReceiptInsert, Repository, TransferUpsert}, + runtime::{consume_request_updates, RequestStreamOutcome}, transfer_state::{TransferDirection, TransferStatus}, CoreEvent, CoreEventSink, VnidropCore, VnidropError, }; @@ -14,6 +21,37 @@ impl CoreEventSink for TestSink { fn on_event(&self, _event: CoreEvent) {} } +#[test] +fn provider_request_stream_distinguishes_success_from_silent_abort() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let (completed_tx, completed_rx) = irpc::channel::mpsc::channel(1); + completed_tx + .send(RequestUpdate::Completed(TransferCompleted { + stats: Box::new(TransferStats { + payload_bytes_sent: 5, + other_bytes_sent: 0, + other_bytes_read: 0, + duration: Duration::ZERO, + }), + })) + .await + .unwrap(); + drop(completed_tx); + assert_eq!( + consume_request_updates(completed_rx, |_| {}).await, + RequestStreamOutcome::TerminalUpdateReceived + ); + + let (aborted_tx, aborted_rx) = irpc::channel::mpsc::channel::(1); + drop(aborted_tx); + assert_eq!( + consume_request_updates(aborted_rx, |_| {}).await, + RequestStreamOutcome::Aborted + ); + }); +} + #[test] fn initializes_and_reports_endpoint() { let temp = tempfile::tempdir().unwrap(); diff --git a/crates/vnidrop/tests/transfer.rs b/crates/vnidrop/tests/transfer.rs index 6544150..680f0d1 100644 --- a/crates/vnidrop/tests/transfer.rs +++ b/crates/vnidrop/tests/transfer.rs @@ -1,7 +1,28 @@ mod support; +use std::time::{Duration, Instant}; + use support::{receive_with_response, share_path, TestNode}; -use vnidrop::VnidropError; +use vnidrop::{CoreEvent, VnidropError}; + +fn wait_for_sender_transfer_event(sender: &TestNode, transfer_id: u64, kind: &str) -> CoreEvent { + let started = Instant::now(); + loop { + if let Some(event) = sender.sink.events().into_iter().find(|event| { + event.transfer_id == Some(transfer_id) + && event.direction.as_deref() == Some("send") + && event.phase == "transfer" + && event.kind == kind + }) { + return event; + } + assert!( + started.elapsed() < Duration::from_secs(5), + "timed out waiting for sender transfer event {kind}" + ); + std::thread::sleep(Duration::from_millis(10)); + } +} #[test] fn transfers_file_between_two_cores() { @@ -50,6 +71,12 @@ fn transfers_file_between_two_cores() { artifacts[0].locator, output_dir.path().join("hello.txt").to_string_lossy() ); + let completed = wait_for_sender_transfer_event(&sender, share.transfer_id, "completed"); + assert!(completed.data_json.contains("\"connection_id\":")); + assert!(completed.data_json.contains("\"request_id\":")); + assert!(completed + .data_json + .contains(receiver.core.status().endpoint_id.as_str())); receiver.core.delete_receive_history().unwrap(); assert_eq!(receiver.core.list_received_artifacts().unwrap(), artifacts);