Reconnect

pub struct Reconnect<M, Target>
where
    M: Service<Target>,
{
    mk_service: M,
    state: State<M::Future, M::Response>,
    target: Target,
    error: Option<M::Error>,
}

#[derive(Debug)]
enum State<F, S> {
    Idle,
    Connecting(F),
    Connected(S),
}

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        loop {
            match &mut self.state {
                State::Idle => {
                    trace!("poll_ready; idle");
                    match self.mk_service.poll_ready(cx) {
                        Poll::Ready(r) => r?,
                        Poll::Pending => {
                            trace!("poll_ready; MakeService not ready");
                            return Poll::Pending;
                        }
                    }

                    let fut = self.mk_service.make_service(self.target.clone());
                    self.state = State::Connecting(fut);
                    continue;
                }
                State::Connecting(ref mut f) => {
                    trace!("poll_ready; connecting");
                    match Pin::new(f).poll(cx) {
                        Poll::Ready(Ok(service)) => {
                            self.state = State::Connected(service);
                        }
                        Poll::Pending => {
                            trace!("poll_ready; not ready");
                            return Poll::Pending;
                        }
                        Poll::Ready(Err(e)) => {
                            trace!("poll_ready; error");
                            self.state = State::Idle;
                            self.error = Some(e);
                            break;
                        }
                    }
                }
                State::Connected(ref mut inner) => {
                    trace!("poll_ready; connected");
                    match inner.poll_ready(cx) {
                        Poll::Ready(Ok(())) => {
                            trace!("poll_ready; ready");
                            return Poll::Ready(Ok(()));
                        }
                        Poll::Pending => {
                            trace!("poll_ready; not ready");
                            return Poll::Pending;
                        }
                        Poll::Ready(Err(_)) => {
                            trace!("poll_ready; error");
                            self.state = State::Idle;
                        }
                    }
                }
            }
        }

        Poll::Ready(Ok(()))
    }

    fn call(&mut self, request: Request) -> Self::Future {
        if let Some(error) = self.error.take() {
            return ResponseFuture::error(error);
        }

        let service = match self.state {
            State::Connected(ref mut service) => service,
            _ => panic!("service not ready; poll_ready must be called first"),
        };

        let fut = service.call(request);
        ResponseFuture::new(fut)
    }

tower = { version = "=0.5.3", features = ["reconnect", "util"] }
tokio = { version = "=1.54.1", features = ["full"] }

use std::future::{ready, Ready};
use std::task::{Context, Poll};
use tower::reconnect::Reconnect;
use tower::{BoxError, Service, ServiceExt};

// A connection. Its poll_ready fails once it has been "killed".
struct Conn {
  id: u32,
  broken: bool,
}

impl Service<&'static str> for Conn {
  type Response = String;
  type Error = BoxError;
  type Future = Ready<Result<String, BoxError>>;

  fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), BoxError>> {
    if self.broken {
      return Poll::Ready(Err("connection reset".into()));
    }
    Poll::Ready(Ok(()))
  }

  fn call(&mut self, req: &'static str) -> Self::Future {
    match req {
      "boom" => ready(Err("bad request".into())),
      "kill" => {
        self.broken = true;
        ready(Ok(format!("conn {} handled {req}", self.id)))
      }
      _ => ready(Ok(format!("conn {} handled {req}", self.id))),
    }
  }
}

// Makes connections; attempts listed in `fail` are refused.
struct Connector {
  attempts: u32,
  fail: Vec<u32>,
}

impl Service<&'static str> for Connector {
  type Response = Conn;
  type Error = BoxError;
  type Future = Ready<Result<Conn, BoxError>>;

  fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), BoxError>> {
    Poll::Ready(Ok(()))
  }

  fn call(&mut self, target: &'static str) -> Self::Future {
    self.attempts += 1;
    let n = self.attempts;
    if self.fail.contains(&n) {
      println!("  connect #{n} to {target}: refused");
      return ready(Err("connection refused".into()));
    }
    println!("  connect #{n} to {target}: ok");
    ready(Ok(Conn { id: n, broken: false }))
  }
}

async fn send(svc: &mut Reconnect<Connector, &'static str>, req: &'static str) {
  let r = svc.ready().await.unwrap().call(req).await;
  println!("{req:>5} -> {:?}", r.map_err(|e| e.to_string()));
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  println!("scenario 1");
  let mut svc = Reconnect::new(Connector { attempts: 0, fail: vec![1] }, "db:5432");
  send(&mut svc, "a").await;
  send(&mut svc, "b").await;
  send(&mut svc, "boom").await;
  send(&mut svc, "c").await;
  send(&mut svc, "kill").await;
  send(&mut svc, "d").await;

  println!("scenario 2: ready() twice after a failed connect");
  let mut svc = Reconnect::new(Connector { attempts: 0, fail: vec![1] }, "db:5432");
  svc.ready().await.unwrap();
  svc.ready().await.unwrap();
  send(&mut svc, "e").await;
  send(&mut svc, "f").await;
}

scenario 1
  connect #1 to db:5432: refused
    a -> Err("connection refused")
  connect #2 to db:5432: ok
    b -> Ok("conn 2 handled b")
 boom -> Err("bad request")
    c -> Ok("conn 2 handled c")
 kill -> Ok("conn 2 handled kill")
  connect #3 to db:5432: ok
    d -> Ok("conn 3 handled d")
scenario 2: ready() twice after a failed connect
  connect #1 to db:5432: refused
  connect #2 to db:5432: ok
    e -> Err("connection refused")
    f -> Ok("conn 2 handled f")