Fix stalled stream detection (#3238)

aws-smithy-types = { path = "../aws-smithy-types", features = ["test-util"] }

futures-util = "0.3.28"

pretty_assertions = "1.4.0"

tokio = { version = "1.25", features = ["macros", "rt", "rt-multi-thread", "test-util"] }

tokio = { version = "1.25", features = ["macros", "rt", "rt-multi-thread", "test-util", "full"] }

tracing-subscriber = { version = "0.3.16", features = ["env-filter"] }

tracing-test = "0.2.1"

hyper_0_14 = { package = "hyper", version = "0.14.27",features = ["client", "server", "tcp", "http1", "http2"] }

[package.metadata.docs.rs]

all-features = true

/// Options for a [`MinimumThroughputBody`].

pub mod options;

pub use throughput::Throughput;

mod throughput;

use aws_smithy_async::rt::sleep::Sleep;

use aws_smithy_runtime_api::shared::IntoShared;

use options::MinimumThroughputBodyOptions;

use std::fmt;

use throughput::{Throughput, ThroughputLogs};

use std::time::SystemTime;

use throughput::ThroughputLogs;

pin_project_lite::pin_project! {

    /// A body-wrapping type that ensures data is being streamed faster than some lower limit.

    }

}

const SIZE_OF_ONE_LOG: usize = std::mem::size_of::<(std::time::SystemTime, u64)>(); // 24 bytes per log

const SIZE_OF_ONE_LOG: usize = std::mem::size_of::<(SystemTime, u64)>(); // 24 bytes per log

const NUMBER_OF_LOGS_IN_ONE_KB: f64 = 1024.0 / SIZE_OF_ONE_LOG as f64;

impl<B> MinimumThroughputBody<B> {

                // Never keep more than 10KB of logs in memory. This currently

                // equates to 426 logs.

                (NUMBER_OF_LOGS_IN_ONE_KB * 10.0) as usize,

                options.minimum_throughput().per_time_elapsed(),

            ),

            async_sleep: async_sleep.into_shared(),

            time_source: time_source.into_shared(),

        mut self: Pin<&mut Self>,

        cx: &mut Context<'_>,

    ) -> Poll<Option<Result<Self::Data, Self::Error>>> {

        // this code is called quite frequently in production—one every millisecond or so when downloading

        // a stream. However, SystemTime::now is on the order of nanoseconds

        let now = self.time_source.now();

        // Attempt to read the data from the inner body, then update the

        // throughput logs.

        let mut this = self.as_mut().project();

        let poll_res = match this.inner.poll_data(cx) {

            Poll::Ready(Some(Ok(bytes))) => {

                tracing::trace!("received data: {}", bytes.len());

                this.throughput_logs.push((now, bytes.len() as u64));

                Poll::Ready(Some(Ok(bytes)))

            }

            Poll::Pending => Poll::Pending,

            Poll::Pending => {

                tracing::trace!("received poll pending");

                this.throughput_logs.push((now, 0));

                Poll::Pending

            }

            // If we've read all the data or an error occurred, then return that result.

            res => return res,

        };

        // Push a start log if we haven't already done so.

        if this.throughput_logs.is_empty() {

            this.throughput_logs.push((now, 0));

        }

        // Check the sleep future to see if it needs refreshing.

        let mut sleep_fut = this.sleep_fut.take().unwrap_or_else(|| {

            this.async_sleep

                .sleep(this.options.minimum_throughput().per_time_elapsed())

        });

        let mut sleep_fut = this

            .sleep_fut

            .take()

            .unwrap_or_else(|| this.async_sleep.sleep(this.options.check_interval()));

        if let Poll::Ready(()) = pin!(&mut sleep_fut).poll(cx) {

            tracing::trace!("sleep future triggered—triggering a wakeup");

            // Whenever the sleep future expires, we replace it.

            sleep_fut = this

                .async_sleep

                .sleep(this.options.minimum_throughput().per_time_elapsed());

            sleep_fut = this.async_sleep.sleep(this.options.check_interval());

            // We also schedule a wake up for current task to ensure that

            // it gets polled at least one more time.

            cx.waker().wake_by_ref();

        };

        this.sleep_fut.replace(sleep_fut);

        let calculated_tpt = match this

            .throughput_logs

            .calculate_throughput(now, this.options.check_window())

        {

            Some(tpt) => tpt,

            None => {

                tracing::trace!("calculated throughput is None!");

                return poll_res;

            }

        };

        tracing::trace!(

            "calculated throughput {:?} (window: {:?})",

            calculated_tpt,

            this.options.check_window()

        );

        // Calculate the current throughput and emit an error if it's too low and

        // the grace period has elapsed.

        let actual_throughput = this.throughput_logs.calculate_throughput(now);

        let is_below_minimum_throughput = actual_throughput

            .map(|t| t < this.options.minimum_throughput())

            .unwrap_or_default();

        let is_below_minimum_throughput = calculated_tpt <= this.options.minimum_throughput();

        if is_below_minimum_throughput {

            // Check the grace period future to see if it needs creating.

            tracing::trace!(

                in_grace_period = this.grace_period_fut.is_some(),

                observed_throughput = ?calculated_tpt,

                minimum_throughput = ?this.options.minimum_throughput(),

                "below minimum throughput"

            );

            let mut grace_period_fut = this

                .grace_period_fut

                .take()

                // The grace period has ended!

                return Poll::Ready(Some(Err(Box::new(Error::ThroughputBelowMinimum {

                    expected: self.options.minimum_throughput(),

                    actual: actual_throughput.unwrap(),

                    actual: calculated_tpt,

                }))));

            };

            this.grace_period_fut.replace(grace_period_fut);

}

