rama_net/stream/layer/throttle/
incoming.rs

1use super::{ThrottleMode, ThrottledIo};
2
3use rama_core::{Layer, Service, io::Io};
4use rama_utils::macros::define_inner_service_accessors;
5
6/// A [`Service`] that wraps a [`Service`]'s input IO [`Stream`] with
7/// a byte-rate throttle. See [`ThrottledIo`].
8///
9/// [`Service`]: rama_core::Service
10/// [`Stream`]: rama_core::io::Io
11#[derive(Debug, Clone)]
12pub struct ThrottleService<S> {
13    inner: S,
14    config: super::ThrottleConfig,
15}
16
17impl<S> ThrottleService<S> {
18    define_inner_service_accessors!();
19}
20
21impl<S, IO> Service<IO> for ThrottleService<S>
22where
23    S: Service<ThrottledIo<IO>>,
24    IO: Io,
25{
26    type Output = S::Output;
27    type Error = S::Error;
28
29    fn serve(
30        &self,
31        stream: IO,
32    ) -> impl Future<Output = Result<Self::Output, Self::Error>> + Send + '_ {
33        self.inner.serve(self.config.wrap(stream))
34    }
35}
36
37/// A [`Layer`] that wraps a [`Service`]'s input IO [`Stream`] with
38/// a byte-rate throttle. See [`ThrottledIo`].
39///
40/// Directions are relative to the wrapped connection: `read` throttles
41/// ingress from the peer (back-pressuring it through transport flow
42/// control), `write` paces egress toward it.
43///
44/// [`Layer`]: rama_core::Layer
45/// [`Service`]: rama_core::Service
46/// [`Stream`]: rama_core::io::Io
47#[derive(Debug, Clone, Default)]
48pub struct ThrottleLayer {
49    config: super::ThrottleConfig,
50}
51
52impl ThrottleLayer {
53    /// Create a new [`ThrottleLayer`] throttling both directions with
54    /// the given [`ThrottleMode`].
55    ///
56    /// [`ThrottleMode::PerConn`] gives each direction its own
57    /// (independent) bucket; [`ThrottleMode::Shared`] spends both
58    /// directions from the same aggregate budget.
59    #[must_use]
60    pub fn symmetric(mode: ThrottleMode) -> Self {
61        Self {
62            config: super::ThrottleConfig {
63                read: Some(mode.clone()),
64                write: Some(mode),
65                quantum: None,
66            },
67        }
68    }
69
70    /// Create a new [`ThrottleLayer`] throttling only the read
71    /// (ingress) direction.
72    #[must_use]
73    pub fn read_only(mode: ThrottleMode) -> Self {
74        Self {
75            config: super::ThrottleConfig {
76                read: Some(mode),
77                write: None,
78                quantum: None,
79            },
80        }
81    }
82
83    /// Create a new [`ThrottleLayer`] throttling only the write
84    /// (egress) direction.
85    #[must_use]
86    pub fn write_only(mode: ThrottleMode) -> Self {
87        Self {
88            config: super::ThrottleConfig {
89                read: None,
90                write: Some(mode),
91                quantum: None,
92            },
93        }
94    }
95
96    /// Create a new [`ThrottleLayer`] with per-direction modes.
97    #[must_use]
98    pub fn new(read: Option<ThrottleMode>, write: Option<ThrottleMode>) -> Self {
99        Self {
100            config: super::ThrottleConfig {
101                read,
102                write,
103                quantum: None,
104            },
105        }
106    }
107
108    rama_utils::macros::generate_set_and_with! {
109        /// Override the grant quantum in bytes: the budget reserved per
110        /// IO operation (clamped to the burst capacity; defaults to a
111        /// tenth of a period worth of bytes, at most 16 KiB).
112        pub fn quantum(mut self, quantum: Option<u64>) -> Self {
113            self.config.quantum = quantum;
114            self
115        }
116    }
117}
118
119impl<S> Layer<S> for ThrottleLayer {
120    type Service = ThrottleService<S>;
121
122    fn layer(&self, inner: S) -> Self::Service {
123        ThrottleService {
124            inner,
125            config: self.config.clone(),
126        }
127    }
128}
rama_net/stream/layer/throttle/incoming.rs

rama_net/stream/layer/throttle/
incoming.rs
