add per-request transfer byte tracking

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: ryansmith3136

Cargo.lock

@@ -60,6 +60,8 @@ nonzero = "0.2.0"
 bytes = "1.9.0"
 serde_html_form = "0.2.7"
 dashmap = "6.1.0"
+http-body = "1"
+pin-project-lite = "0.2"
 time = { version = "0.3", features = ["serde", "formatting"] }
 handlebars = "6.3.2"
 

be/Cargo.toml

@@ -3,7 +3,12 @@ use std::{
     convert::Infallible,
     fmt::{self, Debug},
     net::SocketAddr,
-    sync::{Arc, Mutex},
+    pin::Pin,
+    sync::{
+        atomic::{AtomicU64, Ordering},
+        Arc, Mutex,
+    },
+    task::{Context, Poll},
 };
 
 use axum::{
@@ -26,7 +31,10 @@ use serde_json::{json, Value};
 use tokio::sync::{OwnedSemaphorePermit, Semaphore};
 use url::Url;
 
-use crate::{broadcast, gafe};
+use http_body::Frame;
+use pin_project_lite::pin_project;
+
+use crate::{broadcast, gafe, user_query};
 
 macro_rules! user_error {
     ($e:expr) => {
@@ -127,6 +135,7 @@ pub enum Error {
     User(String),
     Timeout(Option<String>),
     TooManyRequests(Option<String>),
+    TransferExceeded(Option<String>),
 
     Server(Box<dyn std::error::Error + Send + Sync>),
 }
@@ -150,6 +159,10 @@ impl Serialize for Error {
                 state.serialize_field("error", "too_many_requests")?;
                 state.serialize_field("message", &opt_msg)?;
             }
+            Error::TransferExceeded(opt_msg) => {
+                state.serialize_field("error", "transfer_exceeded")?;
+                state.serialize_field("message", &opt_msg)?;
+            }
             Error::Server(err) => {
                 state.serialize_field("error", "server")?;
                 state.serialize_field("message", &err.to_string())?;
@@ -167,6 +180,8 @@ impl std::fmt::Display for Error {
             Error::Timeout(None) => write!(f, "Operation timed out"),
             Error::TooManyRequests(Some(msg)) => write!(f, "Too many requests: {msg}"),
             Error::TooManyRequests(None) => write!(f, "Too many requests"),
+            Error::TransferExceeded(Some(msg)) => write!(f, "Transfer exceeded: {msg}"),
+            Error::TransferExceeded(None) => write!(f, "Transfer limit exceeded"),
             Error::Server(err) => write!(f, "Server error: {err}"),
         }
     }
@@ -219,6 +234,12 @@ impl axum::response::IntoResponse for Error {
                 StatusCode::TOO_MANY_REQUESTS,
                 msg.unwrap_or(String::from("too many requests")),
             ),
+            Self::TransferExceeded(msg) => (
+                StatusCode::TOO_MANY_REQUESTS,
+                msg.unwrap_or(String::from(
+                    "Transfer limit exceeded. Upgrade at: https://www.indexsupply.net",
+                )),
+            ),
             Self::User(msg) => (StatusCode::BAD_REQUEST, msg),
             Self::Server(e) => {
                 tracing::error!(%e, "server-error={:?}", e);
@@ -285,10 +306,21 @@ pub async fn limit(
             "Rate limited. Create or upgrade API Key at: https://www.indexsupply.net",
         ))));
     }
-    match tokio::time::timeout(account_limit.timeout, next.run(request)).await {
-        Ok(response) => Ok(response),
-        Err(_) => Err(Error::Timeout(None)),
-    }
+    let log = request.extensions().get::<user_query::RequestLog>().cloned();
+    let response = match tokio::time::timeout(account_limit.timeout, next.run(request)).await {
+        Ok(response) => response,
+        Err(_) => return Err(Error::Timeout(None)),
+    };
+    let (parts, body) = response.into_parts();
+    let counting = CountingBody {
+        inner: body,
+        count: Arc::new(AtomicU64::new(0)),
+        log,
+    };
+    Ok(axum::http::Response::from_parts(
+        parts,
+        axum::body::Body::new(counting),
+    ))
 }
 
 #[derive(Clone, Copy, Default, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
@@ -496,19 +528,45 @@ pub async fn latency_header(
     Ok(response)
 }
 
-pub async fn content_length_header(
-    request: axum::extract::Request,
-    next: axum::middleware::Next,
-) -> Result<axum::response::Response, Error> {
-    let response = next.run(request).await;
-    let span = tracing::Span::current();
-    response
-        .headers()
-        .get("content-length")
-        .and_then(|cl| cl.to_str().ok())
-        .map(|cl| cl.parse::<u64>().ok())
-        .map(|size| span.record("size", size));
-    Ok(response)
+pin_project! {
+    pub struct CountingBody<B> {
+        #[pin]
+        inner: B,
+        count: Arc<AtomicU64>,
+        log: Option<user_query::RequestLog>,
+    }
+}
+
+impl<B> http_body::Body for CountingBody<B>
+where
+    B: http_body::Body<Data = bytes::Bytes>,
+{
+    type Data = B::Data;
+    type Error = B::Error;
+
+    fn poll_frame(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
+        let this = self.project();
+        match this.inner.poll_frame(cx) {
+            Poll::Ready(Some(Ok(frame))) => {
+                if let Some(data) = frame.data_ref() {
+                    this.count.fetch_add(data.len() as u64, Ordering::Relaxed);
+                }
+                Poll::Ready(Some(Ok(frame)))
+            }
+            Poll::Ready(None) => {
+                let total = this.count.load(Ordering::Relaxed);
+                tracing::Span::current().record("size", total);
+                if let Some(log) = this.log.take() {
+                    log.set_bytes(total);
+                }
+                Poll::Ready(None)
+            }
+            other => other,
+        }
+    }
 }
 
 pub async fn log_fields(

be/src/api.rs

@@ -201,7 +201,6 @@ fn service(config: api::Config) -> IntoMakeServiceWithConnectInfo<Router, Socket
         .layer(axum::middleware::from_fn(api::latency_header))
         .layer(tracing)
         .layer(axum::middleware::from_fn(api::log_fields))
-        .layer(axum::middleware::from_fn(api::content_length_header))
         .layer(HandleErrorLayer::new(api::handle_service_error))
         .load_shed()
         .concurrency_limit(1024)

be/src/main.rs

@@ -15,6 +15,7 @@ pub struct Row {
     pub latency: u16,
     pub status: u32,
     pub qty: u16,
+    pub bytes: u64,
 }
 
 impl Row {
@@ -28,6 +29,7 @@ impl Row {
             latency: 0,
             status: 0,
             qty: 1,
+            bytes: 0,
         }
     }
 }
@@ -72,6 +74,12 @@ impl RequestLog {
         }
     }
 
