use super::bytes::BytesRWTracker;
use rama_core::{Layer, Service, extensions::ExtensionsRef, io::Io};
use rama_utils::macros::define_inner_service_accessors;
#[derive(Debug, Clone)]
pub struct IncomingBytesTrackerService<S> {
inner: S,
}
impl<S> IncomingBytesTrackerService<S> {
pub const fn new(inner: S) -> Self {
Self { inner }
}
define_inner_service_accessors!();
}
impl<S, IO> Service<IO> for IncomingBytesTrackerService<S>
where
S: Service<BytesRWTracker<IO>>,
IO: Io + ExtensionsRef,
{
type Output = S::Output;
type Error = S::Error;
fn serve(
&self,
stream: IO,
) -> impl Future<Output = Result<Self::Output, Self::Error>> + Send + '_ {
let tracked_stream = BytesRWTracker::new(stream);
let handle = tracked_stream.handle();
tracked_stream.extensions().insert(handle);
self.inner.serve(tracked_stream)
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct IncomingBytesTrackerLayer;
impl IncomingBytesTrackerLayer {
#[must_use]
pub const fn new() -> Self {
Self
}
}
impl Default for IncomingBytesTrackerLayer {
fn default() -> Self {
Self::new()
}
}
impl<S> Layer<S> for IncomingBytesTrackerLayer {
type Service = IncomingBytesTrackerService<S>;
fn layer(&self, inner: S) -> Self::Service {
IncomingBytesTrackerService { inner }
}
}