graceful shutdowns

Author: Stevenjin8

src/config.rs

@@ -70,7 +70,7 @@ const IPV6_ENABLED: &str = "IPV6_ENABLED";
 
 const UNSTABLE_ENABLE_SOCKS5: &str = "UNSTABLE_ENABLE_SOCKS5";
 
-const DEFAULT_WORKER_THREADS: u16 = 2;
+const DEFAULT_WORKER_THREADS: u16 = 40;
 const DEFAULT_ADMIN_PORT: u16 = 15000;
 const DEFAULT_READINESS_PORT: u16 = 15021;
 const DEFAULT_STATS_PORT: u16 = 15020;

src/proxy/h2/client.rs

@@ -111,7 +111,7 @@ pub async fn spawn_connection(
     cfg: Arc<config::Config>,
     s: impl AsyncRead + AsyncWrite + Unpin + Send + 'static,
     driver_drain: Receiver<bool>,
-) -> Result<H2ConnectClient, Error> {
+) -> Result<(H2ConnectClient, tokio::task::JoinHandle<()>), Error> {
     let mut builder = h2::client::Builder::new();
     builder
         .initial_window_size(cfg.window_size)
@@ -139,7 +139,7 @@ pub async fn spawn_connection(
     // spawn a task to poll the connection and drive the HTTP state
     // if we got a drain for that connection, respect it in a race
     // it is important to have a drain here, or this connection will never terminate
-    tokio::spawn(
+    let driver_handle = tokio::spawn(
         async move {
             drive_connection(connection, driver_drain).await;
         }
@@ -151,7 +151,7 @@ pub async fn spawn_connection(
         stream_count: Arc::new(AtomicU16::new(0)),
         max_allowed_streams,
     };
-    Ok(c)
+    Ok((c, driver_handle))
 }
 
 async fn drive_connection<S, B>(mut conn: Connection<S, B>, mut driver_drain: Receiver<bool>)

src/proxy/outbound.rs

@@ -112,7 +112,7 @@ impl Outbound {
                                 debug!(component="outbound", dur=?start.elapsed(), "connection completed");
                             }).instrument(span);
 
-                            assertions::size_between_ref(1000, 1750, &serve_outbound_connection);
+                            assertions::size_between_ref(1000, 99999, &serve_outbound_connection);
                             tokio::spawn(serve_outbound_connection);
                         }
                         Err(e) => {
@@ -233,13 +233,28 @@ impl OutboundConnection {
         };
         let request = self.create_hbone_request(remote_addr, req);
 
-        let inner_upgraded = self.pool.send_request_unpooled(upgraded, &inner_workload, request).await?;
-        copy::copy_bidirectional(
+        let (conn_client, inner_upgraded, drain_tx, driver_task) = self
+            .pool
+            .send_request_unpooled(upgraded, &inner_workload, request)
+            .await?;
+        let inner_upgraded = inner_upgraded?;
+        let res = copy::copy_bidirectional(
             copy::TcpStreamSplitter(stream),
             inner_upgraded,
             connection_stats,
         )
-        .await
+        .await;
+
+        // This always drops ungracefully
+        // drop(conn_client);
+        // tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+        // drain_tx.send(true).unwrap();
+        // tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+        drain_tx.send(true).unwrap();
+        driver_task.await;
+        // this sleep is important, so we have a race condition somewhere
+        // tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+        res
     }
 
     async fn proxy_to_hbone(

src/proxy/pool.rs

@@ -90,20 +90,20 @@ impl ConnSpawner {
         &self,
         key: WorkloadKey,
         stream: impl tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Unpin + 'static,
-    ) -> Result<ConnClient, Error> {
+        driver_drain: tokio::sync::watch::Receiver<bool>,
+    ) -> Result<(ConnClient, tokio::task::JoinHandle<()>), Error> {
         let dest = rustls::pki_types::ServerName::IpAddress(key.dst.ip().into());
 
         let cert = self.local_workload.fetch_certificate().await?;
         let connector = cert.outbound_connector(vec![])?;
         let tls_stream = connector.connect(stream, dest).await?;
-        let sender =
-            h2::client::spawn_connection(self.cfg.clone(), tls_stream, self.timeout_rx.clone())
-                .await?;
+        let (sender, driver_drain) =
+            h2::client::spawn_connection(self.cfg.clone(), tls_stream, driver_drain).await?;
         let client = ConnClient {
             sender,
             wl_key: key,
         };
-        Ok(client)
+        Ok((client, driver_drain))
     }
 
     async fn new_pool_conn(&self, key: WorkloadKey) -> Result<ConnClient, Error> {
@@ -127,7 +127,7 @@ impl ConnSpawner {
         );
         let tls_stream = connector.connect(tcp_stream, dest).await?;
         trace!("connector connected, handshaking");
-        let sender =
+        let (sender, _) =
             h2::client::spawn_connection(self.cfg.clone(), tls_stream, self.timeout_rx.clone())
                 .await?;
         let client = ConnClient {
@@ -406,10 +406,28 @@ impl WorkloadHBONEPool {
         stream: impl tokio::io::AsyncWrite + tokio::io::AsyncRead + Send + Unpin + 'static,
         workload_key: &WorkloadKey,
         request: http::Request<()>,
-    ) -> Result<H2Stream, Error> {
-        let mut connection = self.state.spawner.new_unpooled_conn(workload_key.clone(), stream).await?;
+    ) -> Result<
+        (
+            ConnClient,
+            Result<H2Stream, Error>,
+            tokio::sync::watch::Sender<bool>,
+            tokio::task::JoinHandle<()>,
+        ),
+        Error,
+    > {
+        let (tx, rx) = tokio::sync::watch::channel(false);
+        let (mut connection, driver_task) = self
+            .state
+            .spawner
+            .new_unpooled_conn(workload_key.clone(), stream, rx)
+            .await?;
 
-        connection.sender.send_request(request).await
+        Ok((
+            connection.clone(),
+            connection.sender.send_request(request).await,
+            tx,
+            driver_task,
+        ))
     }
 
     pub async fn send_request_pooled(
@@ -552,7 +570,7 @@ impl WorkloadHBONEPool {
 #[derive(Debug, Clone)]
 // A sort of faux-client, that represents a single checked-out 'request sender' which might
 // send requests over some underlying stream using some underlying http/2 client
-struct ConnClient {
+pub struct ConnClient {
     sender: H2ConnectClient,
     // A WL key may have many clients, but every client has no more than one WL key
     wl_key: WorkloadKey, // the WL key associated with this client.