Steer 与 MakeService

pub trait Picker<S, Req> {
    /// Return an index into the iterator of `Service` passed to [`Steer::new`].
    fn pick(&mut self, r: &Req, services: &[S]) -> usize;
}

pub struct Steer<S, F, Req> {
    router: F,
    services: Vec<S>,
    not_ready: VecDeque<usize>,
    _phantom: PhantomData<Req>,
}

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        loop {
            // must wait for *all* services to be ready.
            // this will cause head-of-line blocking unless the underlying services are always ready.
            if self.not_ready.is_empty() {
                return Poll::Ready(Ok(()));
            } else {
                if self.services[self.not_ready[0]]
                    .poll_ready(cx)?
                    .is_pending()
                {
                    return Poll::Pending;
                }

                self.not_ready.pop_front();
            }
        }
    }

    fn call(&mut self, req: Req) -> Self::Future {
        assert!(
            self.not_ready.is_empty(),
            "Steer must wait for all services to be ready. Did you forget to call poll_ready()?"
        );

        let idx = self.router.pick(&req, &self.services[..]);
        let cl = &mut self.services[idx];
        self.not_ready.push_back(idx);
        cl.call(req)
    }

use std::convert::Infallible;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::task::Poll;
use tokio_test::task;
use tower::buffer::Buffer;
use tower::limit::ConcurrencyLimit;
use tower::steer::Steer;
use tower::{Service, ServiceExt, service_fn};
use tower_test::mock;

fn pick<S>(req: &&'static str, _: &[S]) -> usize {
  if req.starts_with("to0") { 0 } else { 1 }
}

