cf_mach/nq_core/body/

counting_body.rs

// Copyright (c) 2023-2024 Cloudflare, Inc.
// Licensed under the BSD-3-Clause license found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause

use std::{sync::Arc, task::Poll, time::Duration};

use hyper::body::{Body, Bytes};
use tokio::sync::mpsc;
use tracing::{debug, error, trace};

use crate::nq_core::{Time, Timestamp};

/// [`BodyEvent`]s are generated by a [`CountingBody`] and describe the number
/// of total bytes seen or that the body has finished.
#[derive(Debug)]
pub enum BodyEvent {
    /// The number of bytes sent by the wrapped body at the given [`Timestamp`].
    ByteCount {
        /// When the event was generated.
        at: Timestamp,
        /// The total number of bytes seen.
        total: usize,
    },
    /// The [`CountingBody`] has finished sending the body it wraps.
    Finished {
        /// When the body finished.
        at: Timestamp,
    },
    /// The transfer terminated early with an error and will produce no further
    /// bytes.
    ///
    /// Emitted either by the [`CountingBody`] itself when the wrapped body
    /// yields an error, or by the client when the request fails or the server
    /// rejects it (e.g. an HTTP 413 on an upload). Consumers must treat this as
    /// terminal: the transfer did *not* complete.
    Failed {
        /// When the failure was observed.
        at: Timestamp,
        /// Human-readable cause, e.g. `"unexpected status 413 Payload Too Large"`.
        reason: String,
    },
}

pin_project_lite::pin_project! {
    #[allow(missing_docs)]
    pub struct CountingBody<B> {
        #[pin]
        inner: B,
        time: Arc<dyn Time>,
        last_sent: Timestamp,
        update_every: Duration,
        total: usize,
        events_tx: mpsc::UnboundedSender<BodyEvent>,
        sent_finished: bool,
    }
}

impl<B> CountingBody<B> {
    /// Create a [`CountingBody`] by wrapping the given body. Updates are sent
    /// every `update_every` duration and timestamps are taken with the given
    /// [`Arc<dyn Time>`].
    pub fn new(
        inner: B,
        update_every: Duration,
        time: Arc<dyn Time>,
    ) -> (Self, mpsc::UnboundedReceiver<BodyEvent>) {
        let (events_tx, events_rx) = mpsc::unbounded_channel();
        let last_sent = time.now();

        events_tx
            .send(BodyEvent::ByteCount {
                at: last_sent,
                total: 0,
            })
            .expect("no data buffered");

        (
            Self {
                inner,
                time,
                last_sent,
                update_every,
                total: 0,
                events_tx,
                sent_finished: false,
            },
            events_rx,
        )
    }

    /// A handle for reporting a failure the body itself cannot observe, such as
    /// an upload rejected by the server with a non-success status.
    ///
    /// The returned sender must be dropped as soon as it is no longer needed.
    /// [`CountingBody`] is otherwise the sole owner of the sender, and
    /// consumers rely on the channel closing when the body is dropped to detect
    /// a transfer that died without reporting anything.
    pub fn sender(&self) -> mpsc::UnboundedSender<BodyEvent> {
        self.events_tx.clone()
    }
}

