tower_load/
constant.rs

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