Skip to content

fix(receive): await Endpoint::close for iroh 0.97 #124

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cbenhagen wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(receive): await Endpoint::close for iroh 0.97 #124

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
177 changes: 95 additions & 82 deletions src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1016,101 +1016,113 @@ async fn receive(args: ReceiveArgs) -> anyhow::Result<()> {
builder = builder.bind_addr(addr)?;
}
let endpoint = builder.bind().await?;
let endpoint_work = endpoint.clone();
let dir_name = format!(".sendme-recv-{}", ticket.hash().to_hex());
let iroh_data_dir = std::env::current_dir()?.join(dir_name);
let db = iroh_blobs::store::fs::FsStore::load(&iroh_data_dir).await?;
let db2 = db.clone();
trace!("load done!");
let fut = async move {
trace!("running");
let mut mp: MultiProgress = MultiProgress::new();
let draw_target = if args.common.no_progress {
ProgressDrawTarget::hidden()
} else {
ProgressDrawTarget::stderr()
};
mp.set_draw_target(draw_target);
let hash_and_format = ticket.hash_and_format();
trace!("computing local");
let local = db.remote().local(hash_and_format).await?;
trace!("local done");
let (stats, total_files, payload_size) = if !local.is_complete() {
trace!("{} not complete", hash_and_format.hash);
let cp = mp.add(make_connect_progress());
let connection = endpoint.connect(addr, iroh_blobs::protocol::ALPN).await?;
cp.finish_and_clear();
let sp = mp.add(make_get_sizes_progress());
let (_hash_seq, sizes) =
get_hash_seq_and_sizes(&connection, &hash_and_format.hash, 1024 * 1024 * 32, None)
.await
.map_err(show_get_error)?;
sp.finish_and_clear();
let total_size = sizes.iter().copied().sum::<u64>();
let payload_size = sizes.iter().skip(2).copied().sum::<u64>();
let total_files = (sizes.len().saturating_sub(1)) as u64;
eprintln!(
"getting collection {} {} files, {}",
print_hash(&ticket.hash(), args.common.format),
total_files,
HumanBytes(payload_size)
);
// print the details of the collection only in verbose mode
if args.common.verbose > 0 {
let res = async {
trace!("running");
let mut mp: MultiProgress = MultiProgress::new();
let draw_target = if args.common.no_progress {
ProgressDrawTarget::hidden()
} else {
ProgressDrawTarget::stderr()
};
mp.set_draw_target(draw_target);
let hash_and_format = ticket.hash_and_format();
trace!("computing local");
let local = db.remote().local(hash_and_format).await?;
trace!("local done");
let (stats, total_files, payload_size) = if !local.is_complete() {
trace!("{} not complete", hash_and_format.hash);
let cp = mp.add(make_connect_progress());
let connection = endpoint_work
.connect(addr, iroh_blobs::protocol::ALPN)
.await?;
cp.finish_and_clear();
let sp = mp.add(make_get_sizes_progress());
let (_hash_seq, sizes) = get_hash_seq_and_sizes(
&connection,
&hash_and_format.hash,
1024 * 1024 * 32,
None,
)
.await
.map_err(show_get_error)?;
sp.finish_and_clear();
let total_size = sizes.iter().copied().sum::<u64>();
let payload_size = sizes.iter().skip(2).copied().sum::<u64>();
let total_files = (sizes.len().saturating_sub(1)) as u64;
eprintln!(
"getting {} blobs in total, {}",
total_files + 1,
HumanBytes(total_size)
"getting collection {} {} files, {}",
print_hash(&ticket.hash(), args.common.format),
total_files,
HumanBytes(payload_size)
);
}
let (tx, rx) = mpsc::channel(32);
let local_size = local.local_bytes();
let get = db.remote().execute_get(connection, local.missing());
let task = tokio::spawn(show_download_progress(
mp.clone(),
rx,
local_size,
total_size,
));
// let mut stream = get.stream();
let mut stats = Stats::default();
let mut stream = get.stream();
while let Some(item) = stream.next().await {
trace!("got item {item:?}");
match item {
GetProgressItem::Progress(offset) => {
tx.send(offset).await.ok();
}
GetProgressItem::Done(value) => {
stats = value;
break;
}
GetProgressItem::Error(cause) => {
anyhow::bail!(show_get_error(cause));
// print the details of the collection only in verbose mode
if args.common.verbose > 0 {
eprintln!(
"getting {} blobs in total, {}",
total_files + 1,
HumanBytes(total_size)
);
}
let (tx, rx) = mpsc::channel(32);
let local_size = local.local_bytes();
let get = db.remote().execute_get(connection, local.missing());
let task = tokio::spawn(show_download_progress(
mp.clone(),
rx,
local_size,
total_size,
));
// let mut stream = get.stream();
let mut stats = Stats::default();
let mut stream = get.stream();
while let Some(item) = stream.next().await {
trace!("got item {item:?}");
match item {
GetProgressItem::Progress(offset) => {
tx.send(offset).await.ok();
}
GetProgressItem::Done(value) => {
stats = value;
break;
}
GetProgressItem::Error(cause) => {
anyhow::bail!(show_get_error(cause));
}
}
}
drop(tx);
task.await.ok();
(stats, total_files, payload_size)
} else {
println!("{} already complete", hash_and_format.hash);
let total_files = local.children().unwrap() - 1;
let payload_bytes = 0; // todo local.sizes().skip(2).map(Option::unwrap).sum::<u64>();
(Stats::default(), total_files, payload_bytes)
};
let collection = Collection::load(hash_and_format.hash, db.as_ref()).await?;
if args.common.verbose > 1 {
for (name, hash) in collection.iter() {
println!(" {} {name}", print_hash(hash, args.common.format));
}
}
drop(tx);
task.await.ok();
(stats, total_files, payload_size)
} else {
println!("{} already complete", hash_and_format.hash);
let total_files = local.children().unwrap() - 1;
let payload_bytes = 0; // todo local.sizes().skip(2).map(Option::unwrap).sum::<u64>();
(Stats::default(), total_files, payload_bytes)
};
let collection = Collection::load(hash_and_format.hash, db.as_ref()).await?;
if args.common.verbose > 1 {
for (name, hash) in collection.iter() {
println!(" {} {name}", print_hash(hash, args.common.format));
}
}
if let Some((name, _)) = collection.iter().next() {
if let Some(first) = name.split('/').next() {
println!("exporting to {first}");
if let Some((name, _)) = collection.iter().next() {
if let Some(first) = name.split('/').next() {
println!("exporting to {first}");
}
}
export(&db, collection, &mut mp).await?;
anyhow::Ok((total_files, payload_size, stats))
}
export(&db, collection, &mut mp).await?;
anyhow::Ok((total_files, payload_size, stats))
.await;
endpoint_work.close().await;
res
};
let (total_files, payload_size, stats) = select! {
x = fut => match x {
Expand All @@ -1123,6 +1135,7 @@ async fn receive(args: ReceiveArgs) -> anyhow::Result<()> {
}
},
_ = tokio::signal::ctrl_c() => {
endpoint.close().await;
db2.shutdown().await?;
std::process::exit(130);
}
Expand Down
42 changes: 42 additions & 0 deletions tests/cli.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,48 @@ fn send_recv_file() {
assert_eq!(tgt_data, data);
}

#[test]
fn receive_closes_endpoint_no_iroh_socket_error() {
let name = "graceful-close.bin";
let data = vec![0xabu8; 64];
let src_dir = tempfile::tempdir().unwrap();
let tgt_dir = tempfile::tempdir().unwrap();
let src_file = src_dir.path().join(name);
std::fs::write(&src_file, &data).unwrap();
let mut send_cmd = duct::cmd(
sendme_bin(),
["send", src_file.as_os_str().to_str().unwrap()],
)
.dir(src_dir.path())
.env_remove("RUST_LOG")
.stderr_to_stdout()
.reader()
.unwrap();
let output = read_ascii_lines(3, &mut send_cmd).unwrap();
let output = String::from_utf8(output).unwrap();
let ticket = output.split_ascii_whitespace().last().unwrap();
let ticket = BlobTicket::from_str(ticket).unwrap();
let receive_output = duct::cmd(sendme_bin(), ["receive", &ticket.to_string()])
.dir(tgt_dir.path())
.env("RUST_LOG", "iroh::socket=error")
.stdout_capture()
.stderr_capture()
.run()
.unwrap();
assert!(receive_output.status.success(), "{receive_output:?}");
let stderr = String::from_utf8_lossy(&receive_output.stderr);
assert!(
!stderr.contains("Endpoint dropped"),
"unexpected iroh shutdown log on stderr: {stderr}"
);
assert!(
!stderr.contains("Aborting ungracefully"),
"unexpected iroh shutdown log on stderr: {stderr}"
);
let tgt_file = tgt_dir.path().join(name);
assert_eq!(std::fs::read(&tgt_file).unwrap(), data);
}

#[test]
fn send_recv_dir() {
fn create_file(base: &Path, i: usize, j: usize, k: usize) -> (PathBuf, Vec<u8>) {
Expand Down