abstract out listen-accept loops

proxy/src/console_redirect_proxy.rs

@@ -1,21 +1,19 @@
 use crate::auth::backend::ConsoleRedirectBackend;
-use crate::config::{ProxyConfig, ProxyProtocolV2};
-use crate::proxy::{
-    prepare_client_connection, run_until_cancelled, ClientRequestError, ErrorSource,
-};
+use crate::config::ProxyConfig;
+use crate::metrics::Protocol;
+use crate::proxy::{prepare_client_connection, transition_connection, ClientRequestError};
 use crate::{
     cancellation::CancellationHandlerMain,
     context::RequestMonitoring,
-    error::ReportableError,
     metrics::{Metrics, NumClientConnectionsGuard},
-    protocol2::read_proxy_protocol,
     proxy::handshake::{handshake, HandshakeData},
 };
 use futures::TryFutureExt;
+use std::net::IpAddr;
 use std::sync::Arc;
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
 use tokio_util::sync::CancellationToken;
-use tracing::{error, info, Instrument};
+use tracing::{info, Instrument};
 
 use crate::proxy::{
     connect_compute::{connect_to_compute, TcpMechanism},
@@ -33,107 +31,53 @@ pub async fn task_main(
         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)?;
+    super::connection_loop(
+        config,
+        listener,
+        cancellation_token,
+        Protocol::Tcp,
+        C {
+            config,
+            backend,
+            cancellation_handler,
+        },
+    )
+    .await
+}
 
-    let connections = tokio_util::task::task_tracker::TaskTracker::new();
+#[derive(Clone)]
+struct C {
+    config: &'static ProxyConfig,
+    backend: &'static ConsoleRedirectBackend,
+    cancellation_handler: Arc<CancellationHandlerMain>,
+}
 
-    while let Some(accept_result) =
-        run_until_cancelled(listener.accept(), &cancellation_token).await
-    {
-        let (socket, peer_addr) = accept_result?;
+impl super::ConnHandler for C {
+    async fn handle(
+        self,
+        session_id: uuid::Uuid,
+        peer_addr: IpAddr,
+        socket: crate::protocol2::ChainRW<tokio::net::TcpStream>,
+        conn_gauge: crate::metrics::NumClientConnectionsGuard<'static>,
+    ) {
+        let ctx = RequestMonitoring::new(session_id, peer_addr, Protocol::Tcp, &self.config.region);
+        let span = ctx.span();
 
-        let conn_gauge = Metrics::get()
-            .proxy
-            .client_connections
-            .guard(crate::metrics::Protocol::Tcp);
+        let startup = Box::pin(
+            handle_client(
+                self.config,
+                self.backend,
+                &ctx,
+                self.cancellation_handler,
+                socket,
+                conn_gauge,
+            )
+            .instrument(span.clone()),
+        );
 
-        let session_id = uuid::Uuid::new_v4();
-        let cancellation_handler = Arc::clone(&cancellation_handler);
-
-        tracing::info!(protocol = "tcp", %session_id, "accepted new TCP connection");
-
-        connections.spawn(async move {
-            let (socket, peer_addr) = match read_proxy_protocol(socket).await {
-                Err(e) => {
-                    error!("per-client task finished with an error: {e:#}");
-                    return;
-                }
-                Ok((_socket, None)) if config.proxy_protocol_v2 == ProxyProtocolV2::Required => {
-                    error!("missing required proxy protocol header");
-                    return;
-                }
-                Ok((_socket, Some(_))) if config.proxy_protocol_v2 == ProxyProtocolV2::Rejected => {
-                    error!("proxy protocol header not supported");
-                    return;
-                }
-                Ok((socket, Some(addr))) => (socket, addr.ip()),
-                Ok((socket, None)) => (socket, peer_addr.ip()),
-            };
-
-            match socket.inner.set_nodelay(true) {
-                Ok(()) => {}
-                Err(e) => {
-                    error!("per-client task finished with an error: failed to set socket option: {e:#}");
-                    return;
-                }
-            };
-
-            let ctx = RequestMonitoring::new(
-                session_id,
-                peer_addr,
-                crate::metrics::Protocol::Tcp,
-                &config.region,
-            );
-            let span = ctx.span();
-
-            let startup = Box::pin(
-                handle_client(
-                    config,
-                    backend,
-                    &ctx,
-                    cancellation_handler,
-                    socket,
-                    conn_gauge,
-                )
-                .instrument(span.clone()),
-            );
-            let res = startup.await;
-
-            match res {
-                Err(e) => {
-                    // todo: log and push to ctx the error kind
-                    ctx.set_error_kind(e.get_error_kind());
-                    error!(parent: &span, "per-client task finished with an error: {e:#}");
-                }
-                Ok(None) => {
-                    ctx.set_success();
-                }
-                Ok(Some(p)) => {
-                    ctx.set_success();
-                    ctx.log_connect();
-                    match p.proxy_pass().instrument(span.clone()).await {
-                        Ok(()) => {}
-                        Err(ErrorSource::Client(e)) => {
-                            error!(parent: &span, "per-client task finished with an IO error from the client: {e:#}");
-                        }
-                        Err(ErrorSource::Compute(e)) => {
-                            error!(parent: &span, "per-client task finished with an IO error from the compute: {e:#}");
-                        }
-                    }
-                }
-            }
-        });
+        let res = startup.await;
+        transition_connection(ctx, res).await;
     }
-
-    connections.close();
-    drop(listener);
-
-    // Drain connections
-    connections.wait().await;
-
-    Ok(())
 }
 
 pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(

proxy/src/lib.rs

@@ -82,13 +82,16 @@
     impl_trait_overcaptures,
 )]
 
