#![allow(unsafe_code)]
#![allow(
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
clippy::cast_possible_wrap
)]
use std::fs::{File, OpenOptions};
use std::io;
use std::os::fd::{AsRawFd, OwnedFd};
use std::os::unix::fs::OpenOptionsExt;
use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::OnceLock;
use io_uring::{opcode, squeue, types, IoUring};
const QD: u32 = 256;
pub const MAX_FILE: u64 = 1 << 20;
pub const MAX_LARGE_FILE: u64 = 64 * 1024 * 1024;
const SPLICE_CHUNK: usize = 1 * 1024 * 1024;
static LARGE_FILE_ATTEMPTS: AtomicU64 = AtomicU64::new(0);
static LARGE_FILE_SUCCESSES: AtomicU64 = AtomicU64::new(0);
#[must_use]
pub fn large_file_stats() -> (u64, u64) {
(
LARGE_FILE_ATTEMPTS.load(Ordering::Relaxed),
LARGE_FILE_SUCCESSES.load(Ordering::Relaxed),
)
}
static SPLICE_SUPPORTED: OnceLock<bool> = OnceLock::new();
#[must_use]
pub fn kernel_supports_splice() -> bool {
*SPLICE_SUPPORTED.get_or_init(probe_kernel_splice)
}
fn probe_kernel_splice() -> bool {
let uname = rustix::system::uname();
let release = uname.release().to_str().unwrap_or("0.0");
let mut parts = release.split(|c: char| !c.is_ascii_digit());
let major: u32 = parts.next().and_then(|s| s.parse().ok()).unwrap_or(0);
let minor: u32 = parts.next().and_then(|s| s.parse().ok()).unwrap_or(0);
major > 5 || (major == 5 && minor >= 19)
}
pub struct Job<'a> {
pub src: &'a Path,
pub tmp: &'a Path,
pub len: u64,
}
#[must_use]
pub fn copy_batch(jobs: &[Job]) -> Vec<io::Result<u64>> {
let mut results: Vec<io::Result<u64>> = Vec::with_capacity(jobs.len());
for _ in jobs {
results.push(Ok(0));
}
let mut ring = match IoUring::new(QD) {
Ok(r) => r,
Err(e) => {
for r in &mut results {
*r = Err(io::Error::other(format!("io_uring init: {e}")));
}
return results;
}
};
let mut small_indices: Vec<usize> = Vec::new();
let mut large_indices: Vec<usize> = Vec::new();
for (i, job) in jobs.iter().enumerate() {
if job.len <= MAX_FILE {
small_indices.push(i);
} else if job.len <= MAX_LARGE_FILE {
large_indices.push(i);
} else {
results[i] = Err(io::Error::new(
io::ErrorKind::Unsupported,
"file too large for uring path",
));
}
}
for chunk in chunk_indices(small_indices.len(), QD as usize) {
let mapped: Vec<usize> = chunk.iter().map(|&idx| small_indices[idx]).collect();
copy_chunk(&mut ring, jobs, &mapped, &mut results);
}
let splice_ok = kernel_supports_splice();
for &job_idx in &large_indices {
if !splice_ok {
results[job_idx] = Err(io::Error::new(
io::ErrorKind::Unsupported,
"kernel too old for io_uring splice; upgrade to ≥5.19",
));
continue;
}
LARGE_FILE_ATTEMPTS.fetch_add(1, Ordering::Relaxed);
match copy_large_file_splice(&mut ring, &jobs[job_idx]) {
Ok(bytes) => {
LARGE_FILE_SUCCESSES.fetch_add(1, Ordering::Relaxed);
results[job_idx] = Ok(bytes);
}
Err(e) => {
results[job_idx] = Err(e);
}
}
}
results
}
pub fn copy_single_large(src: &Path, tmp: &Path, len: u64) -> io::Result<u64> {
if !kernel_supports_splice() {
return Err(io::Error::new(
io::ErrorKind::Unsupported,
"kernel too old for io_uring splice; upgrade to ≥5.19",
));
}
if len > MAX_LARGE_FILE {
return Err(io::Error::new(
io::ErrorKind::Unsupported,
"file too large for uring splice path",
));
}
LARGE_FILE_ATTEMPTS.fetch_add(1, Ordering::Relaxed);
let mut ring = IoUring::new(QD)?;
match copy_large_file_splice(&mut ring, &Job { src, tmp, len }) {
Ok(bytes) => {
LARGE_FILE_SUCCESSES.fetch_add(1, Ordering::Relaxed);
Ok(bytes)
}
Err(e) => Err(e),
}
}
fn chunk_indices(n: usize, size: usize) -> impl Iterator<Item = Vec<usize>> {
(0..n)
.step_by(size)
.map(move |start| (start..(start + size).min(n)).collect())
}
struct Open {
job_idx: usize,
src: File,
tmp: File,
buf: Vec<u8>,
read: usize,
failed: bool,
}
fn copy_chunk(ring: &mut IoUring, jobs: &[Job], chunk: &[usize], results: &mut [io::Result<u64>]) {
let mut opens: Vec<Open> = Vec::with_capacity(chunk.len());
for &job_idx in chunk {
let job = &jobs[job_idx];
if job.len > MAX_FILE {
results[job_idx] = Err(io::Error::new(
io::ErrorKind::Unsupported,
"file too large for uring path",
));
continue;
}
match open_pair(job) {
Ok((src, tmp)) => {
let len = usize::try_from(job.len).unwrap_or(0);
opens.push(Open {
job_idx,
src,
tmp,
buf: vec![0u8; len],
read: 0,
failed: false,
});
}
Err(e) => results[job_idx] = Err(e),
}
}
submit_reads(ring, &mut opens, results);
submit_writes(ring, &mut opens, results);
for o in &opens {
if !o.failed {
results[o.job_idx] = Ok(o.buf.len() as u64);
}
}
}
fn open_pair(job: &Job) -> io::Result<(File, File)> {
let src = File::open(job.src)?;
let tmp = OpenOptions::new()
.write(true)
.create_new(true)
.mode(0o600)
.open(job.tmp)?;
Ok((src, tmp))
}
fn submit_reads(ring: &mut IoUring, opens: &mut [Open], results: &mut [io::Result<u64>]) {
let mut submitted = 0u32;
for (slot, o) in opens.iter().enumerate() {
if o.buf.is_empty() {
continue; }
let sqe = opcode::Read::new(
types::Fd(o.src.as_raw_fd()),
o.buf.as_ptr().cast_mut(),
o.buf.len() as u32,
)
.offset(0)
.build()
.user_data(slot as u64);
if unsafe { ring.submission().push(&sqe) }.is_err() {
break;
}
submitted += 1;
}
reap(ring, submitted, opens, results, true);
}
fn submit_writes(ring: &mut IoUring, opens: &mut [Open], results: &mut [io::Result<u64>]) {
let mut submitted = 0u32;
for (slot, o) in opens.iter().enumerate() {
if o.failed || o.read == 0 {
continue;
}
let sqe = opcode::Write::new(types::Fd(o.tmp.as_raw_fd()), o.buf.as_ptr(), o.read as u32)
.offset(0)
.build()
.user_data(slot as u64);
if unsafe { ring.submission().push(&sqe) }.is_err() {
break;
}
submitted += 1;
}
reap(ring, submitted, opens, results, false);
}
fn reap(
ring: &mut IoUring,
count: u32,
opens: &mut [Open],
results: &mut [io::Result<u64>],
is_read: bool,
) {
if count == 0 {
return;
}
if let Err(e) = ring.submit_and_wait(count as usize) {
for o in opens.iter_mut() {
o.failed = true;
results[o.job_idx] = Err(io::Error::other(format!("uring submit: {e}")));
}
return;
}
let cqes: Vec<(usize, i32)> = ring
.completion()
.map(|c| (c.user_data() as usize, c.result()))
.collect();
for (slot, res) in cqes {
let Some(o) = opens.get_mut(slot) else {
continue;
};
if res < 0 {
o.failed = true;
results[o.job_idx] = Err(io::Error::from_raw_os_error(-res));
} else if is_read {
o.read = res as usize;
if o.read != o.buf.len() {
o.failed = true;
results[o.job_idx] =
Err(io::Error::new(io::ErrorKind::UnexpectedEof, "short read"));
}
} else if res as usize != o.read {
o.failed = true;
results[o.job_idx] = Err(io::Error::new(io::ErrorKind::WriteZero, "short write"));
}
}
}
fn copy_large_file_splice(ring: &mut IoUring, job: &Job) -> io::Result<u64> {
let src_file = File::open(job.src)?;
let dst_file = OpenOptions::new()
.write(true)
.create_new(true)
.mode(0o600)
.open(job.tmp)?;
let src_fd = src_file.as_raw_fd();
let dst_fd = dst_file.as_raw_fd();
let (pipe_read, pipe_write) = rustix::pipe::pipe()?;
let pipe_size = maximize_pipe_size(&pipe_write);
let chunk_size = pipe_size.min(SPLICE_CHUNK);
let num_chunks = ((job.len as usize) + chunk_size - 1) / chunk_size;
let total_sqes = num_chunks * 2;
if total_sqes > QD as usize {
return Err(io::Error::new(
io::ErrorKind::Other,
format!(
"too many splice SQEs needed ({total_sqes}) for {} byte file; \
increase QD or reduce file size",
job.len
),
));
}
let pipe_read_fd = pipe_read.as_raw_fd();
let pipe_write_fd = pipe_write.as_raw_fd();
let mut submitted: u32 = 0;
for chunk_idx in 0..num_chunks {
let offset = (chunk_idx * chunk_size) as u64;
let remaining = job.len.saturating_sub(offset);
let this_chunk = (remaining as usize).min(chunk_size);
if this_chunk == 0 {
break;
}
let sqe1 = opcode::Splice::new(
types::Fd(src_fd),
offset as i64, types::Fd(pipe_write_fd),
-1_i64, this_chunk as u32,
)
.build()
.flags(squeue::Flags::IO_LINK)
.user_data((chunk_idx * 2) as u64);
if unsafe { ring.submission().push(&sqe1) }.is_err() {
return Err(io::Error::other("uring submission queue full (splice src→pipe)"));
}
submitted += 1;
let sqe2 = opcode::Splice::new(
types::Fd(pipe_read_fd),
-1_i64, types::Fd(dst_fd),
offset as i64, this_chunk as u32,
)
.build()
.user_data((chunk_idx * 2 + 1) as u64);
if unsafe { ring.submission().push(&sqe2) }.is_err() {
return Err(io::Error::other("uring submission queue full (splice pipe→dst)"));
}
submitted += 1;
}
if submitted == 0 {
return Ok(0);
}
if let Err(e) = ring.submit_and_wait(submitted as usize) {
return Err(io::Error::other(format!("uring splice submit: {e}")));
}
let mut total_copied: u64 = 0;
let mut first_error: Option<io::Error> = None;
for cqe in ring.completion() {
let res = cqe.result();
if res < 0 {
let err = io::Error::from_raw_os_error(-res);
if first_error.is_none() {
first_error = Some(io::Error::other(format!(
"uring splice SQE failed: {err}; consider upgrading kernel to ≥5.19"
)));
}
} else {
total_copied += res as u64;
}
}
if let Some(e) = first_error {
tracing::warn!(
"io_uring splice failed for {} ({} bytes): {e}; falling back to portable copy",
job.src.display(),
job.len
);
drop(dst_file);
let _ = std::fs::remove_file(job.tmp);
return Err(e);
}
if total_copied != job.len {
tracing::warn!(
"io_uring splice short copy for {}: expected {} bytes, got {total_copied}",
job.src.display(),
job.len
);
drop(dst_file);
let _ = std::fs::remove_file(job.tmp);
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
format!(
"splice short copy: expected {} bytes, got {total_copied}",
job.len
),
));
}
Ok(total_copied)
}
fn maximize_pipe_size(pipe_fd: &OwnedFd) -> usize {
let desired = SPLICE_CHUNK;
if rustix::pipe::fcntl_setpipe_size(pipe_fd, desired).is_ok() {
rustix::pipe::fcntl_getpipe_size(pipe_fd).unwrap_or(desired)
} else {
rustix::pipe::fcntl_getpipe_size(pipe_fd).unwrap_or(65536)
}
}