// These tests use `hyper::body::Body::wrap_stream`

#[cfg(all(test, feature = "connector-hyper-0-14-x"))]

#[cfg(all(test, feature = "connector-hyper-0-14-x", feature = "test-util"))]

mod test {

    use super::{super::Throughput, Error, MinimumThroughputBody};

    use crate::client::http::body::minimum_throughput::options::MinimumThroughputBodyOptions;

    use crate::test_util::capture_test_logs::capture_test_logs;

    use aws_smithy_async::rt::sleep::AsyncSleep;

    use aws_smithy_async::test_util::{instant_time_and_sleep, InstantSleep, ManualTimeSource};

    use aws_smithy_types::body::SdkBody;

    use pretty_assertions::assert_eq;

    use std::convert::Infallible;

    use std::error::Error as StdError;

    use std::future::Future;

    use std::pin::Pin;

    use std::future::{poll_fn, Future};

    use std::pin::{pin, Pin};

    use std::task::{Context, Poll};

    use std::time::{Duration, UNIX_EPOCH};

            eight_byte_per_second_stream_with_minimum_throughput_timeout(minimum_throughput);

        let expected_err = Error::ThroughputBelowMinimum {

            expected: minimum_throughput,

            actual: Throughput::new(8.889, Duration::from_secs(1)),

            actual: Throughput::new(8, Duration::from_secs(1)),

        };

