Updating integ test for TokenBucket time source

Author: landonxjames

aws/sdk/integration-tests/s3/tests/token_bucket_time_source.rs

@@ -5,45 +5,57 @@
 
 #![cfg(feature = "test-util")]
 
+use aws_config::retry::RetryConfig;
+use aws_sdk_s3::config::retry::RetryPartition;
 use aws_sdk_s3::{config::Region, Client, Config};
 use aws_smithy_async::test_util::ManualTimeSource;
 use aws_smithy_async::time::SharedTimeSource;
 use aws_smithy_http_client::test_util::{ReplayEvent, StaticReplayClient};
 use aws_smithy_runtime::client::retries::TokenBucket;
 use aws_smithy_runtime_api::box_error::BoxError;
-use aws_smithy_runtime_api::client::interceptors::context::BeforeTransmitInterceptorContextMut;
 use aws_smithy_runtime_api::client::interceptors::Intercept;
 use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents;
 use aws_smithy_types::body::SdkBody;
 use aws_smithy_types::config_bag::ConfigBag;
-use std::sync::LazyLock;
+use std::sync::Mutex;
+use std::sync::{Arc, LazyLock};
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
 static THE_TIME: LazyLock<SystemTime> =
     LazyLock::new(|| UNIX_EPOCH + Duration::from_secs(12344321));
 
 #[derive(Debug)]
-struct TimeSourceValidationInterceptor;
+struct TimeSourceValidationInterceptor {
+    current_attempt: Arc<Mutex<u32>>,
+}
 
 impl Intercept for TimeSourceValidationInterceptor {
     fn name(&self) -> &'static str {
         "TimeSourceValidationInterceptor"
     }
 
-    fn modify_before_transmit(
+    fn read_before_attempt(
         &self,
-        _context: &mut BeforeTransmitInterceptorContextMut<'_>,
+        _context: &aws_sdk_s3::config::interceptors::BeforeTransmitInterceptorContextRef<'_>,
         _runtime_components: &RuntimeComponents,
         cfg: &mut ConfigBag,
     ) -> Result<(), BoxError> {
         if let Some(token_bucket) = cfg.load::<TokenBucket>() {
-            let token_bucket_time_source = token_bucket.time_source();
-            let token_time = token_bucket_time_source.now();
+            *self.current_attempt.lock().unwrap() += 1;
 
-            assert_eq!(
-                *THE_TIME, token_time,
-                "Token source should match the configured time source"
-            );
+            if *self.current_attempt.lock().unwrap() == 1 {
+                let last_refill = token_bucket
+                    .last_refill_time_secs()
+                    .load(std::sync::atomic::Ordering::Relaxed);
+                assert_eq!(last_refill, 0);
+            } else if *self.current_attempt.lock().unwrap() == 2 {
+                let last_refill = token_bucket
+                    .last_refill_time_secs()
+                    .load(std::sync::atomic::Ordering::Relaxed);
+                assert_eq!(last_refill, 12344321);
+            } else {
+                panic!("No attempts past the second should happen");
+            }
         }
         Ok(())
     }
@@ -54,22 +66,42 @@ async fn test_token_bucket_gets_time_source_from_config() {
     let time_source = ManualTimeSource::new(*THE_TIME);
     let shared_time_source = SharedTimeSource::new(time_source);
 
-    let http_client = StaticReplayClient::new(vec![ReplayEvent::new(
-        http_1x::Request::builder()
-            .uri("https://www.doesntmatter.com")
-            .body(SdkBody::empty())
-            .unwrap(),
-        http_1x::Response::builder()
-            .status(200)
-            .body(SdkBody::from("<ListBucketResult></ListBucketResult>"))
-            .unwrap(),
-    )]);
+    let http_client = StaticReplayClient::new(vec![
+        ReplayEvent::new(
+            http_1x::Request::builder()
+                .uri("https://www.doesntmatter.com")
+                .body(SdkBody::empty())
+                .unwrap(),
+            http_1x::Response::builder()
+                .status(500)
+                .body(SdkBody::from("This was an error"))
+                .unwrap(),
+        ),
+        ReplayEvent::new(
+            http_1x::Request::builder()
+                .uri("https://www.doesntmatter.com")
+                .body(SdkBody::empty())
+                .unwrap(),
+            http_1x::Response::builder()
+                .status(200)
+                .body(SdkBody::from("<ListBucketResult></ListBucketResult>"))
+                .unwrap(),
+        ),
+    ]);
 
     let config = Config::builder()
         .region(Region::new("us-east-1"))
         .http_client(http_client)
         .time_source(shared_time_source)
-        .interceptor(TimeSourceValidationInterceptor)
+        .interceptor(TimeSourceValidationInterceptor {
+            current_attempt: Arc::new(Mutex::new(0)),
+        })
+        .retry_config(RetryConfig::standard())
+        .retry_partition(
+            RetryPartition::custom("test")
+                .token_bucket(TokenBucket::builder().refill_rate(100.0).build())
+                .build(),
+        )
         .build();
 
     let client = Client::from_conf(config);

rust-runtime/aws-smithy-runtime/src/client/retries/token_bucket.rs

@@ -186,6 +186,7 @@ impl TokenBucket {
     /// Refills tokens based on elapsed time since last refill.
     /// This method implements lazy evaluation - tokens are only calculated when accessed.
     /// Uses a single compare-and-swap to ensure only one thread processes each time window.
+    #[inline]
     fn refill_tokens_based_on_time(&self, time_source: &impl TimeSource) {
         if self.refill_rate > 0.0 {
             let current_time_secs = time_source
@@ -274,6 +275,14 @@ impl TokenBucket {
     pub(crate) fn available_permits(&self) -> usize {
         self.semaphore.available_permits()
     }
+
+    /// Only used in tests
+    #[allow(dead_code)]
+    #[doc(hidden)]
+    #[cfg(any(test, feature = "test-util", feature = "legacy-test-util"))]
+    pub fn last_refill_time_secs(&self) -> Arc<AtomicU64> {
+        self.last_refill_time_secs.clone()
+    }
 }
 
 /// Builder for constructing a `TokenBucket`.
@@ -350,7 +359,7 @@ impl TokenBucketBuilder {
 mod tests {
 
     use super::*;
-    use aws_smithy_async::{test_util::ManualTimeSource, time::SharedTimeSource};
+    use aws_smithy_async::test_util::ManualTimeSource;
     use std::{sync::LazyLock, time::UNIX_EPOCH};
 
     static TIME_SOURCE: LazyLock<ManualTimeSource> =
@@ -907,7 +916,7 @@ mod tests {
 
         // Advance time by 10 seconds
         time_source.advance(Duration::from_secs(10));
-        let shared_time_source = SharedTimeSource::new(time_source);
+        let shared_time_source = aws_smithy_async::time::SharedTimeSource::new(time_source);
 
         // Launch 100 threads that all try to refill simultaneously
         let barrier = Arc::new(Barrier::new(100));