fn show<T, E>(p: Poll<Result<T, E>>) -> &'static str {
  match p {
    Poll::Ready(Ok(_)) => "Ready",
    Poll::Ready(Err(_)) => "Err",
    Poll::Pending => "Pending",
  }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  // 1. Backend 0 is not ready, backend 1 is.
  let (s0, mut h0) = mock::pair::<&'static str, &'static str>();
  let (s1, mut h1) = mock::pair::<&'static str, &'static str>();
  h0.allow(0);
  h1.allow(10);
  let mut steer = task::spawn(Steer::new(vec![s0, s1], pick));
  let p = steer.enter(|cx, mut s| s.poll_ready(cx));
  println!("plain:    poll_ready for a request to 1 -> {}", show(p));

  // 2. Same, but each backend sits behind a Buffer.
  let (s0, mut h0) = mock::pair::<&'static str, &'static str>();
  let (s1, mut h1b) = mock::pair::<&'static str, &'static str>();
  h0.allow(0);
  h1b.allow(10);
  let mut steer = Steer::new(vec![Buffer::new(s0, 8), Buffer::new(s1, 8)], pick);
  let fut = steer.ready().await.unwrap().call("to1: hello");
  let (req, send) = h1b.next_request().await.unwrap();
  send.send_response("pong");
  println!("buffered: backend 1 got {req:?}, response {:?}", fut.await.unwrap());
  drop(h1);

  // 3. Clone a Steer that is ready.
  let svc = |n: u32| ConcurrencyLimit::new(service_fn(move |_: &'static str| async move { Ok::<_, Infallible>(n) }), 1);
  let mut steer = Steer::new(vec![svc(0), svc(1)], pick);
  steer.ready().await.unwrap();
  let mut copy = steer.clone();
  std::panic::set_hook(Box::new(|_| {}));
  let r = catch_unwind(AssertUnwindSafe(|| {
    let _ = copy.call("to1: from the clone");
  }));
  let msg = r.err().and_then(|p| {
    p.downcast_ref::<String>()
      .cloned()
      .or_else(|| p.downcast_ref::<&str>().map(|s| s.to_string()))
  });
  println!("clone:    call without poll_ready -> panic {msg:?}");
}

plain:    poll_ready for a request to 1 -> Pending
buffered: backend 1 got "to1: hello", response "pong"
clone:    call without poll_ready -> panic Some("max requests in-flight; poll_ready must be called first")

pub trait MakeService<Target, Request>: Sealed<(Target, Request)> {
    type Response;
    type Error;
    type Service: Service<Request, Response = Self::Response, Error = Self::Error>;
    type MakeError;
    type Future: Future<Output = Result<Self::Service, Self::MakeError>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::MakeError>>;
    fn make_service(&mut self, target: Target) -> Self::Future;
    // ...
}

impl<M, S, Target, Request> MakeService<Target, Request> for M
where
    M: Service<Target, Response = S>,
    S: Service<Request>,
{
    type Response = S::Response;
    type Error = S::Error;
    type Service = S;
    type MakeError = M::Error;
    type Future = M::Future;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::MakeError>> {
        Service::poll_ready(self, cx)
    }

    fn make_service(&mut self, target: Target) -> Self::Future {
        Service::call(self, target)
    }
}

impl<S, T> Service<T> for Shared<S>
where
    S: Clone,
{
    type Response = S;
    type Error = Infallible;
    type Future = SharedFuture<S>;

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

    fn call(&mut self, _target: T) -> Self::Future {
        SharedFuture::new(std::future::ready(Ok(self.service.clone())))
    }
}

use std::convert::Infallible;
use std::future::{Ready, ready};
use std::task::{Context, Poll};
use axum::body::Body;
use axum::extract::Request;
use axum::response::Response;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tower::make::{MakeService, Shared};
use tower::{Service, ServiceExt};

// A service with per-instance state: how many requests this instance has seen.
#[derive(Clone)]
struct Counter {
  count: u32,
}

impl Service<Request> for Counter {
  type Response = Response;
  type Error = Infallible;
  type Future = Ready<Result<Response, Infallible>>;
  fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Infallible>> {
    Poll::Ready(Ok(()))
  }
  fn call(&mut self, _: Request) -> Self::Future {
    self.count += 1;
    ready(Ok(Response::new(Body::from(format!("count={}", self.count)))))
  }
}

async fn body_of(svc: &mut Counter) -> String {
  let rsp = svc.ready().await.unwrap().call(Request::new(Body::empty())).await.unwrap();
  let bytes = axum::body::to_bytes(rsp.into_body(), 64).await.unwrap();
  String::from_utf8(bytes.to_vec()).unwrap()
}

// Send one HTTP/1.1 request on an open connection and return the body.
async fn get(conn: &mut TcpStream) -> String {
  conn.write_all(b"GET / HTTP/1.1\r\nHost: x\r\n\r\n").await.unwrap();
  let mut buf = Vec::new();
  loop {
    let mut chunk = [0u8; 1024];
    let n = conn.read(&mut chunk).await.unwrap();
    buf.extend_from_slice(&chunk[..n]);
    let text = String::from_utf8_lossy(&buf).to_string();
    if let Some(pos) = text.find("\r\n\r\n") {
      let len: usize = text
        .lines()
        .find_map(|l| l.strip_prefix("content-length: "))
        .unwrap()
        .trim()
        .parse()
        .unwrap();
      if text.len() >= pos + 4 + len {
        return text[pos + 4..pos + 4 + len].to_string();
      }
    }
  }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  // 1. Shared by hand: one service per "connection", each a clone of the original.
  let mut make = Shared::new(Counter { count: 0 });
  let mut conn_a = make.make_service(()).await.unwrap();
  let mut conn_b = make.make_service(()).await.unwrap();
  let a1 = body_of(&mut conn_a).await;
  let a2 = body_of(&mut conn_a).await;
  let b1 = body_of(&mut conn_b).await;
  println!("by hand: conn A -> {a1}, {a2}; conn B -> {b1}");

  // 2. The same Shared handed to axum::serve.
  let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
  let addr = listener.local_addr().unwrap();
  tokio::spawn(async move {
    axum::serve(listener, Shared::new(Counter { count: 0 })).await.unwrap();
  });
  let mut a = TcpStream::connect(addr).await.unwrap();
  let mut b = TcpStream::connect(addr).await.unwrap();
  let a1 = get(&mut a).await;
  let a2 = get(&mut a).await;
  let b1 = get(&mut b).await;
  println!("axum:    conn A -> {a1}, {a2}; conn B -> {b1}");
}

by hand: conn A -> count=1, count=2; conn B -> count=1
axum:    conn A -> count=1, count=1; conn B -> count=1