From 469632aa71c0667409ef5df18c811b809f1401f7 Mon Sep 17 00:00:00 2001
From: Sergei Lebedev <slebedev@google.com>
Date: Thu, 17 Jan 2019 20:28:51 +0000
Subject: [PATCH] Changed Prober to gracefully terminate on SIGINT

---
 src/prober.cc            | 12 +++++++++++-
 tests/test_end_to_end.py | 16 ++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/src/prober.cc b/src/prober.cc
index a691c83..f3a5afe 100644
--- a/src/prober.cc
+++ b/src/prober.cc
@@ -22,6 +22,7 @@
 
 #include <algorithm>
 #include <cassert>
+#include <csignal>
 #include <cstring>
 #include <fstream>
 #include <iostream>
@@ -391,6 +392,14 @@ 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;
@@ -398,6 +407,7 @@ int Prober::ProbeLoop(const PyFrob &frobber, std::ostream *out) {
   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_);
diff --git a/tests/test_end_to_end.py b/tests/test_end_to_end.py
index b573539..94872fa 100644
--- a/tests/test_end_to_end.py
+++ b/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