use super::bytes::BytesRWTracker;
use crate::client::{ConnectorService, EstablishedClientConnection};
use rama_core::{Layer, Service, extensions::ExtensionsRef, io::Io};
use rama_utils::macros::define_inner_service_accessors;
#[derive(Debug, Clone)]
pub struct OutgoingBytesTrackerService<S> {
inner: S,
}
impl<S> OutgoingBytesTrackerService<S> {
pub const fn new(inner: S) -> Self {
Self { inner }
}
define_inner_service_accessors!();
}
impl<S, Input> Service<Input> for OutgoingBytesTrackerService<S>
where
S: ConnectorService<Input, Connection: Io + Unpin>,
Input: Send + 'static,
{
type Output = EstablishedClientConnection<BytesRWTracker<S::Connection>, Input>;
type Error = S::Error;
async fn serve(&self, input: Input) -> Result<Self::Output, Self::Error> {
let EstablishedClientConnection { input, conn } = self.inner.connect(input).await?;
let conn = BytesRWTracker::new(conn);
let handle = conn.handle();
conn.extensions().insert(handle);
Ok(EstablishedClientConnection { input, conn })
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct OutgoingBytesTrackerLayer;
impl OutgoingBytesTrackerLayer {
#[must_use]
pub const fn new() -> Self {
Self
}
}
impl Default for OutgoingBytesTrackerLayer {
fn default() -> Self {
Self::new()
}
}
impl<S> Layer<S> for OutgoingBytesTrackerLayer {
type Service = OutgoingBytesTrackerService<S>;
fn layer(&self, inner: S) -> Self::Service {
OutgoingBytesTrackerService { inner }
}
}