From f9ab3aa1acadad02c3dc65672211da91e9158f29 Mon Sep 17 00:00:00 2001
From: Jules Villard <jul@fb.com>
Date: Thu, 19 Jan 2017 08:54:13 -0800
Subject: [PATCH] run epilogues register with `at_exit` upon receiving SIGINT
 as well

Summary:
Epilogue tasks such as closing logs or putting files back were we found them
run automatically at the end of our executables by registering them with
`at_exit`. They do not run if the program is interrupted by a signal. This diff
makes sure they are run when the user stops infer with Ctrl-C (SIGINT).

Reviewed By: cristianoc

Differential Revision: D4435575

fbshipit-source-id: c3ab702
---
 infer/src/backend/PerfStats.ml |  2 +-
 infer/src/base/Logging.ml      |  2 +-
 infer/src/base/Utils.ml        | 17 +++++++++++++++++
 infer/src/base/Utils.mli       |  4 ++++
 infer/src/integration/Maven.ml |  4 ++--
 5 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/infer/src/backend/PerfStats.ml b/infer/src/backend/PerfStats.ml
index 11ad6dc73..1d0b6fa9b 100644
--- a/infer/src/backend/PerfStats.ml
+++ b/infer/src/backend/PerfStats.ml
@@ -141,7 +141,7 @@ let stats () =
   }
 
 let register_report_at_exit file =
-  Pervasives.at_exit (fun () ->
+  Utils.register_epilogue (fun () ->
       try
         let json_stats = to_json (stats ()) in
         try
diff --git a/infer/src/base/Logging.ml b/infer/src/base/Logging.ml
index c97c3ee3c..841814273 100644
--- a/infer/src/base/Logging.ml
+++ b/infer/src/base/Logging.ml
@@ -77,7 +77,7 @@ let create_log_file exe name_prefix outerr =
     dup_formatter file_fmt outerr_fmt
   );
   (* flush files on exit *)
-  Pervasives.at_exit (fun () -> close_log_file (lazy file_fmt) (lazy chan) (lazy file));
+  Utils.register_epilogue (fun () -> close_log_file (lazy file_fmt) (lazy chan) (lazy file));
   (file_fmt, chan, file)
 
 let should_setup_log_files (exe : CLOpt.exe) = match exe with
diff --git a/infer/src/base/Utils.ml b/infer/src/base/Utils.ml
index 9d626692e..315f672c4 100644
--- a/infer/src/base/Utils.ml
+++ b/infer/src/base/Utils.ml
@@ -289,3 +289,20 @@ let compare_versions v1 v2 =
   let lv1  = int_list_of_version v1 in
   let lv2  = int_list_of_version v2 in
   [%compare : int list] lv1 lv2
+
+(* 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);
+    (* Epilogues are registered with [at_exit] so exiting will make them run. *)
+    exit 0 in
+  Signal.Expert.handle Signal.int run_epilogues_on_signal
+)
+
+let register_epilogue f =
+  (* We call `exit` in a bunch of places, so register the epilogues with [at_exit]. *)
+  Pervasives.at_exit f;
+  (* Register signal masking. *)
+  Lazy.force activate_run_epilogues_on_signal
diff --git a/infer/src/base/Utils.mli b/infer/src/base/Utils.mli
index 884cf7730..a0b148a28 100644
--- a/infer/src/base/Utils.mli
+++ b/infer/src/base/Utils.mli
@@ -85,3 +85,7 @@ val suppress_stderr2 : ('a -> 'b -> 'c) -> 'a -> 'b -> 'c
     -1 if v1 is older than v2 and 0 if they are the same version.
     The versions are strings of the shape "n.m.t", the order is lexicographic. *)
 val compare_versions : string -> string -> int
+
+(** 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_epilogue : (unit -> unit) -> unit
diff --git a/infer/src/integration/Maven.ml b/infer/src/integration/Maven.ml
index 6230606c0..4ed6b99df 100644
--- a/infer/src/integration/Maven.ml
+++ b/infer/src/integration/Maven.ml
@@ -132,10 +132,10 @@ let add_profile_to_pom_in_directory dir =
   let infer_pom_path = dir ^/ "pom.xml.infer" in
   add_infer_profile maven_pom_path infer_pom_path;
   Unix.rename ~src:maven_pom_path ~dst:saved_pom_path;
-  Pervasives.at_exit (fun () -> Unix.rename ~src:saved_pom_path ~dst:maven_pom_path);
+  Utils.register_epilogue (fun () -> Unix.rename ~src:saved_pom_path ~dst:maven_pom_path);
   Unix.rename ~src:infer_pom_path ~dst:maven_pom_path;
   if Config.debug_mode || Config.stats_mode then
-    Pervasives.at_exit (fun () -> Unix.rename ~src:maven_pom_path ~dst:infer_pom_path)
+    Utils.register_epilogue (fun () -> Unix.rename ~src:maven_pom_path ~dst:infer_pom_path)
 
 let capture ~prog ~args =
   while not (List.is_empty !pom_worklist); do