From c216b16b0ff5e065d4ac9f1385a1e70ee15b003b Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conrad@neon.tech>
Date: Thu, 5 Oct 2023 07:30:28 +0100
Subject: [PATCH] proxy: fix memory leak (#5472)

## Problem

these JoinSets live for the duration of the process. they might have
many millions of connections spawned on them and they never get cleared.

Fixes #4672

## Summary of changes

Drain the connections as we go
---
 proxy/src/bin/pg_sni_router.rs | 5 +++++
 proxy/src/proxy.rs             | 5 +++++
 2 files changed, 10 insertions(+)

diff --git a/proxy/src/bin/pg_sni_router.rs b/proxy/src/bin/pg_sni_router.rs
index 849af47cfc..6574500c97 100644
--- a/proxy/src/bin/pg_sni_router.rs
+++ b/proxy/src/bin/pg_sni_router.rs
@@ -168,6 +168,11 @@ async fn task_main(
                     .instrument(tracing::info_span!("handle_client", ?session_id))
                 );
             }
+            Some(Err(e)) = connections.join_next(), if !connections.is_empty() => {
+                if !e.is_panic() && !e.is_cancelled() {
+                    warn!("unexpected error from joined connection task: {e:?}");
+                }
+            }
             _ = cancellation_token.cancelled() => {
                 drop(listener);
                 break;
diff --git a/proxy/src/proxy.rs b/proxy/src/proxy.rs
index 71e00ed58f..cef3cea514 100644
--- a/proxy/src/proxy.rs
+++ b/proxy/src/proxy.rs
@@ -130,6 +130,11 @@ pub async fn task_main(
                     }),
                 );
             }
+            Some(Err(e)) = connections.join_next(), if !connections.is_empty() => {
+                if !e.is_panic() && !e.is_cancelled() {
+                    warn!("unexpected error from joined connection task: {e:?}");
+                }
+            }
             _ = cancellation_token.cancelled() => {
                 drop(listener);
                 break;