smelter_aws_fargate/
client.rs

1//! Functionality for building AWS clients
2use std::marker::PhantomData;
3
4use aws_sdk_ecs::{
5    config::retry::ClassifyRetry,
6    error::ProvideErrorMetadata,
7    operation::{
8        describe_tasks::DescribeTasksError, run_task::RunTaskError, stop_task::StopTaskError,
9    },
10};
11use aws_sdk_s3::config::{interceptors::InterceptorContext, retry::RetryAction};
12use aws_smithy_runtime_api::client::retries::classifiers::SharedRetryClassifier;
13
14/// Build a configuration for the ECS client that will retry failed requests, including
15/// rate-limited requests.
16pub fn build_ecs_client_config(sdk_config: &aws_config::SdkConfig) -> aws_sdk_ecs::config::Builder {
17    let mut builder = aws_sdk_ecs::config::Builder::from(sdk_config);
18
19    let retry_config = aws_sdk_ecs::config::retry::RetryConfig::adaptive()
20        .with_initial_backoff(std::time::Duration::from_secs(2))
21        .with_max_attempts(5)
22        .with_max_backoff(std::time::Duration::from_secs(20));
23
24    builder.set_retry_config(Some(retry_config));
25
26    // This should include all the types of operations that the spawner does. Unfortunate that we
27    // need a separate one for every error type but it works out since there's only a few.
28    builder.push_retry_classifier(SharedRetryClassifier::new(RateLimitErrorClassifier::<
29        DescribeTasksError,
30    >::new()));
31    builder.push_retry_classifier(SharedRetryClassifier::new(RateLimitErrorClassifier::<
32        RunTaskError,
33    >::new()));
34    builder.push_retry_classifier(SharedRetryClassifier::new(RateLimitErrorClassifier::<
35        StopTaskError,
36    >::new()));
37
38    builder
39}
40
41/// An error classifier that triggers retries on ThrottlingException, which the SDK does not yet
42/// handle.
43#[derive(Debug, Default)]
44pub struct RateLimitErrorClassifier<E> {
45    _marker: PhantomData<E>,
46}
47
48impl<E> RateLimitErrorClassifier<E> {
49    /// Create a new error RateLimitErrorClassifier
50    pub fn new() -> Self {
51        Self {
52            _marker: PhantomData,
53        }
54    }
55}
56
57impl<E> ClassifyRetry for RateLimitErrorClassifier<E>
58where
59    // Adding a trait bound for ProvideErrorMetadata allows us to inspect the error code.
60    E: std::error::Error + ProvideErrorMetadata + Send + Sync + 'static,
61{
62    fn classify_retry(&self, ctx: &InterceptorContext) -> RetryAction {
63        // Check for a result
64        let output_or_error = ctx.output_or_error();
65        // Check for an error
66        let error = match output_or_error {
67            Some(Ok(_)) | None => return RetryAction::NoActionIndicated,
68            Some(Err(err)) => err,
69        };
70
71        // Downcast the generic error and extract the code
72        let rate_throttle_error = error
73            .as_operation_error()
74            .and_then(|err| err.downcast_ref::<aws_sdk_ecs::error::SdkError<E>>())
75            .and_then(|err| err.as_service_error())
76            .and_then(|err| err.code())
77            .map(|code| code == "ThrottlingException")
78            .unwrap_or(false);
79
80        if rate_throttle_error {
81            // If the error is a ThrottlingException, return that we should retry.
82            RetryAction::throttling_error()
83        } else {
84            // Otherwise, return that no action is indicated i.e. that this classifier doesn't require a retry.
85            // Another classifier may still classify this response as retryable.
86            RetryAction::NoActionIndicated
87        }
88    }
89
90    fn name(&self) -> &'static str {
91        "ThrottlingException Classifier"
92    }
93}
smelter_aws_fargate/client.rs

smelter_aws_fargate/
client.rs
