Minimum throughput body timeouts Pt.1

aws/sdk/integration-tests/s3/Cargo.toml

@@ -36,7 +36,8 @@ futures-util = { version = "0.3.16", default-features = false }
 hdrhistogram = "7.5.2"
 http = "0.2.3"
 http-body = "0.4.5"
-hyper = "0.14.26"
+hyper = { version = "0.14.26", features = ["stream"] }
+once_cell = "1.18.0"
 pretty_assertions = "1.3"
 serde_json = "1"
 smol = "1.2"

aws/sdk/integration-tests/s3/tests/throughput-timeout.rs

@@ -0,0 +1,142 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use aws_sdk_sts::error::DisplayErrorContext;
+use aws_smithy_async::rt::sleep::AsyncSleep;
+use aws_smithy_async::test_util::instant_time_and_sleep;
+use aws_smithy_async::time::SharedTimeSource;
+use aws_smithy_http::body::SdkBody;
+use aws_smithy_http::byte_stream::ByteStream;
+use aws_smithy_runtime::client::http::body::minimum_throughput::MinimumThroughputBody;
+use aws_smithy_runtime::test_util::capture_test_logs::capture_test_logs;
+use aws_smithy_runtime_api::shared::IntoShared;
+use aws_types::sdk_config::SharedAsyncSleep;
+use bytes::Bytes;
+use once_cell::sync::Lazy;
+use std::convert::Infallible;
+use std::time::{Duration, UNIX_EPOCH};
+
+#[should_panic = "minimum throughput was specified at 2 B/s, but throughput of 1.5 B/s was observed"]
+#[tokio::test]
+async fn test_throughput_timeout_less_than() {
+    let _logs = capture_test_logs();
+    let (time_source, sleep) = instant_time_and_sleep(UNIX_EPOCH);
+
+    let shared_sleep = sleep.clone();
+    // Will send ~1 byte per second.
+    let stream = futures_util::stream::unfold(1, move |state| {
+        let sleep = shared_sleep.clone();
+        async move {
+            if state > 255 {
+                None
+            } else {
+                sleep.sleep(Duration::from_secs(1)).await;
+                Some((
+                    Result::<_, Infallible>::Ok(Bytes::from(vec![state as u8])),
+                    state + 1,
+                ))
+            }
+        }
+    });
+    let body = ByteStream::new(SdkBody::from(hyper::body::Body::wrap_stream(stream)));
+    let body = body.map(move |body| {
+        let ts = time_source.clone();
+        // Throw an error if the stream sends less than 2 bytes per second
+        let minimum_throughput = (2u64, Duration::from_secs(1));
+        SdkBody::from_dyn(aws_smithy_http::body::BoxBody::new(
+            MinimumThroughputBody::new(ts, body, minimum_throughput),
+        ))
+    });
+    let res = body.collect().await;
+
+    match res {
+        Ok(_) => panic!("Expected an error due to slow stream but no error occurred."),
+        Err(e) => panic!("{}", DisplayErrorContext(e)),
+    }
+}
+
+const EXPECTED_BYTES: Lazy<Vec<u8>> = Lazy::new(|| (1..=255).map(|i| i as u8).collect::<Vec<_>>());
+
+#[tokio::test]
+async fn test_throughput_timeout_equal_to() {
+    let _logs = capture_test_logs();
+    let (time_source, sleep) = instant_time_and_sleep(UNIX_EPOCH);
+    let time_source: SharedTimeSource = time_source.into_shared();
+    let sleep: SharedAsyncSleep = sleep.into_shared();
+
+    // Will send ~1 byte per second.
+    let stream = futures_util::stream::unfold(1, move |state| {
+        let sleep = sleep.clone();
+        async move {
+            if state > 255 {
+                None
+            } else {
+                sleep.sleep(Duration::from_secs(1)).await;
+                Some((
+                    Result::<_, Infallible>::Ok(Bytes::from(vec![state as u8])),
+                    state + 1,
+                ))
+            }
+        }
+    });
+    let body = ByteStream::new(SdkBody::from(hyper::body::Body::wrap_stream(stream)));
+    let body = body.map(move |body| {
+        let time_source = time_source.clone();
+        // Throw an error if the stream sends less than 1 byte per second
+        let minimum_throughput = (1u64, Duration::from_secs(1));
+        SdkBody::from_dyn(aws_smithy_http::body::BoxBody::new(
+            MinimumThroughputBody::new(time_source, body, minimum_throughput),
+        ))
+    });
+    // assert_eq!(255.0, time_source.seconds_since_unix_epoch());
+    // assert_eq!(Duration::from_secs(255), sleep.total_duration());
+    let res = body
+        .collect()
+        .await
+        .expect("no streaming error occurs because data is sent fast enough")
+        .to_vec();
+    assert_eq!(*EXPECTED_BYTES, res);
+}
+
+#[tokio::test]
+async fn test_throughput_timeout_greater_than() {
+    let _logs = capture_test_logs();
+    let (time_source, sleep) = instant_time_and_sleep(UNIX_EPOCH);
+    let time_source: SharedTimeSource = time_source.into_shared();
+    let sleep: SharedAsyncSleep = sleep.into_shared();
+
+    // Will send ~1 byte per second.
+    let stream = futures_util::stream::unfold(1, move |state| {
+        let sleep = sleep.clone();
+        async move {
+            if state > 255 {
+                None
+            } else {
+                sleep.sleep(Duration::from_secs(1)).await;
+                Some((
+                    Result::<_, Infallible>::Ok(Bytes::from(vec![state as u8])),
+                    state + 1,
+                ))
+            }
+        }
+    });
+    let body = ByteStream::new(SdkBody::from(hyper::body::Body::wrap_stream(stream)));
+    let body = body.map(move |body| {
+        let time_source = time_source.clone();
+        // Throw an error if the stream sends less than 1 byte per 2s
+        let minimum_throughput = (1u64, Duration::from_secs(2));
+        SdkBody::from_dyn(aws_smithy_http::body::BoxBody::new(
+            MinimumThroughputBody::new(time_source, body, minimum_throughput),
+        ))
+    });
+    // assert_eq!(255.0, time_source.seconds_since_unix_epoch());
+    // assert_eq!(Duration::from_secs(255), sleep.total_duration());
+    let res = body
+        .collect()
+        .await
+        .expect("no streaming error occurs because data is sent fast enough")
+        .to_vec();
+    assert_eq!(*EXPECTED_BYTES, res);
+}

