fix(rustproxy-http): report streamed HTTP and WebSocket bytes per chunk for real-time throughput metrics

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

@@ -11,20 +11,17 @@ 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`.
+/// Each chunk is reported to the `MetricsCollector` immediately so that
+/// the throughput tracker (sampled at 1 Hz) reflects real-time data flow.
 ///
 /// 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>,
     source_ip: 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,
     /// Optional connection-level activity tracker. When set, poll_frame updates this
     /// to keep the idle watchdog alive during active body streaming (uploads/downloads).
     connection_activity: Option<Arc<AtomicU64>>,
@@ -52,12 +49,10 @@ impl<B> CountingBody<B> {
     ) -> Self {
         Self {
             inner: Box::pin(inner),
-            counted_bytes: AtomicU64::new(0),
             metrics,
             route_id,
             source_ip,
             direction,
-            reported: false,
             connection_activity: None,
             activity_start: None,
         }
@@ -72,33 +67,18 @@ impl<B> CountingBody<B> {
         self
     }
 
-    /// 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;
-        }
-
+    /// Report a chunk of bytes immediately to the metrics collector.
+    #[inline]
+    fn report_chunk(&self, len: u64) {
         let route_id = self.route_id.as_deref();
         let source_ip = self.source_ip.as_deref();
         match self.direction {
-            Direction::In => self.metrics.record_bytes(bytes, 0, route_id, source_ip),
-            Direction::Out => self.metrics.record_bytes(0, bytes, route_id, source_ip),
+            Direction::In => self.metrics.record_bytes(len, 0, route_id, source_ip),
+            Direction::Out => self.metrics.record_bytes(0, len, route_id, source_ip),
         }
     }
 }
 
-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> {}
 
@@ -118,7 +98,9 @@ where
         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);
+                    let len = data.len() as u64;
+                    // Report bytes immediately so the 1 Hz throughput sampler sees them
+                    this.report_chunk(len);
                     // Keep the connection-level idle watchdog alive during body streaming
                     if let (Some(activity), Some(start)) = (&this.connection_activity, &this.activity_start) {
                         activity.store(start.elapsed().as_millis() as u64, Ordering::Relaxed);
@@ -127,11 +109,7 @@ where
                 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::Ready(None) => Poll::Ready(None),
             Poll::Pending => Poll::Pending,
         }
     }