1use std::fmt;
19use std::future::Future;
20use std::sync::{Arc, Mutex};
21use std::time::Duration;
22
23use rustev_contract::remote::{PhaseExpiry, PhaseLimit, RemoteCode, TransportPhase};
24use serde::Serialize;
25use tokio::sync::Notify;
26use tokio::time::Instant;
27
28use crate::budget::Budget;
29use crate::http::TransportError;
30use crate::wire::SentTracker;
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
34#[serde(deny_unknown_fields)]
35pub struct PhasePolicy {
36 pub resolution_ms: u64,
37 pub connect_ms: u64,
38 pub secure_handshake_ms: u64,
39 pub request_write_ms: u64,
40 pub first_response_byte_ms: u64,
41 pub response_body_ms: u64,
42 pub cancellation_ack_ms: u64,
43}
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub struct ZeroCeiling(pub TransportPhase);
49
50impl fmt::Display for ZeroCeiling {
51 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
52 write!(f, "the {:?} phase ceiling must be positive", self.0)
53 }
54}
55
56impl std::error::Error for ZeroCeiling {}
57
58impl PhasePolicy {
59 pub fn ceiling_ms(&self, phase: TransportPhase) -> u64 {
61 match phase {
62 TransportPhase::Resolution => self.resolution_ms,
63 TransportPhase::Connect => self.connect_ms,
64 TransportPhase::SecureHandshake => self.secure_handshake_ms,
65 TransportPhase::RequestWrite => self.request_write_ms,
66 TransportPhase::FirstResponseByte => self.first_response_byte_ms,
67 TransportPhase::ResponseBody => self.response_body_ms,
68 TransportPhase::CancellationAck => self.cancellation_ack_ms,
69 }
70 }
71
72 pub fn validate(&self) -> Result<(), ZeroCeiling> {
74 match TransportPhase::ALL
75 .into_iter()
76 .find(|p| self.ceiling_ms(*p) == 0)
77 {
78 Some(p) => Err(ZeroCeiling(p)),
79 None => Ok(()),
80 }
81 }
82}
83
84#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize)]
87#[serde(rename_all = "snake_case")]
88pub enum TransportPhases {
89 #[default]
90 Disabled,
91 Enabled(PhasePolicy),
92}
93
94impl TransportPhases {
95 pub fn is_disabled(&self) -> bool {
96 matches!(self, TransportPhases::Disabled)
97 }
98
99 pub fn policy(&self) -> Option<&PhasePolicy> {
100 match self {
101 TransportPhases::Disabled => None,
102 TransportPhases::Enabled(p) => Some(p),
103 }
104 }
105
106 pub fn validate(&self) -> Result<(), ZeroCeiling> {
108 self.policy().map_or(Ok(()), PhasePolicy::validate)
109 }
110
111 pub fn cancel_budget(&self, budget: Budget) -> Budget {
114 match self.policy() {
115 None => budget,
116 Some(p) => budget.min(Budget::from_ms(p.cancellation_ack_ms)),
117 }
118 }
119}
120
121#[derive(Debug, Clone, Copy)]
122struct Current {
123 phase: TransportPhase,
124 started: Instant,
125 allowance: Duration,
126 limit: PhaseLimit,
127 generation: u64,
129}
130
131struct ClockInner {
132 policy: PhasePolicy,
133 budget: Budget,
134 current: Mutex<Current>,
135 changed: Notify,
136 body_taken: std::sync::atomic::AtomicBool,
139}
140
141#[derive(Clone)]
144pub struct PhaseClock(Arc<ClockInner>);
145
146impl fmt::Debug for PhaseClock {
147 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
148 f.debug_tuple("PhaseClock")
149 .field(&self.current().phase)
150 .finish()
151 }
152}
153
154fn lock<T>(m: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
155 m.lock().unwrap_or_else(|e| e.into_inner())
156}
157
158fn allowance(
159 policy: &PhasePolicy,
160 budget: &Budget,
161 phase: TransportPhase,
162) -> (Duration, PhaseLimit) {
163 let ceiling = Duration::from_millis(policy.ceiling_ms(phase));
164 let left = budget.remaining();
165 if ceiling <= left {
166 (ceiling, PhaseLimit::Ceiling)
167 } else {
168 (left, PhaseLimit::Budget)
169 }
170}
171
172fn ms(d: Duration) -> u64 {
173 u64::try_from(d.as_millis()).unwrap_or(u64::MAX)
174}
175
176impl PhaseClock {
177 pub fn start(policy: PhasePolicy, budget: Budget, first: TransportPhase) -> Self {
179 let (allowance, limit) = allowance(&policy, &budget, first);
180 PhaseClock(Arc::new(ClockInner {
181 policy,
182 budget,
183 current: Mutex::new(Current {
184 phase: first,
185 started: Instant::now(),
186 allowance,
187 limit,
188 generation: 0,
189 }),
190 changed: Notify::new(),
191 body_taken: std::sync::atomic::AtomicBool::new(false),
192 }))
193 }
194
195 pub fn enter(&self, phase: TransportPhase) {
200 let mut c = lock(&self.0.current);
201 if phase <= c.phase {
202 return;
203 }
204 let (allowance, limit) = allowance(&self.0.policy, &self.0.budget, phase);
205 *c = Current {
206 phase,
207 started: Instant::now(),
208 allowance,
209 limit,
210 generation: c.generation + 1,
211 };
212 drop(c);
213 self.0.changed.notify_waiters();
214 }
215
216 pub fn phase(&self) -> TransportPhase {
218 self.current().phase
219 }
220
221 fn current(&self) -> Current {
222 *lock(&self.0.current)
223 }
224
225 pub fn body_taken(&self) {
227 self.0
228 .body_taken
229 .store(true, std::sync::atomic::Ordering::SeqCst);
230 }
231
232 pub fn flushed(&self) {
236 if self.0.body_taken.load(std::sync::atomic::Ordering::SeqCst) {
237 self.enter(TransportPhase::FirstResponseByte);
238 }
239 }
240
241 pub fn response_byte(&self) {
243 self.enter(TransportPhase::ResponseBody);
244 }
245
246 fn expiry(&self, c: &Current, sent: bool) -> PhaseExpiry {
247 PhaseExpiry {
248 phase: c.phase,
249 ceiling_ms: self.0.policy.ceiling_ms(c.phase),
250 allowance_ms: ms(c.allowance),
251 elapsed_ms: ms(Instant::now().saturating_duration_since(c.started)),
252 limit: c.limit,
253 bytes_sent: sent,
254 code: if sent {
255 RemoteCode::TransportInterrupted
256 } else {
257 RemoteCode::TransportUnsent
258 },
259 }
260 }
261}
262
263pub async fn phased<T, F>(
269 clock: &PhaseClock,
270 sent: &SentTracker,
271 fut: F,
272) -> Result<T, TransportError>
273where
274 F: Future<Output = Result<T, TransportError>>,
275{
276 tokio::pin!(fut);
277 loop {
278 let changed = clock.0.changed.notified();
279 tokio::pin!(changed);
280 changed.as_mut().enable();
281 let c = clock.current();
282 tokio::select! {
283 biased;
284 r = &mut fut => return r,
285 _ = &mut changed => continue,
286 _ = tokio::time::sleep_until(c.started + c.allowance) => {
287 let now = clock.current();
288 if now.generation != c.generation {
289 continue;
290 }
291 return Err(TransportError::PhaseExpired(clock.expiry(&now, sent.may_have_sent())));
292 }
293 }
294 }
295}
296
297#[cfg(test)]
298mod tests {
299 use super::*;
300
301 fn policy() -> PhasePolicy {
302 PhasePolicy {
303 resolution_ms: 10,
304 connect_ms: 20,
305 secure_handshake_ms: 30,
306 request_write_ms: 40,
307 first_response_byte_ms: 50,
308 response_body_ms: 60,
309 cancellation_ack_ms: 70,
310 }
311 }
312
313 #[test]
314 fn every_ceiling_must_be_positive() {
315 assert!(policy().validate().is_ok());
316 for p in TransportPhase::ALL {
317 let mut bad = policy();
318 match p {
319 TransportPhase::Resolution => bad.resolution_ms = 0,
320 TransportPhase::Connect => bad.connect_ms = 0,
321 TransportPhase::SecureHandshake => bad.secure_handshake_ms = 0,
322 TransportPhase::RequestWrite => bad.request_write_ms = 0,
323 TransportPhase::FirstResponseByte => bad.first_response_byte_ms = 0,
324 TransportPhase::ResponseBody => bad.response_body_ms = 0,
325 TransportPhase::CancellationAck => bad.cancellation_ack_ms = 0,
326 }
327 assert_eq!(bad.validate(), Err(ZeroCeiling(p)));
328 assert_eq!(
329 TransportPhases::Enabled(bad).validate(),
330 Err(ZeroCeiling(p))
331 );
332 }
333 assert!(TransportPhases::Disabled.validate().is_ok());
334 assert!(TransportPhases::default().is_disabled());
335 }
336
337 #[tokio::test(flavor = "current_thread", start_paused = true)]
338 async fn the_cancel_budget_is_cut_to_the_acknowledgement_ceiling() {
339 let b = Budget::from_ms(1_000);
340 assert_eq!(TransportPhases::Disabled.cancel_budget(b), b);
341 let cut = TransportPhases::Enabled(policy()).cancel_budget(b);
342 assert_eq!(cut.remaining_ms(), 70);
343 let short = Budget::from_ms(5);
344 assert_eq!(
345 TransportPhases::Enabled(policy()).cancel_budget(short),
346 short
347 );
348 }
349}