[epilogues] do not rely on `at_exit`

Summary:
For some unexplained reason, some of the functions registered in the Epilogues would sometimes be executed several times. I could not figure out why.

This diff fixes that, but also has more explainable benefits:
- Do not run epilogues registered in the parent in the children. Previously it
  would do so, but probably only if the children registered some epilogue given
  that `at_exit` must be called again once on the child (but the value of the ref
  in `Pervasives` would not have been reset).
- Unified behaviour for early and late epilogues given that we now handle both of these directly

We already have all the control needed to run epilogues when needed: we know
when infer exits, and we know when children processes exit.

Reviewed By: mbouaziz

Differential Revision: D9752046

fbshipit-source-id: 13af40081

infer/src/backend/PerfStats.ml

@@ -356,4 +356,5 @@ let get_reporter stats_type =
 let register_report_at_exit stats_type =
   let relative_path = relative_path_of_stats_type stats_type in
   register_report TimeAndMemory stats_type ;
-  Epilogues.register ~f:(get_reporter stats_type) ("stats reporting in " ^ relative_path)
+  Epilogues.register ~f:(get_reporter stats_type)
+    ~description:("stats reporting in " ^ relative_path)

infer/src/backend/Tasks.ml

@@ -26,6 +26,7 @@ let run_sequentially ~(f : 'a doer) (tasks : 'a list) : unit =
 
 let fork_protect ~f x =
   (* this is needed whenever a new process is started *)
+  Epilogues.reset () ;
   EventLogger.prepare () ;
   L.reset_formatters () ;
   ResultsDatabase.new_database_connection () ;

infer/src/base/Config.ml

@@ -2287,8 +2287,6 @@ let inferconfig_file =
       find (Sys.getcwd ()) |> Option.map ~f:(fun dir -> dir ^/ CommandDoc.inferconfig_file)
 
 
