Merge pull request #7438 from karlmcdowall/head_perf2

head: rework handling of non-seekable files

Author: RenjiSann

src/uu/head/src/head.rs

@@ -17,7 +17,6 @@ use thiserror::Error;
 use uucore::display::Quotable;
 use uucore::error::{FromIo, UError, UResult};
 use uucore::line_ending::LineEnding;
-use uucore::lines::lines;
 use uucore::{format_usage, help_about, help_usage, show};
 
 const BUF_SIZE: usize = 65536;
@@ -37,7 +36,8 @@ mod options {
 
 mod parse;
 mod take;
-use take::take_all_but;
+use take::copy_all_but_n_bytes;
+use take::copy_all_but_n_lines;
 use take::take_lines;
 
 #[derive(Error, Debug)]
@@ -274,14 +274,16 @@ fn read_n_lines(input: &mut impl std::io::BufRead, n: u64, separator: u8) -> std
     let mut reader = take_lines(input, n, separator);
 
     // Write those bytes to `stdout`.
-    let mut stdout = std::io::stdout();
+    let stdout = std::io::stdout();
+    let stdout = stdout.lock();
+    let mut writer = BufWriter::with_capacity(BUF_SIZE, stdout);
 
-    let bytes_written = io::copy(&mut reader, &mut stdout).map_err(wrap_in_stdout_error)?;
+    let bytes_written = io::copy(&mut reader, &mut writer).map_err(wrap_in_stdout_error)?;
 
     // Make sure we finish writing everything to the target before
     // exiting. Otherwise, when Rust is implicitly flushing, any
     // error will be silently ignored.
-    stdout.flush().map_err(wrap_in_stdout_error)?;
+    writer.flush().map_err(wrap_in_stdout_error)?;
 
     Ok(bytes_written)
 }
@@ -296,43 +298,37 @@ fn catch_too_large_numbers_in_backwards_bytes_or_lines(n: u64) -> Option<usize>
     }
 }
 
-fn read_but_last_n_bytes(input: impl std::io::BufRead, n: u64) -> std::io::Result<u64> {
-    let mut bytes_written = 0;
+fn read_but_last_n_bytes(mut input: impl Read, n: u64) -> std::io::Result<u64> {
+    let mut bytes_written: u64 = 0;
     if let Some(n) = catch_too_large_numbers_in_backwards_bytes_or_lines(n) {
         let stdout = std::io::stdout();
-        let stdout = stdout.lock();
-        // Even though stdout is buffered, it will flush on each newline in the
-        // input stream. This can be costly, so add an extra layer of buffering
-        // over the top. This gives a significant speedup (approx 4x).
-        let mut writer = BufWriter::with_capacity(BUF_SIZE, stdout);
-        for byte in take_all_but(input.bytes(), n) {
-            writer.write_all(&[byte?]).map_err(wrap_in_stdout_error)?;
-            bytes_written += 1;
-        }
+        let mut stdout = stdout.lock();
+
+        bytes_written = copy_all_but_n_bytes(&mut input, &mut stdout, n)
+            .map_err(wrap_in_stdout_error)?
+            .try_into()
+            .unwrap();
+
         // Make sure we finish writing everything to the target before
         // exiting. Otherwise, when Rust is implicitly flushing, any
         // error will be silently ignored.
-        writer.flush().map_err(wrap_in_stdout_error)?;
+        stdout.flush().map_err(wrap_in_stdout_error)?;
     }
     Ok(bytes_written)
 }
 
-fn read_but_last_n_lines(
-    input: impl std::io::BufRead,
-    n: u64,
-    separator: u8,
-) -> std::io::Result<u64> {
+fn read_but_last_n_lines(mut input: impl Read, n: u64, separator: u8) -> std::io::Result<u64> {
+    let stdout = std::io::stdout();
+    let mut stdout = stdout.lock();
+    if n == 0 {
+        return io::copy(&mut input, &mut stdout).map_err(wrap_in_stdout_error);
+    }
     let mut bytes_written: u64 = 0;
     if let Some(n) = catch_too_large_numbers_in_backwards_bytes_or_lines(n) {
-        let stdout = std::io::stdout();
-        let mut stdout = stdout.lock();
-
-        for bytes in take_all_but(lines(input, separator), n) {
-            let bytes = bytes?;
-            bytes_written += u64::try_from(bytes.len()).unwrap();
-
-            stdout.write_all(&bytes).map_err(wrap_in_stdout_error)?;
-        }
+        bytes_written = copy_all_but_n_lines(input, &mut stdout, n, separator)
+            .map_err(wrap_in_stdout_error)?
+            .try_into()
+            .unwrap();
         // Make sure we finish writing everything to the target before
         // exiting. Otherwise, when Rust is implicitly flushing, any
         // error will be silently ignored.
@@ -434,10 +430,9 @@ fn head_backwards_without_seek_file(
     input: &mut std::fs::File,
     options: &HeadOptions,
 ) -> std::io::Result<u64> {
-    let reader = std::io::BufReader::with_capacity(BUF_SIZE, &*input);
     match options.mode {
-        Mode::AllButLastBytes(n) => read_but_last_n_bytes(reader, n),
-        Mode::AllButLastLines(n) => read_but_last_n_lines(reader, n, options.line_ending.into()),
+        Mode::AllButLastBytes(n) => read_but_last_n_bytes(input, n),
+        Mode::AllButLastLines(n) => read_but_last_n_lines(input, n, options.line_ending.into()),
         _ => unreachable!(),
     }
 }
@@ -452,28 +447,20 @@ fn head_backwards_on_seekable_file(
             if n >= size {
                 Ok(0)
             } else {
-                read_n_bytes(
-                    &mut std::io::BufReader::with_capacity(BUF_SIZE, input),
-                    size - n,
-                )
+                read_n_bytes(input, size - n)
             }
         }
         Mode::AllButLastLines(n) => {
             let found = find_nth_line_from_end(input, n, options.line_ending.into())?;
-            read_n_bytes(
-                &mut std::io::BufReader::with_capacity(BUF_SIZE, input),
-                found,
-            )
+            read_n_bytes(input, found)
         }
         _ => unreachable!(),
     }
 }
 
 fn head_file(input: &mut std::fs::File, options: &HeadOptions) -> std::io::Result<u64> {
     match options.mode {
-        Mode::FirstBytes(n) => {
-            read_n_bytes(&mut std::io::BufReader::with_capacity(BUF_SIZE, input), n)
-        }
+        Mode::FirstBytes(n) => read_n_bytes(input, n),
         Mode::FirstLines(n) => read_n_lines(
             &mut std::io::BufReader::with_capacity(BUF_SIZE, input),
             n,