From 22e6aef6e776b00ce36f8c7e003f475d9bd18242 Mon Sep 17 00:00:00 2001
From: Fabien Gaud <me@fabiengaud.net>
Date: Thu, 16 Dec 2021 11:34:00 -0800
Subject: [PATCH] net: allow to set linger on `TcpSocket` (#4324)

For now, this is only allowed on TcpStream. This is a problem when one
want to disable lingering (i.e. set it to Duration(0, 0)). Without being
able to set it prior to the connect call, if the connect future is
dropped it would leave sockets in a TIME_WAIT state.

Co-authored-by: Fabien Gaud <fgaud@amazon.com>
---
 tokio/src/net/tcp/socket.rs | 23 +++++++++++++++++++++++
 tokio/tests/tcp_socket.rs   | 14 ++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/tokio/src/net/tcp/socket.rs b/tokio/src/net/tcp/socket.rs
index fb75f75a5b8..3c6870221c2 100644
--- a/tokio/src/net/tcp/socket.rs
+++ b/tokio/src/net/tcp/socket.rs
@@ -8,6 +8,7 @@ use std::net::SocketAddr;
 use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
 #[cfg(windows)]
 use std::os::windows::io::{AsRawSocket, FromRawSocket, IntoRawSocket, RawSocket};
+use std::time::Duration;
 
 cfg_net! {
     /// A TCP socket that has not yet been converted to a `TcpStream` or
@@ -349,6 +350,28 @@ impl TcpSocket {
         self.inner.get_recv_buffer_size()
     }
 
+    /// Sets the linger duration of this socket by setting the SO_LINGER option.
+    ///
+    /// This option controls the action taken when a stream has unsent messages and the stream is
+    /// closed. If SO_LINGER is set, the system shall block the process until it can transmit the
+    /// data or until the time expires.
+    ///
+    /// If SO_LINGER is not specified, and the socket is closed, the system handles the call in a
+    /// way that allows the process to continue as quickly as possible.
+    pub fn set_linger(&self, dur: Option<Duration>) -> io::Result<()> {
+        self.inner.set_linger(dur)
+    }
+
+    /// Reads the linger duration for this socket by getting the `SO_LINGER`
+    /// option.
+    ///
+    /// For more information about this option, see [`set_linger`].
+    ///
+    /// [`set_linger`]: TcpSocket::set_linger
+    pub fn linger(&self) -> io::Result<Option<Duration>> {
+        self.inner.get_linger()
+    }
+
     /// Gets the local address of this socket.
     ///
     /// Will fail on windows if called before `bind`.
diff --git a/tokio/tests/tcp_socket.rs b/tokio/tests/tcp_socket.rs
index 9258864d416..3030416502a 100644
--- a/tokio/tests/tcp_socket.rs
+++ b/tokio/tests/tcp_socket.rs
@@ -1,6 +1,7 @@
 #![warn(rust_2018_idioms)]
 #![cfg(feature = "full")]
 
+use std::time::Duration;
 use tokio::net::TcpSocket;
 use tokio_test::assert_ok;
 
@@ -58,3 +59,16 @@ async fn bind_before_connect() {
     // Accept
     let _ = assert_ok!(srv.accept().await);
 }
+
+#[tokio::test]
+async fn basic_linger() {
+    // Create server
+    let addr = assert_ok!("127.0.0.1:0".parse());
+    let srv = assert_ok!(TcpSocket::new_v4());
+    assert_ok!(srv.bind(addr));
+
+    assert!(srv.linger().unwrap().is_none());
+
+    srv.set_linger(Some(Duration::new(0, 0))).unwrap();
+    assert_eq!(srv.linger().unwrap(), Some(Duration::new(0, 0)));
+}