+    pub fn set_bytes(&self, bytes: u64) {
+        for row in self.0.lock().unwrap().rows.iter_mut() {
+            row.bytes = bytes;
+        }
+    }
+
     async fn insert(self, pool: deadpool_postgres::Pool, status: u16, ip: String) {
         // only if no one else has the log
         if let Ok(log) = Arc::try_unwrap(self.0).map(Mutex::into_inner) {
@@ -125,8 +133,9 @@ pub async fn insert(pool: deadpool_postgres::Pool, row: Row) {
                         latency,
                         status,
                         ip,
-                        qty
-                    ) values ($1, $2, $3, $4, $5, $6, $7, $8)",
+                        qty,
+                        bytes
+                    ) values ($1, $2, $3, $4, $5, $6, $7, $8, $9)",
                 &[
                     &row.api_key,
                     &U64::from(row.chain),
@@ -136,6 +145,7 @@ pub async fn insert(pool: deadpool_postgres::Pool, row: Row) {
                     &(row.status as i16),
                     &row.ip,
                     &(row.qty as i16),
+                    &(row.bytes as i64),
                 ],
             )
             .await;

be/src/user_query.rs

@@ -276,23 +276,29 @@ async fn update_daily_user_queries(
                 q.api_key,
                 k.owner_email,
                 date_trunc('day', q.created_at)::date as day,
-                q.qty
+                q.qty,
+                coalesce(q.bytes, 0) as bytes
               from user_queries q
               join api_keys k on q.api_key = k.secret
               where q.created_at >= date_trunc('day', now())
                 and q.created_at < date_trunc('day', now() + interval '1 day')
             ),
             aggregated as (
-              select owner_email, day, sum(qty)::int8 as total_qty
+              select owner_email, day,
+                sum(qty)::int8 as total_qty,
+                sum(bytes)::int8 as total_bytes
               from one_day
               group by owner_email, day
             ),
             upserted as (
-              insert into daily_user_queries (owner_email, day, n, updated_at)
-              select owner_email, day, total_qty, now()
+              insert into daily_user_queries (owner_email, day, n, bytes, updated_at)
+              select owner_email, day, total_qty, total_bytes, now()
               from aggregated
               on conflict (owner_email, day)
-              do update set n = excluded.n, updated_at = excluded.updated_at
+              do update set
+                n = excluded.n,
+                bytes = excluded.bytes,
+                updated_at = excluded.updated_at
             )
             select count(*)::int8, coalesce(sum(qty), 0)::int8
             from one_day
@@ -321,23 +327,29 @@ async fn update_wl_daily_user_queries(
                 k.provision_key,
                 k.org,
                 date_trunc('day', q.created_at)::date as day,
-                q.qty
+                q.qty,
+                coalesce(q.bytes, 0) as bytes
               from user_queries q
               join wl_api_keys k on q.api_key = k.secret
               where q.created_at >= date_trunc('day', now())
                 and q.created_at < date_trunc('day', now() + interval '1 day')
             ),
             aggregated as (
-              select provision_key, org, day, sum(qty)::int8 as total_qty
+              select provision_key, org, day,
+                sum(qty)::int8 as total_qty,
+                sum(bytes)::int8 as total_bytes
               from one_day
               group by provision_key, org, day
             ),
             upserted as (
-              insert into wl_daily_user_queries (provision_key, org, day, n, updated_at)
-              select provision_key, org, day, total_qty, now()
+              insert into wl_daily_user_queries (provision_key, org, day, n, bytes, updated_at)
+              select provision_key, org, day, total_qty, total_bytes, now()
               from aggregated
               on conflict (provision_key, org, day)
-              do update set n = excluded.n, updated_at = excluded.updated_at
+              do update set
+                n = excluded.n,
+                bytes = excluded.bytes,
+                updated_at = excluded.updated_at
             )
             select count(*)::int8, coalesce(sum(qty), 0)::int8
             from one_day;

fe/src/main.rs

@@ -130,13 +130,15 @@ create table if not exists user_queries(
     status int2,
     qty int2 default 1,
     created_at timestamptz default now(),
-    ip text
+    ip text,
+    bytes int8 default 0
 );
 
 create table if not exists daily_user_queries (
     owner_email text not null,
     day date not null,
     n int8 not null,
+    bytes int8 not null default 0,
     updated_at timestamptz not null default now(),
     primary key (owner_email, day)
 );
@@ -146,6 +148,7 @@ create table if not exists wl_daily_user_queries(
     org text not null,
     day date not null,
     n int8 not null,
+    bytes int8 not null default 0,
     updated_at timestamptz not null default now(),
     primary key (provision_key, org, day)
 );