1 //! A constant [`Load`] implementation.
2 
3 #[cfg(feature = "discover")]
4 use crate::discover::{Change, Discover};
5 #[cfg(feature = "discover")]
6 use futures_core::{ready, Stream};
7 #[cfg(feature = "discover")]
8 use std::pin::Pin;
9 
10 use super::Load;
11 use pin_project_lite::pin_project;
12 use std::task::{Context, Poll};
13 use tower_service::Service;
14 
15 pin_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 
28 impl<T, M: Copy> Constant<T, M> {
29     /// Wraps a `T`-typed service with a constant `M`-typed load metric.
new(inner: T, load: M) -> Self30     pub fn new(inner: T, load: M) -> Self {
31         Self { inner, load }
32     }
33 }
34 
35 impl<T, M: Copy + PartialOrd> Load for Constant<T, M> {
36     type Metric = M;
37 
load(&self) -> M38     fn load(&self) -> M {
39         self.load
40     }
41 }
42 
43 impl<S, M, Request> Service<Request> for Constant<S, M>
44 where
45     S: Service<Request>,
46     M: Copy,
47 {
48     type Response = S::Response;
49     type Error = S::Error;
50     type Future = S::Future;
51 
poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>52     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
53         self.inner.poll_ready(cx)
54     }
55 
call(&mut self, req: Request) -> Self::Future56     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")))]
64 impl<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.
poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>>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 }
81