//! Middleware for shedding load when inner services aren't ready. use std::task::{Context, Poll}; use tower_service::Service; pub mod error; pub mod future; mod layer; use self::future::ResponseFuture; pub use self::layer::LoadShedLayer; /// A [`Service`] that sheds load when the inner service isn't ready. /// /// [`Service`]: crate::Service #[derive(Debug)] pub struct LoadShed { inner: S, is_ready: bool, } // ===== impl LoadShed ===== impl LoadShed { /// Wraps a service in [`LoadShed`] middleware. pub fn new(inner: S) -> Self { LoadShed { inner, is_ready: false, } } } impl Service for LoadShed where S: Service, S::Error: Into, { type Response = S::Response; type Error = crate::BoxError; type Future = ResponseFuture; fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { // We check for readiness here, so that we can know in `call` if // the inner service is overloaded or not. self.is_ready = match self.inner.poll_ready(cx) { Poll::Ready(Err(e)) => return Poll::Ready(Err(e.into())), r => r.is_ready(), }; // But we always report Ready, so that layers above don't wait until // the inner service is ready (the entire point of this layer!) Poll::Ready(Ok(())) } fn call(&mut self, req: Req) -> Self::Future { if self.is_ready { // readiness only counts once, you need to check again! self.is_ready = false; ResponseFuture::called(self.inner.call(req)) } else { ResponseFuture::overloaded() } } } impl Clone for LoadShed { fn clone(&self) -> Self { LoadShed { inner: self.inner.clone(), // new clones shouldn't carry the readiness state, as a cloneable // inner service likely tracks readiness per clone. is_ready: false, } } }