From c8108a4b8490b9c2f29924cf30af34b9f9f5a0ec Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conradludgate@gmail.com>
Date: Thu, 24 Oct 2024 11:55:31 +0100
Subject: [PATCH] make ComputeConnectBackend dyn

---
 proxy/src/auth/backend/mod.rs      |  2 +-
 proxy/src/proxy/connect_compute.rs |  6 ++--
 proxy/src/proxy/wake_compute.rs    |  4 +--
 proxy/src/serverless/backend.rs    | 53 +++++++++++-------------------
 4 files changed, 26 insertions(+), 39 deletions(-)

diff --git a/proxy/src/auth/backend/mod.rs b/proxy/src/auth/backend/mod.rs
index c689bf4e69..81f6eaa6dd 100644
--- a/proxy/src/auth/backend/mod.rs
+++ b/proxy/src/auth/backend/mod.rs
@@ -389,7 +389,7 @@ pub struct ControlPlaneComputeBackend<'a> {
 }
 
 #[async_trait::async_trait]
-impl ComputeConnectBackend for ControlPlaneComputeBackend<'_> {
+impl ComputeConnectBackend for ControlPlaneComputeBackend<'static> {
     async fn wake_compute(
         &self,
         ctx: &RequestMonitoring,
diff --git a/proxy/src/proxy/connect_compute.rs b/proxy/src/proxy/connect_compute.rs
index 659b7afa68..b5a187bca3 100644
--- a/proxy/src/proxy/connect_compute.rs
+++ b/proxy/src/proxy/connect_compute.rs
@@ -56,7 +56,7 @@ pub(crate) trait ConnectMechanism {
 }
 
 #[async_trait]
-pub(crate) trait ComputeConnectBackend {
+pub(crate) trait ComputeConnectBackend: Send + Sync + 'static {
     async fn wake_compute(
         &self,
         ctx: &RequestMonitoring,
@@ -98,10 +98,10 @@ impl ConnectMechanism for TcpMechanism<'_> {
 
 /// Try to connect to the compute node, retrying if necessary.
 #[tracing::instrument(skip_all)]
-pub(crate) async fn connect_to_compute<M: ConnectMechanism, B: ComputeConnectBackend>(
+pub(crate) async fn connect_to_compute<M: ConnectMechanism>(
     ctx: &RequestMonitoring,
     mechanism: &M,
-    user_info: &B,
+    user_info: &dyn ComputeConnectBackend,
     allow_self_signed_compute: bool,
     wake_compute_retry_config: RetryConfig,
     connect_to_compute_retry_config: RetryConfig,
diff --git a/proxy/src/proxy/wake_compute.rs b/proxy/src/proxy/wake_compute.rs
index 4e61094264..69d8111434 100644
--- a/proxy/src/proxy/wake_compute.rs
+++ b/proxy/src/proxy/wake_compute.rs
@@ -11,10 +11,10 @@ use crate::metrics::{
 };
 use crate::proxy::retry::{retry_after, should_retry};
 
-pub(crate) async fn wake_compute<B: ComputeConnectBackend>(
+pub(crate) async fn wake_compute(
     num_retries: &mut u32,
     ctx: &RequestMonitoring,
-    api: &B,
+    api: &dyn ComputeConnectBackend,
     config: RetryConfig,
 ) -> Result<CachedNodeInfo, WakeComputeError> {
     let retry_type = RetryType::WakeCompute;
diff --git a/proxy/src/serverless/backend.rs b/proxy/src/serverless/backend.rs
index b8f89bcfcd..1508de95e7 100644
--- a/proxy/src/serverless/backend.rs
+++ b/proxy/src/serverless/backend.rs
@@ -30,7 +30,7 @@ use crate::control_plane::provider::ApiLockError;
 use crate::control_plane::{Api, CachedNodeInfo};
 use crate::error::{ErrorKind, ReportableError, UserFacingError};
 use crate::intern::EndpointIdInt;
-use crate::proxy::connect_compute::ConnectMechanism;
+use crate::proxy::connect_compute::{ComputeConnectBackend, ConnectMechanism};
 use crate::proxy::retry::{CouldRetry, ShouldRetryWakeCompute};
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::types::{EndpointId, Host};
@@ -190,40 +190,27 @@ impl PoolingBackend {
         tracing::Span::current().record("conn_id", display(conn_id));
         info!(%conn_id, "pool: opening a new connection '{conn_info}'");
 
-        match &self.auth_backend {
+        let api = match &self.auth_backend {
             ServerlessBackend::ControlPlane(cplane) => {
-                crate::proxy::connect_compute::connect_to_compute(
-                    ctx,
-                    &TokioMechanism {
-                        conn_id,
-                        conn_info,
-                        pool: self.pool.clone(),
-                        locks: &self.config.connect_compute_locks,
-                    },
-                    &cplane.attach_to_credentials(keys),
-                    false, // do not allow self signed compute for http flow
-                    self.config.wake_compute_retry_config,
-                    self.config.connect_to_compute_retry_config,
-                )
-                .await
+                &cplane.attach_to_credentials(keys) as &dyn ComputeConnectBackend
             }
-            ServerlessBackend::Local(local_proxy) => {
-                crate::proxy::connect_compute::connect_to_compute(
-                    ctx,
-                    &TokioMechanism {
-                        conn_id,
-                        conn_info,
-                        pool: self.pool.clone(),
-                        locks: &self.config.connect_compute_locks,
-                    },
-                    &**local_proxy,
-                    false, // do not allow self signed compute for http flow
-                    self.config.wake_compute_retry_config,
-                    self.config.connect_to_compute_retry_config,
-                )
-                .await
-            }
-        }
+            ServerlessBackend::Local(local_proxy) => &**local_proxy as &dyn ComputeConnectBackend,
+        };
+
+        crate::proxy::connect_compute::connect_to_compute(
+            ctx,
+            &TokioMechanism {
+                conn_id,
+                conn_info,
+                pool: self.pool.clone(),
+                locks: &self.config.connect_compute_locks,
+            },
+            api,
+            false, // do not allow self signed compute for http flow
+            self.config.wake_compute_retry_config,
+            self.config.connect_to_compute_retry_config,
+        )
+        .await
     }
 
     // Wake up the destination if needed