abstract byte handling

proxy/src/binary/pg_sni_router.rs

@@ -297,7 +297,7 @@ async fn handle_client(
     // Starting from here we only proxy the client's traffic.
     info!("performing the proxy pass...");
 
-    match copy_bidirectional_client_compute(&mut tls_stream, &mut client).await {
+    match copy_bidirectional_client_compute(&mut tls_stream, &mut client, |_, _| {}).await {
         Ok(_) => Ok(()),
         Err(ErrorSource::Client(err)) => Err(err).context("client"),
         Err(ErrorSource::Compute(err)) => Err(err).context("compute"),

proxy/src/binary/proxy.rs

@@ -24,6 +24,7 @@ use crate::config::{
 use crate::context::parquet::ParquetUploadArgs;
 use crate::http::health_server::AppMetrics;
 use crate::metrics::Metrics;
+use crate::proxy::conntrack::ConnectionTracking;
 use crate::rate_limiter::{
     EndpointRateLimiter, LeakyBucketConfig, RateBucketInfo, WakeComputeRateLimiter,
 };
@@ -418,6 +419,8 @@ pub async fn run() -> anyhow::Result<()> {
         64,
     ));
 
+    let conntracking = Arc::new(ConnectionTracking::default());
+
     // client facing tasks. these will exit on error or on cancellation
     // cancellation returns Ok(())
     let mut client_tasks = JoinSet::new();
@@ -431,6 +434,7 @@ pub async fn run() -> anyhow::Result<()> {
                     cancellation_token.clone(),
                     cancellation_handler.clone(),
                     endpoint_rate_limiter.clone(),
+                    conntracking.clone(),
                 ));
             }
 
@@ -453,6 +457,7 @@ pub async fn run() -> anyhow::Result<()> {
                     proxy_listener,
                     cancellation_token.clone(),
                     cancellation_handler.clone(),
+                    conntracking.clone(),
                 ));
             }
         }

proxy/src/console_redirect_proxy.rs

@@ -13,6 +13,7 @@ use crate::error::ReportableError;
 use crate::metrics::{Metrics, NumClientConnectionsGuard};
 use crate::protocol2::{ConnectHeader, ConnectionInfo, read_proxy_protocol};
 use crate::proxy::connect_compute::{TcpMechanism, connect_to_compute};
