ReadyCache

pub struct ReadyCache<K, S, Req>
where
    K: Eq + Hash,
{
    /// A stream of services that are not yet ready.
    pending: FuturesUnordered<Pending<K, S, Req>>,
    /// An index of cancelation handles for pending streams.
    pending_cancel_txs: IndexMap<K, CancelTx>,

    /// Services that have previously become ready. Readiness can become stale,
    /// so a given service should be polled immediately before use.
    ///
    /// The cancelation oneshot is preserved (though unused) while the service is
    /// ready so that it need not be reallocated each time a request is
    /// dispatched.
    ready: IndexMap<K, (S, CancelPair)>,
}

    fn push_pending(&mut self, key: K, svc: S, (cancel_tx, cancel_rx): CancelPair) {
        if let Some(c) = self.pending_cancel_txs.insert(key.clone(), cancel_tx) {
            // If there is already a service for this key, cancel it.
            c.cancel();
        }
        self.pending.push(Pending {
            key: Some(key),
            cancel: Some(cancel_rx),
            ready: Some(svc),
            _pd: std::marker::PhantomData,
        });
    }

    pub fn poll_pending(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), error::Failed<K>>> {
        loop {
            match Pin::new(&mut self.pending).poll_next(cx) {
                Poll::Pending => return Poll::Pending,
                Poll::Ready(None) => return Poll::Ready(Ok(())),
                Poll::Ready(Some(Ok((key, svc, cancel_rx)))) => {
                    trace!("endpoint ready");
                    let cancel_tx = self.pending_cancel_txs.swap_remove(&key);
                    if let Some(cancel_tx) = cancel_tx {
                        // Keep track of the cancelation so that it need not be
                        // recreated after the service is used.
                        self.ready.insert(key, (svc, (cancel_tx, cancel_rx)));
                    } else {
                        assert!(
                            cancel_tx.is_some(),
                            "services that become ready must have a pending cancelation"
                        );
                    }
                }
                Poll::Ready(Some(Err(PendingError::Canceled(_)))) => {
                    debug!("endpoint canceled");
                    // The cancellation for this service was removed in order to
                    // cause this cancellation.
                }
                Poll::Ready(Some(Err(PendingError::Inner(key, e)))) => {
                    let cancel_tx = self.pending_cancel_txs.swap_remove(&key);
                    assert!(
                        cancel_tx.is_some(),
                        "services that return an error must have a pending cancelation"
                    );
                    return Err(error::Failed(key, e.into())).into();
                }
            }
        }
    }

    pub fn evict<Q: Hash + Equivalent<K>>(&mut self, key: &Q) -> bool {
        let canceled = if let Some(c) = self.pending_cancel_txs.swap_remove(key) {
            c.cancel();
            true
        } else {
            false
        };

        self.ready
            .swap_remove_full(key)
            .map(|_| true)
            .unwrap_or(canceled)
    }

/// Creates a cancelation sender and receiver.
///
/// A `tokio::sync::oneshot` is NOT used, as a `Receiver` is not guaranteed to
/// observe results as soon as a `Sender` fires. Using an `AtomicBool` allows
/// the state to be observed as soon as the cancelation is triggered.
fn cancelable() -> CancelPair {
    let cx = Arc::new(Cancel {
        waker: AtomicWaker::new(),
        canceled: AtomicBool::new(false),
    });
    (CancelTx(cx.clone()), CancelRx(cx))
}

        match svc.poll_ready(cx) {
            Poll::Ready(Ok(())) => Ok(true),
            Poll::Pending => {
                // became unready; so move it back there.
                let (key, (svc, cancel)) = self
                    .ready
                    .swap_remove_index(index)
                    .expect("invalid ready index");

                // If a new version of this service has been added to the
                // unready set, don't overwrite it.
                if !self.pending_contains(&key) {
                    self.push_pending(key, svc, cancel);
                }

                Ok(false)
            }
            Poll::Ready(Err(e)) => {
                // failed, so drop it.
                let (key, _) = self
                    .ready
                    .swap_remove_index(index)
                    .expect("invalid ready index");
                Err(error::Failed(key, e.into()))
            }
        }

    pub fn call_ready_index(&mut self, index: usize, req: Req) -> S::Future {
        let (key, (mut svc, cancel)) = self
            .ready
            .swap_remove_index(index)
            .expect("check_ready_index was not called");

        let fut = svc.call(req);

        // If a new version of this service has been added to the
        // unready set, don't overwrite it.
        if !self.pending_contains(&key) {
            self.push_pending(key, svc, cancel);
        }

        fut
    }

