Merge pull request superfly#373 from superfly/gorbak/track-alive-api-conn

Track active TCP streams in the user facing API

Author: gorbak25

Cargo.lock

@@ -85,6 +85,7 @@ uhlc = { version = "0.7", features = ["defmt"] }
 uuid = { version = "1.3.1", features = ["v4", "serde"] }
 webpki = { version = "0.22.0", features = ["std"] }
 http = { version = "0.2.9" }
+lazy_static = "1.5.0"
 
 [profile.release]
 debug = 1

Cargo.toml

@@ -61,6 +61,7 @@ uhlc = { workspace = true }
 uuid = { workspace = true }
 corro-pg = { path = "../corro-pg" }
 indexmap = { workspace = true }
+pin-project-lite = { workspace = true }
 governor.workspace = true
 
 [dev-dependencies]

crates/corro-agent/Cargo.toml

@@ -1,2 +1,3 @@
 pub mod peer;
 pub mod public;
+pub mod utils;

crates/corro-agent/src/api/mod.rs

@@ -5,6 +5,7 @@ use std::{
     time::{Duration, Instant},
 };
 
+use crate::api::utils::CountedBody;
 use antithesis_sdk::assert_sometimes;
 use axum::{extract::ConnectInfo, response::IntoResponse, Extension};
 use bytes::{BufMut, BytesMut};
@@ -18,6 +19,7 @@ use corro_types::{
     base::CrsqlDbVersion,
     broadcast::Timestamp,
     change::{insert_local_changes, InsertChangesInfo, SqliteValue},
+    persistent_gauge,
     schema::{apply_schema, parse_sql},
     sqlite::SqlitePoolError,
 };
