Skip to content

Commit

Permalink
Fixes rust-lang#2748 by filling the streaming buffers fully
Browse files Browse the repository at this point in the history
Diagnosis by Daniel. Code by Robert.
  • Loading branch information
Robert Collins authored and rbtcollins committed May 4, 2021
1 parent 8e632bb commit 094e497
Showing 1 changed file with 23 additions and 11 deletions.
34 changes: 23 additions & 11 deletions src/dist/component/package.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
//! for installing from a directory or tarball to an installation
//! prefix, represented by a `Components` instance.
use std::cmp::min;
use std::collections::{HashMap, HashSet};
use std::fmt;
use std::io::{self, ErrorKind as IOErrorKind, Read};
Expand Down Expand Up @@ -364,16 +365,19 @@ fn unpack_without_first_dir<'a, R: Read>(
continue;
}

struct SenderEntry<'a, 'b, R: std::io::Read> {
sender: Box<dyn FnMut(Vec<u8>) -> bool + 'a>,
entry: tar::Entry<'b, R>,
read: usize,
}

/// true if either no sender_entry was provided, or the incremental file
/// has been fully dispatched.
fn flush_ios<'a, R: std::io::Read, P: AsRef<Path>>(
mut budget: &mut MemoryBudget,
io_executor: &dyn Executor,
mut directories: &mut HashMap<PathBuf, DirStatus>,
mut sender_entry: Option<&mut (
Box<dyn FnMut(Vec<u8>) -> bool + 'a>,
&mut tar::Entry<'_, R>,
)>,
mut sender_entry: Option<&mut SenderEntry<'a, '_, R>>,
full_path: P,
) -> Result<bool> {
let mut result = sender_entry.is_none();
Expand All @@ -384,16 +388,19 @@ fn unpack_without_first_dir<'a, R: Read>(
trigger_children(&*io_executor, &mut directories, &mut budget, op)?;
}
// Maybe stream a file incrementally
if let Some((sender, entry)) = sender_entry.as_mut() {
if let Some(sender) = sender_entry.as_mut() {
if budget.available() as u64 >= IO_CHUNK_SIZE {
let mut v = vec![0; IO_CHUNK_SIZE as usize];
let len = entry.read(&mut v)?;
let chunk_size = min(IO_CHUNK_SIZE, sender.entry.size() - sender.read as u64);
let mut v = vec![0; chunk_size as usize];
let len = v.len();
if len == 0 {
result = true;
} else {
sender.entry.read_exact(&mut v)?;
sender.read += v.len();
}
v.resize(len, 0);
budget.claim_chunk(len);
if !sender(v) {
if !(sender.sender)(v) {
bail!(format!(
"IO receiver for '{}' disconnected",
full_path.as_ref().display()
Expand Down Expand Up @@ -519,8 +526,13 @@ fn unpack_without_first_dir<'a, R: Read>(
}
}

let mut incremental_file_sender = incremental_file_sender
.map(|incremental_file_sender| (incremental_file_sender, &mut entry));
let mut incremental_file_sender = incremental_file_sender.map(|incremental_file_sender| {
(SenderEntry {
sender: incremental_file_sender,
entry,
read: 0,
})
});

// monitor io queue and feed in the content of the file (if needed)
while !flush_ios(
Expand Down

0 comments on commit 094e497

Please sign in to comment.