From 193ac8bd091c84249d3d6ebc431177b566489afb Mon Sep 17 00:00:00 2001 From: Benjamin Gilbert Date: Wed, 1 Jul 2020 17:03:53 -0400 Subject: [PATCH 1/2] osmet: move copy_n and copy_exactly_n out of osmet --- src/io.rs | 69 +++++++++++++++++++++++++++++++++++++++++ src/main.rs | 1 + src/osmet/io_helpers.rs | 53 +------------------------------ src/osmet/mod.rs | 1 + 4 files changed, 72 insertions(+), 52 deletions(-) create mode 100644 src/io.rs diff --git a/src/io.rs b/src/io.rs new file mode 100644 index 000000000..601208b3e --- /dev/null +++ b/src/io.rs @@ -0,0 +1,69 @@ +// Copyright 2019 CoreOS, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use error_chain::bail; +use std::io::{ErrorKind, Read, Write}; + +use crate::errors::*; + +/// This is like `std::io:copy()`, but limits the number of bytes copied over. The `Read` trait has +/// `take()`, but that takes ownership of the reader. We also take a buf to avoid re-initializing a +/// block each time (std::io::copy() gets around this by using MaybeUninit, but that requires using +/// nightly and unsafe functions). +pub fn copy_n( + reader: &mut impl Read, + writer: &mut impl Write, + mut n: u64, + buf: &mut [u8], +) -> Result { + let mut written = 0; + loop { + if n == 0 { + return Ok(written); + } + let bufn = if n < (buf.len() as u64) { + &mut buf[..n as usize] + } else { + &mut buf[..] + }; + let len = match reader.read(bufn) { + Ok(0) => return Ok(written), + Ok(len) => len, + Err(ref e) if e.kind() == ErrorKind::Interrupted => continue, + Err(e) => return Err(e.into()), + }; + assert!(len as u64 <= n); + writer.write_all(&bufn[..len])?; + written += len as u64; + n -= len as u64; + } +} + +/// This is like `copy_n()` but errors if the number of bytes copied is less than expected. +pub fn copy_exactly_n( + reader: &mut impl Read, + writer: &mut impl Write, + n: u64, + buf: &mut [u8], +) -> Result { + let bytes_copied = copy_n(reader, writer, n, buf)?; + if bytes_copied != n { + bail!( + "expected to copy {} bytes but instead copied {} bytes", + n, + bytes_copied + ); + } + Ok(n) +} diff --git a/src/main.rs b/src/main.rs index 2145cb4ec..21e895219 100644 --- a/src/main.rs +++ b/src/main.rs @@ -17,6 +17,7 @@ mod cmdline; mod download; mod errors; mod install; +mod io; mod iso; mod osmet; mod source; diff --git a/src/osmet/io_helpers.rs b/src/osmet/io_helpers.rs index c755e2bab..61472ceb5 100644 --- a/src/osmet/io_helpers.rs +++ b/src/osmet/io_helpers.rs @@ -14,7 +14,7 @@ use std::convert::{TryFrom, TryInto}; use std::fs::OpenOptions; -use std::io::{self, ErrorKind, Read, Write}; +use std::io::{self, Write}; use std::os::unix::io::AsRawFd; use std::path::Path; @@ -41,57 +41,6 @@ impl TryFrom for Sha256Digest { } } -/// This is like `std::io:copy()`, but limits the number of bytes copied over. The `Read` trait has -/// `take()`, but that takes ownership of the reader. We also take a buf to avoid re-initializing a -/// block each time (std::io::copy() gets around this by using MaybeUninit, but that requires using -/// nightly and unsafe functions). -pub fn copy_n( - reader: &mut impl Read, - writer: &mut impl Write, - mut n: u64, - buf: &mut [u8], -) -> Result { - let mut written = 0; - loop { - if n == 0 { - return Ok(written); - } - let bufn = if n < (buf.len() as u64) { - &mut buf[..n as usize] - } else { - &mut buf[..] - }; - let len = match reader.read(bufn) { - Ok(0) => return Ok(written), - Ok(len) => len, - Err(ref e) if e.kind() == ErrorKind::Interrupted => continue, - Err(e) => return Err(e.into()), - }; - assert!(len as u64 <= n); - writer.write_all(&bufn[..len])?; - written += len as u64; - n -= len as u64; - } -} - -/// This is like `copy_n()` but errors if the number of bytes copied is less than expected. -pub fn copy_exactly_n( - reader: &mut impl Read, - writer: &mut impl Write, - n: u64, - buf: &mut [u8], -) -> Result { - let bytes_copied = copy_n(reader, writer, n, buf)?; - if bytes_copied != n { - bail!( - "expected to copy {} bytes but instead copied {} bytes", - n, - bytes_copied - ); - } - Ok(n) -} - // ab/cdef....file --> 0xabcdef... pub fn object_path_to_checksum(path: &Path) -> Result { let chksum2 = path diff --git a/src/osmet/mod.rs b/src/osmet/mod.rs index 5460e1419..af14eb43f 100644 --- a/src/osmet/mod.rs +++ b/src/osmet/mod.rs @@ -37,6 +37,7 @@ use xz2::write::XzEncoder; use crate::blockdev::*; use crate::cmdline::*; use crate::errors::*; +use crate::io::*; mod fiemap; mod file; From 63e9dd2b97468475a1bbfa9328732f00047871e3 Mon Sep 17 00:00:00 2001 From: Benjamin Gilbert Date: Wed, 1 Jul 2020 17:18:15 -0400 Subject: [PATCH 2/2] *: use our own std::io::copy() implementation The std one uses an 8 KiB buffer. Use a larger buffer to amortize syscall overhead. --- src/download.rs | 3 ++- src/install.rs | 3 ++- src/io.rs | 9 +++++++++ src/iso.rs | 3 ++- src/osmet/file.rs | 4 ++-- src/osmet/io_helpers.rs | 4 ++-- src/osmet/mod.rs | 6 +++--- src/osmet/unpacker.rs | 2 +- 8 files changed, 23 insertions(+), 11 deletions(-) diff --git a/src/download.rs b/src/download.rs index e76b400df..08de56703 100644 --- a/src/download.rs +++ b/src/download.rs @@ -18,7 +18,7 @@ use flate2::read::GzDecoder; use nix::unistd::isatty; use progress_streams::ProgressReader; use std::fs::{remove_file, File, OpenOptions}; -use std::io::{copy, stderr, BufRead, BufReader, Read, Seek, SeekFrom, Write}; +use std::io::{stderr, BufRead, BufReader, Read, Seek, SeekFrom, Write}; use std::num::NonZeroU32; use std::os::unix::io::AsRawFd; use std::path::{Path, PathBuf}; @@ -28,6 +28,7 @@ use xz2::read::XzDecoder; use crate::blockdev::detect_formatted_sector_size; use crate::cmdline::*; use crate::errors::*; +use crate::io::*; use crate::source::*; use crate::verify::*; diff --git a/src/install.rs b/src/install.rs index fe51f8334..c85295c02 100644 --- a/src/install.rs +++ b/src/install.rs @@ -15,7 +15,7 @@ use error_chain::{bail, ensure, ChainedError}; use nix::mount; use std::fs::{copy as fscopy, create_dir_all, read_dir, File, OpenOptions}; -use std::io::{copy, Read, Seek, SeekFrom, Write}; +use std::io::{Read, Seek, SeekFrom, Write}; use std::os::unix::fs::FileTypeExt; use std::path::Path; @@ -23,6 +23,7 @@ use crate::blockdev::*; use crate::cmdline::*; use crate::download::*; use crate::errors::*; +use crate::io::*; use crate::source::*; /// Integrity verification hash for an Ignition config. diff --git a/src/io.rs b/src/io.rs index 601208b3e..973b2ee09 100644 --- a/src/io.rs +++ b/src/io.rs @@ -17,6 +17,15 @@ use std::io::{ErrorKind, Read, Write}; use crate::errors::*; +/// This is like `std::io:copy()`, but uses a buffer larger than 8 KiB +/// to amortize syscall overhead. +pub fn copy(reader: &mut impl Read, writer: &mut impl Write) -> Result { + // https://github.com/rust-lang/rust/issues/49921 + // https://github.com/coreutils/coreutils/blob/6a3d2883/src/ioblksize.h + let mut buf = [0u8; 256 * 1024]; + copy_n(reader, writer, std::u64::MAX, &mut buf) +} + /// This is like `std::io:copy()`, but limits the number of bytes copied over. The `Read` trait has /// `take()`, but that takes ownership of the reader. We also take a buf to avoid re-initializing a /// block each time (std::io::copy() gets around this by using MaybeUninit, but that requires using diff --git a/src/iso.rs b/src/iso.rs index 6a3b47451..df45fd41d 100644 --- a/src/iso.rs +++ b/src/iso.rs @@ -18,10 +18,11 @@ use flate2::read::GzDecoder; use flate2::{Compression, GzBuilder}; use std::convert::TryInto; use std::fs::{remove_file, File, OpenOptions}; -use std::io::{copy, stdin, Cursor, Read, Seek, SeekFrom, Write}; +use std::io::{stdin, Cursor, Read, Seek, SeekFrom, Write}; use crate::cmdline::*; use crate::errors::*; +use crate::io::*; const FILENAME: &str = "config.ign"; diff --git a/src/osmet/file.rs b/src/osmet/file.rs index 6c16164be..f92c59885 100644 --- a/src/osmet/file.rs +++ b/src/osmet/file.rs @@ -13,7 +13,7 @@ // limitations under the License. use std::fs::{File, OpenOptions}; -use std::io::{self, BufReader, BufWriter, Read}; +use std::io::{BufReader, BufWriter, Read}; use std::path::Path; use bincode::Options; @@ -90,7 +90,7 @@ pub(super) fn osmet_file_write( .chain_err(|| "failed to serialize osmet")?; // and followed by the xz-compressed packed image - io::copy(&mut xzpacked_image, &mut f)?; + copy(&mut xzpacked_image, &mut f)?; f.into_inner() .chain_err(|| "failed to flush write buffer")? diff --git a/src/osmet/io_helpers.rs b/src/osmet/io_helpers.rs index 61472ceb5..890aa3fc0 100644 --- a/src/osmet/io_helpers.rs +++ b/src/osmet/io_helpers.rs @@ -14,7 +14,7 @@ use std::convert::{TryFrom, TryInto}; use std::fs::OpenOptions; -use std::io::{self, Write}; +use std::io::Write; use std::os::unix::io::AsRawFd; use std::path::Path; @@ -101,7 +101,7 @@ pub fn get_path_digest(path: &Path) -> Result { } let mut hasher = Hasher::new(MessageDigest::sha256()).chain_err(|| "creating SHA256 hasher")?; - io::copy(&mut f, &mut hasher)?; + copy(&mut f, &mut hasher)?; Ok(hasher.try_into()?) } diff --git a/src/osmet/mod.rs b/src/osmet/mod.rs index af14eb43f..5dbfebdc0 100644 --- a/src/osmet/mod.rs +++ b/src/osmet/mod.rs @@ -24,7 +24,7 @@ use std::collections::hash_map::Entry; use std::collections::HashMap; use std::convert::TryInto; use std::fs::{File, OpenOptions}; -use std::io::{self, Seek, SeekFrom, Write}; +use std::io::{Seek, SeekFrom, Write}; use std::os::unix::fs::FileTypeExt; use std::path::{Path, PathBuf}; @@ -162,7 +162,7 @@ pub fn osmet_unpack(config: &OsmetUnpackConfig) -> Result<()> { } let mut unpacker = OsmetUnpacker::new(Path::new(&config.osmet), Path::new(&config.repo))?; - io::copy(&mut unpacker, &mut dev) + copy(&mut unpacker, &mut dev) .chain_err(|| format!("copying to block device {}", &config.device))?; Ok(()) @@ -427,7 +427,7 @@ fn write_packed_image( } // and finally write out the remainder of the disk - io::copy(dev, w).chain_err(|| "copying remainder of disk")?; + copy(dev, w).chain_err(|| "copying remainder of disk")?; Ok(total_bytes_skipped) } diff --git a/src/osmet/unpacker.rs b/src/osmet/unpacker.rs index c70852e14..1d4528202 100644 --- a/src/osmet/unpacker.rs +++ b/src/osmet/unpacker.rs @@ -182,7 +182,7 @@ fn write_unpacked_image( } // and copy the rest - cursor += io::copy(packed_image, w)?; + cursor += copy(packed_image, w)?; Ok(cursor) }