From 1a8f7f55b4ef758947ec68cc3597e90cafa0dc6d Mon Sep 17 00:00:00 2001
From: Sergei Lebedev <slebedev@google.com>
Date: Fri, 18 Jan 2019 10:55:46 +0000
Subject: [PATCH] Allowed tracing a single Python thread

This is useful for running PyFlame from Jupyter Notebook,
where a cell could be isolated to a separate thread prior
to profiling.

See also #163
---
 src/prober.cc            | 16 ++++++++++++++--
 src/prober.h             |  2 ++
 tests/test_end_to_end.py | 25 +++++++++++++++++++++++++
 3 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/src/prober.cc b/src/prober.cc
index a691c83..5b995f4 100644
--- a/src/prober.cc
+++ b/src/prober.cc
@@ -48,6 +48,7 @@ static const char usage_str[] =
      "Common Options:\n"
 #ifdef ENABLE_THREADS
      "  --threads                Enable multi-threading support\n"
+     "  --only=IDENT             Only trace a thread identified by IDENT (requires --threads)\n"
      "  -d, --dump               Dump stacks from all threads (implies --threads)\n"
 #else
      "  -d, --dump               Dump the current interpreter stack\n"
@@ -189,6 +190,7 @@ int Prober::ParseOpts(int argc, char **argv) {
     {"seconds", required_argument, 0, 's'},
 #if ENABLE_THREADS
     {"threads", no_argument, 0, 'L'},
+    {"only", required_argument, 0, 'i'},
 #endif
     {"no-line-numbers", no_argument, 0, 'n'},
     {"output", required_argument, 0, 'o'},
@@ -238,7 +240,10 @@ int Prober::ParseOpts(int argc, char **argv) {
         std::cout << PYFLAME_VERSION_STR << "\n\n" << usage_str;
         return 0;
         break;
-#ifdef ENABLE_THREADS
+#if ENABLE_THREADS
+      case 'i':
+        thread_id_ = std::stoul(optarg);
+        break;
       case 'L':
         enable_threads_ = true;
         break;
@@ -309,6 +314,11 @@ int Prober::ParseOpts(int argc, char **argv) {
     std::cerr << "WARNING: Specifying a PID to trace without -p is deprecated; "
                  "see Pyflame issue #99 for details.\n";
   }
+  if (thread_id_ > 0 && !enable_threads_) {
+    std::cerr << "Option --only requires --threads.\n";
+    std::cerr << usage_str;
+    return 1;
+  }
   interval_ = ToMicroseconds(sample_rate_);
   return -1;
 }
@@ -416,7 +426,9 @@ int Prober::ProbeLoop(const PyFrob &frobber, std::ostream *out) {
       }
 
       for (const auto &thread : threads) {
-        call_stacks.push_back({now, thread.frames()});
+        if (thread_id_ == 0 || thread.id() == thread_id_) {
+          call_stacks.push_back({now, thread.frames()});
+        }
       }
 
       if (check_end && (now + interval_ >= end)) {
diff --git a/src/prober.h b/src/prober.h
index 1087def..4565a97 100644
--- a/src/prober.h
+++ b/src/prober.h
@@ -39,6 +39,7 @@ class Prober {
         include_ts_(false),
         include_line_number_(true),
         enable_threads_(false),
+        thread_id_(0),
         seconds_(1),
         sample_rate_(0.01) {}
   Prober(const Prober &other) = delete;
@@ -63,6 +64,7 @@ class Prober {
   bool include_ts_;
   bool include_line_number_;
   bool enable_threads_;
+  unsigned long thread_id_;
   double seconds_;
   double sample_rate_;
   std::chrono::microseconds interval_;
diff --git a/tests/test_end_to_end.py b/tests/test_end_to_end.py
index b573539..849e730 100644
--- a/tests/test_end_to_end.py
+++ b/tests/test_end_to_end.py
@@ -22,6 +22,7 @@
 import subprocess
 import sys
 import time
+import threading
 
 IDLE_RE = re.compile(r'^\(idle\) \d+$')
 FLAMEGRAPH_RE = re.compile(
@@ -637,3 +638,27 @@ def test_no_line_numbers(dijkstra):
     for line in lines:
         assert_flamegraph(
             line, allow_idle=True, line_re=FLAMEGRAPH_NONUMBER_RE)
+
+
+def test_only():
+    should_stop = False
+
+    def infinite_sleeper():
+        while not should_stop:
+            time.sleep(0.01)
+
+    thread = threading.Thread(target=infinite_sleeper)
+    thread.daemon = True
+    thread.start()
+
+    proc = subprocess.Popen(
+        [path_to_pyflame(), '-p',
+         str(os.getpid()), '--seconds=1',
+         '--threads', '--only=' + str(thread.ident)],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        universal_newlines=True)
+    out, err = communicate(proc)
+    assert not err
+    assert proc.returncode == 0
+    assert 'test_only' not in out