relay_core_lib/proxy/
throttle.rs

1use crate::interceptor::{BoxError, HttpBody};
2use hyper::body::{Body, Bytes, Frame, SizeHint};
3use std::pin::Pin;
4use std::task::{Context, Poll};
5use std::time::Duration;
6use tokio::time::Instant;
7
8/// Wraps a body stream with bandwidth throttling (bytes/sec).
9/// Inserts artificial delays between data frames to ensure
10/// throughput does not exceed the configured rate.
11pub struct ThrottleBody {
12    inner: HttpBody,
13    bytes_per_sec: u64,
14    last_frame_at: Option<Instant>,
15}
16
17impl ThrottleBody {
18    pub fn new(inner: HttpBody, bytes_per_sec: u64) -> Self {
19        Self {
20            inner,
21            bytes_per_sec,
22            last_frame_at: None,
23        }
24    }
25}
26
27impl Body for ThrottleBody {
28    type Data = Bytes;
29    type Error = BoxError;
30
31    fn poll_frame(
32        mut self: Pin<&mut Self>,
33        cx: &mut Context<'_>,
34    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
35        let frame = match Pin::new(&mut self.inner).poll_frame(cx) {
36            Poll::Ready(Some(Ok(frame))) => frame,
37            other => return other,
38        };
39
40        // Calculate per-frame delay based on data size
41        if let Some(data) = frame.data_ref() {
42            let bytes = data.len() as u64;
43            if bytes > 0 && self.bytes_per_sec > 0 {
44                let frame_dur = Duration::from_micros(bytes * 1_000_000 / self.bytes_per_sec);
45                let now = Instant::now();
46
47                if let Some(last) = self.last_frame_at {
48                    let elapsed = now.duration_since(last);
49                    if elapsed < frame_dur {
50                        let remaining = frame_dur - elapsed;
51                        // Since we can't .await in poll_frame, schedule a wake
52                        let waker = cx.waker().clone();
53                        tokio::spawn(async move {
54                            tokio::time::sleep(remaining).await;
55                            waker.wake();
56                        });
57                        return Poll::Pending;
58                    }
59                }
60                self.last_frame_at = Some(now);
61            }
62        }
63
64        Poll::Ready(Some(Ok(frame)))
65    }
66
67    fn is_end_stream(&self) -> bool {
68        self.inner.is_end_stream()
69    }
70
71    fn size_hint(&self) -> SizeHint {
72        self.inner.size_hint()
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79    use bytes::Bytes;
80    use http_body_util::{BodyExt, Full};
81
82    #[tokio::test]
83    async fn test_throttle_body_preserves_data() {
84        let data = Bytes::from("test-body-data");
85        let body: HttpBody = Full::new(data.clone())
86            .map_err(|e| -> BoxError { Box::new(e) })
87            .boxed();
88        // High rate limit — no effective throttling
89        let throttled = ThrottleBody::new(body, 1_000_000);
90        let collected = throttled.collect().await.unwrap().to_bytes();
91        assert_eq!(collected, data);
92    }
93
94    #[tokio::test]
95    async fn test_throttle_body_passthrough_empty() {
96        let body: HttpBody = Full::new(Bytes::new())
97            .map_err(|e| -> BoxError { Box::new(e) })
98            .boxed();
99        let throttled = ThrottleBody::new(body, 1000);
100        let collected = throttled.collect().await.unwrap().to_bytes();
101        assert_eq!(collected.len(), 0);
102    }
103}
relay_core_lib/proxy/throttle.rs

relay_core_lib/proxy/
throttle.rs
