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

This commit is contained in:
2026-02-13 23:18:22 +00:00
parent 37372353d7
commit 1f95d2b6c4
17 changed files with 1109 additions and 228 deletions

View File

@@ -14,6 +14,7 @@ rustproxy-metrics = { workspace = true }
hyper = { workspace = true }
hyper-util = { workspace = true }
regex = { workspace = true }
http-body = { workspace = true }
http-body-util = { workspace = true }
bytes = { workspace = true }
tokio = { workspace = true }

View File

@@ -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()
}
}

View File

@@ -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::*;

View File

@@ -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())
}