1 //! Tower middleware for shedding load when inner services aren't ready. 2 3 use std::task::{Context, Poll}; 4 use tower_service::Service; 5 6 pub mod error; 7 pub mod future; 8 mod layer; 9 10 use self::error::Error; 11 use self::future::ResponseFuture; 12 pub use self::layer::LoadShedLayer; 13 14 /// A `Service` that sheds load when the inner service isn't ready. 15 #[derive(Debug)] 16 pub struct LoadShed<S> { 17 inner: S, 18 is_ready: bool, 19 } 20 21 // ===== impl LoadShed ===== 22 23 impl<S> LoadShed<S> { 24 /// Wraps a service in `LoadShed` middleware. new(inner: S) -> Self25 pub fn new(inner: S) -> Self { 26 LoadShed { 27 inner, 28 is_ready: false, 29 } 30 } 31 } 32 33 impl<S, Req> Service<Req> for LoadShed<S> 34 where 35 S: Service<Req>, 36 S::Error: Into<Error>, 37 { 38 type Response = S::Response; 39 type Error = Error; 40 type Future = ResponseFuture<S::Future>; 41 poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>42 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { 43 // We check for readiness here, so that we can know in `call` if 44 // the inner service is overloaded or not. 45 self.is_ready = match self.inner.poll_ready(cx) { 46 Poll::Ready(Err(e)) => return Poll::Ready(Err(e.into())), 47 r => r.is_ready(), 48 }; 49 50 // But we always report Ready, so that layers above don't wait until 51 // the inner service is ready (the entire point of this layer!) 52 Poll::Ready(Ok(())) 53 } 54 call(&mut self, req: Req) -> Self::Future55 fn call(&mut self, req: Req) -> Self::Future { 56 if self.is_ready { 57 // readiness only counts once, you need to check again! 58 self.is_ready = false; 59 ResponseFuture::called(self.inner.call(req)) 60 } else { 61 ResponseFuture::overloaded() 62 } 63 } 64 } 65 66 impl<S: Clone> Clone for LoadShed<S> { clone(&self) -> Self67 fn clone(&self) -> Self { 68 LoadShed { 69 inner: self.inner.clone(), 70 // new clones shouldn't carry the readiness state, as a cloneable 71 // inner service likely tracks readiness per clone. 72 is_ready: false, 73 } 74 } 75 } 76