use std::{
cell::RefCell,
io::{Error, ErrorKind},
os::unix::io::AsRawFd,
ptr,
};
use crate::{metrics::names, socket::SocketResult};
pub const DEFAULT_SPLICE_PIPE_CAPACITY: usize = 65_536;
static PIPE_CAPACITY_OVERRIDE: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
pub fn set_pipe_capacity(bytes: usize) {
if bytes == 0 {
return;
}
let _ = PIPE_CAPACITY_OVERRIDE.set(bytes);
}
fn requested_pipe_capacity() -> usize {
PIPE_CAPACITY_OVERRIDE
.get()
.copied()
.unwrap_or(DEFAULT_SPLICE_PIPE_CAPACITY)
}
const SPLICE_PIPE_POOL_CAPACITY: usize = 32;
thread_local! {
static SPLICE_PIPE_POOL: RefCell<Vec<SplicePipe>> = const { RefCell::new(Vec::new()) };
}
pub struct SplicePipe {
pub in_pipe: [libc::c_int; 2],
pub out_pipe: [libc::c_int; 2],
pub in_pipe_pending: usize,
pub out_pipe_pending: usize,
pub capacity: usize,
}
impl SplicePipe {
pub fn new() -> Option<Self> {
let pooled = SPLICE_PIPE_POOL
.try_with(|pool| pool.try_borrow_mut().ok().and_then(|mut pool| pool.pop()))
.ok()
.flatten();
if pooled.is_some() {
return pooled;
}
let in_pipe = create_pipe()?;
let out_pipe = match create_pipe() {
Some(p) => p,
None => {
unsafe {
libc::close(in_pipe[0]);
libc::close(in_pipe[1]);
}
return None;
}
};
let requested = requested_pipe_capacity();
let capacity = apply_pipe_capacity(in_pipe[0], out_pipe[0], requested);
Some(SplicePipe {
in_pipe,
out_pipe,
in_pipe_pending: 0,
out_pipe_pending: 0,
capacity,
})
}
}
fn apply_pipe_capacity(in_read: libc::c_int, out_read: libc::c_int, requested: usize) -> usize {
let in_actual = set_and_query_pipe_size(in_read, requested);
let out_actual = set_and_query_pipe_size(out_read, requested);
in_actual.min(out_actual)
}
fn set_and_query_pipe_size(fd: libc::c_int, requested: usize) -> usize {
let set_ret = unsafe { libc::fcntl(fd, libc::F_SETPIPE_SZ, requested as libc::c_int) };
if set_ret == -1 {
let err = Error::last_os_error();
warn!(
"SPLICE\tF_SETPIPE_SZ({}) on pipe fd({}) failed: {:?}; keeping the kernel default. Lower the requested value or raise /proc/sys/fs/pipe-max-size.",
requested, fd, err
);
}
let get_ret = unsafe { libc::fcntl(fd, libc::F_GETPIPE_SZ) };
if get_ret < 0 {
DEFAULT_SPLICE_PIPE_CAPACITY
} else {
get_ret as usize
}
}
impl SplicePipe {
fn is_empty(&self) -> bool {
self.in_pipe_pending == 0
&& self.out_pipe_pending == 0
&& pipe_is_empty(self.in_pipe[0])
&& pipe_is_empty(self.out_pipe[0])
}
}
fn pipe_is_empty(read_end: libc::c_int) -> bool {
let mut queued: libc::c_int = 0;
let ret = unsafe { libc::ioctl(read_end, libc::FIONREAD, &mut queued) };
ret == 0 && queued == 0
}
impl Drop for SplicePipe {
fn drop(&mut self) {
let pooled = SPLICE_PIPE_POOL
.try_with(|pool| match pool.try_borrow_mut() {
Ok(mut pool) if pool.len() < SPLICE_PIPE_POOL_CAPACITY && self.is_empty() => {
pool.push(SplicePipe {
in_pipe: self.in_pipe,
out_pipe: self.out_pipe,
in_pipe_pending: 0,
out_pipe_pending: 0,
capacity: self.capacity,
});
true
}
_ => false,
})
.unwrap_or(false);
if pooled {
return;
}
unsafe {
libc::close(self.in_pipe[0]);
libc::close(self.in_pipe[1]);
libc::close(self.out_pipe[0]);
libc::close(self.out_pipe[1]);
}
}
}
fn create_pipe() -> Option<[libc::c_int; 2]> {
let mut fds: [libc::c_int; 2] = [0; 2];
let ret = unsafe { libc::pipe2(fds.as_mut_ptr(), libc::O_NONBLOCK | libc::O_CLOEXEC) };
if ret == 0 { Some(fds) } else { None }
}
pub fn splice_in(
fd: &dyn AsRawFd,
pipe_write_end: libc::c_int,
len: usize,
) -> (usize, SocketResult) {
let res = unsafe {
libc::splice(
fd.as_raw_fd(),
ptr::null_mut(),
pipe_write_end,
ptr::null_mut(),
len,
libc::SPLICE_F_NONBLOCK | libc::SPLICE_F_MOVE,
)
};
match res {
-1 => {
let err = Error::last_os_error();
match err.kind() {
ErrorKind::WouldBlock => (0, SocketResult::WouldBlock),
_ => {
if matches!(
err.kind(),
ErrorKind::ConnectionReset | ErrorKind::ConnectionAborted
) {
incr!(names::tcp::READ_RESET);
}
error!(
"SPLICE\terr splicing from fd({}) to pipe({}): {:?}",
fd.as_raw_fd(),
pipe_write_end,
err
);
(0, SocketResult::Error)
}
}
}
0 => (0, SocketResult::Closed),
n => (n as usize, SocketResult::Continue),
}
}
pub fn splice_out(
pipe_read_end: libc::c_int,
fd: &dyn AsRawFd,
len: usize,
) -> (usize, SocketResult) {
if len == 0 {
return (0, SocketResult::Continue);
}
let res = unsafe {
libc::splice(
pipe_read_end,
ptr::null_mut(),
fd.as_raw_fd(),
ptr::null_mut(),
len,
libc::SPLICE_F_NONBLOCK | libc::SPLICE_F_MOVE,
)
};
match res {
-1 => {
let err = Error::last_os_error();
match err.kind() {
ErrorKind::WouldBlock => (0, SocketResult::WouldBlock),
_ => {
if matches!(
err.kind(),
ErrorKind::ConnectionReset | ErrorKind::ConnectionAborted
) {
incr!(names::tcp::WRITE_RESET);
}
error!(
"SPLICE\terr splicing from pipe({}) to fd({}): {:?}",
pipe_read_end,
fd.as_raw_fd(),
err
);
(0, SocketResult::Error)
}
}
}
0 => (0, SocketResult::Closed),
n => (n as usize, SocketResult::Continue),
}
}
#[cfg(test)]
mod tests {
use std::{
io::{Read, Write},
net::{TcpListener, TcpStream},
thread,
time::Duration,
};
use super::*;
fn pooled_count() -> usize {
SPLICE_PIPE_POOL.with(|pool| pool.borrow().len())
}
fn queued(read_end: libc::c_int) -> usize {
let mut queued: libc::c_int = 0;
let ret = unsafe { libc::ioctl(read_end, libc::FIONREAD, &mut queued) };
assert_eq!(ret, 0, "FIONREAD must succeed on a pipe");
queued as usize
}
#[test]
fn a_splice_through_a_reset_socket_counts_the_reset() {
use crate::socket::tests::{proxy_count, reset_stream};
let pipe = SplicePipe::new().expect("create splice pipe");
let from = reset_stream();
let before = proxy_count(names::tcp::READ_RESET);
assert_eq!(
splice_in(&from, pipe.in_pipe[1], 4096),
(0, SocketResult::Error)
);
assert_eq!(proxy_count(names::tcp::READ_RESET) - before, 1);
let payload = b"after the reset";
let written =
unsafe { libc::write(pipe.in_pipe[1], payload.as_ptr().cast(), payload.len()) };
assert_eq!(written, payload.len() as isize);
let to = reset_stream();
let before = proxy_count(names::tcp::WRITE_RESET);
assert_eq!(
splice_out(pipe.in_pipe[0], &to, payload.len()),
(0, SocketResult::Error)
);
assert_eq!(proxy_count(names::tcp::WRITE_RESET) - before, 1);
}
#[test]
fn an_empty_splice_pipe_is_reused() {
let pipe = SplicePipe::new().expect("create splice pipe");
let fds = (pipe.in_pipe, pipe.out_pipe);
drop(pipe);
assert_eq!(pooled_count(), 1, "an empty pair must return to the pool");
let reused = SplicePipe::new().expect("take splice pipe");
assert_eq!((reused.in_pipe, reused.out_pipe), fds);
assert_eq!(pooled_count(), 0);
}
#[test]
fn a_splice_pipe_returned_with_bytes_in_flight_is_never_reused() {
let leaked = b"session A";
for backend_to_frontend in [false, true] {
for counted in [true, false] {
let mut pipe = SplicePipe::new().expect("create splice pipe");
let write_end = if backend_to_frontend {
pipe.out_pipe[1]
} else {
pipe.in_pipe[1]
};
let written =
unsafe { libc::write(write_end, leaked.as_ptr().cast(), leaked.len()) };
assert_eq!(written, leaked.len() as isize);
if counted {
if backend_to_frontend {
pipe.out_pipe_pending = leaked.len();
} else {
pipe.in_pipe_pending = leaked.len();
}
}
drop(pipe);
assert_eq!(
pooled_count(),
0,
"a pair holding bytes must not return to the pool \
(backend_to_frontend={backend_to_frontend}, counted={counted})"
);
let next = SplicePipe::new().expect("create splice pipe");
assert_eq!(
(queued(next.in_pipe[0]), queued(next.out_pipe[0])),
(0, 0),
"the next session must get empty pipes \
(backend_to_frontend={backend_to_frontend}, counted={counted})"
);
}
}
}
#[test]
fn the_splice_pipe_pool_is_bounded() {
let pipes: Vec<SplicePipe> = (0..=SPLICE_PIPE_POOL_CAPACITY)
.map(|_| SplicePipe::new().expect("create splice pipe"))
.collect();
drop(pipes);
assert_eq!(pooled_count(), SPLICE_PIPE_POOL_CAPACITY);
}
#[test]
fn splice_roundtrip() {
let proxy_listener = TcpListener::bind("127.0.0.1:0").expect("bind proxy");
let proxy_addr = proxy_listener.local_addr().expect("local_addr");
let pipe = create_pipe().expect("create_pipe");
let pipe_thread = thread::spawn(move || {
let (conn, _) = proxy_listener.accept().expect("accept");
conn.set_read_timeout(Some(Duration::from_secs(2))).ok();
let mut moved = 0usize;
for _ in 0..50 {
let (sz, status) = splice_in(&conn, pipe[1], DEFAULT_SPLICE_PIPE_CAPACITY);
if sz > 0 {
moved = sz;
assert_eq!(status, SocketResult::Continue);
break;
}
thread::sleep(Duration::from_millis(20));
}
assert!(moved > 0, "splice_in moved 0 bytes");
let (sz_out, status_out) = splice_out(pipe[0], &conn, moved);
assert_eq!(sz_out, moved, "splice_out byte count mismatch");
assert_eq!(status_out, SocketResult::Continue);
unsafe {
libc::close(pipe[0]);
libc::close(pipe[1]);
}
});
let mut client = TcpStream::connect(proxy_addr).expect("connect");
client.set_read_timeout(Some(Duration::from_secs(2))).ok();
let payload = b"splice test data";
client.write_all(payload).expect("client write");
let mut buf = [0u8; 128];
let n = client.read(&mut buf).expect("client read");
assert_eq!(&buf[..n], payload);
pipe_thread.join().expect("pipe thread");
}
}