-use std::convert::Infallible;
+use std::{convert::Infallible, future::Future, net::IpAddr};
 
 use anyhow::{bail, Context};
 use intern::{EndpointIdInt, EndpointIdTag, InternId};
-use tokio::task::JoinError;
+use protocol2::{get_client_conn_info, ChainRW};
+use proxy::run_until_cancelled;
+use tokio::{net::TcpStream, task::JoinError};
 use tokio_util::sync::CancellationToken;
-use tracing::warn;
+use tracing::{error, warn};
+use uuid::Uuid;
 
 extern crate hyper0 as hyper;
 
@@ -276,3 +279,81 @@ impl EndpointId {
         ProjectId(self.0.clone())
     }
 }
+
+pub(crate) trait ConnHandler: Clone + Send + 'static {
+    fn handle(
+        self,
+        session_id: Uuid,
+        peer_addr: IpAddr,
+        stream: ChainRW<TcpStream>,
+        conn_gauge: metrics::NumClientConnectionsGuard<'static>,
+    ) -> impl Future<Output = ()> + Send;
+}
+
+/// Accept connections, parse the proxy-protocol v2 header and spawn a tracked connection task.
+pub(crate) async fn connection_loop<C>(
+    config: &'static config::ProxyConfig,
+    listener: tokio::net::TcpListener,
+    cancellation_token: CancellationToken,
+    protocol: metrics::Protocol,
+    conn_handler: C,
+) -> anyhow::Result<()>
+where
+    C: ConnHandler,
+{
+    // When set for the server socket, the keepalive setting
+    // will be inherited by all accepted client sockets.
+    socket2::SockRef::from(&listener).set_keepalive(true)?;
+
+    let connections = tokio_util::task::task_tracker::TaskTracker::new();
+
+    while let Some(accept_result) =
+        run_until_cancelled(listener.accept(), &cancellation_token).await
+    {
+        let (socket, peer_addr) = accept_result?;
+
+        let conn_gauge = metrics::Metrics::get()
+            .proxy
+            .client_connections
+            .guard(protocol);
+
+        let session_id = uuid::Uuid::new_v4();
+        let conn_handler = conn_handler.clone();
+
+        tracing::info!(protocol = protocol.as_str(), %session_id, "accepted new TCP connection");
+
+        connections.spawn(async move {
+            let (socket, peer_addr) = match get_client_conn_info(socket, config.proxy_protocol_v2).await {
+                Err(e) => {
+                    error!("per-client task finished with an error: {e:#}");
+                    return;
+                }
+                Ok((socket, Some(addr))) => (socket, addr),
+                Ok((socket, None)) => (socket, peer_addr.ip()),
+            };
+
+            match socket.inner.set_nodelay(true) {
+                Ok(()) => {}
+                Err(e) => {
+                    error!("per-client task finished with an error: failed to set socket option: {e:#}");
+                    return;
+                }
+            };
+
+            conn_handler.handle(
+                session_id,
+                peer_addr,
+                socket,
+                conn_gauge,
+            ).await;
+        });
+    }
+
+    connections.close();
+    drop(listener);
+
+    // Drain connections
+    connections.wait().await;
+
+    Ok(())
+}

proxy/src/protocol2.rs

