use super::{ThrottleMode, ThrottledIo};
use rama_core::{Layer, Service, io::Io};
use rama_utils::macros::define_inner_service_accessors;
#[derive(Debug, Clone)]
pub struct ThrottleService<S> {
inner: S,
config: super::ThrottleConfig,
}
impl<S> ThrottleService<S> {
define_inner_service_accessors!();
}
impl<S, IO> Service<IO> for ThrottleService<S>
where
S: Service<ThrottledIo<IO>>,
IO: Io,
{
type Output = S::Output;
type Error = S::Error;
fn serve(
&self,
stream: IO,
) -> impl Future<Output = Result<Self::Output, Self::Error>> + Send + '_ {
self.inner.serve(self.config.wrap(stream))
}
}
#[derive(Debug, Clone, Default)]
pub struct ThrottleLayer {
config: super::ThrottleConfig,
}
impl ThrottleLayer {
#[must_use]
pub fn symmetric(mode: ThrottleMode) -> Self {
Self {
config: super::ThrottleConfig {
read: Some(mode.clone()),
write: Some(mode),
quantum: None,
},
}
}
#[must_use]
pub fn read_only(mode: ThrottleMode) -> Self {
Self {
config: super::ThrottleConfig {
read: Some(mode),
write: None,
quantum: None,
},
}
}
#[must_use]
pub fn write_only(mode: ThrottleMode) -> Self {
Self {
config: super::ThrottleConfig {
read: None,
write: Some(mode),
quantum: None,
},
}
}
#[must_use]
pub fn new(read: Option<ThrottleMode>, write: Option<ThrottleMode>) -> Self {
Self {
config: super::ThrottleConfig {
read,
write,
quantum: None,
},
}
}
rama_utils::macros::generate_set_and_with! {
pub fn quantum(mut self, quantum: Option<u64>) -> Self {
self.config.quantum = quantum;
self
}
}
}
impl<S> Layer<S> for ThrottleLayer {
type Service = ThrottleService<S>;
fn layer(&self, inner: S) -> Self::Service {
ThrottleService {
inner,
config: self.config.clone(),
}
}
}