rama_net/stream/layer/throttle/
outgoing.rs

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

rama_net/stream/layer/throttle/
outgoing.rs
