gh-138122: Allow to filter by thread in tachyon's flamegraph (#139216)

Lib/profiling/sampling/collector.py

@@ -30,4 +30,4 @@ def _iter_all_frames(self, stack_frames, skip_idle=False):
                     continue
                 frames = thread_info.frame_info
                 if frames:
-                    yield frames
+                    yield frames, thread_info.thread_id

Lib/profiling/sampling/flamegraph.css

@@ -227,6 +227,65 @@ .controls button.secondary:hover {
   background: #ffcd02;
 }
 
+.thread-filter-wrapper {
+  display: none;
+  align-items: center;
+  margin-left: 16px;
+  background: white;
+  border-radius: 6px;
+  padding: 4px 8px 4px 12px;
+  border: 2px solid #3776ab;
+  transition: all 0.2s ease;
+}
+
+.thread-filter-wrapper:hover {
+  border-color: #2d5aa0;
+  box-shadow: 0 2px 6px rgba(55, 118, 171, 0.2);
+}
+
+.thread-filter-label {
+  color: #3776ab;
+  font-size: 14px;
+  font-weight: 600;
+  margin-right: 8px;
+  display: flex;
+  align-items: center;
+}
+
+.thread-filter-select {
+  background: transparent;
+  color: #2e3338;
+  border: none;
+  padding: 4px 24px 4px 4px;
+  font-size: 14px;
+  font-weight: 600;
+  cursor: pointer;
+  min-width: 120px;
+  font-family: inherit;
+  appearance: none;
+  -webkit-appearance: none;
+  -moz-appearance: none;
+  background-image: url("data:image/svg+xml;charset=UTF-8,%3csvg xmlns='http://www.w3.org/2000/svg' viewBox='0 0 24 24' fill='none' stroke='%233776ab' stroke-width='2' stroke-linecap='round' stroke-linejoin='round'%3e%3cpolyline points='6 9 12 15 18 9'%3e%3c/polyline%3e%3c/svg%3e");
+  background-repeat: no-repeat;
+  background-position: right 4px center;
+  background-size: 16px;
+}
+
+.thread-filter-select:focus {
+  outline: none;
+}
+
+.thread-filter-select:hover {
+  color: #3776ab;
+}
+
+.thread-filter-select option {
+  padding: 8px;
+  background: white;
+  color: #2e3338;
+  font-weight: normal;
+}
+
 #chart {
   width: 100%;
   height: calc(100vh - 160px);

Lib/profiling/sampling/flamegraph.js

@@ -2,6 +2,8 @@ const EMBEDDED_DATA = {{FLAMEGRAPH_DATA}};
 
 // Global string table for resolving string indices
 let stringTable = [];
