Source file shutdown.ml

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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
(* Unit tests are in ../../lib_test/shutdown_tests.ml *)

open Core
open Import

let debug = Debug.shutdown
let todo = ref []

let at_shutdown f =
  let backtrace = Backtrace.get () in
  if debug then Debug.log "at_shutdown" backtrace [%sexp_of: Backtrace.t];
  todo := (backtrace, f) :: !todo
;;

let shutting_down_ref = ref `No
let default_force_ref = ref (fun () -> Clock.after (sec 10.))
let default_force () = !default_force_ref
let set_default_force force = default_force_ref := force
let shutting_down () = !shutting_down_ref

(* Be careful to ensure [shutdown] doesn't raise just because
   stderr is closed *)
let ignore_exn f =
  try f () with
  | _ -> ()
;;

let exit_reliably status =
  match (exit status : Nothing.t) with
  | exception exn ->
    ignore_exn (fun () -> Core.Debug.eprints "Caml.exit raised" exn [%sexp_of: Exn.t]);
    Core.Unix.exit_immediately (if status = 0 then 1 else status)
  | _ -> .
;;

let shutdown ?force status =
  if debug then ignore_exn (fun () -> Debug.log "shutdown" status [%sexp_of: int]);
  match !shutting_down_ref with
  | `Yes status' ->
    if status <> 0 && status' <> 0 && status <> status'
    then
      raise_s
        [%message "shutdown with inconsistent status" (status : int) (status' : int)]
    else if status' = 0 && status <> 0
    then shutting_down_ref := `Yes status
  | `No ->
    shutting_down_ref := `Yes status;
    upon
      (Deferred.all
         (List.map !todo ~f:(fun (backtrace, f) ->
            let%map result = Monitor.try_with_or_error f in
            (match result with
             | Ok () -> ()
             | Error error ->
               ignore_exn (fun () ->
                 Core.Debug.eprints
                   "at_shutdown function raised"
                   (error, backtrace)
                   [%sexp_of: Error.t * Backtrace.t]));
            if debug
            then
              ignore_exn (fun () ->
                Debug.log
                  "one at_shutdown function finished"
                  backtrace
                  [%sexp_of: Backtrace.t]);
            result)))
      (fun results ->
         match shutting_down () with
         | `No -> assert false
         | `Yes status ->
           let status =
             match Or_error.combine_errors_unit results with
             | Ok () -> status
             | Error _ -> if status = 0 then 1 else status
           in
           exit_reliably status);
    let force =
      match force with
      | None -> !default_force_ref ()
      | Some f -> f
    in
    upon force (fun () ->
      ignore_exn (fun () -> Debug.log_string "Shutdown forced.");
      exit_reliably 1)
;;

let shutdown_on_unhandled_exn () =
  Monitor.detach_and_iter_errors Monitor.main ~f:(fun exn ->
    try
      Debug.log "shutting down due to unhandled exception" exn [%sexp_of: exn];
      shutdown 1
    with
    | _ -> ())
;;

let exit ?force status =
  shutdown ?force status;
  Deferred.never ()
;;

let don't_finish_before =
  let proceed_with_shutdown = Ivar.create () in
  let num_waiting = ref 0 in
  let check () = if !num_waiting = 0 then Ivar.fill proceed_with_shutdown () in
  at_shutdown (fun () ->
    check ();
    Ivar.read proceed_with_shutdown);
  fun d ->
    incr num_waiting;
    upon d (fun () ->
      decr num_waiting;
      match shutting_down () with
      | `No -> ()
      | `Yes _ -> check ())
;;