impl<B> Body for CountingBody<B>
where
    B: Body<Data = Bytes>,
    B::Error: std::fmt::Debug,
{
    type Data = B::Data;

    type Error = B::Error;

    #[inline(always)]
    fn poll_frame(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> Poll<Option<Result<hyper::body::Frame<Self::Data>, Self::Error>>> {
        let mut this = self.project();

        // stop the body if there's no event sender.
        if this.events_tx.is_closed() {
            debug!("events_tx is closed, stopping");
            return Poll::Ready(None);
        }

        trace!("polling frame");

        match this.inner.as_mut().poll_frame(cx) {
            Poll::Ready(Some(Ok(frame))) => {
                if let Some(data) = frame.data_ref() {
                    *this.total += data.len();
                }

                let now = this.time.now();

                // We've waited long enough, send an update.
                if now.duration_since(*this.last_sent) >= *this.update_every {
                    let event = BodyEvent::ByteCount {
                        at: now,
                        total: *this.total,
                    };

                    *this.last_sent = now;

                    debug!(?event, "sending event");

                    // We can drop the error here since this is an
                    // increasing counter. The next send will hopefully
                    // capture it.
                    let _ = this.events_tx.send(event);
                }

                if this.inner.is_end_stream() && !*this.sent_finished {
                    debug!(
                        total = *this.total,
                        "body reached end of stream, sending finished event"
                    );
                    let _ = this.events_tx.send(BodyEvent::ByteCount {
                        at: now,
                        total: *this.total,
                    });
                    let _ = this.events_tx.send(BodyEvent::Finished { at: now });
                    *this.sent_finished = true;
                }

                Poll::Ready(Some(Ok(frame)))
            }
            // Stream finished, send the last count
            Poll::Ready(None) => {
                let now = this.time.now();
                let event = BodyEvent::ByteCount {
                    at: now,
                    total: *this.total,
                };

                if !*this.sent_finished {
                    debug!(
                        ?event,
                        total = *this.total,
                        "sending final byte count event"
                    );
                    let _ = this.events_tx.send(event);
                    debug!(at=?now, "sending finished event");
                    let _ = this.events_tx.send(BodyEvent::Finished { at: now });
                    *this.sent_finished = true;
                } else {
                    debug!("already sent finish");
                }

                Poll::Ready(None)
            }
            Poll::Ready(Some(Err(e))) => {
                let now = this.time.now();
                error!(error=?e, "body errored");

                // Report the failure so consumers retire this transfer instead
                // of leaving it looking permanently in-flight. Only emitted
                // once, and never after a `Finished`.
                if !*this.sent_finished {
                    let _ = this.events_tx.send(BodyEvent::Failed {
                        at: now,
                        reason: format!("body error: {e:?}"),
                    });
                    *this.sent_finished = true;
                }

                Poll::Ready(Some(Err(e)))
            }
            Poll::Pending => {
                trace!("body pending");
                Poll::Pending
            }
        }
    }

    fn is_end_stream(&self) -> bool {
        self.inner.is_end_stream()
    }

    fn size_hint(&self) -> hyper::body::SizeHint {
        self.inner.size_hint()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::nq_core::{TokioTime, body::UploadBody};
    use http_body_util::BodyExt;

    // Regression test for the upload completion logic: a fully-consumed body
    // must emit exactly one `Finished` event.
    #[tokio::test]
    async fn upload_body_emits_finished() {
        let size = 512 * 1024;
        let body = UploadBody::new(size);
        let time: Arc<dyn Time> = Arc::new(TokioTime::new());
        let (body, mut events) = CountingBody::new(body, Duration::ZERO, time);

        body.collect().await.unwrap();
        events.close();

        let mut got_finished = false;
        while let Some(ev) = events.recv().await {
            if matches!(ev, BodyEvent::Finished { .. }) {
                assert!(!got_finished, "duplicate Finished");
                got_finished = true;
            }
        }
        assert!(got_finished, "never received Finished");
    }

    /// A body that yields one data frame and then errors, standing in for a
    /// transfer killed mid-flight (reset stream, dropped connection).
    struct ErroringBody {
        sent: bool,
    }

    impl Body for ErroringBody {
        type Data = Bytes;
        type Error = &'static str;

        fn poll_frame(
            mut self: std::pin::Pin<&mut Self>,
            _cx: &mut std::task::Context<'_>,
        ) -> Poll<Option<Result<hyper::body::Frame<Self::Data>, Self::Error>>> {
            if !self.sent {
                self.sent = true;
                return Poll::Ready(Some(Ok(hyper::body::Frame::data(Bytes::from_static(
                    b"hello",
                )))));
            }
            Poll::Ready(Some(Err("stream reset")))
        }
    }

    // A body that dies mid-transfer must report `Failed`, otherwise consumers
    // cannot distinguish it from one that is still running.
    #[tokio::test]
    async fn errored_body_emits_failed() {
        let time: Arc<dyn Time> = Arc::new(TokioTime::new());
        let (body, mut events) =
            CountingBody::new(ErroringBody { sent: false }, Duration::ZERO, time);

        // Drive the body until it errors.
        let _ = body.collect().await;
        events.close();

        let mut failed = None;
        let mut got_finished = false;
        while let Some(ev) = events.recv().await {
            match ev {
                BodyEvent::Failed { reason, .. } => failed = Some(reason),
                BodyEvent::Finished { .. } => got_finished = true,
                BodyEvent::ByteCount { .. } => {}
            }
        }

        assert!(failed.is_some(), "never received Failed");
        assert!(!got_finished, "a failed body must not also report Finished");
    }
}