Skip to content
37 changes: 30 additions & 7 deletions src/uu/sort/src/merge.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ use uucore::error::{FromIo, UResult};
use crate::{
GlobalSettings, Output, SortError,
chunks::{self, Chunk, RecycledChunk},
compare_by, open,
compare_by, fd_soft_limit, open,
tmp_dir::TmpDirWrapper,
};

Expand Down Expand Up @@ -62,6 +62,26 @@ fn replace_output_file_in_input_files(
Ok(())
}

fn effective_merge_batch_size(settings: &GlobalSettings) -> usize {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it needs more comments

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

added

const MIN_BATCH_SIZE: usize = 2;
const RESERVED_STDIO: usize = 3;
const RESERVED_OUTPUT: usize = 1;
const SAFETY_MARGIN: usize = 1;
let mut batch_size = settings.merge_batch_size.max(MIN_BATCH_SIZE);

if let Some(limit) = fd_soft_limit() {
let reserved = RESERVED_STDIO + RESERVED_OUTPUT + SAFETY_MARGIN;
let available_inputs = limit.saturating_sub(reserved);
if available_inputs >= MIN_BATCH_SIZE {
batch_size = batch_size.min(available_inputs);
} else {
batch_size = MIN_BATCH_SIZE;
}
}

batch_size
}

/// Merge pre-sorted `Box<dyn Read>`s.
///
/// If `settings.merge_batch_size` is greater than the length of `files`, intermediate files will be used.
Expand Down Expand Up @@ -94,18 +114,21 @@ pub fn merge_with_file_limit<
output: Output,
tmp_dir: &mut TmpDirWrapper,
) -> UResult<()> {
if files.len() <= settings.merge_batch_size {
let batch_size = effective_merge_batch_size(settings);
debug_assert!(batch_size >= 2);

if files.len() <= batch_size {
let merger = merge_without_limit(files, settings);
merger?.write_all(settings, output)
} else {
let mut temporary_files = vec![];
let mut batch = vec![];
let mut batch = Vec::with_capacity(batch_size);
for file in files {
batch.push(file);
if batch.len() >= settings.merge_batch_size {
assert_eq!(batch.len(), settings.merge_batch_size);
if batch.len() >= batch_size {
assert_eq!(batch.len(), batch_size);
let merger = merge_without_limit(batch.into_iter(), settings)?;
batch = vec![];
batch = Vec::with_capacity(batch_size);

let mut tmp_file =
Tmp::create(tmp_dir.next_file()?, settings.compress_prog.as_deref())?;
Expand All @@ -115,7 +138,7 @@ pub fn merge_with_file_limit<
}
// Merge any remaining files that didn't get merged in a full batch above.
if !batch.is_empty() {
assert!(batch.len() < settings.merge_batch_size);
assert!(batch.len() < batch_size);
let merger = merge_without_limit(batch.into_iter(), settings)?;

let mut tmp_file =
Expand Down
41 changes: 31 additions & 10 deletions src/uu/sort/src/sort.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ use custom_str_cmp::custom_str_cmp;
use ext_sort::ext_sort;
use fnv::FnvHasher;
#[cfg(target_os = "linux")]
use nix::libc::{RLIMIT_NOFILE, getrlimit, rlimit};
use nix::libc;
use numeric_str_cmp::{NumInfo, NumInfoParseSettings, human_numeric_str_cmp, numeric_str_cmp};
use rand::{Rng, rng};
use rayon::prelude::*;
Expand Down Expand Up @@ -1073,17 +1073,30 @@ fn make_sort_mode_arg(mode: &'static str, short: char, help: String) -> Arg {
}

#[cfg(target_os = "linux")]
fn get_rlimit() -> UResult<usize> {
let mut limit = rlimit {
pub(crate) fn fd_soft_limit() -> Option<usize> {
let mut limit = libc::rlimit {
rlim_cur: 0,
rlim_max: 0,
};
match unsafe { getrlimit(RLIMIT_NOFILE, &raw mut limit) } {
0 => Ok(limit.rlim_cur as usize),
_ => Err(UUsageError::new(2, translate!("sort-failed-fetch-rlimit"))),

let result = unsafe { libc::getrlimit(libc::RLIMIT_NOFILE, &raw mut limit) };

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please use use nix::sys::resource::getrlimit;
to avoid the unsafe

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fix it

if result == 0 {
let current = limit.rlim_cur;
if current == libc::RLIM_INFINITY {
None
} else {
usize::try_from(current).ok()
}
} else {
None
}
}

#[cfg(not(target_os = "linux"))]
pub(crate) fn fd_soft_limit() -> Option<usize> {
None
}

const STDIN_FILE: &str = "-";
#[cfg(target_os = "linux")]
const LINUX_BATCH_DIVISOR: usize = 4;
Expand All @@ -1096,12 +1109,12 @@ fn default_merge_batch_size() -> usize {
#[cfg(target_os = "linux")]
{
// Adjust merge batch size dynamically based on available file descriptors.
match get_rlimit() {
Ok(limit) => {
match fd_soft_limit() {
Some(limit) => {
let usable_limit = limit.saturating_div(LINUX_BATCH_DIVISOR);
usable_limit.clamp(LINUX_BATCH_MIN, LINUX_BATCH_MAX)
}
Err(_) => 64,
None => 64,
}
}

Expand Down Expand Up @@ -1279,7 +1292,15 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {

translate!(
"sort-maximum-batch-size-rlimit",
"rlimit" => get_rlimit()?
"rlimit" => {
let Some(rlimit) = fd_soft_limit() else {
return Err(UUsageError::new(
2,
translate!("sort-failed-fetch-rlimit"),
));
};
rlimit
}
)
}
#[cfg(not(target_os = "linux"))]
Expand Down
Loading