P2C 负载均衡

pub struct Balance<D, Req>
where
    D: Discover,
    D::Key: Hash,
{
    discover: D,

    services: ReadyCache<D::Key, D::Service, Req>,
    ready_index: Option<usize>,

    rng: Box<dyn Rng + Send + Sync>,

    _req: PhantomData<Req>,
}

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        // `ready_index` may have already been set by a prior invocation. These
        // updates cannot disturb the order of existing ready services.
        let _ = self.update_pending_from_discover(cx)?;
        self.promote_pending_to_ready(cx);

        loop {
            // If a service has already been selected, ensure that it is ready.
            // This ensures that the underlying service is ready immediately
            // before a request is dispatched to it (i.e. in the same task
            // invocation). If, e.g., a failure detector has changed the state
            // of the service, it may be evicted from the ready set so that
            // another service can be selected.
            if let Some(index) = self.ready_index.take() {
                match self.services.check_ready_index(cx, index) {
                    Ok(true) => {
                        // The service remains ready.
                        self.ready_index = Some(index);
                        return Poll::Ready(Ok(()));
                    }
                    Ok(false) => {
                        // The service is no longer ready. Try to find a new one.
                        trace!("ready service became unavailable");
                    }
                    Err(Failed(_, error)) => {
                        // The ready endpoint failed, so log the error and try
                        // to find a new one.
                        debug!(%error, "endpoint failed");
                    }
                }
            }

            // Select a new service by comparing two at random and using the
            // lesser-loaded service.
            self.ready_index = self.p2c_ready_index();
            if self.ready_index.is_none() {
                debug_assert_eq!(self.services.ready_len(), 0);
                // We have previously registered interest in updates from
                // discover and pending services.
                return Poll::Pending;
            }
        }
    }

        loop {
            match ready!(Pin::new(&mut self.discover).poll_discover(cx))
                .transpose()
                .map_err(|e| error::Discover(e.into()))?
            {
                None => return Poll::Ready(None),
                Some(Change::Remove(key)) => {
                    trace!("remove");
                    self.services.evict(&key);
                }
                Some(Change::Insert(key, svc)) => {
                    trace!("insert");
                    // If this service already existed in the set, it will be
                    // replaced as the new one becomes ready.
                    self.services.push(key, svc);
                }
            }
        }

    fn p2c_ready_index(&mut self) -> Option<usize> {
        match self.services.ready_len() {
            0 => None,
            1 => Some(0),
            len => {
                // Get two distinct random indexes (in a random order) and
                // compare the loads of the service at each index.
                let [aidx, bidx] = sample_floyd2(&mut self.rng, len as u64);
                debug_assert_ne!(aidx, bidx, "random indices must be distinct");

                let aload = self.ready_index_load(aidx as usize);
                let bload = self.ready_index_load(bidx as usize);
                let chosen = if aload <= bload { aidx } else { bidx };
                // ...
                Some(chosen as usize)
            }
        }
    }

pub(crate) fn sample_floyd2<R: Rng>(rng: &mut R, length: u64) -> [u64; 2] {
    debug_assert!(2 <= length);
    let aidx = rng.next_range(0..length - 1);
    let bidx = rng.next_range(0..length);
    let aidx = if aidx == bidx { length - 1 } else { aidx };
    [aidx, bidx]
}

impl<H> Rng for HasherRng<H>
where
    H: BuildHasher,
{
    fn next_u64(&mut self) -> u64 {
        let mut hasher = self.hasher.build_hasher();
        hasher.write_u64(self.counter);
        self.counter = self.counter.wrapping_add(1);
        hasher.finish()
    }
}

use std::convert::Infallible;
use std::future::poll_fn;
use std::task::Poll;
use tower::balance::p2c::Balance;
use tower::discover::ServiceList;
use tower::load::Constant;
use tower::{Service, ServiceExt, service_fn};

type Svc = Constant<tower::util::BoxService<(), usize, Infallible>, u32>;

fn endpoint(id: usize, load: u32) -> Svc {
  let svc = service_fn(move |_: ()| async move { Ok::<_, Infallible>(id) });
  Constant::new(tower::util::BoxService::new(svc), load)
}

async fn poll_once(b: &mut Balance<ServiceList<Vec<Svc>>, ()>) -> &'static str {
  poll_fn(|cx| {
    Poll::Ready(match b.poll_ready(cx) {
      Poll::Ready(Ok(())) => "Ready",
      Poll::Ready(Err(_)) => "Err",
      Poll::Pending => "Pending",
    })
  })
  .await
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  // 0 services: the discover stream ends immediately.
  let mut empty = Balance::new(ServiceList::new(Vec::<Svc>::new()));
  println!("0 services: poll_ready -> {}, len = {}", poll_once(&mut empty).await, empty.len());
  let waited = tokio::time::timeout(std::time::Duration::from_secs(1), empty.ready()).await;
  println!("0 services: ready() after 1s -> {}", if waited.is_err() { "timed out" } else { "returned" });

  // 3 always-ready services with constant loads 0, 1, 10.
  let mut b = Balance::new(ServiceList::new(vec![
    endpoint(0, 0),
    endpoint(1, 1),
    endpoint(2, 10),
  ]));
  let mut hits = [0u32; 3];
  for _ in 0..9000 {
    let id = b.ready().await.unwrap().call(()).await.unwrap();
    hits[id] += 1;
  }
  println!("loads [0, 1, 10], 9000 requests -> hits {:?}", hits);
}

    fn call(&mut self, request: Req) -> Self::Future {
        let index = self.ready_index.take().expect("called before ready");
        self.services
            .call_ready_index(index, request)
            .map_err(Into::into)
    }

