tower/load/
constant.rs

1//! A constant [`Load`] implementation.
2
3#[cfg(feature = "discover")]
4use crate::discover::{Change, Discover};
5#[cfg(feature = "discover")]
6use futures_core::{ready, Stream};
7#[cfg(feature = "discover")]
8use std::pin::Pin;
9
10use super::Load;
11use pin_project_lite::pin_project;
12use std::task::{Context, Poll};
13use tower_service::Service;
14
15pin_project! {
16    #[derive(Debug)]
17    /// Wraps a type so that it implements [`Load`] and returns a constant load metric.
18    ///
19    /// This load estimator is primarily useful for testing.
20    pub struct Constant<T, M> {
21        inner: T,
22        load: M,
23    }
24}
25
26// ===== impl Constant =====
27
28impl<T, M: Copy> Constant<T, M> {
29    /// Wraps a `T`-typed service with a constant `M`-typed load metric.
30    pub fn new(inner: T, load: M) -> Self {
31        Self { inner, load }
32    }
33}
34
35impl<T, M: Copy + PartialOrd> Load for Constant<T, M> {
36    type Metric = M;
37
38    fn load(&self) -> M {
39        self.load
40    }
41}
42
43impl<S, M, Request> Service<Request> for Constant<S, M>
44where
45    S: Service<Request>,
46    M: Copy,
47{
48    type Response = S::Response;
49    type Error = S::Error;
50    type Future = S::Future;
51
52    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
53        self.inner.poll_ready(cx)
54    }
55
56    fn call(&mut self, req: Request) -> Self::Future {
57        self.inner.call(req)
58    }
59}
60
61/// Proxies [`Discover`] such that all changes are wrapped with a constant load.
62#[cfg(feature = "discover")]
63#[cfg_attr(docsrs, doc(cfg(feature = "discover")))]
64impl<D: Discover + Unpin, M: Copy> Stream for Constant<D, M> {
65    type Item = Result<Change<D::Key, Constant<D::Service, M>>, D::Error>;
66
67    /// Yields the next discovery change set.
68    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
69        use self::Change::*;
70
71        let this = self.project();
72        let change = match ready!(Pin::new(this.inner).poll_discover(cx)).transpose()? {
73            None => return Poll::Ready(None),
74            Some(Insert(k, svc)) => Insert(k, Constant::new(svc, *this.load)),
75            Some(Remove(k)) => Remove(k),
76        };
77
78        Poll::Ready(Some(Ok(change)))
79    }
80}