tower/load_shed/
mod.rs

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

tower/load_shed/
mod.rs