// Reproduces tower issue #856 on tower 0.5.3: a Discover removal after
// poll_ready picked a service makes the cached ready_index point elsewhere.
use std::cell::RefCell;
use std::collections::VecDeque;
use std::convert::Infallible;
use std::future::{Ready, ready};
use std::pin::Pin;
use std::rc::Rc;
use std::task::{Context, Poll};
use futures::Stream;
use tower::balance::p2c::Balance;
use tower::discover::Change;
use tower::load::Constant;
use tower::util::rng::Rng;
use tower::Service;

// Always ready; answers with its own id.
struct IdSvc(char);
impl Service<()> for IdSvc {
  type Response = char;
  type Error = Infallible;
  type Future = Ready<Result<char, Infallible>>;
  fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Infallible>> {
    Poll::Ready(Ok(()))
  }
  fn call(&mut self, _: ()) -> Self::Future {
    ready(Ok(self.0))
  }
}

type Endpoint = Constant<IdSvc, u32>;
type Queue = Rc<RefCell<VecDeque<Change<char, Endpoint>>>>;

// A Discover we feed by hand, one change per poll_ready.
struct Disco(Queue);
impl Stream for Disco {
  type Item = Result<Change<char, Endpoint>, Infallible>;
  fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
    match self.0.borrow_mut().pop_front() {
      Some(c) => Poll::Ready(Some(Ok(c))),
      None => Poll::Pending,
    }
  }
}

// Every draw is 0, so sample_floyd2(len) always returns [len - 1, 0].
struct ZeroRng;
impl Rng for ZeroRng {
  fn next_u64(&mut self) -> u64 {
    0
  }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let q: Queue = Rc::new(RefCell::new(VecDeque::new()));
  let mut bal = Balance::from_rng(Disco(q.clone()), ZeroRng);
  let poll = |bal: &mut Balance<Disco, ()>| {
    let waker = futures::task::noop_waker();
    let mut cx = Context::from_waker(&waker);
    match bal.poll_ready(&mut cx) {
      Poll::Ready(Ok(())) => "Ready",
      Poll::Ready(Err(_)) => "Err",
      Poll::Pending => "Pending",
    }
  };

  for (id, load) in [('a', 0), ('b', 1), ('c', 10)] {
    q.borrow_mut().push_back(Change::Insert(id, Constant::new(IdSvc(id), load)));
    println!("insert {id} (load {load:>2}) -> poll_ready {}", poll(&mut bal));
  }
  // `a` was picked by the first poll_ready and is cached as ready_index = 0.
  q.borrow_mut().push_back(Change::Remove('a'));
  println!("remove a          -> poll_ready {}", poll(&mut bal));
  let got = bal.call(()).await.unwrap();
  println!("request went to {got}");

  // What a fresh P2C over the same ready set picks: poll again, call again.
  println!("next poll_ready {}", poll(&mut bal));
  let next = bal.call(()).await.unwrap();
  println!("next request went to {next}");
}

insert a (load  0) -> poll_ready Ready
insert b (load  1) -> poll_ready Ready
insert c (load 10) -> poll_ready Ready
remove a          -> poll_ready Ready
request went to c
next poll_ready Ready
next request went to b

                 Some(Change::Remove(key)) => {
                     trace!("remove");
                     self.services.evict(&key);
+                    // `evict` removes from the ready set with `swap_remove`,
+                    // which can move a different endpoint into the slot a
+                    // previously-selected `ready_index` points at. Discard the
+                    // cached selection so `poll_ready` re-runs P2C over the
+                    // current ready set instead of dispatching to the
+                    // swapped-in endpoint.
+                    self.ready_index = None;
                 }

     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        // `ready_index` may have already been set by a prior invocation. These
-        // updates cannot disturb the order of existing ready services.
+        // `ready_index` may have already been set by a prior invocation. If a
+        // discovered removal evicts a service, `update_pending_from_discover`
+        // clears it, since eviction can reorder the ready set (via
+        // `swap_remove`) and leave the cached index pointing at a different
+        // endpoint.

insert a (load  0) -> poll_ready Ready
insert b (load  1) -> poll_ready Ready
insert c (load 10) -> poll_ready Ready
remove a          -> poll_ready Ready
request went to b
next poll_ready Ready
next request went to b

0 services: poll_ready -> Pending, len = 0
0 services: ready() after 1s -> timed out
loads [0, 1, 10], 9000 requests -> hits [5969, 3031, 0]

0 services: poll_ready -> Pending, len = 0
0 services: ready() after 1s -> timed out
loads [0, 1, 10], 9000 requests -> hits [5984, 3016, 0]

0 services: poll_ready -> Pending, len = 0
0 services: ready() after 1s -> timed out
loads [0, 1, 10], 9000 requests -> hits [5968, 3032, 0]