Skip to main content

cloud_sdk/rate_limit/
decision.rs

1use super::{DelaySeconds, QuotaBuckets, QuotaReset, RetryAfter, WallClockTimestamp};
2
3/// Handling for timestamps at or before the caller's observation time.
4#[derive(Clone, Copy, Debug, Eq, PartialEq)]
5pub enum PastTimestampPolicy {
6    /// Treat a past reset as an immediate zero-second delay.
7    Immediate,
8    /// Reject stale absolute metadata.
9    Reject,
10}
11
12/// Handling for a computed delay above the caller's maximum.
13#[derive(Clone, Copy, Debug, Eq, PartialEq)]
14pub enum ExcessDelayPolicy {
15    /// Clamp the decision to the caller's maximum.
16    Clamp,
17    /// Reject the response metadata.
18    Reject,
19}
20
21/// Conflict policy when `Retry-After` and exhausted quota buckets disagree.
22#[derive(Clone, Copy, Debug, Eq, PartialEq)]
23pub enum DelayConflictPolicy {
24    /// Follow the standard `Retry-After` instruction.
25    RetryAfterPrecedence,
26    /// Use the longest advertised delay.
27    Longest,
28    /// Reject unequal delay instructions.
29    RejectMismatch,
30}
31
32/// Caller-owned quota delay policy.
33#[derive(Clone, Copy, Debug, Eq, PartialEq)]
34pub struct QuotaDelayPolicy {
35    maximum: DelaySeconds,
36    past: PastTimestampPolicy,
37    excess: ExcessDelayPolicy,
38    conflict: DelayConflictPolicy,
39}
40
41impl QuotaDelayPolicy {
42    /// Creates an explicit bounded delay policy.
43    #[must_use]
44    pub const fn new(
45        maximum: DelaySeconds,
46        past: PastTimestampPolicy,
47        excess: ExcessDelayPolicy,
48        conflict: DelayConflictPolicy,
49    ) -> Self {
50        Self {
51            maximum,
52            past,
53            excess,
54            conflict,
55        }
56    }
57
58    /// Returns the caller's maximum accepted delay.
59    #[must_use]
60    pub const fn maximum(self) -> DelaySeconds {
61        self.maximum
62    }
63}
64
65/// Metadata source selected for a delay decision.
66#[derive(Clone, Copy, Debug, Eq, PartialEq)]
67pub enum DelaySource {
68    /// Standard `Retry-After` metadata.
69    RetryAfter,
70    /// One or more exhausted provider quota buckets.
71    ProviderQuota,
72    /// Both sources agreed exactly.
73    Both,
74}
75
76/// Pure bounded delay decision. The caller owns sleeping and clock acquisition.
77#[derive(Clone, Copy, Debug, Eq, PartialEq)]
78pub struct DelayDecision {
79    delay: DelaySeconds,
80    source: DelaySource,
81    clamped: bool,
82}
83
84impl DelayDecision {
85    /// Returns the selected delay.
86    #[must_use]
87    pub const fn delay(self) -> DelaySeconds {
88        self.delay
89    }
90    /// Returns the selected metadata source.
91    #[must_use]
92    pub const fn source(self) -> DelaySource {
93        self.source
94    }
95    /// Reports whether the caller maximum shortened the selected delay.
96    #[must_use]
97    pub const fn was_clamped(self) -> bool {
98        self.clamped
99    }
100}
101
102/// Derives a bounded delay without reading a clock or sleeping.
103pub fn decide_delay(
104    buckets: &QuotaBuckets,
105    retry_after: Option<RetryAfter>,
106    now: WallClockTimestamp,
107    previous_now: Option<WallClockTimestamp>,
108    policy: QuotaDelayPolicy,
109) -> Result<Option<DelayDecision>, DelayDecisionError> {
110    if previous_now.is_some_and(|previous| now < previous) {
111        return Err(DelayDecisionError::ClockRollback);
112    }
113    let quota_delay = quota_delay(buckets, now, policy.past)?;
114    let retry_delay = retry_after
115        .map(|value| retry_delay(value, now, policy.past))
116        .transpose()?;
117    let selected = match (retry_delay, quota_delay) {
118        (None, None) => return Ok(None),
119        (Some(delay), None) => (delay, DelaySource::RetryAfter),
120        (None, Some(delay)) => (delay, DelaySource::ProviderQuota),
121        (Some(retry), Some(quota)) if retry == quota => (retry, DelaySource::Both),
122        (Some(retry), Some(quota)) => match policy.conflict {
123            DelayConflictPolicy::RetryAfterPrecedence => (retry, DelaySource::RetryAfter),
124            DelayConflictPolicy::Longest => (core::cmp::max(retry, quota), DelaySource::Both),
125            DelayConflictPolicy::RejectMismatch => {
126                return Err(DelayDecisionError::ConflictingMetadata);
127            }
128        },
129    };
130    apply_maximum(selected.0, selected.1, policy)
131}
132
133fn quota_delay(
134    buckets: &QuotaBuckets,
135    now: WallClockTimestamp,
136    past: PastTimestampPolicy,
137) -> Result<Option<DelaySeconds>, DelayDecisionError> {
138    let mut selected: Option<DelaySeconds> = None;
139    for bucket in buckets.iter().filter(|bucket| bucket.is_exhausted()) {
140        let delay = match bucket.reset() {
141            QuotaReset::After(delay) => delay,
142            QuotaReset::At(timestamp) => absolute_delay(timestamp, now, past)?,
143            QuotaReset::Unknown => return Err(DelayDecisionError::ExhaustedBucketResetUnknown),
144        };
145        selected = Some(selected.map_or(delay, |current| core::cmp::max(current, delay)));
146    }
147    Ok(selected)
148}
149
150fn retry_delay(
151    retry_after: RetryAfter,
152    now: WallClockTimestamp,
153    past: PastTimestampPolicy,
154) -> Result<DelaySeconds, DelayDecisionError> {
155    match retry_after {
156        RetryAfter::Delay(delay) => Ok(delay),
157        RetryAfter::HttpDate(date) => {
158            let now =
159                i64::try_from(now.get()).map_err(|_| DelayDecisionError::TimestampOverflow)?;
160            if date.epoch_seconds() <= now {
161                return past_delay(past);
162            }
163            let difference = date
164                .epoch_seconds()
165                .checked_sub(now)
166                .and_then(|value| u64::try_from(value).ok())
167                .ok_or(DelayDecisionError::TimestampOverflow)?;
168            Ok(DelaySeconds::new(difference))
169        }
170    }
171}
172
173fn absolute_delay(
174    timestamp: WallClockTimestamp,
175    now: WallClockTimestamp,
176    past: PastTimestampPolicy,
177) -> Result<DelaySeconds, DelayDecisionError> {
178    if timestamp <= now {
179        return past_delay(past);
180    }
181    let delay = timestamp
182        .get()
183        .checked_sub(now.get())
184        .ok_or(DelayDecisionError::TimestampOverflow)?;
185    Ok(DelaySeconds::new(delay))
186}
187
188fn past_delay(past: PastTimestampPolicy) -> Result<DelaySeconds, DelayDecisionError> {
189    match past {
190        PastTimestampPolicy::Immediate => Ok(DelaySeconds::new(0)),
191        PastTimestampPolicy::Reject => Err(DelayDecisionError::PastTimestamp),
192    }
193}
194
195fn apply_maximum(
196    delay: DelaySeconds,
197    source: DelaySource,
198    policy: QuotaDelayPolicy,
199) -> Result<Option<DelayDecision>, DelayDecisionError> {
200    if delay <= policy.maximum {
201        return Ok(Some(DelayDecision {
202            delay,
203            source,
204            clamped: false,
205        }));
206    }
207    match policy.excess {
208        ExcessDelayPolicy::Clamp => Ok(Some(DelayDecision {
209            delay: policy.maximum,
210            source,
211            clamped: true,
212        })),
213        ExcessDelayPolicy::Reject => Err(DelayDecisionError::MaximumExceeded),
214    }
215}
216
217/// Failure to derive an unambiguous bounded delay.
218#[derive(Clone, Copy, Debug, Eq, PartialEq)]
219pub enum DelayDecisionError {
220    /// The supplied wall clock moved behind the previous observation.
221    ClockRollback,
222    /// An exhausted bucket did not expose actionable reset metadata.
223    ExhaustedBucketResetUnknown,
224    /// An absolute timestamp was at or before the current observation.
225    PastTimestamp,
226    /// An absolute date or difference was outside the supported range.
227    TimestampOverflow,
228    /// Retry and provider quota metadata disagreed under strict policy.
229    ConflictingMetadata,
230    /// The selected delay exceeded the caller maximum under reject policy.
231    MaximumExceeded,
232}
233
234impl_static_error!(DelayDecisionError,
235    Self::ClockRollback => "wall clock moved behind the previous observation",
236    Self::ExhaustedBucketResetUnknown => "exhausted quota bucket has no reset instruction",
237    Self::PastTimestamp => "quota reset timestamp is not in the future",
238    Self::TimestampOverflow => "quota reset timestamp cannot be represented",
239    Self::ConflictingMetadata => "Retry-After and provider quota metadata conflict",
240    Self::MaximumExceeded => "quota delay exceeds the caller maximum",
241);