Skip to content

Commit

Permalink
net: add TcpSocket::set_{send, recv}_buffer_size (#1384)
Browse files Browse the repository at this point in the history
This commit adds methods for setting `SO_RECVBUF` and `SO_SNDBUF` to
Mio's `TcpSocket` type. It would be nice to eventually expose these in
Tokio, so adding them to Mio is the first step. 

See tokio-rs/tokio#3082 for details.

Signed-off-by: Eliza Weisman <eliza@buoyant.io>
  • Loading branch information
hawkw authored Nov 2, 2020
1 parent 27fbd5f commit 40c4af7
Show file tree
Hide file tree
Showing 5 changed files with 236 additions and 3 deletions.
58 changes: 58 additions & 0 deletions src/net/tcp/socket.rs
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,64 @@ impl TcpSocket {
sys::tcp::set_linger(self.sys, dur)
}

/// Sets the value of `SO_RCVBUF` on this socket.
pub fn set_recv_buffer_size(&self, size: u32) -> io::Result<()> {
sys::tcp::set_recv_buffer_size(self.sys, size)
}

/// Get the value of `SO_RCVBUF` set on this socket.
///
/// Note that if [`set_recv_buffer_size`] has been called on this socket
/// previously, the value returned by this function may not be the same as
/// the argument provided to `set_recv_buffer_size`. This is for the
/// following reasons:
///
/// * Most operating systems have minimum and maximum allowed sizes for the
/// receive buffer, and will clamp the provided value if it is below the
/// minimum or above the maximum. The minimum and maximum buffer sizes are
/// OS-dependent.
/// * Linux will double the buffer size to account for internal bookkeeping
/// data, and returns the doubled value from `getsockopt(2)`. As per `man
/// 7 socket`:
/// > Sets or gets the maximum socket receive buffer in bytes. The
/// > kernel doubles this value (to allow space for bookkeeping
/// > overhead) when it is set using `setsockopt(2)`, and this doubled
/// > value is returned by `getsockopt(2)`.
///
/// [`set_recv_buffer_size`]: #method.set_recv_buffer_size
pub fn get_recv_buffer_size(&self) -> io::Result<u32> {
sys::tcp::get_recv_buffer_size(self.sys)
}

/// Sets the value of `SO_SNDBUF` on this socket.
pub fn set_send_buffer_size(&self, size: u32) -> io::Result<()> {
sys::tcp::set_send_buffer_size(self.sys, size)
}

/// Get the value of `SO_SNDBUF` set on this socket.
///
/// Note that if [`set_send_buffer_size`] has been called on this socket
/// previously, the value returned by this function may not be the same as
/// the argument provided to `set_send_buffer_size`. This is for the
/// following reasons:
///
/// * Most operating systems have minimum and maximum allowed sizes for the
/// receive buffer, and will clamp the provided value if it is below the
/// minimum or above the maximum. The minimum and maximum buffer sizes are
/// OS-dependent.
/// * Linux will double the buffer size to account for internal bookkeeping
/// data, and returns the doubled value from `getsockopt(2)`. As per `man
/// 7 socket`:
/// > Sets or gets the maximum socket send buffer in bytes. The
/// > kernel doubles this value (to allow space for bookkeeping
/// > overhead) when it is set using `setsockopt(2)`, and this doubled
/// > value is returned by `getsockopt(2)`.
///
/// [`set_send_buffer_size`]: #method.set_send_buffer_size
pub fn get_send_buffer_size(&self) -> io::Result<u32> {
sys::tcp::get_send_buffer_size(self.sys)
}

/// Returns the local address of this socket
///
/// Will return `Err` result in windows if called before calling `bind`
Expand Down
16 changes: 16 additions & 0 deletions src/sys/shell/tcp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,22 @@ pub(crate) fn set_linger(_: TcpSocket, _: Option<Duration>) -> io::Result<()> {
os_required!();
}

pub(crate) fn set_recv_buffer_size(_: TcpSocket, _: u32) -> io::Result<()> {
os_required!();
}

pub(crate) fn get_recv_buffer_size(_: TcpSocket) -> io::Result<u32> {
os_required!();
}

pub(crate) fn set_send_buffer_size(_: TcpSocket, _: u32) -> io::Result<()> {
os_required!();
}

pub(crate) fn get_send_buffer_size(_: TcpSocket) -> io::Result<u32> {
os_required!();
}

pub fn accept(_: &net::TcpListener) -> io::Result<(net::TcpStream, SocketAddr)> {
os_required!();
}
Expand Down
57 changes: 55 additions & 2 deletions src/sys/unix/tcp.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use std::io;
use std::convert::TryInto;
use std::mem;
use std::mem::{size_of, MaybeUninit};
use std::net::{self, SocketAddr};
Expand Down Expand Up @@ -37,8 +38,6 @@ pub(crate) fn connect(socket: TcpSocket, addr: SocketAddr) -> io::Result<net::Tc
}

