1use super::{DelaySeconds, QuotaBuckets, QuotaReset, RetryAfter, WallClockTimestamp};
2
3#[derive(Clone, Copy, Debug, Eq, PartialEq)]
5pub enum PastTimestampPolicy {
6 Immediate,
8 Reject,
10}
11
12#[derive(Clone, Copy, Debug, Eq, PartialEq)]
14pub enum ExcessDelayPolicy {
15 Clamp,
17 Reject,
19}
20
21#[derive(Clone, Copy, Debug, Eq, PartialEq)]
23pub enum DelayConflictPolicy {
24 RetryAfterPrecedence,
26 Longest,
28 RejectMismatch,
30}
31
32#[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 #[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 #[must_use]
60 pub const fn maximum(self) -> DelaySeconds {
61 self.maximum
62 }
63}
64
65#[derive(Clone, Copy, Debug, Eq, PartialEq)]
67pub enum DelaySource {
68 RetryAfter,
70 ProviderQuota,
72 Both,
74}
75
76#[derive(Clone, Copy, Debug, Eq, PartialEq)]
78pub struct DelayDecision {
79 delay: DelaySeconds,
80 source: DelaySource,
81 clamped: bool,
82}
83
84impl DelayDecision {
85 #[must_use]
87 pub const fn delay(self) -> DelaySeconds {
88 self.delay
89 }
90 #[must_use]
92 pub const fn source(self) -> DelaySource {
93 self.source
94 }
95 #[must_use]
97 pub const fn was_clamped(self) -> bool {
98 self.clamped
99 }
100}
101
102pub 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#[derive(Clone, Copy, Debug, Eq, PartialEq)]
219pub enum DelayDecisionError {
220 ClockRollback,
222 ExhaustedBucketResetUnknown,
224 PastTimestamp,
226 TimestampOverflow,
228 ConflictingMetadata,
230 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);