feat(metrics): add real-time throughput sampling and byte-counting metrics

rust/crates/rustproxy-http/src/counting_body.rs

@@ -0,0 +1,122 @@
+//! A body wrapper that counts bytes flowing through and reports them to MetricsCollector.
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::task::{Context, Poll};
+
+use bytes::Bytes;
+use http_body::Frame;
+use rustproxy_metrics::MetricsCollector;
+
+/// Wraps any `http_body::Body` and counts data bytes passing through.
+///
+/// When the body is fully consumed or dropped, accumulated byte counts
+/// are reported to the `MetricsCollector`.
+///
+/// The inner body is pinned on the heap to support `!Unpin` types like `hyper::body::Incoming`.
+pub struct CountingBody<B> {
+    inner: Pin<Box<B>>,
+    counted_bytes: AtomicU64,
+    metrics: Arc<MetricsCollector>,
+    route_id: Option<String>,
+    /// Whether we count bytes as "in" (request body) or "out" (response body).
+    direction: Direction,
+    /// Whether we've already reported the bytes (to avoid double-reporting on drop).
+    reported: bool,
+}
+
+/// Which direction the bytes flow.
+#[derive(Clone, Copy)]
+pub enum Direction {
+    /// Request body: bytes flowing from client → upstream (counted as bytes_in)
+    In,
+    /// Response body: bytes flowing from upstream → client (counted as bytes_out)
+    Out,
+}
+
+impl<B> CountingBody<B> {
+    /// Create a new CountingBody wrapping an inner body.
+    pub fn new(
+        inner: B,
+        metrics: Arc<MetricsCollector>,
+        route_id: Option<String>,
+        direction: Direction,
+    ) -> Self {
+        Self {
+            inner: Box::pin(inner),
+            counted_bytes: AtomicU64::new(0),
+            metrics,
+            route_id,
+            direction,
+            reported: false,
+        }
+    }
+
+    /// Report accumulated bytes to the metrics collector.
+    fn report(&mut self) {
+        if self.reported {
+            return;
+        }
+        self.reported = true;
+
+        let bytes = self.counted_bytes.load(Ordering::Relaxed);
+        if bytes == 0 {
+            return;
+        }
+
+        let route_id = self.route_id.as_deref();
+        match self.direction {
+            Direction::In => self.metrics.record_bytes(bytes, 0, route_id),
+            Direction::Out => self.metrics.record_bytes(0, bytes, route_id),
+        }
+    }
+}
+
+impl<B> Drop for CountingBody<B> {
+    fn drop(&mut self) {
+        self.report();
+    }
+}
+
+// CountingBody is Unpin because inner is Pin<Box<B>> (always Unpin).
+impl<B> Unpin for CountingBody<B> {}
+
+impl<B> http_body::Body for CountingBody<B>
+where
+    B: http_body::Body<Data = Bytes>,
+{
+    type Data = Bytes;
+    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.get_mut();
+
+        match this.inner.as_mut().poll_frame(cx) {
+            Poll::Ready(Some(Ok(frame))) => {
+                if let Some(data) = frame.data_ref() {
+                    this.counted_bytes.fetch_add(data.len() as u64, Ordering::Relaxed);
+                }
+                Poll::Ready(Some(Ok(frame)))
+            }
+            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
+            Poll::Ready(None) => {
+                // Body is fully consumed — report now
+                this.report();
+                Poll::Ready(None)
+            }
+            Poll::Pending => Poll::Pending,
+        }
+    }
+
+    fn is_end_stream(&self) -> bool {
+        self.inner.is_end_stream()
+    }
+
+    fn size_hint(&self) -> http_body::SizeHint {
+        self.inner.size_hint()
+    }
+}

rust/crates/rustproxy-http/src/lib.rs

@@ -3,12 +3,14 @@
 //! Hyper-based HTTP proxy service for RustProxy.
 //! Handles HTTP request parsing, route-based forwarding, and response filtering.
 
+pub mod counting_body;
 pub mod proxy_service;
 pub mod request_filter;
 pub mod response_filter;
 pub mod template;
 pub mod upstream_selector;
 
+pub use counting_body::*;
 pub use proxy_service::*;
 pub use template::*;
 pub use upstream_selector::*;

rust/crates/rustproxy-http/src/proxy_service.rs

@@ -21,6 +21,7 @@ use tracing::{debug, error, info, warn};
 use rustproxy_routing::RouteManager;
 use rustproxy_metrics::MetricsCollector;
 
+use crate::counting_body::{CountingBody, Direction};
 use crate::request_filter::RequestFilter;
 use crate::response_filter::ResponseFilter;
 use crate::upstream_selector::UpstreamSelector;
@@ -345,8 +346,16 @@ impl HttpProxyService {
             }
         }
 
+        // Wrap the request body in CountingBody to track bytes_in
+        let counting_req_body = CountingBody::new(
+            body,
+            Arc::clone(&self.metrics),
+            route_id.map(|s| s.to_string()),
+            Direction::In,
+        );
+
         // Stream the request body through to upstream
-        let upstream_req = upstream_req.body(body).unwrap();
+        let upstream_req = upstream_req.body(counting_req_body).unwrap();
 
         let upstream_response = match sender.send_request(upstream_req).await {
             Ok(resp) => resp,
@@ -401,8 +410,16 @@ impl HttpProxyService {
             }
         }
 
+        // Wrap the request body in CountingBody to track bytes_in
+        let counting_req_body = CountingBody::new(
+            body,
+            Arc::clone(&self.metrics),
+            route_id.map(|s| s.to_string()),
+            Direction::In,
+        );
+
         // Stream the request body through to upstream
-        let upstream_req = upstream_req.body(body).unwrap();
+        let upstream_req = upstream_req.body(counting_req_body).unwrap();
 
         let upstream_response = match sender.send_request(upstream_req).await {
             Ok(resp) => resp,
@@ -417,6 +434,10 @@ impl HttpProxyService {
     }
 
     /// Build the client-facing response from an upstream response, streaming the body.
+    ///
+    /// The response body is wrapped in a `CountingBody` that counts bytes as they
+    /// stream from upstream to client. When the body is fully consumed (or dropped),
+    /// it reports byte counts to the metrics collector and calls `connection_closed`.
     async fn build_streaming_response(
         &self,
         upstream_response: Response<Incoming>,
@@ -433,10 +454,22 @@ impl HttpProxyService {
             ResponseFilter::apply_headers(route, headers, None);
         }
 
+        // Wrap the response body in CountingBody to track bytes_out.
+        // CountingBody will report bytes and we close the connection metric
+        // after the body stream completes (not before it even starts).
+        let counting_body = CountingBody::new(
+            resp_body,
+            Arc::clone(&self.metrics),
+            route_id.map(|s| s.to_string()),
+            Direction::Out,
+        );
+
+        // Close the connection metric now — the HTTP request/response cycle is done
+        // from the proxy's perspective once we hand the streaming body to hyper.
+        // Bytes will still be counted as they flow.
         self.metrics.connection_closed(route_id);
 
-        // Stream the response body directly from upstream to client
-        let body: BoxBody<Bytes, hyper::Error> = BoxBody::new(resp_body);
+        let body: BoxBody<Bytes, hyper::Error> = BoxBody::new(counting_body);
 
         Ok(response.body(body).unwrap())
     }