Skip to content

head: rework handling of non-seekable files #7438

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 18, 2025
Merged

head: rework handling of non-seekable files #7438

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2987531
head: rework handling of non-seekable files
karlmcdowall Feb 21, 2025
e1275f4
Update src/uu/head/src/take.rs
karlmcdowall Mar 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 32 additions & 45 deletions src/uu/head/src/head.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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)]
Expand Down Expand Up @@ -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)
}
Expand All @@ -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.
Expand Down Expand Up @@ -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!(),
}
}
Expand All @@ -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,
Expand Down
Loading
Loading