From 1d1e3dc3736d219cc9c66dff039006a9f5e2c8a3 Mon Sep 17 00:00:00 2001
From: Fabien Gaud <fgaud@amazon.com>
Date: Thu, 16 Dec 2021 10:33:21 -0800
Subject: [PATCH] net: allow to set linger on TcpSocket

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.
---
 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)));
+}