From b1329db495df18431b8a8a47bbf0bb7c213c24c9 Mon Sep 17 00:00:00 2001
From: Stas Kelvich <stas.kelvich@gmail.com>
Date: Fri, 28 Apr 2023 16:02:01 +0300
Subject: [PATCH] fix sigterm handling

---
 proxy/src/bin/pg_sni_router.rs | 29 +++++++++++++++--------------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/proxy/src/bin/pg_sni_router.rs b/proxy/src/bin/pg_sni_router.rs
index 3b8852ecc8..bba2d51caf 100644
--- a/proxy/src/bin/pg_sni_router.rs
+++ b/proxy/src/bin/pg_sni_router.rs
@@ -108,17 +108,19 @@ async fn main() -> anyhow::Result<()> {
     let proxy_listener = TcpListener::bind(proxy_address).await?;
 
     let cancellation_token = CancellationToken::new();
-    let tasks = vec![
-        tokio::spawn(proxy::handle_signals(cancellation_token.clone())),
-        tokio::spawn(task_main(
-            Arc::new(destination),
-            tls_config,
-            proxy_listener,
-            cancellation_token.clone(),
-        )),
-    ];
 
-    let _tasks = futures::future::try_join_all(tasks.into_iter().map(proxy::flatten_err)).await?;
+    let main = proxy::flatten_err(tokio::spawn(task_main(
+        Arc::new(destination),
+        tls_config,
+        proxy_listener,
+        cancellation_token.clone(),
+    )));
+    let signals_task = proxy::flatten_err(tokio::spawn(proxy::handle_signals(cancellation_token)));
+
+    tokio::select! {
+        res = main => { res?; },
+        res = signals_task => { res?; },
+    }
 
     Ok(())
 }
@@ -129,10 +131,6 @@ async fn task_main(
     listener: tokio::net::TcpListener,
     cancellation_token: CancellationToken,
 ) -> anyhow::Result<()> {
-    scopeguard::defer! {
-        info!("proxy has shut down");
-    }
-
     // When set for the server socket, the keepalive setting
     // will be inherited by all accepted client sockets.
     socket2::SockRef::from(&listener).set_keepalive(true)?;
@@ -171,7 +169,9 @@ async fn task_main(
             }
         }
     }
+
     // Drain connections
+    info!("waiting for all client connections to finish");
     while let Some(res) = connections.join_next().await {
         if let Err(e) = res {
             if !e.is_panic() && !e.is_cancelled() {
@@ -179,6 +179,7 @@ async fn task_main(
             }
         }
     }
+    info!("all client connections have finished");
     Ok(())
 }