pub(crate) fn listen(socket: TcpSocket, backlog: u32) -> io::Result<net::TcpListener> {
use std::convert::TryInto;

let backlog = backlog.try_into().unwrap_or(i32::max_value());
syscall!(listen(socket, backlog))?;
Ok(unsafe { net::TcpListener::from_raw_fd(socket) })
Expand Down Expand Up @@ -130,6 +129,60 @@ pub(crate) fn set_linger(socket: TcpSocket, dur: Option<Duration>) -> io::Result
)).map(|_| ())
}

pub(crate) fn set_recv_buffer_size(socket: TcpSocket, size: u32) -> io::Result<()> {
let size = size.try_into().ok().unwrap_or_else(i32::max_value);
syscall!(setsockopt(
socket,
libc::SOL_SOCKET,
libc::SO_RCVBUF,
&size as *const _ as *const libc::c_void,
size_of::<libc::c_int>() as libc::socklen_t
))
.map(|_| ())
}

pub(crate) fn get_recv_buffer_size(socket: TcpSocket) -> io::Result<u32> {
let mut optval: libc::c_int = 0;
let mut optlen = size_of::<libc::c_int>() as libc::socklen_t;

syscall!(getsockopt(
socket,
libc::SOL_SOCKET,
libc::SO_RCVBUF,
&mut optval as *mut _ as *mut _,
&mut optlen,
))?;

Ok(optval as u32)
}

pub(crate) fn set_send_buffer_size(socket: TcpSocket, size: u32) -> io::Result<()> {
let size = size.try_into().ok().unwrap_or_else(i32::max_value);
syscall!(setsockopt(
socket,
libc::SOL_SOCKET,
libc::SO_SNDBUF,
&size as *const _ as *const libc::c_void,
size_of::<libc::c_int>() as libc::socklen_t
))
.map(|_| ())
}

pub(crate) fn get_send_buffer_size(socket: TcpSocket) -> io::Result<u32> {
let mut optval: libc::c_int = 0;
let mut optlen = size_of::<libc::c_int>() as libc::socklen_t;

syscall!(getsockopt(
socket,
libc::SOL_SOCKET,
libc::SO_SNDBUF,
&mut optval as *mut _ as *mut _,
&mut optlen,
))?;

Ok(optval as u32)
}

pub fn accept(listener: &net::TcpListener) -> io::Result<(net::TcpStream, SocketAddr)> {
let mut addr: MaybeUninit<libc::sockaddr_storage> = MaybeUninit::uninit();
let mut length = size_of::<libc::sockaddr_storage>() as libc::socklen_t;
Expand Down
63 changes: 62 additions & 1 deletion src/sys/windows/tcp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use std::io;
use std::mem::size_of;
use std::net::{self, SocketAddr, SocketAddrV4, SocketAddrV6};
use std::time::Duration;
use std::convert::TryInto;
use std::os::windows::io::FromRawSocket;
use std::os::windows::raw::SOCKET as StdSocket; // winapi uses usize, stdlib uses u32/u64.

Expand All @@ -12,7 +13,7 @@ use winapi::shared::ws2ipdef::SOCKADDR_IN6_LH;
use winapi::shared::minwindef::{BOOL, TRUE, FALSE};
use winapi::um::winsock2::{
self, closesocket, linger, setsockopt, getsockopt, getsockname, PF_INET, PF_INET6, SOCKET, SOCKET_ERROR,
SOCK_STREAM, SOL_SOCKET, SO_LINGER, SO_REUSEADDR,
SOCK_STREAM, SOL_SOCKET, SO_LINGER, SO_REUSEADDR, SO_RCVBUF, SO_SNDBUF,
};

use crate::sys::windows::net::{init, new_socket, socket_addr};
Expand Down Expand Up @@ -149,6 +150,66 @@ pub(crate) fn set_linger(socket: TcpSocket, dur: Option<Duration>) -> io::Result
}
}


