Tidy up tower::load

.github/workflows/CI.yml

@@ -12,7 +12,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        rust: [stable, 1.39.0]
+        rust: [stable, 1.40.0]
     steps:
     - uses: actions/checkout@master
     - uses: actions-rs/toolchain@v1
@@ -77,7 +77,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        rust: [stable, beta, nightly, 1.39.0]
+        rust: [stable, beta, nightly, 1.40.0]
     steps:
     - uses: actions/checkout@master
     - uses: actions-rs/toolchain@v1

tower/Cargo.toml

@@ -32,7 +32,7 @@ discover = []
 filter = []
 hedge = ["filter", "futures-util", "hdrhistogram", "tokio/time"]
 limit = ["tokio/time"]
-load = ["discover", "tokio/time"]
+load = ["tokio/time"]
 load-shed = []
 make = ["tokio/io-std"]
 ready-cache = ["futures-util", "indexmap", "tokio/sync"]
@@ -73,4 +73,4 @@ all-features = true
 rustdoc-args = ["--cfg", "docsrs"]
 
 [package.metadata.playground]
-features = ["full"]
+features = ["full"]

tower-balance/examples/demo.rs → tower/examples/tower-balance.rs

@@ -1,21 +1,22 @@
 //! Exercises load balancers with mocked services.
 
-use futures_core::TryStream;
+use futures_core::{Stream, TryStream};
 use futures_util::{stream, stream::StreamExt, stream::TryStreamExt};
 use hdrhistogram::Histogram;
 use pin_project::pin_project;
 use rand::{self, Rng};
+use std::hash::Hash;
 use std::time::Duration;
 use std::{
     pin::Pin,
     task::{Context, Poll},
 };
 use tokio::time::{self, Instant};
+use tower::balance as lb;
+use tower::discover::{Change, Discover};
+use tower::limit::concurrency::ConcurrencyLimit;
+use tower::load;
 use tower::util::ServiceExt;
-use tower_balance as lb;
-use tower_discover::{Change, Discover};
-use tower_limit::concurrency::ConcurrencyLimit;
-use tower_load as load;
 use tower_service::Service;
 
 const REQUESTS: usize = 100_000;
@@ -61,13 +62,15 @@ async fn main() {
         d,
         DEFAULT_RTT,
         decay,
-        load::NoInstrument,
+        load::CompleteOnResponse::default(),
     ));
     run("P2C+PeakEWMA...", pe).await;
 
     let d = gen_disco();
-    let ll =
-        lb::p2c::Balance::from_entropy(load::PendingRequestsDiscover::new(d, load::NoInstrument));
+    let ll = lb::p2c::Balance::from_entropy(load::PendingRequestsDiscover::new(
+        d,
+        load::CompleteOnResponse::default(),
+    ));
     run("P2C+LeastLoaded...", ll).await;
 }
 
@@ -78,20 +81,19 @@ type Key = usize;
 #[pin_project]
 struct Disco<S>(Vec<(Key, S)>);
 
-impl<S> Discover for Disco<S>
+impl<S> Stream for Disco<S>
 where
     S: Service<Req, Response = Rsp, Error = Error>,
 {
-    type Key = Key;
-    type Service = S;
-    type Error = Error;
-    fn poll_discover(
-        self: Pin<&mut Self>,
-        _: &mut Context<'_>,
-    ) -> Poll<Result<Change<Self::Key, Self::Service>, Self::Error>> {
+    type Item = Result<Change<Key, S>, Error>;
+
+    fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         match self.project().0.pop() {
-            Some((k, service)) => Poll::Ready(Ok(Change::Insert(k, service))),
-            None => Poll::Pending,
+            Some((k, service)) => Poll::Ready(Some(Ok(Change::Insert(k, service)))),
+            None => {
+                // there may be more later
+                Poll::Pending
+            }
         }
     }
 }