+let originalData = null;
+let currentThreadFilter = 'all';
 
 // Function to resolve string indices to actual strings
 function resolveString(index) {
@@ -374,6 +376,12 @@ function initFlamegraph() {
     processedData = resolveStringIndices(EMBEDDED_DATA);
   }
 
+  // Store original data for filtering
+  originalData = processedData;
+
+  // Initialize thread filter dropdown
+  initThreadFilter(processedData);
+
   const tooltip = createPythonTooltip(processedData);
   const chart = createFlamegraph(tooltip, processedData.value);
   renderFlamegraph(chart, processedData);
@@ -395,10 +403,26 @@ function populateStats(data) {
   const functionMap = new Map();
 
   function collectFunctions(node) {
-    const filename = resolveString(node.filename);
+    if (!node) return;
-    const funcname = resolveString(node.funcname);
 
-    if (filename && funcname) {
+    let filename = typeof node.filename === 'number' ? resolveString(node.filename) : node.filename;
+    let funcname = typeof node.funcname === 'number' ? resolveString(node.funcname) : node.funcname;
+
+    if (!filename || !funcname) {
+      const nameStr = typeof node.name === 'number' ? resolveString(node.name) : node.name;
+      if (nameStr?.includes('(')) {
+        const match = nameStr.match(/^(.+?)\s*\((.+?):(\d+)\)$/);
+        if (match) {
+          funcname = funcname || match[1];
+          filename = filename || match[2];
+        }
+      }
+    }
+
+    filename = filename || 'unknown';
+    funcname = funcname || 'unknown';
+
+    if (filename !== 'unknown' && funcname !== 'unknown' && node.value > 0) {
       // Calculate direct samples (this node's value minus children's values)
       let childrenValue = 0;
       if (node.children) {
@@ -447,15 +471,17 @@ function populateStats(data) {
   // Populate the 3 cards
   for (let i = 0; i < 3; i++) {
     const num = i + 1;
-    if (i < hotSpots.length) {
+    if (i < hotSpots.length && hotSpots[i]) {
       const hotspot = hotSpots[i];
-      const basename = hotspot.filename.split('/').pop();
+      const filename = hotspot.filename || 'unknown';
-      let funcDisplay = hotspot.funcname;
+      const basename = filename !== 'unknown' ? filename.split('/').pop() : 'unknown';
+      const lineno = hotspot.lineno ?? '?';
+      let funcDisplay = hotspot.funcname || 'unknown';
       if (funcDisplay.length > 35) {
         funcDisplay = funcDisplay.substring(0, 32) + '...';
       }
 
-      document.getElementById(`hotspot-file-${num}`).textContent = `${basename}:${hotspot.lineno}`;
+      document.getElementById(`hotspot-file-${num}`).textContent = `${basename}:${lineno}`;
       document.getElementById(`hotspot-func-${num}`).textContent = funcDisplay;
       document.getElementById(`hotspot-detail-${num}`).textContent = `${hotspot.directPercent.toFixed(1)}% samples (${hotspot.directSamples.toLocaleString()})`;
     } else {
@@ -505,3 +531,102 @@ function clearSearch() {
   }
 }
 
+function initThreadFilter(data) {
+  const threadFilter = document.getElementById('thread-filter');
+  const threadWrapper = document.querySelector('.thread-filter-wrapper');
+
+  if (!threadFilter || !data.threads) {
+    return;
+  }
+
+  // Clear existing options except "All Threads"
+  threadFilter.innerHTML = '<option value="all">All Threads</option>';
+
+  // Add thread options
+  const threads = data.threads || [];
+  threads.forEach(threadId => {
+    const option = document.createElement('option');
+    option.value = threadId;
+    option.textContent = `Thread ${threadId}`;
+    threadFilter.appendChild(option);
+  });
+
+  // Show filter if more than one thread
+  if (threads.length > 1 && threadWrapper) {
+    threadWrapper.style.display = 'inline-flex';
+  }
+}
+
+function filterByThread() {
+  const threadFilter = document.getElementById('thread-filter');
+  if (!threadFilter || !originalData) return;
+
+  const selectedThread = threadFilter.value;
+  currentThreadFilter = selectedThread;
+
+  let filteredData;
+  if (selectedThread === 'all') {
+    // Show all data
+    filteredData = originalData;
+  } else {
+    // Filter data by thread
+    const threadId = parseInt(selectedThread);
+    filteredData = filterDataByThread(originalData, threadId);
+
+    if (filteredData.strings) {
+      stringTable = filteredData.strings;
+      filteredData = resolveStringIndices(filteredData);
+    }
+  }
+
+  // Re-render flamegraph with filtered data
+  const tooltip = createPythonTooltip(filteredData);
+  const chart = createFlamegraph(tooltip, filteredData.value);
+  renderFlamegraph(chart, filteredData);
+}
+
+function filterDataByThread(data, threadId) {
+  function filterNode(node) {
+    if (!node.threads || !node.threads.includes(threadId)) {
+      return null;
+    }
+
+    const filteredNode = {
+      ...node,
+      children: []
+    };
+
+    if (node.children && Array.isArray(node.children)) {
+      filteredNode.children = node.children
+        .map(child => filterNode(child))
+        .filter(child => child !== null);
+    }
+
+    return filteredNode;
+  }
+
+  const filteredRoot = {
+    ...data,
+    children: []
+  };
+
+  if (data.children && Array.isArray(data.children)) {
+    filteredRoot.children = data.children
+      .map(child => filterNode(child))
+      .filter(child => child !== null);
+  }
+
+  function recalculateValue(node) {
+    if (!node.children || node.children.length === 0) {
+      return node.value || 0;
+    }
+    const childrenValue = node.children.reduce((sum, child) => sum + recalculateValue(child), 0);
+    node.value = Math.max(node.value || 0, childrenValue);
+    return node.value;
+  }
+
+  recalculateValue(filteredRoot);
+
+  return filteredRoot;
+}
+

Lib/profiling/sampling/flamegraph_template.html

@@ -63,6 +63,12 @@ <h1>Tachyon Profiler Performance Flamegraph</h1>
         <button onclick="resetZoom()">🏠 Reset Zoom</button>
         <button onclick="exportSVG()" class="secondary">📁 Export SVG</button>
         <button onclick="toggleLegend()">🔥 Heat Map Legend</button>
+        <div class="thread-filter-wrapper">
+          <label class="thread-filter-label">🧵 Thread:</label>
+          <select id="thread-filter" class="thread-filter-select" onchange="filterByThread()">
+            <option value="all">All Threads</option>
+          </select>
+        </div>
       </div>
     </div>
 

Lib/profiling/sampling/pstats_collector.py

@@ -41,7 +41,7 @@ def _process_frames(self, frames):
             self.callers[callee][caller] += 1
 
     def collect(self, stack_frames):
-        for frames in self._iter_all_frames(stack_frames, skip_idle=self.skip_idle):
+        for frames, thread_id in self._iter_all_frames(stack_frames, skip_idle=self.skip_idle):
             self._process_frames(frames)
 
     def export(self, filename):

Lib/profiling/sampling/sample.py

@@ -754,7 +754,7 @@ def main():
         "--mode",
         choices=["wall", "cpu", "gil"],
         default="wall",
-        help="Sampling mode: wall (all threads), cpu (only CPU-running threads), gil (only GIL-holding threads)",
+        help="Sampling mode: wall (all threads), cpu (only CPU-running threads), gil (only GIL-holding threads) (default: wall)",
     )
 
     # Output format selection

Lib/profiling/sampling/stack_collector.py

@@ -15,12 +15,12 @@ def __init__(self, *, skip_idle=False):
         self.skip_idle = skip_idle
 
     def collect(self, stack_frames, skip_idle=False):
-        for frames in self._iter_all_frames(stack_frames, skip_idle=skip_idle):
+        for frames, thread_id in self._iter_all_frames(stack_frames, skip_idle=skip_idle):
             if not frames:
                 continue
-            self.process_frames(frames)
+            self.process_frames(frames, thread_id)
 
-    def process_frames(self, frames):
+    def process_frames(self, frames, thread_id):
         pass
 
 
@@ -29,17 +29,17 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.stack_counter = collections.Counter()
 
-    def process_frames(self, frames):
+    def process_frames(self, frames, thread_id):
         call_tree = tuple(reversed(frames))
-        self.stack_counter[call_tree] += 1
+        self.stack_counter[(call_tree, thread_id)] += 1
 
     def export(self, filename):
         lines = []
-        for call_tree, count in self.stack_counter.items():
+        for (call_tree, thread_id), count in self.stack_counter.items():
             stack_str = ";".join(
                 f"{os.path.basename(f[0])}:{f[2]}:{f[1]}" for f in call_tree
             )
-            lines.append((stack_str, count))
+            lines.append((f"tid:{thread_id};{stack_str}", count))
 
         lines.sort(key=lambda x: (-x[1], x[0]))
 
@@ -53,10 +53,11 @@ class FlamegraphCollector(StackTraceCollector):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.stats = {}
-        self._root = {"samples": 0, "children": {}}
+        self._root = {"samples": 0, "children": {}, "threads": set()}
         self._total_samples = 0
         self._func_intern = {}
         self._string_table = StringTable()
+        self._all_threads = set()
 
     def set_stats(self, sample_interval_usec, duration_sec, sample_rate, error_rate=None):
         """Set profiling statistics to include in flamegraph data."""
@@ -111,6 +112,7 @@ def _convert_to_flamegraph_format(self):
                 "name": self._string_table.intern("No Data"),
                 "value": 0,
                 "children": [],
+                "threads": [],
                 "strings": self._string_table.get_strings()
             }
 
@@ -133,6 +135,7 @@ def convert_children(children, min_samples):
                     "filename": filename_idx,
                     "lineno": func[1],
                     "funcname": funcname_idx,
+                    "threads": sorted(list(node.get("threads", set()))),
                 }
 
                 source = self._get_source_lines(func)
@@ -172,6 +175,7 @@ def convert_children(children, min_samples):
             new_name = f"Program Root: {old_name}"
             main_child["name"] = self._string_table.intern(new_name)
             main_child["stats"] = self.stats
+            main_child["threads"] = sorted(list(self._all_threads))
             main_child["strings"] = self._string_table.get_strings()
             return main_child
 
@@ -180,14 +184,17 @@ def convert_children(children, min_samples):
             "value": total_samples,
             "children": root_children,
             "stats": self.stats,
+            "threads": sorted(list(self._all_threads)),
             "strings": self._string_table.get_strings()
         }
 
-    def process_frames(self, frames):
+    def process_frames(self, frames, thread_id):
         # Reverse to root->leaf
         call_tree = reversed(frames)
         self._root["samples"] += 1
         self._total_samples += 1
+        self._root["threads"].add(thread_id)
+        self._all_threads.add(thread_id)
 
         current = self._root
         for func in call_tree:
@@ -195,9 +202,10 @@ def process_frames(self, frames):
             children = current["children"]
             node = children.get(func)
             if node is None:
-                node = {"samples": 0, "children": {}}
+                node = {"samples": 0, "children": {}, "threads": set()}
                 children[func] = node
             node["samples"] += 1
+            node["threads"].add(thread_id)
             current = node
 
     def _get_source_lines(self, func):

Lib/test/test_profiling/test_sampling_profiler.py

@@ -279,8 +279,9 @@ def test_collapsed_stack_collector_with_empty_and_deep_stacks(self):
         test_frames = [MockInterpreterInfo(0, [MockThreadInfo(1, [("file.py", 10, "func")])])]
         collector.collect(test_frames)
         self.assertEqual(len(collector.stack_counter), 1)
-        ((path,), count), = collector.stack_counter.items()
+        ((path, thread_id), count), = collector.stack_counter.items()
-        self.assertEqual(path, ("file.py", 10, "func"))
+        self.assertEqual(path, (("file.py", 10, "func"),))
+        self.assertEqual(thread_id, 1)
         self.assertEqual(count, 1)
 
         # Test with very deep stack
@@ -289,10 +290,11 @@ def test_collapsed_stack_collector_with_empty_and_deep_stacks(self):
         collector = CollapsedStackCollector()
         collector.collect(test_frames)
         # One aggregated path with 100 frames (reversed)
-        (path_tuple,), = (collector.stack_counter.keys(),)
+        ((path_tuple, thread_id),), = (collector.stack_counter.keys(),)
         self.assertEqual(len(path_tuple), 100)
         self.assertEqual(path_tuple[0], ("file99.py", 99, "func99"))
         self.assertEqual(path_tuple[-1], ("file0.py", 0, "func0"))
+        self.assertEqual(thread_id, 1)
 
     def test_pstats_collector_basic(self):
         """Test basic PstatsCollector functionality."""
@@ -394,9 +396,10 @@ def test_collapsed_stack_collector_basic(self):
 
         # Should store one reversed path
         self.assertEqual(len(collector.stack_counter), 1)
-        (path, count), = collector.stack_counter.items()
+        ((path, thread_id), count), = collector.stack_counter.items()
         expected_tree = (("file.py", 20, "func2"), ("file.py", 10, "func1"))
         self.assertEqual(path, expected_tree)
+        self.assertEqual(thread_id, 1)
         self.assertEqual(count, 1)
 
     def test_collapsed_stack_collector_export(self):
@@ -426,9 +429,9 @@ def test_collapsed_stack_collector_export(self):
         lines = content.strip().split("\n")
         self.assertEqual(len(lines), 2)  # Two unique stacks
 
-        # Check collapsed format: file:func:line;file:func:line count
+        # Check collapsed format: tid:X;file:func:line;file:func:line count
-        stack1_expected = "file.py:func2:20;file.py:func1:10 2"
+        stack1_expected = "tid:1;file.py:func2:20;file.py:func1:10 2"
-        stack2_expected = "other.py:other_func:5 1"
+        stack2_expected = "tid:1;other.py:other_func:5 1"
 
         self.assertIn(stack1_expected, lines)
         self.assertIn(stack2_expected, lines)
@@ -1517,7 +1520,8 @@ def test_collapsed_stack_with_recursion(self):
         self.assertEqual(len(collector.stack_counter), 2)
 
         # First path should be longer (deeper recursion) than the second
-        paths = list(collector.stack_counter.keys())
+        path_tuples = list(collector.stack_counter.keys())
+        paths = [p[0] for p in path_tuples]  # Extract just the call paths
         lengths = [len(p) for p in paths]
         self.assertNotEqual(lengths[0], lengths[1])
 
@@ -1530,7 +1534,7 @@ def test_collapsed_stack_with_recursion(self):
 
         def total_occurrences(func):
             total = 0
-            for path, count in collector.stack_counter.items():
+            for (path, thread_id), count in collector.stack_counter.items():
                 total += sum(1 for f in path if f == func) * count
             return total