Skip to content

Commit

Permalink
tcp_stream: implement keepalive/set_keepalive
Browse files Browse the repository at this point in the history
In order to restore the ability to set keepalive duration
for a TCP stream, a pair of keepalive()/set_keepalive() functions
are implemented on top of mio's keepalive feature.

Refs #3082
  • Loading branch information
psarna committed Dec 14, 2020
1 parent a7833e3 commit f875908
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 0 deletions.
55 changes: 55 additions & 0 deletions tokio/src/net/tcp/stream.rs
Original file line number Diff line number Diff line change
Expand Up @@ -807,6 +807,61 @@ impl TcpStream {
mio_socket.set_linger(dur)
}

/// Reads the keepalive duration for this socket by getting the `SO_KEEPALIVE`
/// option along with more system-specific parameters (e.g. TCP_KEEPALIVE
/// or SIO_KEEPALIVE_VALS).
///
/// For more information about this option, see [`set_keepalive`].
///
/// [`set_keepalive`]: TcpStream::set_keepalive
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
///
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// println!("{:?}", stream.keepalive()?);
/// # Ok(())
/// # }
/// ```
#[cfg_attr(docsrs, doc(cfg(not(target_os = "windows"))))]
#[cfg(not(target_os = "windows"))]
pub fn keepalive(&self) -> io::Result<Option<Duration>> {
let mio_socket = std::mem::ManuallyDrop::new(self.to_mio());
mio_socket.get_keepalive_time()
}

/// Sets the keepalive duration of this socket by setting the SO_KEEPALIVE option
/// along with more system-specific parameters (e.g. TCP_KEEPALIVE or SIO_KEEPALIVE_VALS).
///
/// This option controls whether keep-alive TCP packets should be used
/// for a socket connection and what should be their idle interval.
///
/// # Examples
///
/// ```no_run
/// use tokio::net::TcpStream;
///
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let stream = TcpStream::connect("127.0.0.1:8080").await?;
///
/// stream.set_keepalive(None)?;
/// # Ok(())
/// # }
/// ```
pub fn set_keepalive(&self, dur: Option<Duration>) -> io::Result<()> {
let mio_socket = std::mem::ManuallyDrop::new(self.to_mio());

if let Some(duration) = dur {
mio_socket.set_keepalive_params(mio::net::TcpKeepalive::new().with_time(duration))
} else {
mio_socket.set_keepalive(false)
}
}

fn to_mio(&self) -> mio::net::TcpSocket {
#[cfg(windows)]
{
Expand Down
16 changes: 16 additions & 0 deletions tokio/tests/tcp_stream.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,22 @@ async fn set_linger() {
assert!(stream.linger().unwrap().is_none());
}

#[tokio::test]
#[cfg(not(windows))]
async fn set_keepalive() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();

let stream = TcpStream::connect(listener.local_addr().unwrap())
.await
.unwrap();

assert_ok!(stream.set_keepalive(Some(Duration::from_secs(1))));
assert_eq!(stream.keepalive().unwrap().unwrap().as_secs(), 1);

assert_ok!(stream.set_keepalive(None));
assert!(stream.keepalive().unwrap().is_none());
}

#[tokio::test]
async fn try_read_write() {
const DATA: &[u8] = b"this is some data to write to the socket";
Expand Down

0 comments on commit f875908

Please sign in to comment.