@@ -132,7 +134,7 @@ async fn run<D>(name: &'static str, lb: lb::p2c::Balance<D, Req>)
 where
     D: Discover + Unpin + Send + 'static,
     D::Error: Into<Error>,
-    D::Key: Clone + Send,
+    D::Key: Clone + Send + Hash,
     D::Service: Service<Req, Response = Rsp> + load::Load + Send,
     <D::Service as Service<Req>>::Error: Into<Error>,
     <D::Service as Service<Req>>::Future: Send,

tower/src/load/completion.rs

@@ -0,0 +1,92 @@
+//! Application-specific request completion semantics.
+
+use futures_core::ready;
+use pin_project::pin_project;
+use std::{
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+/// Attaches `H`-typed completion tracker to `V` typed values.
+///
+/// Handles (of type `H`) are intended to be RAII guards that primarily implement `Drop` and update
+/// load metric state as they are dropped. This trait allows implementors to "forward" the handle
+/// to later parts of the request-handling pipeline, so that the handle is only dropped when the
+/// request has truly completed.
+///
+/// This utility allows load metrics to have a protocol-agnostic means to track streams past their
+/// initial response future. For example, if `V` represents an HTTP response type, an
+/// implementation could add `H`-typed handles to each response's extensions to detect when all the
+/// response's extensions have been dropped.
+///
+/// A base `impl<H, V> TrackCompletion<H, V> for CompleteOnResponse` is provided to drop the handle
+/// once the response future is resolved. This is appropriate when a response is discrete and
+/// cannot comprise multiple messages.
+///
+/// In many cases, the `Output` type is simply `V`. However, `TrackCompletion` may alter the type
+/// in order to instrument it appropriately. For example, an HTTP `TrackCompletion` may modify the
+/// body type: so a `TrackCompletion` that takes values of type `http::Response<A>` may output
+/// values of type `http::Response<B>`.
+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;
+}
+
+/// A `TrackCompletion` implementation that considers the request completed when the response
+/// future is resolved.
+#[derive(Clone, Copy, Debug, Default)]
+#[non_exhaustive]
+pub struct CompleteOnResponse;
+
+/// Attaches a `C`-typed completion tracker to the result of an `F`-typed `Future`.
+#[pin_project]
+#[derive(Debug)]
+pub struct TrackCompletionFuture<F, C, H> {
+    #[pin]
+    future: F,
+    handle: Option<H>,
+    completion: C,
+}
+
+// ===== impl InstrumentFuture =====
+
+impl<F, C, H> TrackCompletionFuture<F, C, H> {
+    /// Wraps a future, propagating the tracker into its value if successful.
+    pub fn new(completion: C, handle: H, future: F) -> Self {
+        TrackCompletionFuture {
+            future,
+            completion,
+            handle: Some(handle),
+        }
+    }
+}
+
+impl<F, C, H, T, E> Future for TrackCompletionFuture<F, C, H>
+where
+    F: Future<Output = Result<T, E>>,
+    C: TrackCompletion<H, T>,
+{
+    type Output = Result<C::Output, E>;
+
+    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)))
+    }
+}
+
+// ===== CompleteOnResponse =====
+
+impl<H, V> TrackCompletion<H, V> for CompleteOnResponse {
+    type Output = V;
+
+    fn track_completion(&self, handle: H, value: V) -> V {
+        drop(handle);
+        value
+    }
+}

tower/src/load/constant.rs

@@ -1,17 +1,20 @@
-//! A constant `Load` implementation. Primarily useful for testing.
+//! A constant `Load` implementation.
 
+#[cfg(feature = "discover")]
 use crate::discover::{Change, Discover};
+#[cfg(feature = "discover")]
 use futures_core::{ready, Stream};
-use pin_project::pin_project;
-use std::{
-    pin::Pin,
-    task::{Context, Poll},
-};
-use tower_service::Service;
+#[cfg(feature = "discover")]
+use std::pin::Pin;
 
 use super::Load;
+use pin_project::pin_project;
+use std::task::{Context, Poll};
+use tower_service::Service;
 
-/// Wraps a type so that `Load::load` returns a constant value.
+/// Wraps a type so that it implements `Load` and returns a constant load metric.
+///
+/// This load estimator is primarily useful for testing.
 #[pin_project]
 #[derive(Debug)]
 pub struct Constant<T, M> {
@@ -55,6 +58,7 @@ where
 }
 
 /// Proxies `Discover` such that all changes are wrapped with a constant load.
+#[cfg(feature = "discover")]
 impl<D: Discover + Unpin, M: Copy> Stream for Constant<D, M> {
     type Item = Result<Change<D::Key, Constant<D::Service, M>>, D::Error>;