@@ -2,15 +2,18 @@
 
 use std::{
     io,
-    net::SocketAddr,
+    net::{IpAddr, SocketAddr},
     pin::Pin,
     task::{Context, Poll},
 };
 
+use anyhow::bail;
 use bytes::BytesMut;
 use pin_project_lite::pin_project;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, ReadBuf};
 
+use crate::config::ProxyProtocolV2;
+
 pin_project! {
     /// A chained [`AsyncRead`] with [`AsyncWrite`] passthrough
     pub(crate) struct ChainRW<T> {
@@ -60,7 +63,23 @@ const HEADER: [u8; 12] = [
     0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A,
 ];
 
-pub(crate) async fn read_proxy_protocol<T: AsyncRead + Unpin>(
+pub(crate) async fn get_client_conn_info<T: AsyncRead + Unpin>(
+    socket: T,
+    proxy_protocol_v2: ProxyProtocolV2,
+) -> anyhow::Result<(ChainRW<T>, Option<IpAddr>)> {
+    match read_proxy_protocol(socket).await? {
+        (_socket, None) if proxy_protocol_v2 == ProxyProtocolV2::Required => {
+            bail!("missing required proxy protocol header");
+        }
+        (_socket, Some(_)) if proxy_protocol_v2 == ProxyProtocolV2::Rejected => {
+            bail!("proxy protocol header not supported");
+        }
+        (socket, Some(addr)) => Ok((socket, Some(addr.ip()))),
+        (socket, None) => Ok((socket, None)),
+    }
+}
+
+async fn read_proxy_protocol<T: AsyncRead + Unpin>(
     mut read: T,
 ) -> std::io::Result<(ChainRW<T>, Option<SocketAddr>)> {
     let mut buf = BytesMut::with_capacity(128);

proxy/src/proxy/mod.rs

@@ -10,8 +10,8 @@ pub(crate) mod wake_compute;
 pub use copy_bidirectional::copy_bidirectional_client_compute;
 pub use copy_bidirectional::ErrorSource;
 
-use crate::config::ProxyProtocolV2;
 use crate::control_plane::provider::ControlPlaneBackend;
+use crate::metrics::Protocol;
 use crate::{
     auth,
     cancellation::{self, CancellationHandlerMain},
@@ -20,7 +20,6 @@ use crate::{
     context::RequestMonitoring,
     error::ReportableError,
     metrics::{Metrics, NumClientConnectionsGuard},
-    protocol2::read_proxy_protocol,
     proxy::handshake::{handshake, HandshakeData},
     rate_limiter::EndpointRateLimiter,
     stream::{PqStream, Stream},
@@ -32,6 +31,7 @@ use once_cell::sync::OnceCell;
 use pq_proto::{BeMessage as Be, StartupMessageParams};
 use regex::Regex;
 use smol_str::{format_smolstr, SmolStr};
+use std::net::IpAddr;
 use std::sync::Arc;
 use thiserror::Error;
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
@@ -72,110 +72,91 @@ pub async fn task_main(
         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)?;
+    super::connection_loop(
+        config,
+        listener,
+        cancellation_token,
+        Protocol::Tcp,
+        C {
+            config,
+            auth_backend,
+            cancellation_handler,
+            endpoint_rate_limiter,
+        },
+    )
+    .await
+}
 
-    let connections = tokio_util::task::task_tracker::TaskTracker::new();
+#[derive(Clone)]
+struct C {
+    config: &'static ProxyConfig,
+    auth_backend: &'static ControlPlaneBackend,
+    cancellation_handler: Arc<CancellationHandlerMain>,
+    endpoint_rate_limiter: Arc<EndpointRateLimiter>,
+}
 
-    while let Some(accept_result) =
-        run_until_cancelled(listener.accept(), &cancellation_token).await
-    {
-        let (socket, peer_addr) = accept_result?;
+impl super::ConnHandler for C {
+    async fn handle(
+        self,
+        session_id: uuid::Uuid,
+        peer_addr: IpAddr,
+        socket: crate::protocol2::ChainRW<tokio::net::TcpStream>,
+        conn_gauge: crate::metrics::NumClientConnectionsGuard<'static>,
+    ) {
+        let ctx = RequestMonitoring::new(
+            session_id,
+            peer_addr,
+            crate::metrics::Protocol::Tcp,
+            &self.config.region,
+        );
+        let span = ctx.span();
 
-        let conn_gauge = Metrics::get()
-            .proxy
-            .client_connections
-            .guard(crate::metrics::Protocol::Tcp);
+        let startup = Box::pin(
+            handle_client(
+                self.config,
+                self.auth_backend,
+                &ctx,
+                self.cancellation_handler,
+                socket,
+                ClientMode::Tcp,
+                self.endpoint_rate_limiter,
+                conn_gauge,
+            )
+            .instrument(span.clone()),
+        );
 
-        let session_id = uuid::Uuid::new_v4();
-        let cancellation_handler = Arc::clone(&cancellation_handler);
+        let res = startup.await;
+        transition_connection(ctx, res).await;
+    }
+}
 
-        tracing::info!(protocol = "tcp", %session_id, "accepted new TCP connection");
-        let endpoint_rate_limiter2 = endpoint_rate_limiter.clone();
-
-        connections.spawn(async move {
-            let (socket, peer_addr) = match read_proxy_protocol(socket).await {
-                Err(e) => {
-                    error!("per-client task finished with an error: {e:#}");
-                    return;
-                }
-                Ok((_socket, None)) if config.proxy_protocol_v2 == ProxyProtocolV2::Required => {
-                    error!("missing required proxy protocol header");
-                    return;
-                }
-                Ok((_socket, Some(_))) if config.proxy_protocol_v2 == ProxyProtocolV2::Rejected => {
-                    error!("proxy protocol header not supported");
-                    return;
-                }
-                Ok((socket, Some(addr))) => (socket, addr.ip()),
-                Ok((socket, None)) => (socket, peer_addr.ip()),
-            };
-
-            match socket.inner.set_nodelay(true) {
+pub(crate) async fn transition_connection<S: AsyncRead + AsyncWrite + Unpin>(
+    ctx: RequestMonitoring,
+    res: Result<Option<ProxyPassthrough<S>>, ClientRequestError>,
+) {
+    let span = ctx.span();
+    match res {
+        Err(e) => {
+            ctx.set_error_kind(e.get_error_kind());
+            error!(parent: &span, "per-client task finished with an error: {e:#}");
+        }
+        Ok(None) => {
+            ctx.set_success();
+        }
+        Ok(Some(p)) => {
+            ctx.set_success();
+            ctx.log_connect();
+            match p.proxy_pass().instrument(span.clone()).await {
                 Ok(()) => {}
-                Err(e) => {
-                    error!("per-client task finished with an error: failed to set socket option: {e:#}");
-                    return;
+                Err(ErrorSource::Client(e)) => {
+                    error!(parent: &span, "per-client task finished with an IO error from the client: {e:#}");
                 }
-            };
-
-            let ctx = RequestMonitoring::new(
-                session_id,
-                peer_addr,
-                crate::metrics::Protocol::Tcp,
-                &config.region,
-            );
-            let span = ctx.span();
-
-            let startup = Box::pin(
-                handle_client(
-                    config,
-                    auth_backend,
-                    &ctx,
-                    cancellation_handler,
-                    socket,
-                    ClientMode::Tcp,
-                    endpoint_rate_limiter2,
-                    conn_gauge,
-                )
-                .instrument(span.clone()),
-            );
-            let res = startup.await;
-
-            match res {
-                Err(e) => {
-                    // todo: log and push to ctx the error kind
-                    ctx.set_error_kind(e.get_error_kind());
-                    error!(parent: &span, "per-client task finished with an error: {e:#}");
-                }
-                Ok(None) => {
-                    ctx.set_success();
-                }
-                Ok(Some(p)) => {
-                    ctx.set_success();
-                    ctx.log_connect();
-                    match p.proxy_pass().instrument(span.clone()).await {
-                        Ok(()) => {}
-                        Err(ErrorSource::Client(e)) => {
-                            error!(parent: &span, "per-client task finished with an IO error from the client: {e:#}");
-                        }
-                        Err(ErrorSource::Compute(e)) => {
-                            error!(parent: &span, "per-client task finished with an IO error from the compute: {e:#}");
-                        }
-                    }
+                Err(ErrorSource::Compute(e)) => {
+                    error!(parent: &span, "per-client task finished with an IO error from the compute: {e:#}");
                 }
             }
-        });
+        }
     }
-
-    connections.close();
-    drop(listener);
-
-    // Drain connections
-    connections.wait().await;
-
-    Ok(())
 }
 
 pub(crate) enum ClientMode {

proxy/src/proxy/tests/mod.rs

@@ -10,13 +10,14 @@ use super::*;
 use crate::auth::backend::{
     ComputeCredentialKeys, ComputeCredentials, ComputeUserInfo, TestBackend,
 };
-use crate::config::{CertResolver, RetryConfig};
+use crate::config::{CertResolver, ProxyProtocolV2, RetryConfig};
 use crate::control_plane::messages::{ControlPlaneError, Details, MetricsAuxInfo, Status};
 use crate::control_plane::provider::{
     CachedAllowedIps, CachedRoleSecret, ControlPlaneBackend, NodeInfoCache,
 };
 use crate::control_plane::{self, CachedNodeInfo, NodeInfo};
 use crate::error::ErrorKind;
+use crate::protocol2::get_client_conn_info;
 use crate::{sasl, scram, BranchId, EndpointId, ProjectId};
 use anyhow::{bail, Context};
 use async_trait::async_trait;
@@ -177,7 +178,7 @@ async fn dummy_proxy(
     tls: Option<TlsConfig>,
     auth: impl TestAuth + Send,
 ) -> anyhow::Result<()> {
-    let (client, _) = read_proxy_protocol(client).await?;
+    let (client, _) = get_client_conn_info(client, ProxyProtocolV2::Supported).await?;
     let mut stream =
         match handshake(&RequestMonitoring::test(), client, tls.as_ref(), false).await? {
             HandshakeData::Startup(stream, _) => stream,

proxy/src/serverless/mod.rs

@@ -16,7 +16,6 @@ use atomic_take::AtomicTake;
 use bytes::Bytes;
 pub use conn_pool::GlobalConnPoolOptions;
 
-use anyhow::Context;
 use futures::future::{select, Either};
 use futures::TryFutureExt;
 use http::{Method, Response, StatusCode};
@@ -36,19 +35,18 @@ use crate::auth::ServerlessBackend;
 use crate::cancellation::CancellationHandlerMain;
 use crate::config::ProxyConfig;
 use crate::context::RequestMonitoring;
-use crate::metrics::Metrics;
-use crate::protocol2::{read_proxy_protocol, ChainRW};
-use crate::proxy::run_until_cancelled;
+use crate::metrics::{Metrics, Protocol};
+use crate::protocol2::ChainRW;
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::serverless::backend::PoolingBackend;
 use crate::serverless::http_util::{api_error_into_response, json_response};
 
-use std::net::{IpAddr, SocketAddr};
+use std::net::IpAddr;
 use std::pin::{pin, Pin};
 use std::sync::Arc;
 use tokio::net::{TcpListener, TcpStream};
 use tokio_util::sync::CancellationToken;
-use tracing::{error, info, warn, Instrument};
+use tracing::{error, info, instrument, warn, Instrument};
 use utils::http::error::ApiError;
 
 pub(crate) const SERVERLESS_DRIVER_SNI: &str = "api";
@@ -125,81 +123,100 @@ pub async fn task_main(
         }
     };
 
-    let connections = tokio_util::task::task_tracker::TaskTracker::new();
-    connections.close(); // allows `connections.wait to complete`
+    let requests = TaskTracker::new();
+    requests.close(); // allows `requests.wait to complete`
 
-    while let Some(res) = run_until_cancelled(ws_listener.accept(), &cancellation_token).await {
-        let (conn, peer_addr) = res.context("could not accept TCP stream")?;
-        if let Err(e) = conn.set_nodelay(true) {
-            tracing::error!("could not set nodelay: {e}");
-            continue;
-        }
-        let conn_id = uuid::Uuid::new_v4();
-        let http_conn_span = tracing::info_span!("http_conn", ?conn_id);
+    crate::connection_loop(
+        config,
+        ws_listener,
+        cancellation_token.clone(),
+        Protocol::Http,
+        C {
+            config,
+            backend,
+            cancellation_handler,
+            endpoint_rate_limiter,
+            tls_acceptor,
+            requests: requests.clone(),
+            cancellation_token,
+        },
+    )
+    .await?;
 
+    requests.wait().await;
+
+    Ok(())
+}
+
+#[derive(Clone)]
+struct C {
+    config: &'static ProxyConfig,
+    backend: Arc<PoolingBackend>,
+    cancellation_handler: Arc<CancellationHandlerMain>,
+    endpoint_rate_limiter: Arc<EndpointRateLimiter>,
+    tls_acceptor: Arc<dyn MaybeTlsAcceptor>,
+    requests: TaskTracker,
+    cancellation_token: CancellationToken,
+}
+
+impl super::ConnHandler for C {
+    #[instrument(name = "http_conn", skip_all, fields(conn_id))]
+    async fn handle(
+        self,
+        conn_id: uuid::Uuid,
+        peer_addr: IpAddr,
+        stream: ChainRW<TcpStream>,
+        conn_gauge: crate::metrics::NumClientConnectionsGuard<'static>,
+    ) {
+        // try and close an old HTTP connection.
+        // picked at random
         let n_connections = Metrics::get()
             .proxy
             .client_connections
             .sample(crate::metrics::Protocol::Http);
-        tracing::trace!(?n_connections, threshold = ?config.http_config.client_conn_threshold, "check");
-        if n_connections > config.http_config.client_conn_threshold {
+        tracing::trace!(?n_connections, threshold = ?self.config.http_config.client_conn_threshold, "check");
+        if n_connections > self.config.http_config.client_conn_threshold {
             tracing::trace!("attempting to cancel a random connection");
-            if let Some(token) = config.http_config.cancel_set.take() {
+            if let Some(token) = self.config.http_config.cancel_set.take() {
                 tracing::debug!("cancelling a random connection");
                 token.cancel();
             }
         }
 
-        let conn_token = cancellation_token.child_token();
-        let tls_acceptor = tls_acceptor.clone();
-        let backend = backend.clone();
-        let connections2 = connections.clone();
-        let cancellation_handler = cancellation_handler.clone();
-        let endpoint_rate_limiter = endpoint_rate_limiter.clone();
-        connections.spawn(
-            async move {
-                let conn_token2 = conn_token.clone();
-                let _cancel_guard = config.http_config.cancel_set.insert(conn_id, conn_token2);
+        let conn_token = self.cancellation_token.child_token();
+        let _cancel_guard = self
+            .config
+            .http_config
+            .cancel_set
+            .insert(conn_id, conn_token.clone());
 
-                let session_id = uuid::Uuid::new_v4();
+        let startup_result = Box::pin(connection_startup(
+            self.config,
+            self.tls_acceptor,
+            conn_id,
+            stream,
+            peer_addr,
+        ))
+        .await;
+        let Some((conn, peer_addr)) = startup_result else {
+            return;
+        };
 
-                let _gauge = Metrics::get()
-                    .proxy
-                    .client_connections
-                    .guard(crate::metrics::Protocol::Http);
+        Box::pin(connection_handler(
+            self.config,
+            self.backend,
+            self.requests,
+            self.cancellation_handler,
+            self.endpoint_rate_limiter,
+            conn_token,
+            conn,
+            peer_addr,
+            conn_id,
+        ))
+        .await;
 
-                let startup_result = Box::pin(connection_startup(
-                    config,
-                    tls_acceptor,
-                    session_id,
-                    conn,
-                    peer_addr,
-                ))
-                .await;
-                let Some((conn, peer_addr)) = startup_result else {
-                    return;
-                };
-
-                Box::pin(connection_handler(
-                    config,
-                    backend,
-                    connections2,
-                    cancellation_handler,
-                    endpoint_rate_limiter,
-                    conn_token,
-                    conn,
-                    peer_addr,
-                    session_id,
-                ))
-                .await;
-            }
-            .instrument(http_conn_span),
-        );
+        drop(conn_gauge);
     }
-
-    connections.wait().await;
-
-    Ok(())
 }
 
 pub(crate) trait AsyncReadWrite: AsyncRead + AsyncWrite + Send + 'static {}
@@ -227,26 +244,14 @@ impl MaybeTlsAcceptor for NoTls {
     }
 }
 
-/// Handles the TCP startup lifecycle.
-/// 1. Parses PROXY protocol V2
-/// 2. Handles TLS handshake
+/// Handles the TLS startup handshake.
 async fn connection_startup(
     config: &ProxyConfig,
     tls_acceptor: Arc<dyn MaybeTlsAcceptor>,
     session_id: uuid::Uuid,
-    conn: TcpStream,
-    peer_addr: SocketAddr,
+    conn: ChainRW<TcpStream>,
+    peer_addr: IpAddr,
 ) -> Option<(AsyncRW, IpAddr)> {
-    // handle PROXY protocol
-    let (conn, peer) = match read_proxy_protocol(conn).await {
-        Ok(c) => c,
-        Err(e) => {
-            tracing::error!(?session_id, %peer_addr, "failed to accept TCP connection: invalid PROXY protocol V2 header: {e:#}");
-            return None;
-        }
-    };
-
-    let peer_addr = peer.unwrap_or(peer_addr).ip();
     let has_private_peer_addr = match peer_addr {
         IpAddr::V4(ip) => ip.is_private(),
         IpAddr::V6(_) => false,