Reapply drop PG connection when client is dropped

Author: martykulma

src/postgres-util/src/tunnel.rs

@@ -15,11 +15,10 @@ use std::time::Duration;
 use mz_ore::future::{InTask, OreFutureExt};
 use mz_ore::netio::DUMMY_DNS_PORT;
 use mz_ore::option::OptionExt;
-use mz_ore::task;
+use mz_ore::task::{self, AbortOnDropHandle};
 use mz_repr::CatalogItemId;
 use mz_ssh_util::tunnel::{SshTimeoutConfig, SshTunnelConfig};
 use mz_ssh_util::tunnel_manager::SshTunnelManager;
-use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::net::TcpStream as TokioTcpStream;
 use tokio_postgres::config::{Host, ReplicationMode};
 use tokio_postgres::tls::MakeTlsConnect;
@@ -66,26 +65,13 @@ pub const DEFAULT_SNAPSHOT_STATEMENT_TIMEOUT: Duration = Duration::ZERO;
 pub struct Client {
     inner: tokio_postgres::Client,
     server_version: Option<String>,
+    // Holds a handle to the task with the connection to ensure that when
+    // the client is dropped, the task can be aborted to close the connection.
+    // This is also useful for maintaining the lifetimes of dependent object (e.g. ssh tunnel).
+    _connection_handle: AbortOnDropHandle<()>,
 }
 
 impl Client {
-    fn new<S, T>(
-        client: tokio_postgres::Client,
-        connection: &tokio_postgres::Connection<S, T>,
-    ) -> Client
-    where
-        S: AsyncRead + AsyncWrite + Unpin,
-        T: AsyncRead + AsyncWrite + Unpin,
-    {
-        let server_version = connection
-            .parameter("server_version")
-            .map(|v| v.to_string());
-        Client {
-            inner: client,
-            server_version,
-        }
-    }
-
     /// Reports the value of the `server_version` parameter reported by the
     /// server.
     pub fn server_version(&self) -> Option<&str> {
@@ -247,8 +233,19 @@ impl Config {
                 let (client, connection) = async move { postgres_config.connect(tls).await }
                     .run_in_task_if(self.in_task, || "pg_connect".to_string())
                     .await?;
-                let client = Client::new(client, &connection);
-                task::spawn(|| task_name, connection);
+
+                let client = Client {
+                    inner: client,
+                    server_version: connection
+                        .parameter("server_version")
+                        .map(|v| v.to_string()),
+                    _connection_handle: task::spawn(|| task_name, async {
+                        if let Err(e) = connection.await {
+                            warn!("postgres direct connection failed: {e}");
+                        }
+                    })
+                    .abort_on_drop(),
+                };
                 Ok(client)
             }
             TunnelConfig::Ssh { config } => {
@@ -279,14 +276,20 @@ impl Config {
                     async move { postgres_config.connect_raw(tcp_stream, tls).await }
                         .run_in_task_if(self.in_task, || "pg_connect".to_string())
                         .await?;
-                let client = Client::new(client, &connection);
-                task::spawn(|| task_name, async {
-                    let _tunnel = tunnel; // Keep SSH tunnel alive for duration of connection.
 
-                    if let Err(e) = connection.await {
-                        warn!("postgres connection failed: {e}");
-                    }
-                });
+                let client = Client {
+                    inner: client,
+                    server_version: connection
+                        .parameter("server_version")
+                        .map(|v| v.to_string()),
+                    _connection_handle: task::spawn(|| task_name, async {
+                        let _tunnel = tunnel; // Keep SSH tunnel alive for duration of connection.
+                        if let Err(e) = connection.await {
+                            warn!("postgres via SSH tunnel connection failed: {e}");
+                        }
+                    })
+                    .abort_on_drop(),
+                };
                 Ok(client)
             }
             TunnelConfig::AwsPrivatelink { connection_id } => {
@@ -322,8 +325,19 @@ impl Config {
                 let (client, connection) = async move { postgres_config.connect(tls).await }
                     .run_in_task_if(self.in_task, || "pg_connect".to_string())
                     .await?;
-                let client = Client::new(client, &connection);
-                task::spawn(|| task_name, connection);
+
+                let client = Client {
+                    inner: client,
+                    server_version: connection
+                        .parameter("server_version")
+                        .map(|v| v.to_string()),
+                    _connection_handle: task::spawn(|| task_name, async {
+                        if let Err(e) = connection.await {
+                            warn!("postgres AWS link connection failed: {e}");
+                        }
+                    })
+                    .abort_on_drop(),
+                };
                 Ok(client)
             }
         }

src/storage/src/source/postgres/replication.rs

@@ -777,6 +777,9 @@ async fn raw_stream<'a>(
         // Ensure we don't pre-drop the task
         let _max_lsn_task_handle = max_lsn_task_handle;
 
+        // ensure we don't drop the replication client!
+        let _replication_client = replication_client;
+
         let mut uppers = pin!(uppers);
         let mut last_committed_upper = resume_lsn;