devkit/retry.ml at 76a0117c2eba638ea9051ef56c2d77ec4eb889e5 · ahrefs/devkit · GitHub

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
open Prelude

let log = Log.from "retry"

(** Clamped random exponential backoff for retry *)
let exp_backoff_pause ?max_delay attempt =
  let sleep = Random.int (attempt + 1) in
  let pause = 2. ** float sleep -. 1. in
  Option.map_default (min pause) pause max_delay

let wait_pause'' ?(ignore_should_exit = false) poll pause =
  log #info "will wait for %s" (Time.duration_str pause);
  let need_stamp = Time.now () +. pause in
  let rec loop () =
    match Time.now () with
    | now when now < need_stamp && (ignore_should_exit || Daemon.should_run ()) ->
      Nix.sleep (min 2. (need_stamp -. now));
      poll ();
      loop ()
    | _ -> ()
  in
  loop ()

let wait_pause' ?ignore_should_exit = wait_pause'' ?ignore_should_exit id
let wait_pause ?ignore_should_exit master = wait_pause'' ?ignore_should_exit master#poll


(* Exponentially increasing sleep pause depending of the number of attempts already made.
   Since the loop is not managed by this function, the number of attempts already made must be provided.
*)
let backoff_log ~exn ~name attempt =
  function
  | None -> log #warn ~exn "%s: aborting after %d max_retries" name attempt
  | Some pause ->
    log #warn ~exn "%s: will retry in %s (try #%d)" name (Time.duration_str pause) attempt

let exp_backoff ?(f_retry=id) ~exn ~name ?max_retries ~max_delay attempt =
  match attempt, max_retries with
  | n, Some max_retries when n > max_retries ->
    backoff_log ~exn ~name attempt None;
    Lwt.fail exn
  | attempt, _ ->
    let pause = exp_backoff_pause ~max_delay attempt in
    f_retry ();
    backoff_log ~exn ~name attempt (Some pause);
    let%lwt () = Lwt_unix.sleep pause in
    Lwt.return (attempt + 1)

let backoff_log_result to_string error ~name attempt =
  function
    | None -> log #warn "%s: aborting after %d max_retries %s" name attempt (to_string error)
  | Some pause ->
    log #warn "%s: will retry in %s (try #%d) %s" name (Time.duration_str pause) attempt (to_string error)

let exp_backoff_result ?(f_retry=id) to_string error ~name ?max_retries ~max_delay attempt =
  match attempt, max_retries with
  | n, Some max_retries when n > max_retries ->
    backoff_log_result to_string error ~name attempt None;
    Lwt.return_error error
  | attempt, _ ->
    let pause = exp_backoff_pause ~max_delay attempt in
    f_retry ();
    backoff_log_result to_string error ~name attempt (Some pause);
    let%lwt () = Lwt_unix.sleep pause in
    Lwt.return_ok (attempt + 1)

let with_exp_backoff ~name ?f_retry ?max_retries ~max_delay f =
  let rec loop f attempt =
    try%lwt
      f ()
    with
    | Daemon.ShouldExit | Lwt.Canceled as exn ->
      backoff_log ~exn ~name attempt None;
      Lwt.fail exn
    | exn ->
      let%lwt attempt = exp_backoff ?f_retry ~exn ~name ?max_retries ~max_delay attempt in
      loop f attempt
  in
  loop f 1

let exp_backoff_blocking ~master ~exn ~name ?max_retries ~max_delay attempt =
  match attempt, max_retries with
  | n, Some max_retries when n > max_retries ->
    backoff_log ~exn ~name attempt None;
    raise exn
  | attempt, _ ->
    let pause = exp_backoff_pause ~max_delay attempt in
    backoff_log ~exn ~name attempt (Some pause);
    wait_pause master pause;
    attempt + 1

let exp_backoff_blocking_no_poll ~exn ~name ?max_retries ~max_delay attempt =
  match attempt, max_retries with
  | n, Some max_retries when n > max_retries ->
    backoff_log ~exn ~name attempt None;
    raise exn
  | attempt, _ ->
    let pause = exp_backoff_pause ~max_delay attempt in
    backoff_log ~exn ~name attempt (Some pause);
    wait_pause' pause;
    attempt + 1