cdricms/vnidrop
1
1
Fork 0
mirror of https://github.com/sudosylabs/vnidrop.git synced 2026-08-05 02:29:55 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(core): abort send when provider stream closes

Browse Source
This commit is contained in:
Hammed Abass
2026-07-24 00:11:58 +02:00
parent 4074f4bee8
commit 0448137d84
4 changed files with 144 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/vnidrop/src/runtime/mod.rs
View File

@@ -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},

107
crates/vnidrop/src/runtime/provider.rs
View File

@@ -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<RequestUpdate>,
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<Self>,
@@ -302,7 +327,7 @@ impl CoreInner {
transfer_id: u64,
connection_id: u64,
request_id: u64,
mut rx: irpc::channel::mpsc::Receiver<RequestUpdate>,
rx: irpc::channel::mpsc::Receiver<RequestUpdate>,
) {
// 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,
}),
);
}
});
}

42
crates/vnidrop/src/tests/runtime.rs
View File

@@ -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::<RequestUpdate>(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();

29
crates/vnidrop/tests/transfer.rs
View File

@@ -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);