1use std::any::Any;
5use std::borrow::Cow;
6use std::sync::Arc;
7use std::time::Duration;
8
9use crate::http::{Method, StatusCode};
10use async_trait::async_trait;
11use url::Url;
12
13use crate::error::Error;
14
15#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct OperationInfo {
23 pub service: Cow<'static, str>,
25 pub operation: Cow<'static, str>,
27 pub resource_type: Cow<'static, str>,
29 pub is_mutation: bool,
31 pub resource_id: Option<String>,
33}
34
35#[derive(Debug, Clone)]
38pub struct RequestInfo {
39 pub method: Method,
41 pub url: Url,
43 pub attempt: u32,
45}
46
47#[derive(Debug)]
49pub struct RequestResult<'a> {
50 pub status: Option<StatusCode>,
52 pub duration: Duration,
54 pub error: Option<&'a Error>,
56 pub from_cache: bool,
59 pub retryable: bool,
61 pub retry_after: Option<u64>,
63}
64
65impl<'a> RequestResult<'a> {
66 pub(crate) fn failed(
68 status: Option<StatusCode>,
69 duration: Duration,
70 error: &'a Error,
71 retryable: bool,
72 retry_after: Option<u64>,
73 ) -> RequestResult<'a> {
74 RequestResult {
75 status,
76 duration,
77 error: Some(error),
78 from_cache: false,
79 retryable,
80 retry_after,
81 }
82 }
83}
84
85pub type OperationState = Option<Box<dyn Any + Send>>;
87
88#[async_trait]
123pub trait Hooks: Send + Sync {
124 async fn on_operation_gate(&self, _op: &OperationInfo) -> Result<(), Error> {
133 Ok(())
134 }
135
136 fn on_operation_start(&self, _op: &OperationInfo) -> OperationState {
139 None
140 }
141
142 fn on_operation_abandoned(&self, _op: &OperationInfo) {}
146
147 fn on_operation_end(
150 &self,
151 _op: &OperationInfo,
152 _state: OperationState,
153 _outcome: Result<(), &Error>,
154 _duration: Duration,
155 ) {
156 }
157
158 fn on_request_start(&self, _info: &RequestInfo) {}
160
161 fn on_request_end(&self, _info: &RequestInfo, _result: &RequestResult<'_>) {}
163
164 fn on_retry(&self, _info: &RequestInfo, _next_attempt: u32, _cause: &Error) {}
167
168 fn is_noop(&self) -> bool {
171 false
172 }
173}
174
175#[async_trait]
176impl<H: Hooks + ?Sized> Hooks for Arc<H> {
177 async fn on_operation_gate(&self, op: &OperationInfo) -> Result<(), Error> {
178 (**self).on_operation_gate(op).await
179 }
180
181 fn on_operation_start(&self, op: &OperationInfo) -> OperationState {
182 (**self).on_operation_start(op)
183 }
184
185 fn on_operation_abandoned(&self, op: &OperationInfo) {
186 (**self).on_operation_abandoned(op);
187 }
188
189 fn on_operation_end(
190 &self,
191 op: &OperationInfo,
192 state: OperationState,
193 outcome: Result<(), &Error>,
194 duration: Duration,
195 ) {
196 (**self).on_operation_end(op, state, outcome, duration);
197 }
198
199 fn on_request_start(&self, info: &RequestInfo) {
200 (**self).on_request_start(info);
201 }
202
203 fn on_request_end(&self, info: &RequestInfo, result: &RequestResult<'_>) {
204 (**self).on_request_end(info, result);
205 }
206
207 fn on_retry(&self, info: &RequestInfo, next_attempt: u32, cause: &Error) {
208 (**self).on_retry(info, next_attempt, cause);
209 }
210
211 fn is_noop(&self) -> bool {
212 (**self).is_noop()
213 }
214}
215
216#[derive(Debug, Clone, Copy, Default)]
219pub struct NoopHooks;
220
221impl Hooks for NoopHooks {
222 fn is_noop(&self) -> bool {
223 true
224 }
225}
226
227pub struct ChainHooks {
230 hooks: Vec<Arc<dyn Hooks>>,
231}
232
233impl ChainHooks {
234 pub fn of(hooks: Vec<Arc<dyn Hooks>>) -> Arc<dyn Hooks> {
238 let mut installed: Vec<Arc<dyn Hooks>> =
239 hooks.into_iter().filter(|hook| !hook.is_noop()).collect();
240 if installed.is_empty() {
241 Arc::new(NoopHooks)
242 } else if installed.len() == 1 {
243 installed.remove(0)
244 } else {
245 Arc::new(ChainHooks { hooks: installed })
246 }
247 }
248}
249
250#[async_trait]
251impl Hooks for ChainHooks {
252 async fn on_operation_gate(&self, op: &OperationInfo) -> Result<(), Error> {
256 for (admitted, hook) in self.hooks.iter().enumerate() {
257 if let Err(refusal) = hook.on_operation_gate(op).await {
258 for earlier in self.hooks[..admitted].iter().rev() {
259 earlier.on_operation_abandoned(op);
260 }
261 return Err(refusal);
262 }
263 }
264 Ok(())
265 }
266
267 fn on_operation_abandoned(&self, op: &OperationInfo) {
268 for hook in self.hooks.iter().rev() {
269 hook.on_operation_abandoned(op);
270 }
271 }
272
273 fn on_operation_start(&self, op: &OperationInfo) -> OperationState {
276 let states: Vec<OperationState> = self
277 .hooks
278 .iter()
279 .map(|hook| hook.on_operation_start(op))
280 .collect();
281 Some(Box::new(states))
282 }
283
284 fn on_operation_end(
285 &self,
286 op: &OperationInfo,
287 state: OperationState,
288 outcome: Result<(), &Error>,
289 duration: Duration,
290 ) {
291 let mut states = member_states(state);
292 for hook in self.hooks.iter().rev() {
293 hook.on_operation_end(op, states.pop().flatten(), outcome, duration);
294 }
295 }
296
297 fn on_request_start(&self, info: &RequestInfo) {
298 for hook in &self.hooks {
299 hook.on_request_start(info);
300 }
301 }
302
303 fn on_request_end(&self, info: &RequestInfo, result: &RequestResult<'_>) {
304 for hook in self.hooks.iter().rev() {
305 hook.on_request_end(info, result);
306 }
307 }
308
309 fn on_retry(&self, info: &RequestInfo, next_attempt: u32, cause: &Error) {
310 for hook in &self.hooks {
311 hook.on_retry(info, next_attempt, cause);
312 }
313 }
314}
315
316fn member_states(state: OperationState) -> Vec<OperationState> {
317 match state.and_then(|state| state.downcast::<Vec<OperationState>>().ok()) {
318 Some(states) => *states,
319 None => Vec::new(),
320 }
321}
322
323#[cfg(test)]
324#[allow(clippy::unwrap_used)]
325mod tests {
326 use std::sync::Mutex;
327
328 use super::*;
329 use crate::ErrorCode;
330
331 #[tokio::test]
332 async fn every_noop_callback_is_safe_to_call() {
333 let hooks = NoopHooks;
334 let op = operation_info();
335 let info = request_info();
336
337 hooks.on_operation_gate(&op).await.unwrap();
338 let state = hooks.on_operation_start(&op);
339 assert!(state.is_none());
340 hooks.on_request_start(&info);
341 hooks.on_request_end(&info, &request_result());
342 hooks.on_retry(&info, 2, &Error::usage("nothing"));
343 hooks.on_operation_end(&op, state, Ok(()), Duration::from_secs(1));
344
345 assert!(hooks.is_noop());
346 }
347
348 #[test]
349 fn a_chain_runs_forwards_and_unwinds_backwards() {
350 let log = Log::new();
351 let chain = ChainHooks::of(vec![log.recorder("first"), log.recorder("second")]);
352 let op = operation_info();
353
354 let state = chain.on_operation_start(&op);
355 chain.on_request_start(&request_info());
356 chain.on_request_end(&request_info(), &request_result());
357 chain.on_retry(&request_info(), 2, &Error::usage("nothing"));
358 chain.on_operation_end(&op, state, Ok(()), Duration::from_secs(1));
359
360 assert_eq!(
361 log.entries(),
362 [
363 "first: start Svc.Do",
364 "second: start Svc.Do",
365 "first: request start 1",
366 "second: request start 1",
367 "second: request end 200",
368 "first: request end 200",
369 "first: retry 2",
370 "second: retry 2",
371 "second: end Svc.Do carrying second",
372 "first: end Svc.Do carrying first",
373 ]
374 );
375 }
376
377 #[test]
378 fn a_chain_of_one_is_that_hook() {
379 let log = Log::new();
380 let recorder = log.recorder("only");
381
382 let chain = ChainHooks::of(vec![recorder.clone(), Arc::new(NoopHooks)]);
383
384 assert!(Arc::ptr_eq(&chain, &recorder));
385 }
386
387 #[test]
388 fn a_chain_of_nothing_but_noops_is_a_noop() {
389 assert!(ChainHooks::of(vec![Arc::new(NoopHooks), Arc::new(NoopHooks)]).is_noop());
390 assert!(ChainHooks::of(Vec::new()).is_noop());
391 }
392
393 #[tokio::test]
394 async fn a_chain_answers_the_first_refusal() {
395 let log = Log::new();
396 let chain = ChainHooks::of(vec![
397 log.recorder("first"),
398 Arc::new(Refusing),
399 log.recorder("third"),
400 ]);
401
402 let refused = chain
403 .on_operation_gate(&operation_info())
404 .await
405 .unwrap_err();
406
407 assert_eq!(refused.code(), ErrorCode::Usage);
408 assert_eq!(refused.message(), "blocked");
409 assert_eq!(log.entries(), ["first: gate Svc.Do"]);
410 }
411
412 struct Log {
414 entries: Arc<Mutex<Vec<String>>>,
415 }
416
417 impl Log {
418 fn new() -> Log {
419 Log {
420 entries: Arc::new(Mutex::new(Vec::new())),
421 }
422 }
423
424 fn recorder(&self, name: &'static str) -> Arc<dyn Hooks> {
425 Arc::new(Recorder {
426 name,
427 entries: self.entries.clone(),
428 })
429 }
430
431 fn entries(&self) -> Vec<String> {
432 self.entries.lock().unwrap().clone()
433 }
434 }
435
436 struct Recorder {
437 name: &'static str,
438 entries: Arc<Mutex<Vec<String>>>,
439 }
440
441 impl Recorder {
442 fn record(&self, event: &str) {
443 self.entries
444 .lock()
445 .unwrap()
446 .push(format!("{}: {event}", self.name));
447 }
448 }
449
450 #[async_trait]
451 impl Hooks for Recorder {
452 async fn on_operation_gate(&self, op: &OperationInfo) -> Result<(), Error> {
453 self.record(&format!("gate {}.{}", op.service, op.operation));
454 Ok(())
455 }
456
457 fn on_operation_start(&self, op: &OperationInfo) -> OperationState {
458 self.record(&format!("start {}.{}", op.service, op.operation));
459 Some(Box::new(self.name.to_string()))
460 }
461
462 fn on_operation_end(
463 &self,
464 op: &OperationInfo,
465 state: OperationState,
466 _outcome: Result<(), &Error>,
467 _duration: Duration,
468 ) {
469 let carried = match state.and_then(|state| state.downcast::<String>().ok()) {
470 Some(name) => *name,
471 None => "nothing".to_string(),
472 };
473 self.record(&format!(
474 "end {}.{} carrying {carried}",
475 op.service, op.operation
476 ));
477 }
478
479 fn on_request_start(&self, info: &RequestInfo) {
480 self.record(&format!("request start {}", info.attempt));
481 }
482
483 fn on_request_end(&self, _info: &RequestInfo, result: &RequestResult<'_>) {
484 self.record(&format!("request end {}", result.status.unwrap().as_u16()));
485 }
486
487 fn on_retry(&self, _info: &RequestInfo, next_attempt: u32, _cause: &Error) {
488 self.record(&format!("retry {next_attempt}"));
489 }
490 }
491
492 struct Refusing;
493
494 #[async_trait]
495 impl Hooks for Refusing {
496 async fn on_operation_gate(&self, _op: &OperationInfo) -> Result<(), Error> {
497 Err(Error::usage("blocked"))
498 }
499 }
500
501 fn operation_info() -> OperationInfo {
502 OperationInfo {
503 service: Cow::Borrowed("Svc"),
504 operation: Cow::Borrowed("Do"),
505 resource_type: Cow::Borrowed("thing"),
506 is_mutation: false,
507 resource_id: None,
508 }
509 }
510
511 fn request_info() -> RequestInfo {
512 RequestInfo {
513 method: Method::GET,
514 url: Url::parse("https://fizzy.example/999/boards.json").unwrap(),
515 attempt: 1,
516 }
517 }
518
519 fn request_result() -> RequestResult<'static> {
520 RequestResult {
521 status: Some(StatusCode::OK),
522 duration: Duration::from_millis(3),
523 error: None,
524 from_cache: false,
525 retryable: false,
526 retry_after: None,
527 }
528 }
529}