feat(s2n-quic-dc): accept linger parameter instead of always setting it

aws · Feb 17, 2025 · 215dc72 · 215dc72
1 parent 00e3371
commit 215dc72
Show file tree

Hide file tree

Showing 6 changed files with 30 additions and 6 deletions.
diff --git a/dc/s2n-quic-dc/src/stream/client/tokio.rs b/dc/s2n-quic-dc/src/stream/client/tokio.rs
@@ -11,7 +11,7 @@ use crate::{
         socket::Protocol,
     },
 };
-use std::{io, net::SocketAddr};
+use std::{io, net::SocketAddr, time::Duration};
 use tokio::net::TcpStream;
 
 /// Connects using the UDP transport layer
@@ -54,6 +54,7 @@ pub async fn connect_tcp<H, Sub>(
     acceptor_addr: SocketAddr,
     env: &Environment<Sub>,
     subscriber: Sub,
+    linger: Option<Duration>,
 ) -> io::Result<Stream<Sub>>
 where
     H: core::future::Future<Output = io::Result<secret::map::Peer>>,
@@ -64,7 +65,10 @@ where
 
     // Make sure TCP_NODELAY is set
     let _ = socket.set_nodelay(true);
-    let _ = socket.set_linger(Some(core::time::Duration::ZERO));
+
+    if linger.is_some() {
+        let _ = socket.set_linger(linger);
+    }
 
     // if the acceptor_ip isn't known, then ask the socket to resolve it for us
     let peer_addr = if acceptor_addr.ip().is_unspecified() {

diff --git a/dc/s2n-quic-dc/src/stream/server/tokio/tcp.rs b/dc/s2n-quic-dc/src/stream/server/tokio/tcp.rs
@@ -9,6 +9,7 @@ use crate::{
 };
 use core::{future::poll_fn, task::Poll};
 use s2n_quic_core::{inet::SocketAddress, time::Clock};
+use std::time::Duration;
 use tokio::net::TcpListener;
 use tracing::debug;
 
@@ -26,6 +27,7 @@ where
     secrets: secret::Map,
     backlog: usize,
     accept_flavor: accept::Flavor,
+    linger: Option<Duration>,
     subscriber: Sub,
 }
 
@@ -42,6 +44,7 @@ where
         secrets: &secret::Map,
         backlog: usize,
         accept_flavor: accept::Flavor,
+        linger: Option<Duration>,
         subscriber: Sub,
     ) -> Self {
         let acceptor = Self {
@@ -51,6 +54,7 @@ where
             secrets: secrets.clone(),
             backlog,
             accept_flavor,
+            linger,
             subscriber,
         };
 
@@ -98,6 +102,7 @@ where
                 workers.insert(
                     remote_address,
                     socket,
+                    self.linger,
                     &mut context,
                     subscriber_ctx,
                     &publisher,

diff --git a/dc/s2n-quic-dc/src/stream/server/tokio/tcp/manager.rs b/dc/s2n-quic-dc/src/stream/server/tokio/tcp/manager.rs
@@ -153,6 +153,7 @@ where
         &mut self,
         remote_address: SocketAddress,
         stream: W::Stream,
+        linger: Option<Duration>,
         cx: &mut W::Context,
         connection_context: W::ConnectionContext,
         publisher: &Pub,
@@ -179,6 +180,7 @@ where
         self.inner.workers[idx].worker.replace(
             remote_address,
             stream,
+            linger,
             connection_context,
             publisher,
             clock,
@@ -377,6 +379,7 @@ pub trait Worker {
         &mut self,
         remote_address: SocketAddress,
         stream: Self::Stream,
+        linger: Option<Duration>,
         connection_context: Self::ConnectionContext,
         publisher: &Pub,
         clock: &C,

diff --git a/dc/s2n-quic-dc/src/stream/server/tokio/tcp/manager/tests.rs b/dc/s2n-quic-dc/src/stream/server/tokio/tcp/manager/tests.rs
@@ -64,6 +64,7 @@ impl super::Worker for Worker {
         &mut self,
         _remote_address: SocketAddress,
         _stream: Self::Stream,
+        _linger: Option<Duration>,
         _connection_context: Self::ConnectionContext,
         _publisher: &Pub,
         clock: &C,
@@ -160,6 +161,7 @@ impl Harness {
         self.manager.insert(
             SocketAddress::default(),
             (),
+            None,
             &mut (),
             (),
             &publisher(&self.subscriber, &self.clock),

diff --git a/dc/s2n-quic-dc/src/stream/server/tokio/tcp/worker.rs b/dc/s2n-quic-dc/src/stream/server/tokio/tcp/worker.rs
@@ -98,6 +98,7 @@ where
         &mut self,
         remote_address: SocketAddress,
         stream: TcpStream,
+        linger: Option<Duration>,
         subscriber_ctx: Self::ConnectionContext,
         publisher: &Pub,
         clock: &C,
@@ -107,7 +108,10 @@ where
     {
         // Make sure TCP_NODELAY is set
         let _ = stream.set_nodelay(true);
-        let _ = stream.set_linger(Some(Duration::ZERO));
+
+        if linger.is_some() {
+            let _ = stream.set_linger(linger);
+        }
 
         let now = clock.get_time();
 

diff --git a/dc/s2n-quic-dc/src/stream/testing.rs b/dc/s2n-quic-dc/src/stream/testing.rs
@@ -68,8 +68,14 @@ impl Client {
 
         match server.protocol {
             Protocol::Tcp => {
-                stream_client::connect_tcp(handshake, server.local_addr, &self.env, subscriber)
-                    .await
+                stream_client::connect_tcp(
+                    handshake,
+                    server.local_addr,
+                    &self.env,
+                    subscriber,
+                    None,
+                )
+                .await
             }
             Protocol::Udp => {
                 stream_client::connect_udp(handshake, server.local_addr, &self.env, subscriber)
@@ -291,7 +297,7 @@ pub mod server {
                     let socket = tokio::net::TcpListener::from_std(socket).unwrap();
 
                     let acceptor = stream_server::tcp::Acceptor::new(
-                        0, socket, &sender, &env, &map, backlog, flavor, subscriber,
+                        0, socket, &sender, &env, &map, backlog, flavor, None, subscriber,
                     );
                     let acceptor = drop_handle_receiver.wrap(acceptor.run());
                     let acceptor = acceptor.instrument(tracing::info_span!("tcp"));