pub(crate) fn set_recv_buffer_size(socket: TcpSocket, size: u32) -> io::Result<()> {
let size = size.try_into().ok().unwrap_or_else(i32::max_value);
match unsafe { setsockopt(
socket,
SOL_SOCKET,
SO_RCVBUF,
&size as *const _ as *const c_char,
size_of::<c_int>() as c_int
) } {
SOCKET_ERROR => Err(io::Error::last_os_error()),
_ => Ok(()),
}
}

pub(crate) fn get_recv_buffer_size(socket: TcpSocket) -> io::Result<u32> {
let mut optval: c_int = 0;
let mut optlen = size_of::<c_int>() as c_int;
match unsafe { getsockopt(
socket,
SOL_SOCKET,
SO_RCVBUF,
&mut optval as *mut _ as *mut _,
&mut optlen as *mut _,
) } {
SOCKET_ERROR => Err(io::Error::last_os_error()),
_ => Ok(optval as u32),
}
}

pub(crate) fn set_send_buffer_size(socket: TcpSocket, size: u32) -> io::Result<()> {
let size = size.try_into().ok().unwrap_or_else(i32::max_value);
match unsafe { setsockopt(
socket,
SOL_SOCKET,
SO_SNDBUF,
&size as *const _ as *const c_char,
size_of::<c_int>() as c_int
) } {
SOCKET_ERROR => Err(io::Error::last_os_error()),
_ => Ok(()),
}
}

pub(crate) fn get_send_buffer_size(socket: TcpSocket) -> io::Result<u32> {
let mut optval: c_int = 0;
let mut optlen = size_of::<c_int>() as c_int;
match unsafe { getsockopt(
socket,
SOL_SOCKET,
SO_SNDBUF,
&mut optval as *mut _ as *mut _,
&mut optlen as *mut _,
) } {
SOCKET_ERROR => Err(io::Error::last_os_error()),
_ => Ok(optval as u32),
}
}


pub(crate) fn accept(listener: &net::TcpListener) -> io::Result<(net::TcpStream, SocketAddr)> {
// The non-blocking state of `listener` is inherited. See
// https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-accept#remarks.
Expand Down
45 changes: 45 additions & 0 deletions tests/tcp_socket.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
#![cfg(all(feature = "os-poll", feature = "tcp"))]

use mio::net::TcpSocket;
use std::io;

#[test]
fn is_send_and_sync() {
Expand Down Expand Up @@ -56,3 +57,47 @@ fn get_localaddr() {

let _ = socket.listen(128).unwrap();
}

#[test]
fn send_buffer_size_roundtrips() {
test_buffer_sizes(
TcpSocket::set_send_buffer_size,
TcpSocket::get_send_buffer_size,
)
}

#[test]
fn recv_buffer_size_roundtrips() {
test_buffer_sizes(
TcpSocket::set_recv_buffer_size,
TcpSocket::get_recv_buffer_size,
)
}

// Helper for testing send/recv buffer size.
fn test_buffer_sizes(
set: impl Fn(&TcpSocket, u32) -> io::Result<()>,
get: impl Fn(&TcpSocket) -> io::Result<u32>,
) {
let test = |size: u32| {
println!("testing buffer size: {}", size);
let socket = TcpSocket::new_v4().unwrap();
set(&socket, size).unwrap();
// Note that this doesn't assert that the values are equal: on Linux,
// the kernel doubles the requested buffer size, and returns the doubled
// value from `getsockopt`. As per `man socket(7)`:
// > Sets or gets the maximum socket send buffer in bytes. The
// > kernel doubles this value (to allow space for bookkeeping
// > overhead) when it is set using setsockopt(2), and this doubled
// > value is returned by getsockopt(2).
//
// Additionally, the buffer size may be clamped above a minimum value,
// and this minimum value is OS-dependent.
let actual = get(&socket).unwrap();
assert!(actual >= size, "\tactual: {}\n\texpected: {}", actual, size);
};

test(256);
test(4096);
test(65512);
}

0 comments on commit 40c4af7

Please sign in to comment.