Ensure all forwarded TCP streams keep the proxy running (#786)

The proxy is not intended to shutdown until all in-flight responses and
TCP streams are complete. This is enforced by the `ServeHttp`
middleware, which confusingly handles some TCP forwarding in addition to
HTTP server initialization. But there are HTTP streams that are not
instrumented by this middleware and will therefore be abruptly
interrupted when a shutdown signal is received.

This change modifie the `ServeHttp` middleware to only handle HTTP
streams. a `NewOptional` middleware is introduced to conditionally build
a TCP stack when no HTTP type has been detected; and a `drain::Retain`
middleware is added to prevent the drain signal from being released
until TCP streams are completed.

Author: olix0r

Cargo.lock

@@ -1104,9 +1104,11 @@ version = "0.1.0"
 dependencies = [
  "futures 0.3.5",
  "linkerd2-error",
+ "linkerd2-stack",
  "pin-project 0.4.22",
  "tokio 0.3.5",
  "tokio-test 0.3.0",
+ "tower",
 ]
 
 [[package]]

linkerd/app/inbound/src/lib.rs

@@ -106,7 +106,11 @@ impl Config {
         // Forwards TCP streams that cannot be decoded as HTTP.
         let tcp_forward = svc::stack(tcp_connect)
             .push_make_thunk()
-            .push_on_response(svc::layer::mk(tcp::Forward::new))
+            .push_on_response(
+                svc::layers()
+                    .push(svc::layer::mk(tcp::Forward::new))
+                    .push(drain::Retain::layer(drain.clone())),
+            )
             .instrument(|_: &_| debug_span!("tcp"))
             .into_inner();
 
@@ -380,23 +384,21 @@ impl Config {
             .check_new_service::<(http::Version, TcpAccept), http::Request<_>>()
             .into_inner();
 
-        svc::stack(http::NewServeHttp::new(
-            h2_settings,
-            http_server,
-            svc::stack(tcp_forward)
-                .push_map_target(TcpEndpoint::from)
-                .into_inner(),
-            drain,
-        ))
-        .check_new_clone::<(Option<http::Version>, TcpAccept)>()
-        .push_cache(cache_max_idle_age)
-        .push(transport::NewDetectService::layer(
-            transport::detect::DetectTimeout::new(
-                detect_protocol_timeout,
-                http::DetectHttp::default(),
-            ),
-        ))
-        .into_inner()
+        svc::stack(http::NewServeHttp::new(h2_settings, http_server, drain))
+            .push(svc::stack::NewOptional::layer(
+                svc::stack(tcp_forward)
+                    .push_map_target(TcpEndpoint::from)
+                    .into_inner(),
+            ))
+            .check_new_clone::<(Option<http::Version>, TcpAccept)>()
+            .push_cache(cache_max_idle_age)
+            .push(transport::NewDetectService::layer(
+                transport::detect::DetectTimeout::new(
+                    detect_protocol_timeout,
+                    http::DetectHttp::default(),
+                ),
+            ))
+            .into_inner()
     }
 
     pub fn build_tls_accept<D, A, F, B>(

linkerd/app/outbound/src/ingress.rs

@@ -130,9 +130,8 @@ where
         ))
         .into_inner();
 
-    svc::stack(http::NewServeHttp::new(h2_settings, http, tcp, drain))
-        .check_new_service::<(Option<http::Version>, tcp::Accept), io::PrefixedIo<transport::metrics::SensorIo<I>>>()
-        .check_new_clone::<(Option<http::Version>, tcp::Accept)>()
+    svc::stack(http::NewServeHttp::new(h2_settings, http, drain))
+        .push(svc::stack::NewOptional::layer(tcp))
         .push_cache(cache_max_idle_age)
         .push(transport::NewDetectService::layer(
             transport::detect::DetectTimeout::new(

linkerd/app/outbound/src/server.rs

@@ -56,7 +56,8 @@ where
     P::Future: Unpin + Send,
     P::Error: Send,
 {
-    let tcp_balance = tcp::balance::stack(&config.proxy, tcp_connect.clone(), resolve);
+    let tcp_balance =
+        tcp::balance::stack(&config.proxy, tcp_connect.clone(), resolve, drain.clone());
     let accept = accept_stack(
         config,
         profiles,
@@ -206,36 +207,30 @@ where
 
     // Load balances TCP streams that cannot be decoded as HTTP.
     let tcp_balance = svc::stack(tcp_balance)
-    .push_map_target(tcp::Concrete::from)
-    .push(profiles::split::layer())
-    .check_new_service::<tcp::Logical, transport::io::PrefixedIo<transport::metrics::SensorIo<I>>>()
-    .push_switch(tcp::Logical::should_resolve, tcp_forward)
-    .push_on_response(
-        svc::layers()
-            .push_failfast(dispatch_timeout)
-            .push_spawn_buffer(buffer_capacity),
-    )
-    .instrument(|_: &_| debug_span!("tcp"))
-    .check_new_service::<tcp::Logical, transport::io::PrefixedIo<transport::metrics::SensorIo<I>>>()
-    .into_inner();
+        .push_map_target(tcp::Concrete::from)
+        .push(profiles::split::layer())
+        .check_new_service::<tcp::Logical, transport::io::PrefixedIo<transport::metrics::SensorIo<I>>>()
+        .push_switch(tcp::Logical::should_resolve, tcp_forward)
+        .push_on_response(
+            svc::layers()
+                .push_failfast(dispatch_timeout)
+                .push_spawn_buffer(buffer_capacity),
+        )
+        .instrument(|_: &_| debug_span!("tcp"))
+        .check_new_service::<tcp::Logical, transport::io::PrefixedIo<transport::metrics::SensorIo<I>>>()
+        .into_inner();
 
-    let http = svc::stack(http::NewServeHttp::new(
-        h2_settings,
-        http_server,
-        tcp_balance,
-        drain,
-    ))
-    .check_new_clone::<(Option<http::Version>, tcp::Logical)>()
-    .check_new_service::<(Option<http::Version>, tcp::Logical), transport::io::PrefixedIo<transport::metrics::SensorIo<I>>>()
-    .push_cache(cache_max_idle_age)
-    .push(transport::NewDetectService::layer(
-        transport::detect::DetectTimeout::new(
-            detect_protocol_timeout,
-            http::DetectHttp::default(),
-        ),
-    ))
-    .check_new_service::<tcp::Logical, transport::metrics::SensorIo<I>>()
-    .into_inner();
+    let http = svc::stack(http::NewServeHttp::new(h2_settings, http_server, drain))
+        .push(svc::stack::NewOptional::layer(tcp_balance))
+        .push_cache(cache_max_idle_age)
+        .push(transport::NewDetectService::layer(
+            transport::detect::DetectTimeout::new(
+                detect_protocol_timeout,
+                http::DetectHttp::default(),
+            ),
+        ))
+        .check_new_service::<tcp::Logical, transport::metrics::SensorIo<I>>()
+        .into_inner();
 
     let tcp = svc::stack(tcp::connect::forward(tcp_connect))
         .push_map_target(tcp::Endpoint::from_logical(

linkerd/app/outbound/src/tcp/balance.rs

@@ -2,6 +2,7 @@ use super::{Concrete, Endpoint};
 use crate::resolve;
 use linkerd2_app_core::{
     config::ProxyConfig,
+    drain,
     proxy::{api_resolve::Metadata, core::Resolve, tcp},
     svc,
     transport::io,
@@ -14,6 +15,7 @@ pub fn stack<C, R, I>(
     config: &ProxyConfig,
     connect: C,
     resolve: R,
+    drain: drain::Watch,
 ) -> impl svc::NewService<
     Concrete,
     Service = impl tower::Service<
@@ -51,7 +53,8 @@ where
                     crate::EWMA_DEFAULT_RTT,
                     crate::EWMA_DECAY,
                 ))
-                .push(svc::layer::mk(tcp::Forward::new)),
+                .push(svc::layer::mk(tcp::Forward::new))
+                .push(drain::Retain::layer(drain)),
         )
         .check_make_service::<Concrete, I>()
         .into_new_service()

linkerd/app/outbound/src/tcp/tests.rs

@@ -69,7 +69,8 @@ async fn plaintext_tcp() {
     );
 
     // Build the outbound TCP balancer stack.
-    let forward = super::balance::stack(&cfg.proxy, connect, resolver).new_service(concrete);
+    let (_, drain) = drain::channel();
+    let forward = super::balance::stack(&cfg.proxy, connect, resolver, drain).new_service(concrete);
 
     forward
         .oneshot(client_io)
@@ -150,7 +151,8 @@ async fn tls_when_hinted() {
     client_io.read(b"hello").write(b"world");
 
     // Build the outbound TCP balancer stack.
-    let mut balance = super::balance::stack(&cfg.proxy, connect, resolver);
+    let (_, drain) = drain::channel();
+    let mut balance = super::balance::stack(&cfg.proxy, connect, resolver, drain);
 
     let plain = balance
         .new_service(plain_concrete)

linkerd/drain/Cargo.toml

@@ -6,10 +6,12 @@ edition = "2018"
 publish = false
 
 [dependencies]
+futures = "0.3"
 linkerd2-error = { path = "../error" }
-tokio = {version = "0.3", features = ["sync", "macros", "stream"]}
+linkerd2-stack = { path = "../stack" }
 pin-project = "0.4"
-futures = "0.3"
+tokio = { version = "0.3", features = ["sync", "macros", "stream"]}
+tower = { version = "0.4", default-features = false }
 
 [dev-dependencies]
-tokio-test = "0.3"
+tokio-test = "0.3"

linkerd/drain/src/lib.rs

@@ -1,9 +1,15 @@
 #![deny(warnings, rust_2018_idioms)]
+
+mod retain;
+
+pub use crate::retain::Retain;
 use linkerd2_error::Never;
 use pin_project::pin_project;
-use std::future::Future;
-use std::pin::Pin;
-use std::task::{Context, Poll};
+use std::{
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+};
 use tokio::{
     stream::Stream,
     sync::{mpsc, watch},

linkerd/drain/src/retain.rs

@@ -0,0 +1,49 @@
+use crate::Watch;
+use linkerd2_stack::layer;
+use std::{
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+/// Holds a drain::Watch for as long as a request is pending.
+#[derive(Clone, Debug)]
+pub struct Retain<S> {
+    inner: S,
+    drain: Watch,
+}
+
+// === impl Retain ===
+
+impl<S> Retain<S> {
+    pub fn new(drain: Watch, inner: S) -> Self {
+        Self { drain, inner }
+    }
+
+    pub fn layer(drain: Watch) -> impl layer::Layer<S, Service = Self> + Clone {
+        layer::mk(move |inner| Self::new(drain.clone(), inner))
+    }
+}
+
+impl<Req, S> tower::Service<Req> for Retain<S>
+where
+    S: tower::Service<Req>,
+    S::Future: Send + 'static,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = Pin<Box<dyn Future<Output = Result<S::Response, 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 {
+        Box::pin(
+            self.drain
+                .clone()
+                .ignore_signal()
+                .release_after(self.inner.call(req)),
+        )
+    }
+}