tower_load_shed/
lib.rs

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

tower_load_shed/
lib.rs