rust-runtime/aws-smithy-runtime/src/client/http.rs

@@ -35,3 +35,6 @@ pub fn default_http_client_plugin() -> SharedRuntimePlugin {
         );
     SharedRuntimePlugin::new(plugin)
 }
+
+/// HTTP body and body-wrapper types
+pub mod body;

rust-runtime/aws-smithy-runtime/src/client/http/body.rs

@@ -0,0 +1,6 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+pub mod minimum_throughput;

rust-runtime/aws-smithy-runtime/src/client/http/body/minimum_throughput.rs

@@ -0,0 +1,149 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+//! A body-wrapping type that ensures data is being streamed faster than some lower limit.
+//!
+//! If data is being streamed too slowly, this body type will emit an error next time it's polled.
+
+use aws_smithy_async::time::{SharedTimeSource, TimeSource};
+use aws_smithy_runtime_api::shared::IntoShared;
+use bytes::Buf;
+use http::HeaderMap;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use std::time::{Duration, SystemTime};
+
+pin_project_lite::pin_project! {
+    /// A body-wrapper that will ensure that the wrapped body is emitting bytes faster than some
+    /// `minimum_throughput`.
+    pub struct MinimumThroughputBody<InnerBody> {
+        #[pin]
+        inner: InnerBody,
+        // A record of when and how much data was read
+        throughput_logs: Vec<(SystemTime, u64)>,
+        // The minimum acceptable throughput. If the amount of data per unit of time returned is
+        // less that this, an error will be returned instead.
+        minimum_throughput: (u64, Duration),
+        time_source: SharedTimeSource,
+    }
+}
+
+impl<T: http_body::Body> MinimumThroughputBody<T> {
+    /// Given an HTTP body and a minimum throughput, create a new `MinimumThroughputBody`.
+    pub fn new(
+        time_source: impl TimeSource + 'static,
+        body: T,
+        minimum_throughput: (u64, Duration),
+    ) -> Self {
+        Self {
+            inner: body,
+            throughput_logs: Vec::new(),
+            minimum_throughput,
+            time_source: time_source.into_shared(),
+        }
+    }
+}
+
+impl<T> http_body::Body for MinimumThroughputBody<T>
+where
+    T: http_body::Body<Data = bytes::Bytes, Error = Box<dyn std::error::Error + Send + Sync>>,
+{
+    type Data = T::Data;
+    type Error = T::Error;
+
+    fn poll_data(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Self::Data, Self::Error>>> {
+        let this = self.project();
+
+        let poll_res = this.inner.poll_data(cx);
+
+        if let Poll::Ready(Some(Ok(ref data))) = poll_res {
+            this.throughput_logs
+                .push((this.time_source.now(), data.remaining() as u64));
+        };
+
+        let mut logs = this.throughput_logs.iter();
+        if let Some((first_instant, first_bytes)) = logs.next() {
+            let now = this.time_source.now();
+            let time_elapsed_since_first_poll = now
+                .duration_since(*first_instant)
+                .map_err(|err| Box::new(Error::TimeTravel(err.into())))?;
+            let mut total_bytes_read = *first_bytes;
+
+            while let Some((_, bytes_read)) = logs.next() {
+                total_bytes_read += bytes_read;
+            }
+
+            let minimum_bytes_per_second =
+                this.minimum_throughput.0 as f64 / this.minimum_throughput.1.as_secs_f64();
+            let actual_bytes_per_second =
+                total_bytes_read as f64 / time_elapsed_since_first_poll.as_secs_f64();
+
+            // oh no, too slow!
+            if actual_bytes_per_second < minimum_bytes_per_second {
+                return Poll::Ready(Some(Err(Box::new(Error::ThroughputBelowMinimum {
+                    expected: this.minimum_throughput.clone(),
+                    actual: (total_bytes_read, time_elapsed_since_first_poll),
+                }))));
+            }
+        };
+
+        poll_res
+    }
+
+    fn poll_trailers(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<Option<HeaderMap>, Self::Error>> {
+        self.project().inner.poll_trailers(cx)
+    }
+}
+
+#[derive(Debug)]
+enum Error {
+    ThroughputBelowMinimum {
+        expected: (u64, Duration),
+        actual: (u64, Duration),
+    },
+    TimeTravel(Box<dyn std::error::Error + Send + Sync>),
+}
+
+impl std::fmt::Display for Error {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::ThroughputBelowMinimum { expected, actual } => {
+                let expected = format_throughput(expected);
+                let actual = format_throughput(actual);
+                write!(
+                    f,
+                    "minimum throughput was specified at {}, but throughput of {} was observed",
+                    expected, actual
+                )
+            }
+            Self::TimeTravel(_) => write!(
+                f,
+                "negative time has elapsed while reading the inner body, this is a bug"
+            ),
+        }
+    }
+}
+
+impl std::error::Error for Error {}
+
+/// Format a given throughput as human-readable bytes per second
+fn format_throughput(throughput: &(u64, Duration)) -> String {
+    let b = throughput.0 as f64;
+    let d = throughput.1.as_secs_f64();
+    // The default float formatting behavior will ensure the a number like 2.000 is rendered as 2
+    // while a number like 0.9982107441748642 will be rendered as 0.9982107441748642. This
+    // multiplication and division will truncate a float to have a precision of no greater than 3.
+    // For example, 0.9982107441748642 would become 0.999. This will fail for very large floats
+    // but should suffice for the numbers we're dealing with.
+    let bytes_per_second = ((b / d) * 1000.0).round() / 1000.0;
+
+    format!("{bytes_per_second} B/s")
+}