        match res.await {

            Ok(_) => {

    #[tokio::test]

    async fn test_throughput_timeout_less_than() {

        let minimum_throughput = Throughput::new_bytes_per_second(9.0);

        let minimum_throughput = Throughput::new_bytes_per_second(9);

        expect_error(minimum_throughput).await;

    }

    #[tokio::test]

    async fn test_throughput_timeout_equal_to() {

        let minimum_throughput = Throughput::new(32.0, Duration::from_secs(4));

        let (_guard, _) = capture_test_logs();

        // a tiny bit less. To capture 0-throughput properly, we need to allow 0 to be 0

        let minimum_throughput = Throughput::new(31, Duration::from_secs(4));

        expect_success(minimum_throughput).await;

    }

    #[tokio::test]

    async fn test_throughput_timeout_greater_than() {

        let minimum_throughput = Throughput::new(20.0, Duration::from_secs(3));

        let minimum_throughput = Throughput::new(20, Duration::from_secs(3));

        expect_success(minimum_throughput).await;

    }

    // A multiplier for the sine wave amplitude; Chosen arbitrarily.

    const BYTE_COUNT_UPPER_LIMIT: f64 = 100.0;

    const BYTE_COUNT_UPPER_LIMIT: u64 = 1000;

    fn create_shrinking_sine_wave_stream(

    /// emits 1000B/S for 5 seconds then suddenly stops

    fn sudden_stop(

        async_sleep: impl AsyncSleep + Clone,

    ) -> impl futures_util::Stream<Item = Result<Bytes, Infallible>> {

        let sleep_dur = Duration::from_millis(50);

        fastrand::seed(0);

        futures_util::stream::unfold(1, move |i| {

            let async_sleep = async_sleep.clone();

            async move {

                if i > 255 {

                    None

                let number_seconds = (i * sleep_dur).as_secs_f64();

                async_sleep.sleep(sleep_dur).await;

                if number_seconds > 5.0 {

                    Some((Result::<Bytes, Infallible>::Ok(Bytes::new()), i + 1))

                } else {

                    let byte_count = (i as f64).sin().floor().abs() + 1.0 / (i as f64 + 1.0);

                    let byte_count = (byte_count * BYTE_COUNT_UPPER_LIMIT) as u64;

                    let mut bytes = BytesMut::new();

                    bytes.put_u8(i as u8);

                    if byte_count > 0 {

                        for _ in 0..byte_count {

                    let bytes_per_segment =

                        (BYTE_COUNT_UPPER_LIMIT as f64) * sleep_dur.as_secs_f64();

                    for _ in 0..bytes_per_segment as usize {

                        bytes.put_u8(0)

                    }

                    }

                    async_sleep.sleep(Duration::from_secs(1)).await;

                    Some((Result::<Bytes, Infallible>::Ok(bytes.into()), i + 1))

                }

            }

    }

    #[tokio::test]

    async fn test_throughput_timeout_shrinking_sine_wave() {

    async fn test_stalled_stream_detection() {

        test_suddenly_stopping_stream(0, Duration::from_secs(6)).await

    }

    #[tokio::test]

    async fn test_slow_stream_detection() {

        test_suddenly_stopping_stream(BYTE_COUNT_UPPER_LIMIT / 2, Duration::from_secs_f64(5.50))

            .await

    }

    #[tokio::test]

    async fn test_check_interval() {

        let (_guard, _) = capture_test_logs();

        let (ts, sleep) = instant_time_and_sleep(UNIX_EPOCH);

        let mut body = MinimumThroughputBody::new(

            ts,

            sleep.clone(),

            NeverBody,

            MinimumThroughputBodyOptions::builder()

                .check_interval(Duration::from_millis(1234))

                .grace_period(Duration::from_millis(456))

                .build(),

        );

        let mut body = pin!(body);

        let _ = poll_fn(|cx| body.as_mut().poll_data(cx)).await;

        assert_eq!(

            sleep.logs(),

            vec![

                // sleep, by second sleep we know we have no data, then the grace period

                Duration::from_millis(1234),

                Duration::from_millis(1234),

                Duration::from_millis(456)

            ]

        );

    }

    async fn test_suddenly_stopping_stream(throughput_limit: u64, time_until_timeout: Duration) {

        let (_guard, _) = capture_test_logs();

        let options = MinimumThroughputBodyOptions::builder()

            // Minimum throughput per second will be approx. half of the BYTE_COUNT_UPPER_LIMIT.

            .minimum_throughput(Throughput::new_bytes_per_second(

                BYTE_COUNT_UPPER_LIMIT / 2.0 + 2.0,

            ))

            .minimum_throughput(Throughput::new_bytes_per_second(throughput_limit))

            .build();

        let (time_source, async_sleep) = instant_time_and_sleep(UNIX_EPOCH);

        let time_clone = time_source.clone();

        let stream = create_shrinking_sine_wave_stream(async_sleep.clone());

        let stream = sudden_stop(async_sleep.clone());

        let body = ByteStream::new(SdkBody::from_body_0_4(hyper_0_14::body::Body::wrap_stream(

            stream,

        )));

        match res.await {

            Ok(_res) => {

                assert_eq!(255.0, time_source.seconds_since_unix_epoch());

                assert_eq!(Duration::from_secs(255), async_sleep.total_duration());

                panic!("stream should have timed out");

            }

            Err(err) => panic!(

                "test stopped after {:?} due to {}",

            Err(err) => {

                dbg!(err);

                assert_eq!(

                    async_sleep.total_duration(),

                DisplayErrorContext(err.source().unwrap())

            ),

                    time_until_timeout,

                    "With throughput limit {:?} expected timeout after {:?} (stream starts sending 0's at 5 seconds.",

                    throughput_limit, time_until_timeout

                );

            }

        }

    }

}

    /// If this is set to a positive value, whenever throughput is below the minimum throughput,

    /// a timer is started. If the timer expires before throughput rises above the minimum,

    /// an error is emitted.

    ///

    /// It is recommended to set this to a small value (e.g. 200ms) to avoid issues during

    /// stream-startup.

    grace_period: Duration,

    /// The interval at which the throughput is checked.

    check_interval: Duration,

    /// The period of time to consider when computing the throughput

    ///

    /// This SHOULD be longer than the check interval, or stuck-streams may evade detection.

    check_window: Duration,

}

impl MinimumThroughputBodyOptions {

        self.minimum_throughput

    }

    pub(crate) fn check_window(&self) -> Duration {

        self.check_window

    }

    /// The rate at which the throughput is checked.

    ///

    /// The actual rate throughput is checked may be higher than this value,

            minimum_throughput: DEFAULT_MINIMUM_THROUGHPUT,

            grace_period: DEFAULT_GRACE_PERIOD,

            check_interval: DEFAULT_CHECK_INTERVAL,

            check_window: DEFAULT_CHECK_WINDOW,

        }

    }

}

    grace_period: Option<Duration>,

}

const DEFAULT_CHECK_INTERVAL: Duration = Duration::from_secs(1);

const DEFAULT_CHECK_INTERVAL: Duration = Duration::from_millis(500);

const DEFAULT_GRACE_PERIOD: Duration = Duration::from_secs(0);

const DEFAULT_MINIMUM_THROUGHPUT: Throughput = Throughput {

    bytes_read: 1.0,

    bytes_read: 1,

    per_time_elapsed: Duration::from_secs(1),

};

const DEFAULT_CHECK_WINDOW: Duration = Duration::from_secs(1);

impl MinimumThroughputBodyOptionsBuilder {

    /// Create a new `MinimumThroughputBodyOptionsBuilder`.

    pub fn new() -> Self {

                .minimum_throughput

                .unwrap_or(DEFAULT_MINIMUM_THROUGHPUT),

            check_interval: self.check_interval.unwrap_or(DEFAULT_CHECK_INTERVAL),

            check_window: DEFAULT_CHECK_WINDOW,

        }

    }

}

            grace_period: value.grace_period(),

            minimum_throughput: DEFAULT_MINIMUM_THROUGHPUT,

            check_interval: DEFAULT_CHECK_INTERVAL,

            check_window: DEFAULT_CHECK_WINDOW,

        }

    }

}