Load 与三种负载度量

pub trait Load {
    /// A comparable load metric.
    ///
    /// Lesser values indicate that the service is less loaded, and should be preferred for new
    /// requests over another service with a higher value.
    type Metric: PartialOrd;

    /// Estimate the service's current load.
    fn load(&self) -> Self::Metric;
}

pub trait TrackCompletion<H, V>: Clone {
    /// The instrumented value type.
    type Output;

    /// Attaches a `H`-typed handle to a `V`-typed value.
    fn track_completion(&self, handle: H, value: V) -> Self::Output;
}

impl<H, V> TrackCompletion<H, V> for CompleteOnResponse {
    type Output = V;

    fn track_completion(&self, handle: H, value: V) -> V {
        drop(handle);
        value
    }
}

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        let rsp = ready!(this.future.poll(cx))?;
        let h = this.handle.take().expect("handle");
        Poll::Ready(Ok(this.completion.track_completion(h, rsp)))
    }

#[derive(Debug)]
pub struct PendingRequests<S, C = CompleteOnResponse> {
    service: S,
    ref_count: RefCount,
    completion: C,
}

/// Shared between instances of [`PendingRequests`] and [`Handle`] to track active references.
#[derive(Clone, Debug, Default)]
struct RefCount(Arc<()>);

    fn load(&self) -> Count {
        // Count the number of references that aren't `self`.
        Count(self.ref_count.ref_count() - 1)
    }

tower = { version = "=0.5.3", features = ["load", "util"] }
tokio = { version = "=1.54.1", features = ["macros", "rt"] }

use std::{future::{poll_fn, Future}, task::Poll};
use tower::{
  load::{
    completion::TrackCompletion,
    pending_requests::{Handle, PendingRequests},
    CompleteOnResponse, Load,
  },
  service_fn, Service,
};

// keeps the handle inside the "response", like a streaming body would
#[derive(Clone)]
struct KeepUntilBodyDropped;

impl<V> TrackCompletion<Handle, V> for KeepUntilBodyDropped {
  type Output = (V, Handle);
  fn track_completion(&self, handle: Handle, value: V) -> (V, Handle) {
    (value, handle)
  }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let svc = service_fn(|x: i32| async move {
    if x < 0 { Err("negative") } else { Ok(x) }
  });

  let mut p = PendingRequests::new(svc, CompleteOnResponse::default());
  println!("1 new:                    {:?}", p.load());
  let f1 = p.call(1);
  let f2 = p.call(2);
  println!("2 two futures alive:      {:?}", p.load());
  let _ = f1.await;
  println!("3 f1 awaited:             {:?}", p.load());
  drop(f2);
  println!("4 f2 dropped unpolled:    {:?}", p.load());

  let mut f3 = Box::pin(p.call(-1));
  let r = poll_fn(|cx| Poll::Ready(f3.as_mut().poll(cx))).await;
  println!("5 f3 polled -> {r:?}");
  println!("  f3 still alive:         {:?}", p.load());
  drop(f3);
  println!("6 f3 dropped:             {:?}", p.load());

  let mut q = PendingRequests::new(svc, KeepUntilBodyDropped);
  let (v, body) = q.call(4).await.unwrap();
  println!("7 response {v} received:  {:?}", q.load());
  drop(body);
  println!("8 body dropped:           {:?}", q.load());
}

1 new:                    Count(0)
2 two futures alive:      Count(2)
3 f1 awaited:             Count(1)
4 f2 dropped unpolled:    Count(0)
5 f3 polled -> Ready(Err("negative"))
  f3 still alive:         Count(1)
6 f3 dropped:             Count(0)
7 response 4 received:  Count(1)
8 body dropped:           Count(0)