@@ -472,7 +474,9 @@ pub async fn api_v1_queries(
     axum::extract::Query(params): axum::extract::Query<TimeoutParams>,
     axum::extract::Json(stmt): axum::extract::Json<Statement>,
 ) -> impl IntoResponse {
-    let (mut tx, body) = hyper::Body::channel();
+    let (mut tx, body) = CountedBody::channel(
+        persistent_gauge!("corro.api.active.streams", "source" => "queries", "protocol" => "http"),
+    );
 
     counter!("corro.api.queries.count").increment(1);
     // TODO: timeout on data send instead of infinitely waiting for channel space.

crates/corro-agent/src/api/public/mod.rs

@@ -1,8 +1,10 @@
 use std::{collections::HashMap, io::Write, sync::Arc, time::Duration};
 
+use crate::api::utils::CountedBody;
 use axum::{http::StatusCode, response::IntoResponse, Extension};
 use bytes::{BufMut, Bytes, BytesMut};
 use compact_str::{format_compact, ToCompactString};
+use corro_types::persistent_gauge;
 use corro_types::updates::Handle;
 use corro_types::{
     agent::Agent,
@@ -50,7 +52,7 @@ async fn sub_by_id(
     params: SubParams,
     bcast_cache: &SharedMatcherBroadcastCache,
     tripwire: Tripwire,
-) -> hyper::Response<hyper::Body> {
+) -> impl IntoResponse {
     let matcher_rx = bcast_cache.read().await.get(&id).and_then(|tx| {
         subs.get(&id).map(|matcher| {
             debug!("found matcher by id {id}");
@@ -99,7 +101,9 @@ async fn sub_by_id(
     let query_hash = matcher.hash().to_owned();
     tokio::spawn(catch_up_sub(matcher, params, rx, evt_tx));
 
-    let (tx, body) = hyper::Body::channel();
+    let (tx, body) = CountedBody::channel(
+        persistent_gauge!("corro.api.active.streams", "source" => "subscriptions", "protocol" => "http"),
+    );
 
     tokio::spawn(forward_bytes_to_body_sender(id, evt_rx, tx, tripwire));
 
@@ -302,7 +306,7 @@ impl MatcherUpsertError {
     }
 }
 
-impl From<MatcherUpsertError> for hyper::Response<hyper::Body> {
+impl From<MatcherUpsertError> for hyper::Response<CountedBody<hyper::Body>> {
     fn from(value: MatcherUpsertError) -> Self {
         hyper::Response::builder()
             .status(value.status_code())
@@ -676,7 +680,7 @@ pub async fn api_v1_subs(
 ) -> impl IntoResponse {
     let stmt = match expand_sql(&agent, &stmt).await {
         Ok(stmt) => stmt,
-        Err(e) => return hyper::Response::<hyper::Body>::from(e),
+        Err(e) => return hyper::Response::<CountedBody<hyper::Body>>::from(e),
     };
 
     info!("Received subscription request for query: {stmt}");
@@ -695,10 +699,14 @@ pub async fn api_v1_subs(
 
     let (handle, maybe_created) = match upsert_res {
         Ok(res) => res,
-        Err(e) => return hyper::Response::<hyper::Body>::from(MatcherUpsertError::from(e)),
+        Err(e) => {
+            return hyper::Response::<CountedBody<hyper::Body>>::from(MatcherUpsertError::from(e))
+        }
     };
 
-    let (tx, body) = hyper::Body::channel();
+    let (tx, body) = CountedBody::channel(
+        persistent_gauge!("corro.api.active.streams", "source" => "subscriptions", "protocol" => "http"),
+    );
     let (forward_tx, forward_rx) = mpsc::channel(10240);
 
     tokio::spawn(forward_bytes_to_body_sender(
@@ -720,7 +728,7 @@ pub async fn api_v1_subs(
     .await
     {
         Ok(id) => id,
-        Err(e) => return hyper::Response::<hyper::Body>::from(e),
+        Err(e) => return hyper::Response::<CountedBody<hyper::Body>>::from(e),
     };
 
     hyper::Response::builder()

crates/corro-agent/src/api/public/pubsub.rs

@@ -7,6 +7,7 @@ use compact_str::ToCompactString;
 use corro_types::{
     agent::Agent,
     api::NotifyEvent,
+    persistent_gauge,
     updates::{Handle, UpdateCreated, UpdateHandle, UpdatesManager},
 };
 use futures::future::poll_fn;
@@ -18,7 +19,7 @@ use tracing::{debug, info, warn};
 use tripwire::Tripwire;
 use uuid::Uuid;
 
-use crate::api::public::pubsub::MatcherUpsertError;
+use crate::api::{public::pubsub::MatcherUpsertError, utils::CountedBody};
 
 pub type UpdateBroadcastCache = HashMap<Uuid, broadcast::Sender<Bytes>>;
 pub type SharedUpdateBroadcastCache = Arc<TokioRwLock<UpdateBroadcastCache>>;
@@ -48,16 +49,20 @@ pub async fn api_v1_updates(
 
     let (handle, maybe_created) = match upsert_res {
         Ok(res) => res,
-        Err(e) => return hyper::Response::<hyper::Body>::from(MatcherUpsertError::from(e)),
+        Err(e) => {
+            return hyper::Response::<CountedBody<hyper::Body>>::from(MatcherUpsertError::from(e))
+        }
     };
 
-    let (tx, body) = hyper::Body::channel();
+    let (tx, body) = CountedBody::channel(
+        persistent_gauge!("corro.api.active.streams", "source" => "updates", "protocol" => "http"),
+    );
     // let (forward_tx, forward_rx) = mpsc::channel(10240);
 
     let (update_id, sub_rx) =
         match upsert_update(handle.clone(), maybe_created, updates, &mut bcast_write).await {
             Ok(id) => id,
-            Err(e) => return hyper::Response::<hyper::Body>::from(e),
+            Err(e) => return hyper::Response::<CountedBody<hyper::Body>>::from(e),
         };
 
     tokio::spawn(forward_update_bytes_to_body_sender(

crates/corro-agent/src/api/public/update.rs

@@ -0,0 +1,85 @@
+use corro_types::gauge::PersistentGauge;
+use hyper::body::{Body, HttpBody, Sender, SizeHint};
+use pin_project_lite::pin_project;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+pin_project! {
+    pub struct CountedBody<B: HttpBody> {
+        #[pin]
+        body: B,
+        gauge: Option<PersistentGauge>,
+    }
+
+    impl<B: HttpBody> PinnedDrop for CountedBody<B> {
+        fn drop(this: Pin<&mut Self>) {
+            if let Some(gauge) = &this.gauge {
+                gauge.decrement(1.0);
+            }
+        }
+    }
+}
+
+impl<B: HttpBody> CountedBody<B> {
+    fn new(body: B, gauge: Option<PersistentGauge>) -> Self {
+        if let Some(gauge) = &gauge {
+            gauge.increment(1.0);
+        }
+        Self { body, gauge }
+    }
+}
+
+impl CountedBody<Body> {
+    // Channel bodies need to be counted as they can be long lived
+    pub fn channel(gauge: PersistentGauge) -> (Sender, Self) {
+        let (tx, body) = hyper::Body::channel();
+        (tx, Self::new(body, Some(gauge)))
+    }
+}
+
+impl<B: HttpBody> HttpBody for CountedBody<B> {
+    type Data = B::Data;
+    type Error = B::Error;
+
+    fn poll_data(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Self::Data, Self::Error>>> {
+        let this = self.project();
+        this.body.poll_data(cx)
+    }
+
+    fn poll_trailers(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<Option<hyper::HeaderMap>, Self::Error>> {
+        let this = self.project();
+        this.body.poll_trailers(cx)
+    }
+
+    fn is_end_stream(&self) -> bool {
+        self.body.is_end_stream()
+    }
+    fn size_hint(&self) -> SizeHint {
+        self.body.size_hint()
+    }
+}
+
+// If the underlying body can be constructed from some simple type
+// then we can implement From<T> for CountedBody<B>
+// No need to track their count as they are short lived
+pub trait SimpleBody {}
+impl SimpleBody for Vec<u8> {}
+impl SimpleBody for &'static [u8] {}
+impl SimpleBody for String {}
+impl SimpleBody for &'static str {}
+
+impl<B, T> From<T> for CountedBody<B>
+where
+    B: HttpBody + From<T>,
+    T: SimpleBody,
+{
+    fn from(value: T) -> Self {
+        Self::new(value.into(), None)
+    }
+}

crates/corro-agent/src/api/utils/mod.rs

@@ -33,6 +33,7 @@ sqlparser = { version = "0.39.0" }
 chrono = { version = "0.4.31" }
 socket2 = { version = "0.5" }
 tokio-rustls = "0.24.1"
+pin-project-lite = { workspace = true }
 
 [dev-dependencies]
 corro-tests = { path = "../corro-tests" }

crates/corro-pg/Cargo.toml

@@ -1,4 +1,5 @@
 pub mod sql_state;
+pub mod utils;
 mod vtab;
 
 use std::{
@@ -19,6 +20,7 @@ use corro_types::{
     broadcast::{broadcast_changes, Timestamp},
     change::{insert_local_changes, InsertChangesInfo},
     config::PgConfig,
+    persistent_gauge,
     schema::{parse_sql, Column, Schema, SchemaError, SqliteType, Table},
     sqlite::CrConn,
 };
@@ -72,6 +74,7 @@ use tripwire::{Outcome, PreemptibleFutureExt, TimeoutFutureExt, Tripwire};
 
 use crate::{
     sql_state::SqlState,
+    utils::CountedTcpStream,
     vtab::{
         pg_class::PgClassTable,
         pg_database::{PgDatabase, PgDatabaseTable},
@@ -553,30 +556,36 @@ pub async fn start(
     let server = TcpListener::bind(pg.bind_addr).await?;
     let (tls_acceptor, ssl_required) = setup_tls(pg).await?;
     let local_addr = server.local_addr()?;
+    let conn_gauge = persistent_gauge!("corro.api.active.streams",
+    "source" => "postgres",
+    "protocol" => "pg",
+    "readonly" => readonly.to_string(),
+    );
 
     tokio::spawn(async move {
         loop {
-            let (mut conn, remote_addr) = match server.accept().preemptible(&mut tripwire).await {
+            let (tcp_conn, remote_addr) = match server.accept().preemptible(&mut tripwire).await {
                 Outcome::Completed(res) => res?,
                 Outcome::Preempted(_) => break,
             };
+            let mut conn = CountedTcpStream::wrap(tcp_conn, conn_gauge.clone());
             let tls_acceptor = tls_acceptor.clone();
             debug!("Accepted a PostgreSQL connection (from: {remote_addr})");
 
-            counter!("corro.api.connection.count", "protocol" => "pg").increment(1);
+            counter!("corro.api.connection.count", "protocol" => "pg", "readonly" => readonly.to_string()).increment(1);
 
             let agent = agent.clone();
             tokio::spawn(async move {
-                conn.set_nodelay(true)?;
+                conn.stream.set_nodelay(true)?;
                 {
-                    let sock = SockRef::from(&conn);
+                    let sock = SockRef::from(&conn.stream);
                     let ka = TcpKeepalive::new()
                         .with_time(Duration::from_secs(10))
                         .with_interval(Duration::from_secs(10))
                         .with_retries(4);
                     sock.set_tcp_keepalive(&ka)?;
                 }
-                let is_sslrequest = peek_for_sslrequest(&mut conn).await?;
+                let is_sslrequest = peek_for_sslrequest(&mut conn.stream).await?;
 
                 // reject non-ssl connections if ssl is required (client cert auth)
                 if ssl_required && !is_sslrequest {
@@ -586,7 +595,7 @@ pub async fn start(
 
                 let (mut framed, secured) = match (tls_acceptor, is_sslrequest) {
                     (Some(tls_acceptor), true) => {
-                        conn.write_all(b"S").await?;
+                        conn.stream.write_all(b"S").await?;
                         let tls_conn = tls_acceptor.accept(conn).await?;
                         (
                             Framed::new(
@@ -600,7 +609,7 @@ pub async fn start(
                     }
                     (_, is_sslreq) => {
                         if is_sslreq {
-                            conn.write_all(b"N").await?;
+                            conn.stream.write_all(b"N").await?;
                         }
                         (
                             Framed::new(