tower = { version = "=0.5.3", features = ["ready-cache"] }

use std::{
  cell::Cell,
  convert::Infallible,
  future::{ready, Ready},
  rc::Rc,
  task::{Context, Poll, Waker},
};
use tower::{ready_cache::ReadyCache, Service};

struct Node {
  name: &'static str,
  ready: Rc<Cell<bool>>,
}

fn node(name: &'static str, is_ready: bool) -> (Node, Rc<Cell<bool>>) {
  let flag = Rc::new(Cell::new(is_ready));
  (Node { name, ready: flag.clone() }, flag)
}

impl Service<()> for Node {
  type Response = &'static str;
  type Error = Infallible;
  type Future = Ready<Result<&'static str, Infallible>>;
  fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Infallible>> {
    if self.ready.get() { Poll::Ready(Ok(())) } else { Poll::Pending }
  }
  fn call(&mut self, _: ()) -> Self::Future {
    ready(Ok(self.name))
  }
}

impl Drop for Node {
  fn drop(&mut self) {
    println!("    (drop {})", self.name);
  }
}

type Cache = ReadyCache<&'static str, Node, ()>;

fn show(step: &str, c: &Cache) {
  let ready: Vec<String> = c.iter_ready().enumerate().map(|(i, (k, s))| format!("{i}:{k}={}", s.name)).collect();
  println!("{step:<28} ready [{}]  pending {}", ready.join(" "), c.pending_len());
}

fn poll_pending(c: &mut Cache) {
  let mut cx = Context::from_waker(Waker::noop());
  let _ = c.poll_pending(&mut cx);
}

fn main() {
  let mut cx = Context::from_waker(Waker::noop());
  let mut c = Cache::default();
  for k in ["a", "b", "c", "d"] {
    c.push(k, node(k, true).0);
  }
  show("1 push a b c d", &c);
  poll_pending(&mut c);
  show("2 poll_pending", &c);

  let _ = c.call_ready_index(0, ());
  show("3 call_ready_index(0)", &c);
  poll_pending(&mut c);
  show("4 poll_pending", &c);

  let i = c.get_ready(&"b").unwrap().0;
  println!("  cached index of b = {i}");
  c.evict(&"b");
  show("5 evict(b)", &c);
  let ok = c.check_ready_index(&mut cx, i);
  let who = c.get_ready_index(i).map(|(k, _)| *k);
  println!("  check_ready_index({i}) = {ok:?}, now holds {who:?}");

  let (e, _) = node("e", false);
  c.push("e", e);
  show("6 push e (not ready)", &c);
  println!("  evict(e) = {}", c.evict(&"e"));
  show("7 after evict(e)", &c);
  poll_pending(&mut c);
  show("8 poll_pending", &c);

  let (d2, flag_d2) = node("d2", true);
  c.push("d", d2);
  show("9 push d again (d2)", &c);
  poll_pending(&mut c);
  show("10 poll_pending", &c);

  flag_d2.set(false);
  let ok = c.check_ready_index(&mut cx, 0);
  println!("  d2 turns unready; check_ready_index(0) = {ok:?}");
  show("11", &c);
}

1 push a b c d               ready []  pending 4
2 poll_pending               ready [0:a=a 1:b=b 2:c=c 3:d=d]  pending 0
3 call_ready_index(0)        ready [0:d=d 1:b=b 2:c=c]  pending 1
4 poll_pending               ready [0:d=d 1:b=b 2:c=c 3:a=a]  pending 0
  cached index of b = 1
    (drop b)
5 evict(b)                   ready [0:d=d 1:a=a 2:c=c]  pending 0
  check_ready_index(1) = Ok(true), now holds Some("a")
6 push e (not ready)         ready [0:d=d 1:a=a 2:c=c]  pending 1
  evict(e) = true
7 after evict(e)             ready [0:d=d 1:a=a 2:c=c]  pending 1
    (drop e)
8 poll_pending               ready [0:d=d 1:a=a 2:c=c]  pending 0
9 push d again (d2)          ready [0:d=d 1:a=a 2:c=c]  pending 1
    (drop d)
10 poll_pending              ready [0:d=d2 1:a=a 2:c=c]  pending 0
  d2 turns unready; check_ready_index(0) = Ok(false)
11                           ready [0:c=c 1:a=a]  pending 1
    (drop d2)
    (drop c)
    (drop a)