-let register_late_epilogue = Epilogues.register_late
-
 let post_parsing_initialization command_opt =
   if CommandLineOption.is_originator then (
     (* let subprocesses know where the toplevel process' results dir is *)
@@ -2381,7 +2379,7 @@ let post_parsing_initialization command_opt =
         "Run the command again with `--keep-going` to try and ignore this error.\n" ;
     let exitcode = L.exit_code_of_exception exn in
     L.log_uncaught_exception exn ~exitcode ;
-    Epilogues.late () ;
+    Epilogues.run () ;
     Pervasives.exit exitcode
   in
   Caml.Printexc.set_uncaught_exception_handler uncaught_exception_handler ;

infer/src/base/Config.mli

@@ -688,7 +688,5 @@ val print_usage_exit : unit -> 'a
 
 (** {2 Miscellanous} *)
 
-val register_late_epilogue : (unit -> unit) -> unit
-
 val java_package_is_external : string -> bool
 (** Check if a Java package is external to the repository *)

infer/src/base/Epilogues.ml

@@ -7,39 +7,44 @@
 open! IStd
 module F = Format
 
+let early_callback = ref (fun () -> ())
+
 let late_callback = ref (fun () -> ())
 
-let register_late f =
-  let g = !late_callback in
-  late_callback := fun () -> f () ; g ()
+let register callback_ref ~f ~description =
+  let f_no_exn () =
+    try f () with exn ->
+      F.eprintf "%a: Error while running epilogue \"%s\":@ %a.@ Powering through...@." Pid.pp
+        (Unix.getpid ()) description Exn.pp exn
+  in
+  let g = !callback_ref in
+  callback_ref := fun () -> f_no_exn () ; g ()
+
+
+let register_early ~f ~description = register early_callback ~f ~description
+
+let register_late ~f ~description = register late_callback ~f ~description
 
+let early () = !early_callback ()
 
 let late () = !late_callback ()
 
-(* Run the epilogues when we get SIGINT (Control-C). We do not want to mask SIGINT unless at least
-   one epilogue has been registered, so make this value lazy. *)
-let activate_run_epilogues_on_signal =
-  lazy
-    (let run_epilogues_on_signal s =
-       F.eprintf "*** %s: Caught %s, time to die@."
-         (Filename.basename Sys.executable_name)
-         (Signal.to_string s) ;
-       (* Invoke the callback that runs at the end of uncaught_exception_handler *)
-       late () ;
-       (* Epilogues are registered with [at_exit] so exiting will make them run. *)
-       Pervasives.exit 0
-     in
-     Signal.Expert.handle Signal.int run_epilogues_on_signal)
-
-
-let register ~f desc =
-  let f_no_exn () =
-    if not !ProcessPoolState.in_child then
-      try f () with exn ->
-        F.eprintf "Error while running epilogue \"%s\":@ %a.@ Powering through...@." desc Exn.pp
-          exn
+let run () = early () ; late ()
+
+(* Run the epilogues when we get SIGINT (Control-C). *)
+let () =
+  let run_epilogues_on_signal s =
+    F.eprintf "*** %s: Caught %s, time to die@."
+      (Filename.basename Sys.executable_name)
+      (Signal.to_string s) ;
+    run ()
   in
-  (* We call `exit` in a bunch of places, so register the epilogues with [at_exit]. *)
-  Pervasives.at_exit f_no_exn ;
-  (* Register signal masking. *)
-  Lazy.force activate_run_epilogues_on_signal
+  Signal.Expert.handle Signal.int run_epilogues_on_signal
+
+
+let reset () =
+  (early_callback := fun () -> ()) ;
+  late_callback := fun () -> ()
+
+
+let register = register_early

infer/src/base/Epilogues.mli

@@ -7,10 +7,15 @@
 
 open! IStd
 
-val register : f:(unit -> unit) -> string -> unit
+val register : f:(unit -> unit) -> description:string -> unit
 (** Register a function to run when the program exits or is interrupted. Registered functions are
     run in the reverse order in which they were registered. *)
 
-val register_late : (unit -> unit) -> unit
+val register_late : f:(unit -> unit) -> description:string -> unit
+(** Register a function to run when the program exits or is interrupted. Registered functions are
+   run in the reverse order in which they were registered, but *after* the ones registered with
+   {!register}. *)
+
+val run : unit -> unit
 
-val late : unit -> unit
+val reset : unit -> unit

infer/src/base/EventLogger.ml

@@ -47,7 +47,7 @@ module IO = struct
     List.iter log_files ~f:dump_file_to_stdout
 
 
-  let () = Config.register_late_epilogue close
+  let () = Epilogues.register_late ~f:close ~description:"closing EventLogger file"
 end
 
 module Random_id : sig

infer/src/base/Logging.ml

@@ -166,7 +166,7 @@ let close_logs () =
       F.pp_print_flush file_fmt () ; Out_channel.close chan )
 
 
-let () = Epilogues.register ~f:close_logs "flushing logs and closing log file"
+let () = Epilogues.register ~f:close_logs ~description:"flushing logs and closing log file"
 
 let log ~to_console ?(to_file = true) (lazy formatters) =
   match (to_console, to_file) with

infer/src/base/ProcessPool.ml

@@ -223,6 +223,7 @@ let fork_child ~child_prelude ~slot (updates_r, updates_w) ~f =
       child_loop ~slot send_to_parent receive_from_parent ~f ;
       Out_channel.close updates_oc ;
       In_channel.close orders_ic ;
+      Epilogues.run () ;
       Pervasives.exit 0
   | `In_the_parent pid ->
       let[@warning "-26"] to_child_r = Unix.close to_child_r in

infer/src/base/ResultsDatabase.ml

@@ -161,4 +161,4 @@ let db_canonicalize () =
   SqliteUtils.exec db ~log:"running VACUUM" ~stmt:"VACUUM"
 
 
-let () = Config.register_late_epilogue db_close
+let () = Epilogues.register_late ~f:db_close ~description:"closing database connection"

infer/src/base/Utils.ml

@@ -349,8 +349,8 @@ let rec rmtree name =
 let better_hash x = Marshal.to_string x [Marshal.No_sharing] |> Caml.Digest.string
 
 let unlink_file_on_exit temp_file =
-  "Cleaning temporary file " ^ temp_file
-  |> Epilogues.register ~f:(fun () -> try Unix.unlink temp_file with _ -> ())
+  let description = "Cleaning temporary file " ^ temp_file in
+  Epilogues.register ~description ~f:(fun () -> try Unix.unlink temp_file with _ -> ())
 
 
 (** drop at most one layer of well-balanced first and last characters satisfying [drop] from the

infer/src/integration/Maven.ml

@@ -146,12 +146,12 @@ let add_profile_to_pom_in_directory dir =
   Unix.rename ~src:maven_pom_path ~dst:saved_pom_path ;
   Epilogues.register
     ~f:(fun () -> Unix.rename ~src:saved_pom_path ~dst:maven_pom_path)
-    "restoring Maven's pom.xml to its original state" ;
+    ~description:"restoring Maven's pom.xml to its original state" ;
   Unix.rename ~src:infer_pom_path ~dst:maven_pom_path ;
   if Config.debug_mode then
     Epilogues.register
       ~f:(fun () -> Unix.rename ~src:maven_pom_path ~dst:infer_pom_path)
-      "saving infer's pom.xml"
+      ~description:"saving infer's pom.xml"
 
 
 let capture ~prog ~args =