From 56eb08735be451a3351543bf5c98d0e51c9fa036 Mon Sep 17 00:00:00 2001
From: Ben Kimock <kimockb@gmail.com>
Date: Fri, 19 May 2023 13:49:24 -0400
Subject: [PATCH] Use a signal handler to observe ctrl+c and cleanly drop the
 measureme profiler

---
 src/tools/miri/Cargo.lock                | 29 ++++++++++++++++++++++++
 src/tools/miri/Cargo.toml                |  1 +
 src/tools/miri/src/concurrency/thread.rs | 17 ++++++++++++++
 src/tools/miri/src/machine.rs            |  9 ++++++++
 4 files changed, 56 insertions(+)

diff --git a/src/tools/miri/Cargo.lock b/src/tools/miri/Cargo.lock
index 737423a2cd1..7bedb5d48f6 100644
--- a/src/tools/miri/Cargo.lock
+++ b/src/tools/miri/Cargo.lock
@@ -183,6 +183,16 @@ dependencies = [
  "cfg-if",
 ]
 
+[[package]]
+name = "ctrlc"
+version = "3.2.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bbcf33c2a618cbe41ee43ae6e9f2e48368cd9f9db2896f10167d8d762679f639"
+dependencies = [
+ "nix",
+ "windows-sys 0.45.0",
+]
+
 [[package]]
 name = "diff"
 version = "0.1.13"
@@ -421,6 +431,7 @@ name = "miri"
 version = "0.1.0"
 dependencies = [
  "colored",
+ "ctrlc",
  "env_logger",
  "getrandom",
  "lazy_static",
@@ -437,6 +448,18 @@ dependencies = [
  "ui_test",
 ]
 
+[[package]]
+name = "nix"
+version = "0.26.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bfdda3d196821d6af13126e40375cdf7da646a96114af134d5f417a9a1dc8e1a"
+dependencies = [
+ "bitflags",
+ "cfg-if",
+ "libc",
+ "static_assertions",
+]
+
 [[package]]
 name = "object"
 version = "0.30.3"
@@ -713,6 +736,12 @@ version = "1.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a507befe795404456341dfab10cef66ead4c041f62b8b11bbb92bffe5d0953e0"
 
+[[package]]
+name = "static_assertions"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
+
 [[package]]
 name = "syn"
 version = "2.0.15"
diff --git a/src/tools/miri/Cargo.toml b/src/tools/miri/Cargo.toml
index 5987b0df8d6..683f99ca4de 100644
--- a/src/tools/miri/Cargo.toml
+++ b/src/tools/miri/Cargo.toml
@@ -29,6 +29,7 @@ smallvec = "1.7"
 # for more information.
 rustc-workspace-hack = "1.0.0"
 measureme = "10.0.0"
+ctrlc = "3.2.5"
 
 [target.'cfg(unix)'.dependencies]
 libc = "0.2"
diff --git a/src/tools/miri/src/concurrency/thread.rs b/src/tools/miri/src/concurrency/thread.rs
index e9bbae4d504..d85cac7bcbb 100644
--- a/src/tools/miri/src/concurrency/thread.rs
+++ b/src/tools/miri/src/concurrency/thread.rs
@@ -3,6 +3,7 @@
 use std::cell::RefCell;
 use std::collections::hash_map::Entry;
 use std::num::TryFromIntError;
+use std::sync::atomic::{AtomicBool, Ordering::Relaxed};
 use std::task::Poll;
 use std::time::{Duration, SystemTime};
 
@@ -1012,8 +1013,24 @@ fn unregister_timeout_callback_if_exists(&mut self, thread: ThreadId) {
     /// Run the core interpreter loop. Returns only when an interrupt occurs (an error or program
     /// termination).
     fn run_threads(&mut self) -> InterpResult<'tcx, !> {
+        static SIGNALED: AtomicBool = AtomicBool::new(false);
+        ctrlc::set_handler(move || {
+            // Indicate that we have ben signaled to stop. If we were already signaled, exit
+            // immediately. In our interpreter loop we try to consult this value often, but if for
+            // whatever reason we don't get to that check or the cleanup we do upon finding that
+            // this bool has become true takes a long time, the exit here will promptly exit the
+            // process on the second Ctrl-C.
+            if SIGNALED.swap(true, Relaxed) {
+                std::process::exit(1);
+            }
+        })
+        .unwrap();
         let this = self.eval_context_mut();
         loop {
+            if SIGNALED.load(Relaxed) {
+                this.machine.handle_abnormal_termination();
+                std::process::exit(1);
+            }
             match this.machine.threads.schedule(&this.machine.clock)? {
                 SchedulingAction::ExecuteStep => {
                     if !this.step()? {
diff --git a/src/tools/miri/src/machine.rs b/src/tools/miri/src/machine.rs
index 32717a0d28b..29f518fe58b 100644
--- a/src/tools/miri/src/machine.rs
+++ b/src/tools/miri/src/machine.rs
@@ -713,6 +713,15 @@ pub(crate) fn is_local(&self, frame: &FrameInfo<'_>) -> bool {
         let def_id = frame.instance.def_id();
         def_id.is_local() || self.local_crates.contains(&def_id.krate)
     }
+
+    /// Called when the interpreter is going to shut down abnormally, such as due to a Ctrl-C.
+    pub(crate) fn handle_abnormal_termination(&mut self) {
+        // All strings in the profile data are stored in a single string table which is not
+        // written to disk until the profiler is dropped. If the interpreter exits without dropping
+        // the profiler, it is not possible to interpret the profile data and all measureme tools
+        // will panic when given the file.
+        drop(self.profiler.take());
+    }
 }
 
 impl VisitTags for MiriMachine<'_, '_> {