aws_smithy_runtime_api/client/
retries.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0
 */

//! Retry handling and token bucket.
//!
//! This code defines when and how failed requests should be retried. It also defines the behavior
//! used to limit the rate that requests are sent.

pub mod classifiers;

use crate::box_error::BoxError;
use crate::client::interceptors::context::InterceptorContext;
use crate::client::runtime_components::sealed::ValidateConfig;
use crate::client::runtime_components::RuntimeComponents;
use aws_smithy_types::config_bag::{ConfigBag, Storable, StoreReplace};
use std::fmt;
use std::sync::Arc;
use std::time::Duration;

use crate::impl_shared_conversions;
pub use aws_smithy_types::retry::ErrorKind;
#[cfg(feature = "test-util")]
pub use test_util::AlwaysRetry;

#[derive(Debug, Clone, PartialEq, Eq)]
/// An answer to the question "should I make a request attempt?"
pub enum ShouldAttempt {
    /// Yes, an attempt should be made
    Yes,
    /// No, no attempt should be made
    No,
    /// Yes, an attempt should be made, but only after the given amount of time has passed
    YesAfterDelay(Duration),
}

#[cfg(feature = "test-util")]
impl ShouldAttempt {
    /// Returns the delay duration if this is a `YesAfterDelay` variant.
    pub fn expect_delay(self) -> Duration {
        match self {
            ShouldAttempt::YesAfterDelay(delay) => delay,
            _ => panic!("Expected this to be the `YesAfterDelay` variant but it was the `{self:?}` variant instead"),
        }
    }

    /// If this isn't a `No` variant, panic.
    pub fn expect_no(self) {
        if ShouldAttempt::No == self {
            return;
        }

        panic!("Expected this to be the `No` variant but it was the `{self:?}` variant instead");
    }
}

impl_shared_conversions!(convert SharedRetryStrategy from RetryStrategy using SharedRetryStrategy::new);

/// Decider for whether or not to attempt a request, and when.
///
/// The orchestrator consults the retry strategy every time before making a request.
/// This includes the initial request, and any retry attempts thereafter. The
/// orchestrator will retry indefinitely (until success) if the retry strategy
/// always returns `ShouldAttempt::Yes` from `should_attempt_retry`.
pub trait RetryStrategy: Send + Sync + fmt::Debug {
    /// Decides if the initial attempt should be made.
    fn should_attempt_initial_request(
        &self,
        runtime_components: &RuntimeComponents,
        cfg: &ConfigBag,
    ) -> Result<ShouldAttempt, BoxError>;

    /// Decides if a retry should be done.
    ///
    /// The previous attempt's output or error are provided in the
    /// [`InterceptorContext`] when this is called.
    ///
    /// `ShouldAttempt::YesAfterDelay` can be used to add a backoff time.
    fn should_attempt_retry(
        &self,
        context: &InterceptorContext,
        runtime_components: &RuntimeComponents,
        cfg: &ConfigBag,
    ) -> Result<ShouldAttempt, BoxError>;
}

/// A shared retry strategy.
#[derive(Clone, Debug)]
pub struct SharedRetryStrategy(Arc<dyn RetryStrategy>);

impl SharedRetryStrategy {
    /// Creates a new [`SharedRetryStrategy`] from a retry strategy.
    pub fn new(retry_strategy: impl RetryStrategy + 'static) -> Self {
        Self(Arc::new(retry_strategy))
    }
}

impl RetryStrategy for SharedRetryStrategy {
    fn should_attempt_initial_request(
        &self,
        runtime_components: &RuntimeComponents,
        cfg: &ConfigBag,
    ) -> Result<ShouldAttempt, BoxError> {
        self.0
            .should_attempt_initial_request(runtime_components, cfg)
    }

    fn should_attempt_retry(
        &self,
        context: &InterceptorContext,
        runtime_components: &RuntimeComponents,
        cfg: &ConfigBag,
    ) -> Result<ShouldAttempt, BoxError> {
        self.0
            .should_attempt_retry(context, runtime_components, cfg)
    }
}

impl ValidateConfig for SharedRetryStrategy {}

/// A type to track the number of requests sent by the orchestrator for a given operation.
///
/// `RequestAttempts` is added to the `ConfigBag` by the orchestrator,
/// and holds the current attempt number.
#[derive(Debug, Clone, Copy)]
pub struct RequestAttempts {
    attempts: u32,
}

impl RequestAttempts {
    /// Creates a new [`RequestAttempts`] with the given number of attempts.
    pub fn new(attempts: u32) -> Self {
        Self { attempts }
    }

    /// Returns the number of attempts.
    pub fn attempts(&self) -> u32 {
        self.attempts
    }
}

impl From<u32> for RequestAttempts {
    fn from(attempts: u32) -> Self {
        Self::new(attempts)
    }
}

impl From<RequestAttempts> for u32 {
    fn from(value: RequestAttempts) -> Self {
        value.attempts()
    }
}

impl Storable for RequestAttempts {
    type Storer = StoreReplace<Self>;
}

#[cfg(feature = "test-util")]
mod test_util {
    use super::ErrorKind;
    use crate::client::interceptors::context::InterceptorContext;
    use crate::client::retries::classifiers::{ClassifyRetry, RetryAction};

    /// A retry classifier for testing purposes. This classifier always returns
    /// `Some(RetryAction::Error(ErrorKind))` where `ErrorKind` is the value provided when creating
    /// this classifier.
    #[derive(Debug)]
    pub struct AlwaysRetry(pub ErrorKind);

    impl ClassifyRetry for AlwaysRetry {
        fn classify_retry(&self, error: &InterceptorContext) -> RetryAction {
            tracing::debug!("Retrying error {:?} as an {:?}", error, self.0);
            RetryAction::retryable_error(self.0)
        }

        fn name(&self) -> &'static str {
            "Always Retry"
        }
    }
}