1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

// This file is part of the uutils coreutils package.
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.
use super::{CatResult, FdReadable, InputHandle};

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
/// function `splice()` is used to move data between two file descriptors
/// without copying between kernel and user spaces. This results in a large
/// speedup.
///
/// 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: FdReadable>(
    handle: &mut InputHandle<R>,
    write_fd: &impl AsRawFd,
) -> CatResult<bool> {
    let (pipe_rd, pipe_wr) = pipe()?;

    loop {
        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 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.as_raw_fd(), write_fd.as_raw_fd(), n)?;
                    return Ok(true);
                }
            }
            Err(_) => {
                return Ok(true);
            }
        }
    }
}

/// 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];
    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(())
}