retry: retry requests with <=64KB bodies (#1017)

Linkerd's proxy does not currently support retries on requests with
payloads. This precludes Linkerd from retrying gRPC requests, which is a
substantial limitation.

This PR adds support for retries on requests with bodies, if and only if
the request has a `Content-length` header and the content length is <=
64 KB.

In order to retry requests with payloads, we need to buffer the body
data for that request so that it can be sent again if a retry is
necessary. This is implemented by wrapping profile requests in a new
`Body` type which lazily buffers each chunk of data polled from the
inner `Body`. The buffered data is shared with a clone of the request,
and when the original body is dropped, ownership of the buffered data is
transferred to the clone. When the cloned request is sent, polling its
body will yield the buffered data.

If the server returns an error _before_ an entire streaming body has
been read, the replay body will continue reading from the initial
request body after playing back the buffered portion.

Data is buffered by calling `Buf::copy_to_bytes` on each chunk of data.
Although we call this method on an arbitrary `Buf` type, all data chunks
in the proxy are actually `Bytes`, and their `copy_to_bytes` methods are
therefore a cheap reference-count bump on the `Bytes`. After calling
`copy_to_bytes`, we can clone the returned `Bytes` and store it in a
vector. This allows us to buffer the body without actually copying the
bytes --- we just increase the reference count on the original buffer.

This buffering strategy also has the advantage of allowing us to write
out the entire buffered body in one big `writev` call. Because we store
the buffered body as a list of distinct buffers for each chunk, we can
expand the buffered body to a large number of scatter-gather buffers in
`Buf::bytes_vectored`. This should make replaying the body more
efficient, as we don't have to make a separate `write` call for each
chunk.

I've also added several tests for the new buffering body. In particular,
there are tests for a number of potential edge cases, including:

- when a retry is started before the entire initial body has been read
  (i.e. if the server returns an error before the request completes),
- when a retry body is cloned multiple times, including if the client's
  body has not yet completed,
- dropping clones prior to completion

Closes linkerd/linkerd2#6130.
Signed-off-by: Eliza Weisman <[email protected]>

Author: hawkw

Cargo.lock

@@ -1262,9 +1262,17 @@ dependencies = [
 name = "linkerd-retry"
 version = "0.1.0"
 dependencies = [
+ "bytes",
+ "futures",
+ "http",
+ "http-body",
+ "hyper",
  "linkerd-error",
  "linkerd-stack",
+ "linkerd-tracing",
+ "parking_lot",
  "pin-project",
+ "tokio",
  "tower",
  "tracing",
 ]

linkerd/app/core/src/retry.rs

@@ -1,18 +1,19 @@
 use super::classify;
 use super::dst::Route;
-// use super::handle_time;
 use super::http_metrics::retries::Handle;
 use super::metrics::HttpRouteRetry;
 use crate::profiles;
 use futures::future;
-use hyper::body::HttpBody;
+use http_body::Body;
 use linkerd_http_classify::{Classify, ClassifyEos, ClassifyResponse};
 use linkerd_retry::NewRetryLayer;
 use linkerd_stack::Param;
 use std::marker::PhantomData;
 use std::sync::Arc;
 use tower::retry::budget::Budget;
 
+pub use linkerd_retry::*;
+
 pub fn layer(metrics: HttpRouteRetry) -> NewRetryLayer<NewRetry> {
     NewRetryLayer::new(NewRetry::new(metrics))
 }
@@ -22,18 +23,21 @@ pub trait CloneRequest<Req> {
 }
 
 #[derive(Clone, Debug)]
-pub struct NewRetry<C = ()> {
+pub struct NewRetry<C = WithBody> {
     metrics: HttpRouteRetry,
     _clone_request: PhantomData<C>,
 }
 
-pub struct Retry<C = ()> {
+pub struct Retry<C = WithBody> {
     metrics: Handle,
     budget: Arc<Budget>,
     response_classes: profiles::http::ResponseClasses,
     _clone_request: PhantomData<C>,
 }
 
+#[derive(Clone, Debug)]
+pub struct WithBody;
+
 impl NewRetry {
     pub fn new(metrics: HttpRouteRetry) -> Self {
         Self {
@@ -116,23 +120,47 @@ impl<C> Clone for Retry<C> {
     }
 }
 
-impl<B: Default + HttpBody> CloneRequest<http::Request<B>> for () {
-    fn clone_request(req: &http::Request<B>) -> Option<http::Request<B>> {
-        if !req.body().is_end_stream() {
+// === impl WithBody ===
+
+impl WithBody {
+    /// Allow buffering requests up to 64 kb
+    pub const MAX_BUFFERED_BYTES: usize = 64 * 1024;
+}
+
+impl<B: Body + Unpin> CloneRequest<http::Request<ReplayBody<B>>> for WithBody {
+    fn clone_request(req: &http::Request<ReplayBody<B>>) -> Option<http::Request<ReplayBody<B>>> {
+        let content_length = |req: &http::Request<_>| {
+            req.headers()
+                .get(http::header::CONTENT_LENGTH)
+                .and_then(|value| value.to_str().ok()?.parse::<usize>().ok())
+        };
+
+        // Requests without bodies can always be retried, as we will not need to
+        // buffer the body. If the request *does* have a body, retry it if and
+        // only if the request contains a `content-length` header and the
+        // content length is >= 64 kb.
+        let has_body = !req.body().is_end_stream();
+        if has_body && content_length(&req).unwrap_or(usize::MAX) > Self::MAX_BUFFERED_BYTES {
+            tracing::trace!(
+                req.has_body = has_body,
+                req.content_length = ?content_length(&req),
+                "not retryable",
+            );
             return None;
         }
 
-        let mut clone = http::Request::new(B::default());
+        tracing::trace!(
+            req.has_body = has_body,
+            req.content_length = ?content_length(&req),
+            "retryable",
+        );
+
+        let mut clone = http::Request::new(req.body().clone());
         *clone.method_mut() = req.method().clone();
         *clone.uri_mut() = req.uri().clone();
         *clone.headers_mut() = req.headers().clone();
         *clone.version_mut() = req.version();
 
-        // // Count retries toward the request's total handle time.
-        // if let Some(ext) = req.extensions().get::<handle_time::Tracker>() {
-        //     clone.extensions_mut().insert(ext.clone());
-        // }
-
         Some(clone)
     }
 }

linkerd/app/inbound/fuzz/Cargo.lock

@@ -1120,8 +1120,13 @@ dependencies = [
 name = "linkerd-retry"
 version = "0.1.0"
 dependencies = [
+ "bytes",
+ "futures",
+ "http",
+ "http-body",
  "linkerd-error",
  "linkerd-stack",
+ "parking_lot",
  "pin-project",
  "tower",
  "tracing",

linkerd/app/integration/src/client.rs

@@ -11,7 +11,7 @@ use tracing::instrument::Instrument;
 use webpki::{DNSName, DNSNameRef};
 
 type ClientError = hyper::Error;
-type Request = http::Request<Bytes>;
+type Request = http::Request<hyper::Body>;
 type Response = http::Response<hyper::Body>;
 type Sender = mpsc::UnboundedSender<(Request, oneshot::Sender<Result<Response, ClientError>>)>;
 
@@ -138,7 +138,7 @@ impl Client {
         &self,
         builder: http::request::Builder,
     ) -> impl Future<Output = Result<Response, ClientError>> + Send + Sync + 'static {
-        self.send_req(builder.body(Bytes::new()).unwrap())
+        self.send_req(builder.body(Bytes::new().into()).unwrap())
     }
 
     pub async fn request_body(&self, req: Request) -> Response {
@@ -175,7 +175,7 @@ impl Client {
         }
         tracing::debug!(headers = ?req.headers(), "request");
         let (tx, rx) = oneshot::channel();
-        let _ = self.tx.send((req, tx));
+        let _ = self.tx.send((req.map(Into::into), tx));
         async { rx.await.expect("request cancelled") }.in_current_span()
     }
 

linkerd/app/integration/src/tap.rs

@@ -215,6 +215,10 @@ where
             };
             http::Request::from_parts(parts, body)
         });
-        Box::pin(self.0.send_req(req).map_err(|err| err.to_string()))
+        Box::pin(
+            self.0
+                .send_req(req.map(Into::into))
+                .map_err(|err| err.to_string()),
+        )
     }
 }

linkerd/app/integration/src/tests/profiles.rs

@@ -43,17 +43,23 @@ macro_rules! profile_test {
                     .body("slept".into())
                     .unwrap()
             })
-            .route_fn("/0.5",  move |_req| {
-                if counter.fetch_add(1, Ordering::Relaxed) % 2 == 0 {
-                    Response::builder()
-                        .status(533)
-                        .body("nope".into())
-                        .unwrap()
-                } else {
-                    Response::builder()
-                        .status(200)
-                        .body("retried".into())
-                        .unwrap()
+            .route_async("/0.5",  move |req| {
+                let fail = counter.fetch_add(1, Ordering::Relaxed) % 2 == 0;
+                async move {
+                    // Read the entire body before responding, so that the
+                    // client doesn't fail when writing it out.
+                    let _body = hyper::body::aggregate(req.into_body()).await;
+                    Ok::<_, Error>(if fail {
+                        Response::builder()
+                            .status(533)
+                            .body("nope".into())
+                            .unwrap()
+                    } else {
+                        Response::builder()
+                            .status(200)
+                            .body("retried".into())
+                            .unwrap()
+                    })
                 }
             })
             .route_fn("/0.5/sleep",  move |_req| {
@@ -174,6 +180,48 @@ async fn retry_uses_budget() {
     }
 }
 
+#[tokio::test]
+async fn retry_with_small_post_body() {
+    profile_test! {
+        routes: [
+            controller::route()
+                .request_any()
+                .response_failure(500..600)
+                .retryable(true)
+        ],
+        budget: Some(controller::retry_budget(Duration::from_secs(10), 0.1, 1)),
+        with_client: |client: client::Client| async move {
+            let req = client.request_builder("/0.5")
+                .method("POST")
+                .body("req has a body".into())
+                .unwrap();
+            let res = client.request_body(req).await;
+            assert_eq!(res.status(), 200);
+        }
+    }
+}
+
+#[tokio::test]
+async fn retry_with_small_put_body() {
+    profile_test! {
+        routes: [
+            controller::route()
+                .request_any()
+                .response_failure(500..600)
+                .retryable(true)
+        ],
+        budget: Some(controller::retry_budget(Duration::from_secs(10), 0.1, 1)),
+        with_client: |client: client::Client| async move {
+            let req = client.request_builder("/0.5")
+                .method("PUT")
+                .body("req has a body".into())
+                .unwrap();
+            let res = client.request_body(req).await;
+            assert_eq!(res.status(), 200);
+        }
+    }
+}
+
 #[tokio::test]
 async fn does_not_retry_if_request_does_not_match() {
     profile_test! {
@@ -211,7 +259,7 @@ async fn does_not_retry_if_earlier_response_class_is_success() {
 }
 
 #[tokio::test]
-async fn does_not_retry_if_request_has_body() {
+async fn does_not_retry_if_body_is_too_long() {
     profile_test! {
         routes: [
             controller::route()
@@ -223,14 +271,40 @@ async fn does_not_retry_if_request_has_body() {
         with_client: |client: client::Client| async move {
             let req = client.request_builder("/0.5")
                 .method("POST")
-                .body("req has a body".into())
+                .body(hyper::Body::from(&[1u8; 64 * 1024 + 1][..]))
                 .unwrap();
             let res = client.request_body(req).await;
             assert_eq!(res.status(), 533);
         }
     }
 }
 
+#[tokio::test]
+async fn does_not_retry_if_body_lacks_known_length() {
+    profile_test! {
+        routes: [
+            controller::route()
+                .request_any()
+                .response_failure(500..600)
+                .retryable(true)
+        ],
+        budget: Some(controller::retry_budget(Duration::from_secs(10), 0.1, 1)),
+        with_client: |client: client::Client| async move {
+            let (mut tx, body) = hyper::body::Body::channel();
+            let req = client.request_builder("/0.5")
+                .method("POST")
+                .body(body)
+                .unwrap();
+            let res = tokio::spawn(async move { client.request_body(req).await });
+            let _ = tx.send_data(Bytes::from_static(b"hello"));
+            let _ = tx.send_data(Bytes::from_static(b"world"));
+            drop(tx);
+            let res = res.await.unwrap();
+            assert_eq!(res.status(), 533);
+        }
+    }
+}
+
 #[tokio::test]
 async fn does_not_retry_if_missing_retry_budget() {
     profile_test! {

linkerd/app/outbound/src/http/logical.rs

@@ -15,7 +15,7 @@ use tracing::debug_span;
 impl<E> Outbound<E> {
     pub fn push_http_logical<B, ESvc, R>(self, resolve: R) -> Outbound<svc::BoxNewHttp<Logical, B>>
     where
-        B: http::HttpBody<Error = Error> + std::fmt::Debug + Default + Send + 'static,
+        B: http::HttpBody<Error = Error> + std::fmt::Debug + Default + Unpin + Send + 'static,
         B::Data: Send + 'static,
         E: svc::NewService<Endpoint, Service = ESvc> + Clone + Send + Sync + 'static,
         ESvc: svc::Service<http::Request<http::BoxBody>, Response = http::Response<http::BoxBody>>
@@ -122,6 +122,7 @@ impl<E> Outbound<E> {
                     )
                     // Sets an optional retry policy.
                     .push(retry::layer(rt.metrics.http_route_retry.clone()))
+                    .push_on_response(retry::replay::layer())
                     // Sets an optional request timeout.
                     .push(http::MakeTimeoutLayer::default())
                     // Records per-route metrics.

linkerd/retry/Cargo.toml

@@ -7,8 +7,19 @@ edition = "2018"
 publish = false
 
 [dependencies]
+bytes = "1"
+futures = { version = "0.3", default-features = false }
+http-body = "0.4"
+http = "0.2"
 linkerd-error = { path = "../error" }
 linkerd-stack = { path = "../stack" }
+pin-project = "1"
+parking_lot = "0.11"
 tower = { version = "0.4.7", default-features = false, features = ["retry", "util"] }
 tracing = "0.1.23"
-pin-project = "1"
+
+[dev-dependencies]
+hyper = "0.14"
+linkerd-tracing = { path = "../tracing", features = ["ansi"] }
+tokio = { version = "1", features = ["macros", "rt"]}
+

linkerd/retry/src/lib.rs

@@ -12,6 +12,9 @@ pub use tower::retry::{budget::Budget, Policy};
 use tower::util::{Oneshot, ServiceExt};
 use tracing::trace;
 
+pub mod replay;
+pub use self::replay::ReplayBody;
+
 /// A strategy for obtaining per-target retry polices.
 pub trait NewPolicy<T> {
     type Policy;