1
Fork 0
mirror of https://github.com/RGBCube/uutils-coreutils synced 2025-07-28 11:37:44 +00:00

Merge pull request #2654 from blyxxyz/dedup-splice

Move common pipe and splice functions into uucore
This commit is contained in:
Sylvestre Ledru 2021-09-12 09:23:05 +02:00 committed by GitHub
commit 447d6f2b61
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
12 changed files with 135 additions and 142 deletions

1
Cargo.lock generated
View file

@ -3275,6 +3275,7 @@ dependencies = [
"getopts",
"lazy_static",
"libc",
"nix 0.20.0",
"once_cell",
"termion",
"thiserror",

View file

@ -18,7 +18,7 @@ path = "src/cat.rs"
clap = { version = "2.33", features = ["wrap_help"] }
thiserror = "1.0"
atty = "0.2"
uucore = { version=">=0.0.9", package="uucore", path="../../uucore", features=["fs"] }
uucore = { version=">=0.0.9", package="uucore", path="../../uucore", features=["fs", "pipes"] }
uucore_procs = { version=">=0.0.6", package="uucore_procs", path="../../uucore_procs" }
[target.'cfg(unix)'.dependencies]

View file

@ -29,8 +29,6 @@ use std::os::unix::io::AsRawFd;
/// Linux splice support
#[cfg(any(target_os = "linux", target_os = "android"))]
mod splice;
#[cfg(any(target_os = "linux", target_os = "android"))]
use std::os::unix::io::RawFd;
/// Unix domain socket support
#[cfg(unix)]
@ -137,10 +135,18 @@ struct OutputState {
one_blank_kept: bool,
}
#[cfg(unix)]
trait FdReadable: Read + AsRawFd {}
#[cfg(not(unix))]
trait FdReadable: Read {}
#[cfg(unix)]
impl<T> FdReadable for T where T: Read + AsRawFd {}
#[cfg(not(unix))]
impl<T> FdReadable for T where T: Read {}
/// Represents an open file handle, stream, or other device
struct InputHandle<R: Read> {
#[cfg(any(target_os = "linux", target_os = "android"))]
file_descriptor: RawFd,
struct InputHandle<R: FdReadable> {
reader: R,
is_interactive: bool,
}
@ -297,7 +303,7 @@ pub fn uu_app() -> App<'static, 'static> {
)
}
fn cat_handle<R: Read>(
fn cat_handle<R: FdReadable>(
handle: &mut InputHandle<R>,
options: &OutputOptions,
state: &mut OutputState,
@ -319,8 +325,6 @@ fn cat_path(
if path == "-" {
let stdin = io::stdin();
let mut handle = InputHandle {
#[cfg(any(target_os = "linux", target_os = "android"))]
file_descriptor: stdin.as_raw_fd(),
reader: stdin,
is_interactive: atty::is(atty::Stream::Stdin),
};
@ -333,8 +337,6 @@ fn cat_path(
let socket = UnixStream::connect(path)?;
socket.shutdown(Shutdown::Write)?;
let mut handle = InputHandle {
#[cfg(any(target_os = "linux", target_os = "android"))]
file_descriptor: socket.as_raw_fd(),
reader: socket,
is_interactive: false,
};
@ -347,8 +349,6 @@ fn cat_path(
return Err(CatError::OutputIsInput);
}
let mut handle = InputHandle {
#[cfg(any(target_os = "linux", target_os = "android"))]
file_descriptor: file.as_raw_fd(),
reader: file,
is_interactive: false,
};
@ -437,14 +437,14 @@ fn get_input_type(path: &str) -> CatResult<InputType> {
/// Writes handle to stdout with no configuration. This allows a
/// simple memory copy.
fn write_fast<R: Read>(handle: &mut InputHandle<R>) -> CatResult<()> {
fn write_fast<R: FdReadable>(handle: &mut InputHandle<R>) -> CatResult<()> {
let stdout = io::stdout();
let mut stdout_lock = stdout.lock();
#[cfg(any(target_os = "linux", target_os = "android"))]
{
// If we're on Linux or Android, try to use the splice() system call
// for faster writing. If it works, we're done.
if !splice::write_fast_using_splice(handle, stdout_lock.as_raw_fd())? {
if !splice::write_fast_using_splice(handle, &stdout_lock)? {
return Ok(());
}
}
@ -462,7 +462,7 @@ fn write_fast<R: Read>(handle: &mut InputHandle<R>) -> CatResult<()> {
/// Outputs file contents to stdout in a line-by-line fashion,
/// propagating any errors that might occur.
fn write_lines<R: Read>(
fn write_lines<R: FdReadable>(
handle: &mut InputHandle<R>,
options: &OutputOptions,
state: &mut OutputState,

View file

@ -1,11 +1,11 @@
use super::{CatResult, InputHandle};
use super::{CatResult, FdReadable, InputHandle};
use nix::fcntl::{splice, SpliceFFlags};
use nix::unistd::{self, pipe};
use std::fs::File;
use std::io::Read;
use std::os::unix::io::{FromRawFd, RawFd};
use nix::unistd;
use std::os::unix::io::{AsRawFd, RawFd};
use uucore::pipes::{pipe, splice, splice_exact};
const SPLICE_SIZE: usize = 1024 * 128;
const BUF_SIZE: usize = 1024 * 16;
/// This function is called from `write_fast()` on Linux and Android. The
@ -16,36 +16,25 @@ const BUF_SIZE: usize = 1024 * 16;
/// The `bool` in the result value indicates if we need to fall back to normal
/// copying or not. False means we don't have to.
#[inline]
pub(super) fn write_fast_using_splice<R: Read>(
pub(super) fn write_fast_using_splice<R: FdReadable>(
handle: &mut InputHandle<R>,
write_fd: RawFd,
write_fd: &impl AsRawFd,
) -> CatResult<bool> {
let (pipe_rd, pipe_wr) = pipe()?;
// Ensure the pipe is closed when the function returns.
// SAFETY: The file descriptors do not have other owners.
let _handles = unsafe { (File::from_raw_fd(pipe_rd), File::from_raw_fd(pipe_wr)) };
loop {
match splice(
handle.file_descriptor,
None,
pipe_wr,
None,
BUF_SIZE,
SpliceFFlags::empty(),
) {
match splice(&handle.reader, &pipe_wr, SPLICE_SIZE) {
Ok(n) => {
if n == 0 {
return Ok(false);
}
if splice_exact(pipe_rd, write_fd, n).is_err() {
if splice_exact(&pipe_rd, write_fd, n).is_err() {
// If the first splice manages to copy to the intermediate
// pipe, but the second splice to stdout fails for some reason
// we can recover by copying the data that we have from the
// intermediate pipe to stdout using normal read/write. Then
// we tell the caller to fall back.
copy_exact(pipe_rd, write_fd, n)?;
copy_exact(pipe_rd.as_raw_fd(), write_fd.as_raw_fd(), n)?;
return Ok(true);
}
}
@ -56,35 +45,23 @@ pub(super) fn write_fast_using_splice<R: Read>(
}
}
/// Splice wrapper which handles short writes.
#[inline]
fn splice_exact(read_fd: RawFd, write_fd: RawFd, num_bytes: usize) -> nix::Result<()> {
let mut left = num_bytes;
loop {
let written = splice(read_fd, None, write_fd, None, left, SpliceFFlags::empty())?;
left -= written;
if left == 0 {
break;
}
}
Ok(())
}
/// Caller must ensure that `num_bytes <= BUF_SIZE`, otherwise this function
/// will panic. The way we use this function in `write_fast_using_splice`
/// above is safe because `splice` is set to write at most `BUF_SIZE` to the
/// pipe.
#[inline]
/// Move exactly `num_bytes` bytes from `read_fd` to `write_fd`.
///
/// Panics if not enough bytes can be read.
fn copy_exact(read_fd: RawFd, write_fd: RawFd, num_bytes: usize) -> nix::Result<()> {
let mut left = num_bytes;
let mut buf = [0; BUF_SIZE];
loop {
let read = unistd::read(read_fd, &mut buf[..left])?;
let written = unistd::write(write_fd, &buf[..read])?;
left -= written;
if left == 0 {
break;
while left > 0 {
let read = unistd::read(read_fd, &mut buf)?;
assert_ne!(read, 0, "unexpected end of pipe");
let mut written = 0;
while written < read {
match unistd::write(write_fd, &buf[written..read])? {
0 => panic!(),
n => written += n,
}
}
left -= read;
}
Ok(())
}

View file

@ -16,7 +16,7 @@ path = "src/wc.rs"
[dependencies]
clap = { version = "2.33", features = ["wrap_help"] }
uucore = { version=">=0.0.9", package="uucore", path="../../uucore" }
uucore = { version=">=0.0.9", package="uucore", path="../../uucore", features=["pipes"] }
uucore_procs = { version=">=0.0.6", package="uucore_procs", path="../../uucore_procs" }
bytecount = "0.6.2"
utf-8 = "0.7.6"

View file

@ -3,7 +3,7 @@ use crate::word_count::WordCount;
use super::WordCountable;
#[cfg(any(target_os = "linux", target_os = "android"))]
use std::fs::{File, OpenOptions};
use std::fs::OpenOptions;
use std::io::{self, ErrorKind, Read};
#[cfg(unix)]
@ -11,34 +11,17 @@ use libc::S_IFREG;
#[cfg(unix)]
use nix::sys::stat;
#[cfg(any(target_os = "linux", target_os = "android"))]
use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
use std::os::unix::io::AsRawFd;
#[cfg(any(target_os = "linux", target_os = "android"))]
use libc::S_IFIFO;
#[cfg(any(target_os = "linux", target_os = "android"))]
use nix::fcntl::{splice, SpliceFFlags};
#[cfg(any(target_os = "linux", target_os = "android"))]
use nix::unistd::pipe;
use uucore::pipes::{pipe, splice, splice_exact};
const BUF_SIZE: usize = 16 * 1024;
#[cfg(any(target_os = "linux", target_os = "android"))]
const SPLICE_SIZE: usize = 128 * 1024;
/// Splice wrapper which handles short writes
#[cfg(any(target_os = "linux", target_os = "android"))]
#[inline]
fn splice_exact(read_fd: RawFd, write_fd: RawFd, num_bytes: usize) -> nix::Result<()> {
let mut left = num_bytes;
loop {
let written = splice(read_fd, None, write_fd, None, left, SpliceFFlags::empty())?;
left -= written;
if left == 0 {
break;
}
}
Ok(())
}
/// This is a Linux-specific function to count the number of bytes using the
/// `splice` system call, which is faster than using `read`.
///
@ -46,13 +29,14 @@ fn splice_exact(read_fd: RawFd, write_fd: RawFd, num_bytes: usize) -> nix::Resul
/// caller will fall back to a simpler method.
#[inline]
#[cfg(any(target_os = "linux", target_os = "android"))]
fn count_bytes_using_splice(fd: RawFd) -> Result<usize, usize> {
fn count_bytes_using_splice(fd: &impl AsRawFd) -> Result<usize, usize> {
let null_file = OpenOptions::new()
.write(true)
.open("/dev/null")
.map_err(|_| 0_usize)?;
let null = null_file.as_raw_fd();
let null_rdev = stat::fstat(null).map_err(|_| 0_usize)?.st_rdev;
let null_rdev = stat::fstat(null_file.as_raw_fd())
.map_err(|_| 0_usize)?
.st_rdev;
if (stat::major(null_rdev), stat::minor(null_rdev)) != (1, 3) {
// This is not a proper /dev/null, writing to it is probably bad
// Bit of an edge case, but it has been known to happen
@ -60,17 +44,13 @@ fn count_bytes_using_splice(fd: RawFd) -> Result<usize, usize> {
}
let (pipe_rd, pipe_wr) = pipe().map_err(|_| 0_usize)?;
// Ensure the pipe is closed when the function returns.
// SAFETY: The file descriptors do not have other owners.
let _handles = unsafe { (File::from_raw_fd(pipe_rd), File::from_raw_fd(pipe_wr)) };
let mut byte_count = 0;
loop {
match splice(fd, None, pipe_wr, None, SPLICE_SIZE, SpliceFFlags::empty()) {
match splice(fd, &pipe_wr, SPLICE_SIZE) {
Ok(0) => break,
Ok(res) => {
byte_count += res;
if splice_exact(pipe_rd, null, res).is_err() {
if splice_exact(&pipe_rd, &null_file, res).is_err() {
return Err(byte_count);
}
}
@ -106,7 +86,7 @@ pub(crate) fn count_bytes_fast<T: WordCountable>(handle: &mut T) -> (usize, Opti
// Else, if we're on Linux and our file is a FIFO pipe
// (or stdin), we use splice to count the number of bytes.
if (stat.st_mode & S_IFIFO) != 0 {
match count_bytes_using_splice(fd) {
match count_bytes_using_splice(handle) {
Ok(n) => return (n, None),
Err(n) => byte_count = n,
}

View file

@ -16,7 +16,7 @@ path = "src/yes.rs"
[dependencies]
clap = { version = "2.33", features = ["wrap_help"] }
uucore = { version=">=0.0.9", package="uucore", path="../../uucore" }
uucore = { version=">=0.0.9", package="uucore", path="../../uucore", features=["pipes"] }
uucore_procs = { version=">=0.0.6", package="uucore_procs", path="../../uucore_procs" }
[target.'cfg(any(target_os = "linux", target_os = "android"))'.dependencies]

View file

@ -16,18 +16,11 @@
//! make any effort to rescue data from the pipe if splice() fails, we can
//! just fall back and start over from the beginning.
use std::{
fs::File,
io,
os::unix::io::{AsRawFd, FromRawFd},
};
use std::{io, os::unix::io::AsRawFd};
use nix::{
errno::Errno,
fcntl::SpliceFFlags,
libc::S_IFIFO,
sys::{stat::fstat, uio::IoVec},
};
use nix::{errno::Errno, libc::S_IFIFO, sys::stat::fstat};
use uucore::pipes::{pipe, splice_exact, vmsplice};
pub(crate) fn splice_data(bytes: &[u8], out: &impl AsRawFd) -> Result<()> {
let is_pipe = fstat(out.as_raw_fd())?.st_mode & S_IFIFO != 0;
@ -36,7 +29,7 @@ pub(crate) fn splice_data(bytes: &[u8], out: &impl AsRawFd) -> Result<()> {
loop {
let mut bytes = bytes;
while !bytes.is_empty() {
let len = vmsplice(out, bytes)?;
let len = vmsplice(out, bytes).map_err(maybe_unsupported)?;
bytes = &bytes[len..];
}
}
@ -45,14 +38,8 @@ pub(crate) fn splice_data(bytes: &[u8], out: &impl AsRawFd) -> Result<()> {
loop {
let mut bytes = bytes;
while !bytes.is_empty() {
let len = vmsplice(&write, bytes)?;
let mut remaining = len;
while remaining > 0 {
match splice(&read, out, remaining)? {
0 => panic!("Unexpected end of pipe"),
n => remaining -= n,
};
}
let len = vmsplice(&write, bytes).map_err(maybe_unsupported)?;
splice_exact(&read, out, len).map_err(maybe_unsupported)?;
bytes = &bytes[len..];
}
}
@ -81,30 +68,3 @@ fn maybe_unsupported(error: nix::Error) -> Error {
_ => error.into(),
}
}
fn splice(source: &impl AsRawFd, target: &impl AsRawFd, len: usize) -> Result<usize> {
nix::fcntl::splice(
source.as_raw_fd(),
None,
target.as_raw_fd(),
None,
len,
SpliceFFlags::empty(),
)
.map_err(maybe_unsupported)
}
fn vmsplice(target: &impl AsRawFd, bytes: &[u8]) -> Result<usize> {
nix::fcntl::vmsplice(
target.as_raw_fd(),
&[IoVec::from_slice(bytes)],
SpliceFFlags::empty(),
)
.map_err(maybe_unsupported)
}
fn pipe() -> nix::Result<(File, File)> {
let (read, write) = nix::unistd::pipe()?;
// SAFETY: The file descriptors do not have other owners.
unsafe { Ok((File::from_raw_fd(read), File::from_raw_fd(write))) }
}

View file

@ -31,6 +31,7 @@ data-encoding-macro = { version="0.1.12", optional=true }
z85 = { version="3.0.3", optional=true }
libc = { version="0.2.15", optional=true }
once_cell = "1.8.0"
nix = { version="0.20", optional=true }
[dev-dependencies]
clap = "2.33.3"
@ -57,3 +58,4 @@ signals = []
utf8 = []
utmpx = ["time", "libc", "dns-lookup"]
wide = []
pipes = ["nix"]

View file

@ -19,6 +19,8 @@ pub mod mode;
pub mod entries;
#[cfg(all(unix, feature = "perms"))]
pub mod perms;
#[cfg(all(unix, feature = "pipes"))]
pub mod pipes;
#[cfg(all(unix, feature = "process"))]
pub mod process;

View file

@ -0,0 +1,69 @@
/// Thin pipe-related wrappers around functions from the `nix` crate.
use std::fs::File;
#[cfg(any(target_os = "linux", target_os = "android"))]
use std::os::unix::io::AsRawFd;
use std::os::unix::io::FromRawFd;
#[cfg(any(target_os = "linux", target_os = "android"))]
use nix::{fcntl::SpliceFFlags, sys::uio::IoVec};
pub use nix::{Error, Result};
/// A wrapper around [`nix::unistd::Pipe`] that ensures the pipe is cleaned up.
///
/// Returns two `File` objects: everything written to the second can be read
/// from the first.
pub fn pipe() -> Result<(File, File)> {
let (read, write) = nix::unistd::pipe()?;
// SAFETY: The file descriptors do not have other owners.
unsafe { Ok((File::from_raw_fd(read), File::from_raw_fd(write))) }
}
/// Less noisy wrapper around [`nix::fcntl::splice`].
///
/// Up to `len` bytes are moved from `source` to `target`. Returns the number
/// of successfully moved bytes.
///
/// At least one of `source` and `target` must be some sort of pipe.
/// To get around this requirement, consider splicing from your source into
/// a [`pipe`] and then from the pipe into your target (with `splice_exact`):
/// this is still very efficient.
#[cfg(any(target_os = "linux", target_os = "android"))]
pub fn splice(source: &impl AsRawFd, target: &impl AsRawFd, len: usize) -> Result<usize> {
nix::fcntl::splice(
source.as_raw_fd(),
None,
target.as_raw_fd(),
None,
len,
SpliceFFlags::empty(),
)
}
/// Splice wrapper which fully finishes the write.
///
/// Exactly `len` bytes are moved from `source` into `target`.
///
/// Panics if `source` runs out of data before `len` bytes have been moved.
#[cfg(any(target_os = "linux", target_os = "android"))]
pub fn splice_exact(source: &impl AsRawFd, target: &impl AsRawFd, len: usize) -> Result<()> {
let mut left = len;
while left != 0 {
let written = splice(source, target, left)?;
assert_ne!(written, 0, "unexpected end of data");
left -= written;
}
Ok(())
}
/// Copy data from `bytes` into `target`, which must be a pipe.
///
/// Returns the number of successfully copied bytes.
#[cfg(any(target_os = "linux", target_os = "android"))]
pub fn vmsplice(target: &impl AsRawFd, bytes: &[u8]) -> Result<usize> {
nix::fcntl::vmsplice(
target.as_raw_fd(),
&[IoVec::from_slice(bytes)],
SpliceFFlags::empty(),
)
}

View file

@ -49,6 +49,8 @@ pub use crate::features::mode;
pub use crate::features::entries;
#[cfg(all(unix, feature = "perms"))]
pub use crate::features::perms;
#[cfg(all(unix, feature = "pipes"))]
pub use crate::features::pipes;
#[cfg(all(unix, feature = "process"))]
pub use crate::features::process;
#[cfg(all(unix, not(target_os = "fuchsia"), feature = "signals"))]