1 use super::Balance;
2 use crate::discover::Discover;
3 use futures_core::ready;
4 use pin_project_lite::pin_project;
5 use std::hash::Hash;
6 use std::marker::PhantomData;
7 use std::{
8     fmt,
9     future::Future,
10     pin::Pin,
11     task::{Context, Poll},
12 };
13 use tower_service::Service;
14 
15 /// Constructs load balancers over dynamic service sets produced by a wrapped "inner" service.
16 ///
17 /// This is effectively an implementation of [`MakeService`] except that it forwards the service
18 /// descriptors (`Target`) to an inner service (`S`), and expects that service to produce a
19 /// service set in the form of a [`Discover`]. It then wraps the service set in a [`Balance`]
20 /// before returning it as the "made" service.
21 ///
22 /// See the [module-level documentation](crate::balance) for details on load balancing.
23 ///
24 /// [`MakeService`]: crate::MakeService
25 /// [`Discover`]: crate::discover::Discover
26 /// [`Balance`]: crate::balance::p2c::Balance
27 pub struct MakeBalance<S, Req> {
28     inner: S,
29     _marker: PhantomData<fn(Req)>,
30 }
31 
32 pin_project! {
33     /// A [`Balance`] in the making.
34     ///
35     /// [`Balance`]: crate::balance::p2c::Balance
36     pub struct MakeFuture<F, Req> {
37         #[pin]
38         inner: F,
39         _marker: PhantomData<fn(Req)>,
40     }
41 }
42 
43 impl<S, Req> MakeBalance<S, Req> {
44     /// Build balancers using operating system entropy.
new(make_discover: S) -> Self45     pub fn new(make_discover: S) -> Self {
46         Self {
47             inner: make_discover,
48             _marker: PhantomData,
49         }
50     }
51 }
52 
53 impl<S, Req> Clone for MakeBalance<S, Req>
54 where
55     S: Clone,
56 {
clone(&self) -> Self57     fn clone(&self) -> Self {
58         Self {
59             inner: self.inner.clone(),
60             _marker: PhantomData,
61         }
62     }
63 }
64 
65 impl<S, Target, Req> Service<Target> for MakeBalance<S, Req>
66 where
67     S: Service<Target>,
68     S::Response: Discover,
69     <S::Response as Discover>::Key: Hash,
70     <S::Response as Discover>::Service: Service<Req>,
71     <<S::Response as Discover>::Service as Service<Req>>::Error: Into<crate::BoxError>,
72 {
73     type Response = Balance<S::Response, Req>;
74     type Error = S::Error;
75     type Future = MakeFuture<S::Future, Req>;
76 
poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>77     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
78         self.inner.poll_ready(cx)
79     }
80 
call(&mut self, target: Target) -> Self::Future81     fn call(&mut self, target: Target) -> Self::Future {
82         MakeFuture {
83             inner: self.inner.call(target),
84             _marker: PhantomData,
85         }
86     }
87 }
88 
89 impl<S, Req> fmt::Debug for MakeBalance<S, Req>
90 where
91     S: fmt::Debug,
92 {
fmt(&self, f: &mut fmt::Formatter) -> fmt::Result93     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
94         let Self { inner, _marker } = self;
95         f.debug_struct("MakeBalance").field("inner", inner).finish()
96     }
97 }
98 
99 impl<F, T, E, Req> Future for MakeFuture<F, Req>
100 where
101     F: Future<Output = Result<T, E>>,
102     T: Discover,
103     <T as Discover>::Key: Hash,
104     <T as Discover>::Service: Service<Req>,
105     <<T as Discover>::Service as Service<Req>>::Error: Into<crate::BoxError>,
106 {
107     type Output = Result<Balance<T, Req>, E>;
108 
poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>109     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
110         let this = self.project();
111         let inner = ready!(this.inner.poll(cx))?;
112         let svc = Balance::new(inner);
113         Poll::Ready(Ok(svc))
114     }
115 }
116 
117 impl<F, Req> fmt::Debug for MakeFuture<F, Req>
118 where
119     F: fmt::Debug,
120 {
fmt(&self, f: &mut fmt::Formatter) -> fmt::Result121     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
122         let Self { inner, _marker } = self;
123         f.debug_struct("MakeFuture").field("inner", inner).finish()
124     }
125 }
126