Auto merge of rust-lang#2899 - saethlin:measureme-cleanup, r=RalfJung

Try to ensure measureme output is valid when the interpreter is interrupted

Dropping our `measureme::Profiler` will have the profiler flush its output buffer and write out its string table. If the profiler's `Drop` impl does not run, the output file is not usable.

This approach has zero runtime overhead.

Author: bors

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"

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"

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 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
     /// 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()? {

src/tools/miri/src/machine.rs

@@ -713,6 +713,15 @@ impl<'mir, 'tcx> MiriMachine<'mir, 'tcx> {
         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<'_, '_> {