1mod future;
16
17pub use future::ResponseFuture;
18
19use crate::make::MakeService;
20use std::fmt;
21use std::{
22 future::Future,
23 pin::Pin,
24 task::{Context, Poll},
25};
26use tower_service::Service;
27use tracing::trace;
28
29pub struct Reconnect<M, Target>
31where
32 M: Service<Target>,
33{
34 mk_service: M,
35 state: State<M::Future, M::Response>,
36 target: Target,
37 error: Option<M::Error>,
38}
39
40#[derive(Debug)]
41enum State<F, S> {
42 Idle,
43 Connecting(F),
44 Connected(S),
45}
46
47impl<M, Target> Reconnect<M, Target>
48where
49 M: Service<Target>,
50{
51 pub const fn new(mk_service: M, target: Target) -> Self {
53 Reconnect {
54 mk_service,
55 state: State::Idle,
56 target,
57 error: None,
58 }
59 }
60
61 pub const fn with_connection(init_conn: M::Response, mk_service: M, target: Target) -> Self {
63 Reconnect {
64 mk_service,
65 state: State::Connected(init_conn),
66 target,
67 error: None,
68 }
69 }
70}
71
72impl<M, Target, S, Request> Service<Request> for Reconnect<M, Target>
73where
74 M: Service<Target, Response = S>,
75 S: Service<Request>,
76 M::Future: Unpin,
77 crate::BoxError: From<M::Error> + From<S::Error>,
78 Target: Clone,
79{
80 type Response = S::Response;
81 type Error = crate::BoxError;
82 type Future = ResponseFuture<S::Future, M::Error>;
83
84 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
85 loop {
86 match &mut self.state {
87 State::Idle => {
88 trace!("poll_ready; idle");
89 match self.mk_service.poll_ready(cx) {
90 Poll::Ready(r) => r?,
91 Poll::Pending => {
92 trace!("poll_ready; MakeService not ready");
93 return Poll::Pending;
94 }
95 }
96
97 let fut = self.mk_service.make_service(self.target.clone());
98 self.state = State::Connecting(fut);
99 continue;
100 }
101 State::Connecting(ref mut f) => {
102 trace!("poll_ready; connecting");
103 match Pin::new(f).poll(cx) {
104 Poll::Ready(Ok(service)) => {
105 self.state = State::Connected(service);
106 }
107 Poll::Pending => {
108 trace!("poll_ready; not ready");
109 return Poll::Pending;
110 }
111 Poll::Ready(Err(e)) => {
112 trace!("poll_ready; error");
113 self.state = State::Idle;
114 self.error = Some(e);
115 break;
116 }
117 }
118 }
119 State::Connected(ref mut inner) => {
120 trace!("poll_ready; connected");
121 match inner.poll_ready(cx) {
122 Poll::Ready(Ok(())) => {
123 trace!("poll_ready; ready");
124 return Poll::Ready(Ok(()));
125 }
126 Poll::Pending => {
127 trace!("poll_ready; not ready");
128 return Poll::Pending;
129 }
130 Poll::Ready(Err(_)) => {
131 trace!("poll_ready; error");
132 self.state = State::Idle;
133 }
134 }
135 }
136 }
137 }
138
139 Poll::Ready(Ok(()))
140 }
141
142 fn call(&mut self, request: Request) -> Self::Future {
143 if let Some(error) = self.error.take() {
144 return ResponseFuture::error(error);
145 }
146
147 let service = match self.state {
148 State::Connected(ref mut service) => service,
149 _ => panic!("service not ready; poll_ready must be called first"),
150 };
151
152 let fut = service.call(request);
153 ResponseFuture::new(fut)
154 }
155}
156
157impl<M, Target> fmt::Debug for Reconnect<M, Target>
158where
159 M: Service<Target> + fmt::Debug,
160 M::Future: fmt::Debug,
161 M::Response: fmt::Debug,
162 Target: fmt::Debug,
163{
164 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
165 fmt.debug_struct("Reconnect")
166 .field("mk_service", &self.mk_service)
167 .field("state", &self.state)
168 .field("target", &self.target)
169 .finish()
170 }
171}