feat(http/body-eos): introduce BodyWithEosFn<B, F> (#4246)

this commit introduces a piece of `http_body::Body` middleware.

`BodyWithEosFn<B, F>` wraps an inner `B`-typed body, and invokes an
`F`-typed callback when the end of the stream is reached.

detecting the end of a request body or response body stream is a tricky,
reoccuring problem in different pieces of our `tower` middleware.
streams can end with a `None`, a `Some(Ok(_))` trailers frame, a
`Some(Err(_))` error, can be dropped prior to completion, or a
`Some(Ok(_))` data frame that leaves the body reporting a `true` hint
from `Body::is_end_stream()`.

this is very easy to do incorrectly.

this piece of body middleware abstracts over all of these state
transitions, and allows callers to instead provide a callback that will
be invoked when the stream finishes.

Signed-off-by: katelyn martin <[email protected]>

Author: cratelyn

Cargo.lock

@@ -1710,6 +1710,21 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "linkerd-http-body-eos"
+version = "0.1.0"
+dependencies = [
+ "bytes",
+ "futures",
+ "http",
+ "http-body",
+ "http-body-util",
+ "linkerd-error",
+ "linkerd-mock-http-body",
+ "pin-project",
+ "tokio",
+]
+
 [[package]]
 name = "linkerd-http-box"
 version = "0.1.0"

Cargo.toml

@@ -24,6 +24,7 @@ members = [
     "linkerd/error-respond",
     "linkerd/exp-backoff",
     "linkerd/http/access-log",
+    "linkerd/http/body-eos",
     "linkerd/http/box",
     "linkerd/http/classify",
     "linkerd/http/detect",

linkerd/http/body-eos/Cargo.toml

@@ -0,0 +1,25 @@
+[package]
+name = "linkerd-http-body-eos"
+version = { workspace = true }
+authors = { workspace = true }
+license = { workspace = true }
+edition = { workspace = true }
+publish = { workspace = true }
+description = """
+Body middleware for processing the ends of streams.
+"""
+
+[dependencies]
+futures = { version = "0.3", default-features = false }
+http = { workspace = true }
+http-body = { workspace = true }
+pin-project = "1"
+
+linkerd-error = { path = "../../error" }
+
+[dev-dependencies]
+bytes = { workspace = true }
+http-body-util = { workspace = true }
+tokio = { version = "1", features = ["macros", "test-util"] }
+
+linkerd-mock-http-body = { path = "../../mock/http-body" }

linkerd/http/body-eos/src/lib.rs

@@ -0,0 +1,132 @@
+//! [`Body`] middleware that calls a function when the body ends.
+
+use http::HeaderMap;
+use http_body::{Body, Frame};
+use linkerd_error::Error;
+use std::{
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+#[cfg(test)]
+mod tests;
+
+/// A [`Body`] that calls a function when the body ends.
+#[pin_project::pin_project(PinnedDrop, project = BodyWithEosFnProj)]
+pub struct BodyWithEosFn<B, F>
+where
+    B: Body,
+    F: FnOnce(EosRef<<B as Body>::Error>),
+{
+    #[pin]
+    inner: B,
+    callback: Option<F>,
+}
+
+/// A reference to the end of a [`Body`] stream.
+pub enum EosRef<'a, E = Error> {
+    None,
+    Trailers(&'a HeaderMap),
+    Error(&'a E),
+    Cancelled,
+}
+
+// === imple BodyWithEosFn ===
+
+impl<B, F> BodyWithEosFn<B, F>
+where
+    B: Body,
+    F: FnOnce(EosRef<<B as Body>::Error>),
+{
+    /// Returns a new [`BodyWithEosFn<B, F>`].
+    pub fn new(body: B, f: F) -> Self {
+        let callback = if body.is_end_stream() {
+            // If the body is empty, invoke the callback immediately.
+            f(EosRef::None);
+            None
+        } else {
+            // Otherwise, hold the callback until the end of the stream is reached.
+            Some(f)
+        };
+
+        Self {
+            inner: body,
+            callback,
+        }
+    }
+
+    /// Returns an [`EosRef<'a, E>`] view of the end-of-stream, if applicable.
+    fn eos_ref<'frame, T, E>(
+        frame: &'frame Option<Result<Frame<T>, E>>,
+        body: &B,
+    ) -> Option<EosRef<'frame, E>> {
+        let frame = match frame {
+            Some(Ok(f)) => f,
+            // Errors indicate the end of the stream.
+            Some(Err(error)) => return Some(EosRef::Error(error)),
+            // If nothing was yielded, the end of the stream has been reached.
+            None => return Some(EosRef::None),
+        };
+
+        if let Some(trls) = frame.trailers_ref() {
+            // A trailers frame indicates the end of the stream.
+            Some(EosRef::Trailers(trls))
+        } else {
+            // `is_end_stream()` hints that the end of the stream has been reached.
+            body.is_end_stream().then_some(EosRef::None)
+        }
+    }
+}
+
+impl<B, F> http_body::Body for BodyWithEosFn<B, F>
+where
+    B: Body,
+    F: FnOnce(EosRef<<B as Body>::Error>),
+{
+    type Data = <B as Body>::Data;
+    type Error = <B as Body>::Error;
+
+    fn poll_frame(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
+        let BodyWithEosFnProj {
+            mut inner,
+            callback,
+        } = self.project();
+
+        // Poll the inner body for the next frame.
+        let poll = inner.as_mut().poll_frame(cx);
+        let frame = futures::ready!(poll);
+
+        // Invoke the callback if we have reached the end of the stream.
+        if let Some(eos) = Self::eos_ref(&frame, &inner) {
+            if let Some(callback) = callback.take() {
+                callback(eos);
+            }
+        }
+
+        Poll::Ready(frame)
+    }
+
+    fn is_end_stream(&self) -> bool {
+        let Self { inner: _, callback } = self;
+
+        callback.is_none()
+    }
+}
+
+#[pin_project::pinned_drop]
+impl<B, F> PinnedDrop for BodyWithEosFn<B, F>
+where
+    B: Body,
+    F: FnOnce(EosRef<<B as Body>::Error>),
+{
+    fn drop(self: Pin<&mut Self>) {
+        let BodyWithEosFnProj { inner: _, callback } = self.project();
+        if let Some(callback) = callback.take() {
+            // Invoke the callback if the body was dropped before finishing.
+            callback(EosRef::Cancelled)
+        };
+    }
+}