Changed Prober to gracefully terminate on SIGINT

src/prober.cc

@@ -22,6 +22,7 @@
 
 #include <algorithm>
 #include <cassert>
+#include <csignal>
 #include <cstring>
 #include <fstream>
 #include <iostream>
@@ -391,13 +392,22 @@ int Prober::Run(const PyFrob &frobber) {
   return dump_ ? DumpStacks(frobber, output) : ProbeLoop(frobber, output);
 }
 
+static volatile std::sig_atomic_t should_stop = 0;
+
+void SIGINT_handler(int sig) {
+  if (sig == SIGINT) {
+    should_stop = 1;
+  }
+}
+
 // Main loop to probe the Python process.
 int Prober::ProbeLoop(const PyFrob &frobber, std::ostream *out) {
   std::vector<FrameTS> call_stacks;
   int return_code = 0;
   size_t idle_count = 0;
   size_t failed_count = 0;
   bool check_end = seconds_ >= 0;
+  std::signal(SIGINT, SIGINT_handler);
   auto end = std::chrono::system_clock::now() + ToMicroseconds(seconds_);
   for (;;) {
     auto now = std::chrono::system_clock::now();
@@ -419,7 +429,7 @@ int Prober::ProbeLoop(const PyFrob &frobber, std::ostream *out) {
         call_stacks.push_back({now, thread.frames()});
       }
 
-      if (check_end && (now + interval_ >= end)) {
+      if (should_stop || (check_end && (now + interval_ >= end))) {
         break;
       }
       PtraceCont(pid_);

tests/test_end_to_end.py

@@ -19,6 +19,7 @@
 import platform
 import pytest
 import re
+import signal
 import subprocess
 import sys
 import time
@@ -637,3 +638,18 @@ def test_no_line_numbers(dijkstra):
     for line in lines:
         assert_flamegraph(
             line, allow_idle=True, line_re=FLAMEGRAPH_NONUMBER_RE)
+
+
+def test_sigint(dijkstra):
+    proc = subprocess.Popen(
+        [path_to_pyflame(), '-p',
+         str(dijkstra.pid), "--seconds=-1"],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        universal_newlines=True)
+    time.sleep(0.1)
+    proc.send_signal(signal.SIGINT)
+    out, err = communicate(proc)
+    assert not err
+    assert proc.returncode == 0
+    assert out