#[derive(Debug)]
pub struct PeakEwma<S, C = CompleteOnResponse> {
    service: S,
    decay_ns: f64,
    rtt_estimate: Arc<Mutex<RttEstimate>>,
    completion: C,
}

    fn load(&self) -> Self::Metric {
        let pending = Arc::strong_count(&self.rtt_estimate) as u32 - 1;

        // Update the RTT estimate to account for decay since the last update.
        // If an estimate has not been established, a default is provided
        let estimate = self.update_estimate();

        let cost = Cost(estimate * f64::from(pending + 1));
        // ...
        cost
    }

    fn decay(&mut self, decay_ns: f64) -> f64 {
        // Updates with a 0 duration so that the estimate decays towards 0.
        let now = Instant::now();
        self.update(now, now, decay_ns)
    }

        let rtt = nanos(recv_at.saturating_duration_since(sent_at));

        let now = Instant::now();
        // ...
        self.rtt_ns = if self.rtt_ns < rtt {
            // For Peak-EWMA, always use the worst-case (peak) value as the estimate for
            // subsequent requests.
            // ...
            rtt
        } else {
            // When an RTT is observed that is less than the estimated RTT, we decay the
            // prior estimate according to how much time has elapsed since the last
            // update. The inverse of the decay is used to scale the estimate towards the
            // observed RTT value.
            let elapsed = nanos(now.saturating_duration_since(self.update_at));
            let decay = (-elapsed / decay_ns).exp();
            let recency = 1.0 - decay;
            let next_estimate = (self.rtt_ns * decay) + (rtt * recency);
            // ...
            next_estimate
        };
        self.update_at = now;

        self.rtt_ns

tower = { version = "=0.5.3", features = ["load", "util"] }
tokio = { version = "=1.54.1", features = ["macros", "rt", "time", "test-util"] }

use std::time::Duration;
use tokio::time::{advance, sleep};
use tower::{
  load::{CompleteOnResponse, Load, PeakEwma},
  service_fn, Service,
};

// Cost is opaque: read the f64 out of its Debug output, in ms
fn ms<L: Load>(svc: &L) -> String
where
  L::Metric: std::fmt::Debug,
{
  let s = format!("{:?}", svc.load());
  let ns: f64 = s.trim_start_matches("Cost(").trim_end_matches(')').parse().unwrap();
  format!("{:8.3} ms", ns / 1e6)
}

#[tokio::main(flavor = "current_thread", start_paused = true)]
async fn main() {
  // the request is "how long to take"; a negative value fails after that long
  let svc = service_fn(|ms: i64| async move {
    sleep(Duration::from_millis(ms.unsigned_abs())).await;
    if ms < 0 { Err("failed") } else { Ok(()) }
  });
  let new = || PeakEwma::new(svc, Duration::from_millis(100), 1e9, CompleteOnResponse::default());

  let mut a = new();
  println!("1 fresh, default_rtt 100ms:     {}", ms(&a));
  advance(Duration::from_millis(500)).await;
  println!("2 idle 500ms:                   {}", ms(&a));
  let f = a.call(300);
  println!("3 one request in flight:        {}", ms(&a));
  f.await.unwrap();
  println!("4 it took 300ms:                {}", ms(&a));
  advance(Duration::from_secs(1)).await;
  a.call(10).await.unwrap();
  println!("5 idle 1s, then a 10ms reply:   {}", ms(&a));

  // same history, but load() is read every 100ms while idle
  let mut b = new();
  advance(Duration::from_millis(500)).await;
  let _ = b.load();
  b.call(300).await.unwrap();
  let _ = b.load();
  for _ in 0..10 {
    advance(Duration::from_millis(100)).await;
    let _ = b.load();
  }
  b.call(10).await.unwrap();
  println!("6 same, load() every 100ms:     {}", ms(&b));

  // a request dropped after 2s counts as a 2s round trip
  let f = a.call(5_000);
  advance(Duration::from_secs(2)).await;
  drop(f);
  println!("7 a call dropped after 2s:      {}", ms(&a));

  // ten rounds of 150ms: c answers in 50ms, d fails after 1ms
  let (mut c, mut d) = (new(), new());
  for _ in 0..10 {
    let _ = c.call(50).await;
    advance(Duration::from_millis(50)).await;
    let _ = d.call(-1).await;
    advance(Duration::from_millis(49)).await;
  }
  println!("8 c, ten 50ms replies:          {}", ms(&c));
  println!("  d, ten 1ms errors:            {}", ms(&d));
}

1 fresh, default_rtt 100ms:      100.000 ms
2 idle 500ms:                     60.653 ms
3 one request in flight:         121.306 ms
4 it took 300ms:                 300.000 ms
5 idle 1s, then a 10ms reply:    115.624 ms
6 same, load() every 100ms:      109.365 ms
7 a call dropped after 2s:      2000.000 ms
8 c, ten 50ms replies:            56.398 ms
  d, ten 1ms errors:              23.042 ms