Ensure that services are held as long they are being used (#767)

Middlewares, especially the cache, may want to use RAII to detect when a
service is idle, but the TCP server and HTTP routers drop services as
soon as the request is dispatched.

This change modifies the TCP server to hold the TCP stack until all
work on that connection is complete. It also introduces a new
`http::Retain` middleware that clones its inner service into response
bodies.

This is necessary for upcoming cache eviction changes to address
linkerd/linkerd2#5334.

Author: olix0r

linkerd/app/core/src/serve.rs

@@ -3,7 +3,7 @@ use futures::prelude::*;
 use linkerd2_error::Error;
 use linkerd2_proxy_transport::listen::Addrs;
 use tower::util::ServiceExt;
-use tracing::{debug, info, info_span};
+use tracing::{debug, info, info_span, warn};
 use tracing_futures::Instrument;
 
 /// Spawns a task that binds an `L`-typed listener with an `A`-typed
@@ -28,7 +28,7 @@ where
             match listen.next().await {
                 None => return Ok(()),
                 Some(conn) => {
-                    // If the listener returned an error, complete the task
+                    // If the listener returned an error, complete the task.
                     let (addrs, io) = conn?;
 
                     // The local addr should be instrumented from the listener's context.
@@ -43,9 +43,20 @@ where
                     // Dispatch all of the work for a given connection onto a connection-specific task.
                     tokio::spawn(
                         async move {
-                            match accept.oneshot(io).err_into::<Error>().await {
-                                Ok(()) => debug!("Connection closed"),
-                                Err(error) => info!(%error, "Connection closed"),
+                            match accept.ready_oneshot().err_into::<Error>().await {
+                                Ok(mut accept) => {
+                                    match accept.call(io).err_into::<Error>().await {
+                                        Ok(()) => debug!("Connection closed"),
+                                        Err(error) => info!(%error, "Connection closed"),
+                                    }
+                                    // Hold the service until the connection is
+                                    // complete. This helps tie any inner cache
+                                    // lifetimes to the services they return.
+                                    drop(accept);
+                                }
+                                Err(error) => {
+                                    warn!(%error, "Server failed to become ready");
+                                }
                             }
                         }
                         .instrument(span),

linkerd/app/inbound/src/lib.rs

@@ -164,7 +164,8 @@ impl Config {
             Response = http::Response<http::boxed::BoxBody>,
             Error = Error,
             Future = impl Send,
-        > + Unpin
+        > + Clone
+                      + Unpin
                       + Send,
     > + Unpin
            + Clone
@@ -320,7 +321,8 @@ impl Config {
                 http::Request<http::boxed::BoxBody>,
                 Response = http::Response<http::boxed::BoxBody>,
                 Error = Error,
-            > + Send
+            > + Clone
+            + Send
             + 'static,
         S::Future: Send,
     {

linkerd/proxy/http/src/lib.rs

@@ -16,6 +16,7 @@ pub mod insert;
 pub mod normalize_uri;
 pub mod orig_proto;
 pub mod override_authority;
+mod retain;
 mod server;
 pub mod strip_header;
 pub mod timeout;
@@ -25,8 +26,9 @@ mod version;
 
 pub use self::{
     client_handle::{ClientHandle, SetClientHandle},
-    glue::HyperServerSvc,
+    glue::{HyperServerSvc, UpgradeBody},
     override_authority::CanOverrideAuthority,
+    retain::Retain,
     server::NewServeHttp,
     timeout::MakeTimeoutLayer,
     version::Version,

linkerd/proxy/http/src/retain.rs

@@ -0,0 +1,127 @@
+//! Provides a middleware that holds an inner service as long as responses are
+//! being processed.
+
+use linkerd2_stack::layer;
+use pin_project::pin_project;
+use std::{
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+/// Retains an inner service as long as its HTTP responses are being processsed
+/// so that RAII-guarded resources are held. This is mostly intended to support
+/// cache eviction.
+#[derive(Debug)]
+pub struct Retain<S, B> {
+    inner: S,
+    _marker: std::marker::PhantomData<fn() -> B>,
+}
+
+/// Wraps a `B` typed HTTP body to ensure that a `T` typed instance is held until
+/// the body is fully processed.
+#[pin_project]
+#[derive(Debug)]
+pub struct RetainBody<T, B> {
+    #[pin]
+    inner: B,
+
+    _retain: Option<T>,
+}
+
+// === impl Retain ===
+
+impl<S, B> Retain<S, B> {
+    pub fn new(inner: S) -> Self {
+        Self {
+            inner,
+            _marker: std::marker::PhantomData,
+        }
+    }
+
+    pub fn layer() -> impl layer::Layer<S, Service = Self> + Copy + Clone {
+        layer::mk(Self::new)
+    }
+}
+
+impl<S: Clone, B> Clone for Retain<S, B> {
+    fn clone(&self) -> Self {
+        Self {
+            inner: self.inner.clone(),
+            _marker: std::marker::PhantomData,
+        }
+    }
+}
+
+impl<S, Req, RspB> tower::Service<Req> for Retain<S, RspB>
+where
+    S: tower::Service<Req, Response = http::Response<RspB>> + Clone + Send + 'static,
+    S::Future: Send + 'static,
+{
+    type Response = http::Response<RetainBody<S, RspB>>;
+    type Error = S::Error;
+    type Future = Pin<
+        Box<
+            dyn Future<Output = Result<http::Response<RetainBody<S, RspB>>, S::Error>>
+                + Send
+                + 'static,
+        >,
+    >;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, req: Req) -> Self::Future {
+        let fut = self.inner.call(req);
+
+        // Retain a handle to the inner service until the response is completely
+        // processed.
+        let service = self.inner.clone();
+        Box::pin(async move {
+            let rsp = fut.await?;
+            Ok(rsp.map(move |inner| RetainBody {
+                inner,
+                _retain: Some(service),
+            }))
+        })
+    }
+}
+
+// === impl RetainBody ===
+
+impl<T, B: Default> Default for RetainBody<T, B> {
+    fn default() -> Self {
+        Self {
+            inner: B::default(),
+            _retain: None,
+        }
+    }
+}
+
+impl<T, B: http_body::Body> http_body::Body for RetainBody<T, B> {
+    type Data = B::Data;
+    type Error = B::Error;
+
+    fn is_end_stream(&self) -> bool {
+        self.inner.is_end_stream()
+    }
+
+    fn poll_data(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<B::Data, B::Error>>> {
+        self.project().inner.poll_data(cx)
+    }
+
+    fn poll_trailers(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<Option<http::HeaderMap<http::HeaderValue>>, B::Error>> {
+        self.project().inner.poll_trailers(cx)
+    }
+
+    fn size_hint(&self) -> http_body::SizeHint {
+        self.inner.size_hint()
+    }
+}