+use crate::proxy::conntrack::ConnectionTracking;
 use crate::proxy::handshake::{HandshakeData, handshake};
 use crate::proxy::passthrough::ProxyPassthrough;
 use crate::proxy::{
@@ -25,6 +26,7 @@ pub async fn task_main(
     listener: tokio::net::TcpListener,
     cancellation_token: CancellationToken,
     cancellation_handler: Arc<CancellationHandler>,
+    conntracking: Arc<ConnectionTracking>,
 ) -> anyhow::Result<()> {
     scopeguard::defer! {
         info!("proxy has shut down");
@@ -50,6 +52,7 @@ pub async fn task_main(
         let session_id = uuid::Uuid::new_v4();
         let cancellation_handler = Arc::clone(&cancellation_handler);
         let cancellations = cancellations.clone();
+        let conntracking = Arc::clone(&conntracking);
 
         debug!(protocol = "tcp", %session_id, "accepted new TCP connection");
 
@@ -111,6 +114,7 @@ pub async fn task_main(
                 socket,
                 conn_gauge,
                 cancellations,
+                conntracking,
             )
             .instrument(ctx.span())
             .boxed()
@@ -167,6 +171,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
     stream: S,
     conn_gauge: NumClientConnectionsGuard<'static>,
     cancellations: tokio_util::task::task_tracker::TaskTracker,
+    conntracking: Arc<ConnectionTracking>,
 ) -> Result<Option<ProxyPassthrough<S>>, ClientRequestError> {
     debug!(
         protocol = %ctx.protocol(),
@@ -264,6 +269,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
         compute: node,
         session_id: ctx.session_id(),
         cancel: session,
+        conntracking,
         _req: request_gauge,
         _conn: conn_gauge,
     }))

proxy/src/metrics.rs

@@ -200,8 +200,10 @@ pub enum HttpDirection {
 #[derive(FixedCardinalityLabel, Copy, Clone)]
 #[label(singleton = "direction")]
 pub enum Direction {
-    Tx,
-    Rx,
+    #[label(rename = "tx")]
+    ComputeToClient,
+    #[label(rename = "rx")]
+    ClientToCompute,
 }
 
 #[derive(FixedCardinalityLabel, Clone, Copy, Debug)]

proxy/src/proxy/conntrack.rs

@@ -0,0 +1,564 @@
+use std::fmt;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::time::SystemTime;
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct ConnId(usize);
+
+#[derive(Default)]
+pub struct ConnectionTracking {
+    conns: clashmap::ClashMap<ConnId, (ConnectionState, SystemTime)>,
+}
+
+impl ConnectionTracking {
+    pub fn new_tracker(self: &Arc<Self>) -> ConnectionTracker<Conn> {
+        let conn_id = self.new_conn_id();
+        ConnectionTracker::new(Conn {
+            conn_id,
+            tracking: Arc::clone(self),
+        })
+    }
+
+    fn new_conn_id(&self) -> ConnId {
+        static NEXT_ID: AtomicUsize = AtomicUsize::new(0);
+        let id = ConnId(NEXT_ID.fetch_add(1, Ordering::Relaxed));
+        self.conns
+            .insert(id, (ConnectionState::Idle, SystemTime::now()));
+        id
+    }
+
+    fn update(&self, conn_id: ConnId, new_state: ConnectionState) {
+        let new_timestamp = SystemTime::now();
+        let old_state = self.conns.insert(conn_id, (new_state, new_timestamp));
+
+        if let Some((old_state, _old_timestamp)) = old_state {
+            tracing::debug!(?conn_id, %old_state, %new_state, "conntrack: update");
+        } else {
+            tracing::debug!(?conn_id, %new_state, "conntrack: update");
+        }
+    }
+
+    fn remove(&self, conn_id: ConnId) {
+        if let Some((_, (old_state, _old_timestamp))) = self.conns.remove(&conn_id) {
+            tracing::debug!(?conn_id, %old_state, "conntrack: remove");
+        }
+    }
+}
+
+pub struct Conn {
+    conn_id: ConnId,
+    tracking: Arc<ConnectionTracking>,
+}
+
+impl StateChangeObserver for Conn {
+    fn change(&mut self, _old_state: ConnectionState, new_state: ConnectionState) {
+        match new_state {
+            ConnectionState::Init
+            | ConnectionState::Idle
+            | ConnectionState::Transaction
+            | ConnectionState::Busy
+            | ConnectionState::Unknown => self.tracking.update(self.conn_id, new_state),
+            ConnectionState::Closed => self.tracking.remove(self.conn_id),
+        }
+    }
+}
+
+/// Called by `ConnectionTracker` whenever the `ConnectionState` changed.
+pub trait StateChangeObserver {
+    /// Called iff the connection's state changed.
+    fn change(&mut self, old_state: ConnectionState, new_state: ConnectionState);
+}
+
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
+#[repr(u8)]
+pub enum ConnectionState {
+    #[default]
+    Init = 0,
+    Idle = 1,
+    Transaction = 2,
+    Busy = 3,
+    Closed = 4,
+    Unknown = 5,
+}
+
+impl fmt::Display for ConnectionState {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match *self {
+            ConnectionState::Init => f.write_str("init"),
+            ConnectionState::Idle => f.write_str("idle"),
+            ConnectionState::Transaction => f.write_str("transaction"),
+            ConnectionState::Busy => f.write_str("busy"),
+            ConnectionState::Closed => f.write_str("closed"),
+            ConnectionState::Unknown => f.write_str("unknown"),
+        }
+    }
+}
+
+/// Tracks the `ConnectionState` of a connection by inspecting the frontend and
+/// backend stream and reacting to specific messages. Used in combination with
+/// two `TrackedStream`s.
+pub struct ConnectionTracker<SCO: StateChangeObserver> {
+    state: ConnectionState,
+    observer: SCO,
+}
+
+impl<SCO: StateChangeObserver> Drop for ConnectionTracker<SCO> {
+    fn drop(&mut self) {
+        self.observer.change(self.state, ConnectionState::Closed);
+    }
+}
+
+impl<SCO: StateChangeObserver> ConnectionTracker<SCO> {
+    pub fn new(observer: SCO) -> Self {
+        ConnectionTracker {
+            state: ConnectionState::default(),
+            observer,
+        }
+    }
+
+    pub fn frontend_message_tag(&mut self, tag: Tag) {
+        self.update_state(|old_state| Self::state_from_frontend_tag(old_state, tag));
+    }
+
+    pub fn backend_message_tag(&mut self, tag: Tag) {
+        self.update_state(|old_state| Self::state_from_backend_tag(old_state, tag));
+    }
+
+    fn update_state(&mut self, new_state_fn: impl FnOnce(ConnectionState) -> ConnectionState) {
+        let old_state = self.state;
+        let new_state = new_state_fn(old_state);
+        if old_state != new_state {
+            self.observer.change(old_state, new_state);
+            self.state = new_state;
+        }
+    }
+
+    fn state_from_frontend_tag(_old_state: ConnectionState, fe_tag: Tag) -> ConnectionState {
+        // Most activity from the client puts connection into busy state.
+        // Only the server can put a connection back into idle state.
+        match fe_tag {
+            Tag::Start | Tag::ReadyForQuery(_) | Tag::Message(_) => ConnectionState::Busy,
+            Tag::End => ConnectionState::Closed,
+            Tag::Lost => ConnectionState::Unknown,
+        }
+    }
+
+    fn state_from_backend_tag(old_state: ConnectionState, be_tag: Tag) -> ConnectionState {
+        match be_tag {
+            // Check for RFQ and put connection into idle or idle in transaction state.
+            Tag::ReadyForQuery(b'I') => ConnectionState::Idle,
+            Tag::ReadyForQuery(b'T') => ConnectionState::Transaction,
+            Tag::ReadyForQuery(b'E') => ConnectionState::Transaction,
+            // We can't put a connection into idle state for unknown RFQ status.
+            Tag::ReadyForQuery(_) => ConnectionState::Unknown,
+            // Ignore out-fo message from the server.
+            Tag::NOTICE | Tag::NOTIFICATION_RESPONSE | Tag::PARAMETER_STATUS => old_state,
+            // All other activity from server puts connection into busy state.
+            Tag::Start | Tag::Message(_) => ConnectionState::Busy,
+
+            Tag::End => ConnectionState::Closed,
+            Tag::Lost => ConnectionState::Unknown,
+        }
+    }
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum Tag {
+    Message(u8),
+    ReadyForQuery(u8),
+    Start,
+    End,
+    Lost,
+}
+
+impl Tag {
+    const READY_FOR_QUERY: Tag = Tag::Message(b'Z');
+    const NOTICE: Tag = Tag::Message(b'N');
+    const NOTIFICATION_RESPONSE: Tag = Tag::Message(b'A');
+    const PARAMETER_STATUS: Tag = Tag::Message(b'S');
+}
+
+impl fmt::Display for Tag {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match *self {
+            Tag::Start => f.write_str("start"),
+            Tag::End => f.write_str("end"),
+            Tag::Lost => f.write_str("lost"),
+            Tag::Message(tag) => write!(f, "'{}'", tag as char),
+            Tag::ReadyForQuery(status) => write!(f, "ReadyForQuery:'{}'", status as char),
+        }
+    }
+}
+
+pub trait TagObserver {
+    fn observe(&mut self, tag: Tag);
+}
+
+impl<F: FnMut(Tag)> TagObserver for F {
+    fn observe(&mut self, tag: Tag) {
+        (self)(tag);
+    }
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub(super) enum StreamScannerState {
+    #[allow(dead_code)]
+    /// Initial state when no message has been read and we are looling for a
+    /// message without a tag.
+    Start,
+    /// Read a message tag.
+    Tag,
+    /// Read the length bytes and calculate the total length.
+    Length {
+        tag: Tag,
+        /// Number of bytes missing to know the full length of the message: 0..=4
+        length_bytes_missing: usize,
+        /// Total length of the message (without tag) that is calculated as we
+        /// read the bytes for the length.
+        calculated_length: usize,
+    },
+    /// Read (= skip) the payload.
+    Payload {
+        tag: Tag,
+        /// If this is the first time payload bytes are read. Important when
+        /// inspecting specific messages, like ReadyForQuery.
+        first: bool,
+        /// Number of payload bytes left to read before looking for a new tag.
+        bytes_to_skip: usize,
+    },
+    /// Stream was terminated.
+    End,
+    /// Stream ended up in a lost state. We only stop tracking the stream, not
+    /// interrupt it.
+    Lost,
+}
+
+impl StreamScannerState {
+    pub(super) fn scan_bytes<TO: TagObserver>(&mut self, mut buf: &[u8], observer: &mut TO) {
+        use StreamScannerState as S;
+
+        if matches!(*self, S::End | S::Lost) {
+            return;
+        }
+        if buf.is_empty() {
+            match *self {
+                S::Start | S::Tag => {
+                    observer.observe(Tag::End);
+                    *self = S::End;
+                    return;
+                }
+                S::Length { .. } | S::Payload { .. } => {
+                    observer.observe(Tag::Lost);
+                    *self = S::Lost;
+                    return;
+                }
+                S::End | S::Lost => unreachable!(),
+            }
+        }
+
+        while !buf.is_empty() {
+            match *self {
+                S::Start => {
+                    *self = S::Length {
+                        tag: Tag::Start,
+                        length_bytes_missing: 4,
+                        calculated_length: 0,
+                    };
+                }
+
+                S::Tag => {
+                    let tag = buf.first().copied().expect("buf not empty");
+                    buf = &buf[1..];
+
+                    *self = S::Length {
+                        tag: Tag::Message(tag),
+                        length_bytes_missing: 4,
+                        calculated_length: 0,
+                    };
+                }
+
+                S::Length {
+                    tag,
+                    mut length_bytes_missing,
+                    mut calculated_length,
+                } => {
+                    let consume = length_bytes_missing.min(buf.len());
+
+                    let (length_bytes, remainder) = buf.split_at(consume);
+                    for b in length_bytes {
+                        calculated_length <<= 8;
+                        calculated_length |= *b as usize;
+                    }
+                    buf = remainder;
+
+                    length_bytes_missing -= consume;
+                    if length_bytes_missing == 0 {
+                        let Some(bytes_to_skip) = calculated_length.checked_sub(4) else {
+                            observer.observe(Tag::Lost);
+                            *self = S::Lost;
+                            return;
+                        };
+
+                        if bytes_to_skip == 0 {
+                            observer.observe(tag);
+                            *self = S::Tag;
+                        } else {
+                            *self = S::Payload {
+                                tag,
+                                first: true,
+                                bytes_to_skip,
+                            };
+                        }
+                    } else {
+                        *self = S::Length {
+                            tag,
+                            length_bytes_missing,
+                            calculated_length,
+                        };
+                    }
+                }
+
+                S::Payload {
+                    tag,
+                    first,
+                    mut bytes_to_skip,
+                } => {
+                    let consume = bytes_to_skip.min(buf.len());
+                    bytes_to_skip -= consume;
+                    if bytes_to_skip == 0 {
+                        if tag == Tag::READY_FOR_QUERY && first && consume == 1 {
+                            let status = buf.first().copied().expect("buf not empty");
+                            observer.observe(Tag::ReadyForQuery(status));
+                        } else {
+                            observer.observe(tag);
+                        }
+                        *self = S::Tag;
+                    } else {
+                        *self = S::Payload {
+                            tag,
+                            first: false,
+                            bytes_to_skip,
+                        };
+                    }
+                    buf = &buf[consume..];
+                }
+
+                S::End | S::Lost => unreachable!(),
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::cell::RefCell;
+    use std::io;
+    use std::pin::Pin;
+    use std::rc::Rc;
+    use std::sync::{Arc, Mutex};
+    use std::task::{Context, Poll};
+
+    use tokio::io::{AsyncRead, ReadBuf};
+    use tokio::io::{AsyncReadExt, BufReader};
+
+    use super::*;
+
+    #[test]
+    fn test_stream_scanner() {
+        let tags = Rc::new(RefCell::new(Vec::new()));
+        let observer_tags = tags.clone();
+        let mut observer = move |tag| {
+            observer_tags.borrow_mut().push(tag);
+        };
+        let mut state = StreamScannerState::Start;
+
+        state.scan_bytes(&[0, 0], &mut observer);
+        assert_eq!(tags.borrow().as_slice(), &[]);
+        assert_eq!(
+            state,
+            StreamScannerState::Length {
+                tag: Tag::Start,
+                length_bytes_missing: 2,
+                calculated_length: 0,
+            }
+        );
+
+        state.scan_bytes(&[0x01, 0x01, 0x00], &mut observer);
+        assert_eq!(tags.borrow().as_slice(), &[]);
+        assert_eq!(
+            state,
+            StreamScannerState::Payload {
+                tag: Tag::Start,
+                first: false,
+                bytes_to_skip: 0x00000101 - 4 - 1,
+            }
+        );
+
+        state.scan_bytes(vec![0; 0x00000101 - 4 - 1 - 1].as_slice(), &mut observer);
+        assert_eq!(tags.borrow().as_slice(), &[]);
+        assert_eq!(
+            state,
+            StreamScannerState::Payload {
+                tag: Tag::Start,
+                first: false,
+                bytes_to_skip: 1,
+            }
+        );
+
+        state.scan_bytes(&[0x00, b'A', 0x00, 0x00, 0x00, 0x08], &mut observer);
+        assert_eq!(tags.borrow().as_slice(), &[Tag::Start]);
+        assert_eq!(
+            state,
+            StreamScannerState::Payload {
+                tag: Tag::Message(b'A'),
+                first: true,
+                bytes_to_skip: 4,
+            }
+        );
+
+        state.scan_bytes(&[0, 0, 0, 0], &mut observer);
+        assert_eq!(tags.borrow().as_slice(), &[Tag::Start, Tag::Message(b'A')]);
+        assert_eq!(state, StreamScannerState::Tag);
+
+        state.scan_bytes(&[b'Z', 0x00, 0x00, 0x00, 0x05, b'T'], &mut observer);
+        assert_eq!(
+            tags.borrow().as_slice(),
+            &[Tag::Start, Tag::Message(b'A'), Tag::ReadyForQuery(b'T')]
+        );
+        assert_eq!(state, StreamScannerState::Tag);
+
+        state.scan_bytes(&[], &mut observer);
+        assert_eq!(
+            tags.borrow().as_slice(),
+            &[
+                Tag::Start,
+                Tag::Message(b'A'),
+                Tag::ReadyForQuery(b'T'),
+                Tag::End
+            ]
+        );
+        assert_eq!(state, StreamScannerState::End);
+    }
+
+    #[tokio::test]
+    async fn test_connection_tracker() {
+        let transitions: Arc<Mutex<Vec<(ConnectionState, ConnectionState)>>> = Arc::default();
+        struct Observer(Arc<Mutex<Vec<(ConnectionState, ConnectionState)>>>);
+        impl StateChangeObserver for Observer {
+            fn change(&mut self, old_state: ConnectionState, new_state: ConnectionState) {
+                self.0.lock().unwrap().push((old_state, new_state));
+            }
+        }
+        let mut tracker = ConnectionTracker::new(Observer(transitions.clone()));
+
+        let stream = BufReader::new(
+            &[
+                0, 0, 0, 4, // Init
+                b'Z', 0, 0, 0, 5, b'I', // Init -> Idle
+                b'x', 0, 0, 0, 4, // Idle -> Busy
+                b'Z', 0, 0, 0, 5, b'I', // Busy -> Idle
+            ][..],
+        );
+        // AsyncRead
+        let mut stream = TrackedStream::new(stream, |tag| tracker.backend_message_tag(tag));
+
+        let mut readbuf = [0; 2];
+        let n = stream.read_exact(&mut readbuf).await.unwrap();
+        assert_eq!(n, 2);
+        assert_eq!(&readbuf, &[0, 0,]);
+        assert!(transitions.lock().unwrap().is_empty());
+
+        let mut readbuf = [0; 2];
+        let n = stream.read_exact(&mut readbuf).await.unwrap();
+        assert_eq!(n, 2);
+        assert_eq!(&readbuf, &[0, 4]);
+        assert_eq!(
+            transitions.lock().unwrap().as_slice(),
+            &[(ConnectionState::Init, ConnectionState::Busy)]
+        );
+
+        let mut readbuf = [0; 6];
+        let n = stream.read_exact(&mut readbuf).await.unwrap();
+        assert_eq!(n, 6);
+        assert_eq!(&readbuf, &[b'Z', 0, 0, 0, 5, b'I']);
+        assert_eq!(
+            transitions.lock().unwrap().as_slice(),
+            &[
+                (ConnectionState::Init, ConnectionState::Busy),
+                (ConnectionState::Busy, ConnectionState::Idle),
+            ]
+        );
+
+        let mut readbuf = [0; 5];
+        let n = stream.read_exact(&mut readbuf).await.unwrap();
+        assert_eq!(n, 5);
+        assert_eq!(&readbuf, &[b'x', 0, 0, 0, 4]);
+        assert_eq!(
+            transitions.lock().unwrap().as_slice(),
+            &[
+                (ConnectionState::Init, ConnectionState::Busy),
+                (ConnectionState::Busy, ConnectionState::Idle),
+                (ConnectionState::Idle, ConnectionState::Busy),
+            ]
+        );
+
+        let mut readbuf = [0; 6];
+        let n = stream.read_exact(&mut readbuf).await.unwrap();
+        assert_eq!(n, 6);
+        assert_eq!(&readbuf, &[b'Z', 0, 0, 0, 5, b'I']);
+        assert_eq!(
+            transitions.lock().unwrap().as_slice(),
+            &[
+                (ConnectionState::Init, ConnectionState::Busy),
+                (ConnectionState::Busy, ConnectionState::Idle),
+                (ConnectionState::Idle, ConnectionState::Busy),
+                (ConnectionState::Busy, ConnectionState::Idle),
+            ]
+        );
+    }
+
+    pub struct TrackedStream<S, TO> {
+        stream: S,
+        observer: TO,
+        state: StreamScannerState,
+    }
+
+    impl<S: Unpin, TO> Unpin for TrackedStream<S, TO> {}
+
+    impl<S: AsyncRead + Unpin, TO: TagObserver> TrackedStream<S, TO> {
+        pub const fn new(stream: S, observer: TO) -> Self {
+            TrackedStream {
+                stream,
+                observer,
+                state: StreamScannerState::Start,
+            }
+        }
+    }
+
+    impl<S: AsyncRead + Unpin, TO: TagObserver> AsyncRead for TrackedStream<S, TO> {
+        #[inline]
+        fn poll_read(
+            self: Pin<&mut Self>,
+            cx: &mut Context<'_>,
+            buf: &mut ReadBuf<'_>,
+        ) -> Poll<io::Result<()>> {
+            let Self {
+                stream,
+                observer,
+                state,
+            } = Pin::into_inner(self);
+
+            let old_len = buf.filled().len();
+            match Pin::new(stream).poll_read(cx, buf) {
+                Poll::Ready(Ok(())) => {
+                    let new_len = buf.filled().len();
+                    state.scan_bytes(&buf.filled()[old_len..new_len], observer);
+                    Poll::Ready(Ok(()))
+                }
+                Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
+                Poll::Pending => Poll::Pending,
+            }
+        }
+    }
+}

proxy/src/proxy/copy_bidirectional.rs

@@ -6,6 +6,8 @@ use std::task::{Context, Poll, ready};
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 use tracing::info;
 
+use crate::metrics::Direction;
+
 #[derive(Debug)]
 enum TransferState {
     Running(CopyBuffer),
@@ -43,6 +45,8 @@ pub enum ErrorSource {
 fn transfer_one_direction<A, B>(
     cx: &mut Context<'_>,
     state: &mut TransferState,
+    direction: Direction,
+    conn_tracker: &mut impl for<'a> FnMut(Direction, &'a [u8]),
     r: &mut A,
     w: &mut B,
 ) -> Poll<Result<u64, ErrorDirection>>
@@ -55,7 +59,8 @@ where
     loop {
         match state {
             TransferState::Running(buf) => {
-                let count = ready!(buf.poll_copy(cx, r.as_mut(), w.as_mut()))?;
+                let count =
+                    ready!(buf.poll_copy(cx, direction, conn_tracker, r.as_mut(), w.as_mut()))?;
                 *state = TransferState::ShuttingDown(count);
             }
             TransferState::ShuttingDown(count) => {
@@ -71,6 +76,7 @@ where
 pub async fn copy_bidirectional_client_compute<Client, Compute>(
     client: &mut Client,
     compute: &mut Compute,
+    mut conn_tracker: impl for<'a> FnMut(Direction, &'a [u8]),
 ) -> Result<(u64, u64), ErrorSource>
 where
     Client: AsyncRead + AsyncWrite + Unpin + ?Sized,
@@ -80,12 +86,24 @@ where
     let mut compute_to_client = TransferState::Running(CopyBuffer::new());
 
     poll_fn(|cx| {
-        let mut client_to_compute_result =
-            transfer_one_direction(cx, &mut client_to_compute, client, compute)
-                .map_err(ErrorSource::from_client)?;
-        let mut compute_to_client_result =
-            transfer_one_direction(cx, &mut compute_to_client, compute, client)
-                .map_err(ErrorSource::from_compute)?;
+        let mut client_to_compute_result = transfer_one_direction(
+            cx,
+            &mut client_to_compute,
+            Direction::ClientToCompute,
+            &mut conn_tracker,
+            client,
+            compute,
+        )
+        .map_err(ErrorSource::from_client)?;
+        let mut compute_to_client_result = transfer_one_direction(
+            cx,
+            &mut compute_to_client,
+            Direction::ComputeToClient,
+            &mut conn_tracker,
+            compute,
+            client,
+        )
+        .map_err(ErrorSource::from_compute)?;
 
         // TODO: 1 info log, with a enum label for close direction.
 
@@ -95,9 +113,15 @@ where
                 info!("Compute is done, terminate client");
                 // Initiate shutdown
                 client_to_compute = TransferState::ShuttingDown(buf.amt);
-                client_to_compute_result =
-                    transfer_one_direction(cx, &mut client_to_compute, client, compute)
-                        .map_err(ErrorSource::from_client)?;
+                client_to_compute_result = transfer_one_direction(
+                    cx,
+                    &mut client_to_compute,
+                    Direction::ClientToCompute,
+                    &mut conn_tracker,
+                    client,
+                    compute,
+                )
+                .map_err(ErrorSource::from_client)?;
             }
         }
 
@@ -107,9 +131,15 @@ where
                 info!("Client is done, terminate compute");
                 // Initiate shutdown
                 compute_to_client = TransferState::ShuttingDown(buf.amt);
-                compute_to_client_result =
-                    transfer_one_direction(cx, &mut compute_to_client, compute, client)
-                        .map_err(ErrorSource::from_compute)?;
+                compute_to_client_result = transfer_one_direction(
+                    cx,
+                    &mut compute_to_client,
+                    Direction::ComputeToClient,
+                    &mut conn_tracker,
+                    compute,
+                    client,
+                )
+                .map_err(ErrorSource::from_compute)?;
             }
         }
 
@@ -148,6 +178,8 @@ impl CopyBuffer {
     fn poll_fill_buf<R>(
         &mut self,
         cx: &mut Context<'_>,
+        direction: Direction,
+        conn_tracker: &mut impl for<'a> FnMut(Direction, &'a [u8]),
         reader: Pin<&mut R>,
     ) -> Poll<io::Result<()>>
     where
@@ -158,6 +190,8 @@ impl CopyBuffer {
         buf.set_filled(me.cap);
 
         let res = reader.poll_read(cx, &mut buf);
+        conn_tracker(direction, &buf.filled()[me.cap..]);
+
         if let Poll::Ready(Ok(())) = res {
             let filled_len = buf.filled().len();
             me.read_done = me.cap == filled_len;
@@ -169,6 +203,8 @@ impl CopyBuffer {
     fn poll_write_buf<R, W>(
         &mut self,
         cx: &mut Context<'_>,
+        direction: Direction,
+        conn_tracker: &mut impl for<'a> FnMut(Direction, &'a [u8]),
         mut reader: Pin<&mut R>,
         mut writer: Pin<&mut W>,
     ) -> Poll<Result<usize, ErrorDirection>>
@@ -182,7 +218,8 @@ impl CopyBuffer {
                 // Top up the buffer towards full if we can read a bit more
                 // data - this should improve the chances of a large write
                 if !me.read_done && me.cap < me.buf.len() {
-                    ready!(me.poll_fill_buf(cx, reader.as_mut())).map_err(ErrorDirection::Read)?;
+                    ready!(me.poll_fill_buf(cx, direction, conn_tracker, reader.as_mut()))
+                        .map_err(ErrorDirection::Read)?;
                 }
                 Poll::Pending
             }
@@ -193,6 +230,8 @@ impl CopyBuffer {
     pub(super) fn poll_copy<R, W>(
         &mut self,
         cx: &mut Context<'_>,
+        direction: Direction,
+        conn_tracker: &mut impl for<'a> FnMut(Direction, &'a [u8]),
         mut reader: Pin<&mut R>,
         mut writer: Pin<&mut W>,
     ) -> Poll<Result<u64, ErrorDirection>>
@@ -204,7 +243,7 @@ impl CopyBuffer {
             // If there is some space left in our buffer, then we try to read some
             // data to continue, thus maximizing the chances of a large write.
             if self.cap < self.buf.len() && !self.read_done {
-                match self.poll_fill_buf(cx, reader.as_mut()) {
+                match self.poll_fill_buf(cx, direction, conn_tracker, reader.as_mut()) {
                     Poll::Ready(Ok(())) => (),
                     Poll::Ready(Err(err)) => return Poll::Ready(Err(ErrorDirection::Read(err))),
                     Poll::Pending => {
@@ -227,7 +266,13 @@ impl CopyBuffer {
 
             // If our buffer has some data, let's write it out!
             while self.pos < self.cap {
-                let i = ready!(self.poll_write_buf(cx, reader.as_mut(), writer.as_mut()))?;
+                let i = ready!(self.poll_write_buf(
+                    cx,
+                    direction,
+                    conn_tracker,
+                    reader.as_mut(),
+                    writer.as_mut()
+                ))?;
                 if i == 0 {
                     return Poll::Ready(Err(ErrorDirection::Write(io::Error::new(
                         io::ErrorKind::WriteZero,
@@ -278,9 +323,10 @@ mod tests {
         compute_client.write_all(b"Neon").await.unwrap();
         compute_client.shutdown().await.unwrap();
 
-        let result = copy_bidirectional_client_compute(&mut client_proxy, &mut compute_proxy)
-            .await
-            .unwrap();
+        let result =
+            copy_bidirectional_client_compute(&mut client_proxy, &mut compute_proxy, |_, _| {})
+                .await
+                .unwrap();
 
         // Assert correct transferred amounts
         let (client_to_compute_count, compute_to_client_count) = result;
@@ -301,9 +347,10 @@ mod tests {
             .await
             .unwrap();
 
-        let result = copy_bidirectional_client_compute(&mut client_proxy, &mut compute_proxy)
-            .await
-            .unwrap();
+        let result =
+            copy_bidirectional_client_compute(&mut client_proxy, &mut compute_proxy, |_, _| {})
+                .await
+                .unwrap();
 
         // Assert correct transferred amounts
         let (client_to_compute_count, compute_to_client_count) = result;

proxy/src/proxy/mod.rs

@@ -2,6 +2,7 @@
 mod tests;
 
 pub(crate) mod connect_compute;
+pub mod conntrack;
 mod copy_bidirectional;
 pub(crate) mod handshake;
 pub(crate) mod passthrough;
@@ -30,6 +31,7 @@ use crate::context::RequestContext;
 use crate::error::ReportableError;
 use crate::metrics::{Metrics, NumClientConnectionsGuard};
 use crate::protocol2::{ConnectHeader, ConnectionInfo, ConnectionInfoExtra, read_proxy_protocol};
+use crate::proxy::conntrack::ConnectionTracking;
 use crate::proxy::handshake::{HandshakeData, handshake};
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::stream::{PqStream, Stream};
@@ -60,6 +62,7 @@ pub async fn task_main(
     cancellation_token: CancellationToken,
     cancellation_handler: Arc<CancellationHandler>,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
+    conntracking: Arc<ConnectionTracking>,
 ) -> anyhow::Result<()> {
     scopeguard::defer! {
         info!("proxy has shut down");
@@ -85,6 +88,7 @@ pub async fn task_main(
         let session_id = uuid::Uuid::new_v4();
         let cancellation_handler = Arc::clone(&cancellation_handler);
         let cancellations = cancellations.clone();
+        let conntracking = Arc::clone(&conntracking);
 
         debug!(protocol = "tcp", %session_id, "accepted new TCP connection");
         let endpoint_rate_limiter2 = endpoint_rate_limiter.clone();
@@ -149,6 +153,7 @@ pub async fn task_main(
                 endpoint_rate_limiter2,
                 conn_gauge,
                 cancellations,
+                conntracking,
             )
             .instrument(ctx.span())
             .boxed()
@@ -268,6 +273,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     conn_gauge: NumClientConnectionsGuard<'static>,
     cancellations: tokio_util::task::task_tracker::TaskTracker,
+    conntracking: Arc<ConnectionTracking>,
 ) -> Result<Option<ProxyPassthrough<S>>, ClientRequestError> {
     debug!(
         protocol = %ctx.protocol(),
@@ -409,6 +415,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
         compute: node,
         session_id: ctx.session_id(),
         cancel: session,
+        conntracking,
         _req: request_gauge,
         _conn: conn_gauge,
     }))

proxy/src/proxy/passthrough.rs

@@ -1,7 +1,8 @@
+use std::sync::Arc;
+
 use smol_str::SmolStr;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tracing::debug;
-use utils::measured_stream::MeasuredStream;
 
 use super::copy_bidirectional::ErrorSource;
 use crate::cancellation;
@@ -9,16 +10,19 @@ use crate::compute::PostgresConnection;
 use crate::config::ComputeConfig;
 use crate::control_plane::messages::MetricsAuxInfo;
 use crate::metrics::{Direction, Metrics, NumClientConnectionsGuard, NumConnectionRequestsGuard};
+use crate::proxy::conntrack::{ConnectionTracking, StreamScannerState};
+use crate::proxy::copy_bidirectional::copy_bidirectional_client_compute;
 use crate::stream::Stream;
 use crate::usage_metrics::{Ids, MetricCounterRecorder, USAGE_METRICS};
 
 /// Forward bytes in both directions (client <-> compute).
 #[tracing::instrument(skip_all)]
 pub(crate) async fn proxy_pass(
-    client: impl AsyncRead + AsyncWrite + Unpin,
-    compute: impl AsyncRead + AsyncWrite + Unpin,
+    mut client: impl AsyncRead + AsyncWrite + Unpin,
+    mut compute: impl AsyncRead + AsyncWrite + Unpin,
     aux: MetricsAuxInfo,
     private_link_id: Option<SmolStr>,
+    conntracking: &Arc<ConnectionTracking>,
 ) -> Result<(), ErrorSource> {
     // we will report ingress at a later date
     let usage_tx = USAGE_METRICS.register(Ids {
@@ -27,35 +31,35 @@ pub(crate) async fn proxy_pass(
         private_link_id,
     });
 
+    let mut conn_tracker = conntracking.new_tracker();
+
     let metrics = &Metrics::get().proxy.io_bytes;
-    let m_sent = metrics.with_labels(Direction::Tx);
-    let mut client = MeasuredStream::new(
-        client,
-        |_| {},
-        |cnt| {
-            // Number of bytes we sent to the client (outbound).
-            metrics.get_metric(m_sent).inc_by(cnt as u64);
-            usage_tx.record_egress(cnt as u64);
-        },
-    );
+    let m_sent = metrics.with_labels(Direction::ComputeToClient);
+    let m_recv = metrics.with_labels(Direction::ClientToCompute);
 
-    let m_recv = metrics.with_labels(Direction::Rx);
-    let mut compute = MeasuredStream::new(
-        compute,
-        |_| {},
-        |cnt| {
-            // Number of bytes the client sent to the compute node (inbound).
-            metrics.get_metric(m_recv).inc_by(cnt as u64);
-            usage_tx.record_ingress(cnt as u64);
-        },
-    );
+    let mut client_to_compute = StreamScannerState::Tag;
+    let mut compute_to_client = StreamScannerState::Tag;
 
-    // Starting from here we only proxy the client's traffic.
     debug!("performing the proxy pass...");
-    let _ = crate::proxy::copy_bidirectional::copy_bidirectional_client_compute(
-        &mut client,
-        &mut compute,
-    )
+
+    let _ = copy_bidirectional_client_compute(&mut client, &mut compute, |direction, bytes| {
+        match direction {
+            Direction::ClientToCompute => {
+                client_to_compute
+                    .scan_bytes(bytes, &mut |tag| conn_tracker.frontend_message_tag(tag));
+
+                metrics.get_metric(m_recv).inc_by(bytes.len() as u64);
+                usage_tx.record_ingress(bytes.len() as u64);
+            }
+            Direction::ComputeToClient => {
+                compute_to_client
+                    .scan_bytes(bytes, &mut |tag| conn_tracker.backend_message_tag(tag));
+
+                metrics.get_metric(m_sent).inc_by(bytes.len() as u64);
+                usage_tx.record_egress(bytes.len() as u64);
+            }
+        }
+    })
     .await?;
 
     Ok(())
@@ -68,6 +72,7 @@ pub(crate) struct ProxyPassthrough<S> {
     pub(crate) session_id: uuid::Uuid,
     pub(crate) private_link_id: Option<SmolStr>,
     pub(crate) cancel: cancellation::Session,
+    pub(crate) conntracking: Arc<ConnectionTracking>,
 
     pub(crate) _req: NumConnectionRequestsGuard<'static>,
     pub(crate) _conn: NumClientConnectionsGuard<'static>,
@@ -83,6 +88,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin> ProxyPassthrough<S> {
             self.compute.stream,
             self.aux,
             self.private_link_id,
+            &self.conntracking,
         )
         .await;
         if let Err(err) = self

proxy/src/serverless/mod.rs

@@ -50,6 +50,7 @@ use crate::context::RequestContext;
 use crate::ext::TaskExt;
 use crate::metrics::Metrics;
 use crate::protocol2::{ChainRW, ConnectHeader, ConnectionInfo, read_proxy_protocol};
+use crate::proxy::conntrack::ConnectionTracking;
 use crate::proxy::run_until_cancelled;
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::serverless::backend::PoolingBackend;
@@ -124,6 +125,9 @@ pub async fn task_main(
     connections.close(); // allows `connections.wait to complete`
 
     let cancellations = tokio_util::task::task_tracker::TaskTracker::new();
+
+    let conntracking = Arc::new(ConnectionTracking::default());
+
     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) {
@@ -153,6 +157,8 @@ pub async fn task_main(
         let cancellation_handler = cancellation_handler.clone();
         let endpoint_rate_limiter = endpoint_rate_limiter.clone();
         let cancellations = cancellations.clone();
+        let conntracking = Arc::clone(&conntracking);
+
         connections.spawn(
             async move {
                 let conn_token2 = conn_token.clone();
@@ -185,6 +191,7 @@ pub async fn task_main(
                     cancellation_handler,
                     endpoint_rate_limiter,
                     conn_token,
+                    conntracking,
                     conn,
                     conn_info,
                     session_id,
@@ -309,6 +316,7 @@ async fn connection_handler(
     cancellation_handler: Arc<CancellationHandler>,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     cancellation_token: CancellationToken,
+    conntracking: Arc<ConnectionTracking>,
     conn: AsyncRW,
     conn_info: ConnectionInfo,
     session_id: uuid::Uuid,
@@ -347,6 +355,7 @@ async fn connection_handler(
             // `request_handler` is not cancel safe. It expects to be cancelled only at specific times.
             // By spawning the future, we ensure it never gets cancelled until it decides to.
             let cancellations = cancellations.clone();
+            let conntracking = Arc::clone(&conntracking);
             let handler = connections.spawn(
                 request_handler(
                     req,
@@ -359,6 +368,7 @@ async fn connection_handler(
                     http_request_token,
                     endpoint_rate_limiter.clone(),
                     cancellations,
+                    conntracking,
                 )
                 .in_current_span()
                 .map_ok_or_else(api_error_into_response, |r| r),
@@ -407,6 +417,7 @@ async fn request_handler(
     http_cancellation_token: CancellationToken,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     cancellations: TaskTracker,
+    conntracking: Arc<ConnectionTracking>,
 ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, ApiError> {
     let host = request
         .headers()
@@ -452,6 +463,7 @@ async fn request_handler(
                     endpoint_rate_limiter,
                     host,
                     cancellations,
+                    conntracking,
                 )
                 .await
                 {

proxy/src/serverless/websocket.rs

@@ -17,6 +17,7 @@ use crate::config::ProxyConfig;
 use crate::context::RequestContext;
 use crate::error::ReportableError;
 use crate::metrics::Metrics;
+use crate::proxy::conntrack::ConnectionTracking;
 use crate::proxy::{ClientMode, ErrorSource, handle_client};
 use crate::rate_limiter::EndpointRateLimiter;
 
@@ -133,6 +134,7 @@ pub(crate) async fn serve_websocket(
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     hostname: Option<String>,
     cancellations: tokio_util::task::task_tracker::TaskTracker,
+    conntracking: Arc<ConnectionTracking>,
 ) -> anyhow::Result<()> {
     let websocket = websocket.await?;
     let websocket = WebSocketServer::after_handshake(TokioIo::new(websocket));
@@ -152,6 +154,7 @@ pub(crate) async fn serve_websocket(
         endpoint_rate_limiter,
         conn_gauge,
         cancellations,
+        conntracking,
     ))
     .await;