Refactor TokenBucket to not carry a TimeSource

Author: landonxjames

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientGenerator.kt

@@ -278,7 +278,6 @@ private fun baseClientRuntimePluginsFn(
                             #{DefaultPluginParams}::new()
                                 .with_retry_partition_name(default_retry_partition)
                                 .with_behavior_version(config.behavior_version.expect(${behaviorVersionError.dq()}))
-                                .with_time_source(config.runtime_components.time_source().unwrap_or_default())
                         ))
                         // user config
                         .with_client_plugin(

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

@@ -15,7 +15,7 @@ use crate::client::retries::strategy::standard::TokenBucketProvider;
 use crate::client::retries::strategy::StandardRetryStrategy;
 use crate::client::retries::RetryPartition;
 use aws_smithy_async::rt::sleep::default_async_sleep;
-use aws_smithy_async::time::{SharedTimeSource, SystemTimeSource, TimeSource};
+use aws_smithy_async::time::SystemTimeSource;
 use aws_smithy_runtime_api::box_error::BoxError;
 use aws_smithy_runtime_api::client::behavior_version::BehaviorVersion;
 use aws_smithy_runtime_api::client::http::SharedHttpClient;
@@ -126,7 +126,6 @@ pub fn default_time_source_plugin() -> Option<SharedRuntimePlugin> {
 }
 
 /// Runtime plugin that sets the default retry strategy, config (disabled), and partition.
-#[deprecated = "Use default_retry_config_plugin_v2 to get a TokenBucket that respects the user provided TimeSource."]
 pub fn default_retry_config_plugin(
     default_partition_name: impl Into<Cow<'static, str>>,
 ) -> Option<SharedRuntimePlugin> {
@@ -138,43 +137,7 @@ pub fn default_retry_config_plugin(
                 .with_config_validator(SharedConfigValidator::base_client_config_fn(
                     validate_retry_config,
                 ))
-                .with_interceptor(TokenBucketProvider::new(
-                    retry_partition.clone(),
-                    SharedTimeSource::default(), // Replicates previous behavior
-                ))
-        })
-        .with_config(layer("default_retry_config", |layer| {
-            layer.store_put(RetryConfig::disabled());
-            layer.store_put(retry_partition);
-        }))
-        .into_shared(),
-    )
-}
-
-/// Runtime plugin that sets the default retry strategy, config (disabled), and partition.
-pub fn default_retry_config_plugin_v2(
-    default_plugin_params: &DefaultPluginParams,
-) -> Option<SharedRuntimePlugin> {
-    let retry_partition = RetryPartition::new(
-        default_plugin_params
-            .retry_partition_name()
-            .clone()
-            .expect("retry_partition_name is required"),
-    );
-    Some(
-        default_plugin("default_retry_config_plugin", |components| {
-            components
-                .with_retry_strategy(Some(StandardRetryStrategy::new()))
-                .with_config_validator(SharedConfigValidator::base_client_config_fn(
-                    validate_retry_config,
-                ))
-                .with_interceptor(TokenBucketProvider::new(
-                    retry_partition.clone(),
-                    default_plugin_params
-                        .time_source
-                        .clone()
-                        .unwrap_or_default(),
-                ))
+                .with_interceptor(TokenBucketProvider::new(retry_partition.clone()))
         })
         .with_config(layer("default_retry_config", |layer| {
             layer.store_put(RetryConfig::disabled());
@@ -325,7 +288,6 @@ fn validate_stalled_stream_protection_config(
 pub struct DefaultPluginParams {
     retry_partition_name: Option<Cow<'static, str>>,
     behavior_version: Option<BehaviorVersion>,
-    time_source: Option<SharedTimeSource>,
 }
 
 impl DefaultPluginParams {
@@ -340,32 +302,11 @@ impl DefaultPluginParams {
         self
     }
 
-    /// Gets the retry partition name.
-    pub fn retry_partition_name(&self) -> &Option<Cow<'static, str>> {
-        &self.retry_partition_name
-    }
-
     /// Sets the behavior major version.
     pub fn with_behavior_version(mut self, version: BehaviorVersion) -> Self {
         self.behavior_version = Some(version);
         self
     }
-
-    /// Gets the behavior major version.
-    pub fn behavior_version(&self) -> &Option<BehaviorVersion> {
-        &self.behavior_version
-    }
-
-    /// Sets the time_source.
-    pub fn with_time_source(mut self, time_source: impl TimeSource + 'static) -> Self {
-        self.time_source = Some(SharedTimeSource::new(time_source));
-        self
-    }
-
-    /// Gets the time_source.
-    pub fn time_source(&self) -> &Option<SharedTimeSource> {
-        &self.time_source
-    }
 }
 
 /// All default plugins.
@@ -379,7 +320,11 @@ pub fn default_plugins(
     [
         default_http_client_plugin_v2(behavior_version),
         default_identity_cache_plugin(),
-        default_retry_config_plugin_v2(&params),
+        default_retry_config_plugin(
+            params
+                .retry_partition_name
+                .expect("retry_partition_name is required"),
+        ),
         default_sleep_impl_plugin(),
         default_time_source_plugin(),
         default_timeout_config_plugin(),

rust-runtime/aws-smithy-runtime/src/client/retries/strategy/standard.rs

@@ -6,7 +6,6 @@
 use std::sync::Mutex;
 use std::time::{Duration, SystemTime};
 
-use aws_smithy_async::time::SharedTimeSource;
 use tokio::sync::OwnedSemaphorePermit;
 use tracing::{debug, trace};
 
@@ -257,7 +256,10 @@ impl RetryStrategy for StandardRetryStrategy {
 
         //  acquire permit for retry
         let error_kind = error_kind.expect("result was classified retryable");
-        match token_bucket.acquire(&error_kind) {
+        match token_bucket.acquire(
+            &error_kind,
+            &runtime_components.time_source().unwrap_or_default(),
+        ) {
             Some(permit) => self.set_retry_permit(permit),
             None => {
                 debug!("attempt #{request_attempts} failed with {error_kind:?}; However, not enough retry quota is available for another attempt so no retry will be attempted.");
@@ -372,12 +374,8 @@ impl TokenBucketProvider {
     ///
     /// NOTE: This partition should be the one used for every operation on a client
     /// unless config is overridden.
-    pub(crate) fn new(default_partition: RetryPartition, time_source: SharedTimeSource) -> Self {
-        let token_bucket = TOKEN_BUCKET.get_or_init(default_partition.clone(), || {
-            let mut tb = TokenBucket::default();
-            tb.update_time_source(time_source);
-            tb
-        });
+    pub(crate) fn new(default_partition: RetryPartition) -> Self {
+        let token_bucket = TOKEN_BUCKET.get_or_init_default(default_partition.clone());
         Self {
             default_partition,
             token_bucket,
@@ -393,7 +391,7 @@ impl Intercept for TokenBucketProvider {
     fn modify_before_retry_loop(
         &self,
         _context: &mut BeforeTransmitInterceptorContextMut<'_>,
-        runtime_components: &RuntimeComponents,
+        _runtime_components: &RuntimeComponents,
         cfg: &mut ConfigBag,
     ) -> Result<(), BoxError> {
         let retry_partition = cfg.load::<RetryPartition>().expect("set in default config");
@@ -407,11 +405,7 @@ impl Intercept for TokenBucketProvider {
                     // avoid contention on the global lock
                     self.token_bucket.clone()
                 } else {
-                    TOKEN_BUCKET.get_or_init(retry_partition.clone(), || {
-                        let mut tb = TokenBucket::default();
-                        tb.update_time_source(runtime_components.time_source().unwrap_or_default());
-                        tb
-                    })
+                    TOKEN_BUCKET.get_or_init_default(retry_partition.clone())
                 }
             }
             RetryPartitionInner::Custom { token_bucket, .. } => token_bucket.clone(),