gh-154086: Fix flamegraph thread sample counts

Author: pablogsal

Lib/profiling/sampling/_flamegraph_assets/flamegraph.js

@@ -99,13 +99,23 @@ function getDisplayName(moduleName, filename) {
   return filename;
 }
 
-function selectFlamegraphData() {
-  const baseData = isShowingElided ? elidedFlamegraphData : normalData;
+function selectFlamegraphData(selectedThreadId = null) {
+  let baseData = isShowingElided ? elidedFlamegraphData : normalData;
+
+  if (selectedThreadId !== null) {
+    baseData = filterDataByThread(baseData, selectedThreadId);
+  }
 
   if (!isInverted) {
     return baseData;
   }
 
+  // Thread-filtered trees have different values, so invert them after filtering
+  // instead of using the cached all-thread tree.
+  if (selectedThreadId !== null) {
+    return generateInvertedFlamegraph(baseData);
+  }
+
   if (isShowingElided) {
     if (!invertedElidedData) {
       invertedElidedData = generateInvertedFlamegraph(baseData);
@@ -120,12 +130,11 @@ function selectFlamegraphData() {
 }
 
 function updateFlamegraphView() {
-  const selectedData = selectFlamegraphData();
   const selectedThreadId = currentThreadFilter !== 'all' ? parseInt(currentThreadFilter, 10) : null;
-  const filteredData = selectedThreadId !== null ? filterDataByThread(selectedData, selectedThreadId) : selectedData;
-  const tooltip = createPythonTooltip(filteredData);
-  const chart = createFlamegraph(tooltip, filteredData.value, filteredData);
-  renderFlamegraph(chart, filteredData);
+  const selectedData = selectFlamegraphData(selectedThreadId);
+  const tooltip = createPythonTooltip(selectedData);
+  const chart = createFlamegraph(tooltip, selectedData.value, selectedData);
+  renderFlamegraph(chart, selectedData);
   populateThreadStats(selectedData, selectedThreadId);
 }
 
@@ -1209,7 +1218,7 @@ function initThreadFilter(data) {
   const threadFilter = document.getElementById('thread-filter');
   const threadSection = document.getElementById('thread-section');
 
-  if (!threadFilter || !data.threads) return;
+  if (!threadFilter || !data.threads || data.stats?.is_differential) return;
 
   threadFilter.innerHTML = '<option value="all">All Threads</option>';
 
@@ -1238,11 +1247,17 @@ function filterByThread() {
 
 function filterDataByThread(data, threadId) {
   function filterNode(node) {
-    if (!node.threads || !node.threads.includes(threadId)) {
+    const threadValues = node.thread_values?.[threadId];
+    if (!threadValues) {
       return null;
     }
 
-    const filteredNode = { ...node, children: [] };
+    const filteredNode = {
+      ...node,
+      value: threadValues[0],
+      self: threadValues[1],
+      children: []
+    };
 
     if (node.children && Array.isArray(node.children)) {
       filteredNode.children = node.children
@@ -1253,25 +1268,7 @@ function filterDataByThread(data, threadId) {
     return filteredNode;
   }
 
-  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;
-  }
-
-  const filteredRoot = { ...data, children: [] };
-
-  if (data.children && Array.isArray(data.children)) {
-    filteredRoot.children = data.children
-      .map(child => filterNode(child))
-      .filter(child => child !== null);
-  }
-
-  recalculateValue(filteredRoot);
-  return filteredRoot;
+  return filterNode(data);
 }
 
 // ============================================================================

Lib/profiling/sampling/stack_collector.py

@@ -71,7 +71,13 @@ class FlamegraphCollector(StackTraceCollector):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.stats = {}
-        self._root = {"samples": 0, "children": {}, "threads": set()}
+        self._root = {
+            "samples": 0,
+            "children": {},
+            "threads": set(),
+            "thread_samples": collections.Counter(),
+            "thread_self": collections.Counter(),
+        }
         self._total_samples = 0
         self._sample_count = 0  # Track actual number of samples (not thread traces)
         self._func_intern = {}
@@ -243,6 +249,15 @@ def convert_children(children, min_samples, path_info):
                     "lineno": func[1],
                     "funcname": funcname_idx,
                     "threads": sorted(list(node.get("threads", set()))),
+                    "thread_values": {
+                        thread_id: [
+                            samples,
+                            node["thread_self"].get(thread_id, 0),
+                        ]
+                        for thread_id, samples in sorted(
+                            node["thread_samples"].items()
+                        )
+                    },
                 }
 
                 source = self._get_source_lines(func)
@@ -340,6 +355,12 @@ def convert_children(children, min_samples, path_info):
                 "per_thread_stats": per_thread_stats_with_pct
             },
             "threads": sorted(list(self._all_threads)),
+            "thread_values": {
+                thread_id: [samples, 0]
+                for thread_id, samples in sorted(
+                    self._root["thread_samples"].items()
+                )
+            },
             "strings": self._string_table.get_strings(),
             "opcode_mapping": opcode_mapping
         }
@@ -356,6 +377,7 @@ def process_frames(self, frames, thread_id, weight=1):
         """
         # Reverse to root->leaf order for tree building
         self._root["samples"] += weight
+        self._root["thread_samples"][thread_id] += weight
         self._total_samples += weight
         self._root["threads"].add(thread_id)
         self._all_threads.add(thread_id)
@@ -368,9 +390,18 @@ def process_frames(self, frames, thread_id, weight=1):
 
             node = current["children"].get(func)
             if node is None:
-                node = {"samples": 0, "children": {}, "threads": set(), "opcodes": collections.Counter(), "self": 0}
+                node = {
+                    "samples": 0,
+                    "children": {},
+                    "threads": set(),
+                    "thread_samples": collections.Counter(),
+                    "thread_self": collections.Counter(),
+                    "opcodes": collections.Counter(),
+                    "self": 0,
+                }
                 current["children"][func] = node
             node["samples"] += weight
+            node["thread_samples"][thread_id] += weight
             node["threads"].add(thread_id)
 
             if opcode is not None:
@@ -380,6 +411,7 @@ def process_frames(self, frames, thread_id, weight=1):
 
         if current is not self._root:
             current["self"] += weight
+            current["thread_self"][thread_id] += weight
 
     def _get_source_lines(self, func):
         filename, lineno, _ = func

Lib/test/test_profiling/test_sampling_profiler/test_collectors.py

@@ -1336,6 +1336,27 @@ def test_flamegraph_collector_json_structure_includes_stats(self):
             self.assertIn("gc_pct", thread_data)
             self.assertIn("total", thread_data)
 
+    def test_flamegraph_nodes_include_per_thread_values(self):
+        collector = FlamegraphCollector(sample_interval_usec=1000)
+        root = MockFrameInfo("app.py", 1, "main")
+        collector.process_frames(
+            [MockFrameInfo("app.py", 10, "worker_a"), root],
+            thread_id=1,
+            weight=2,
+        )
+        collector.process_frames(
+            [MockFrameInfo("app.py", 20, "worker_b"), root],
+            thread_id=2,
+            weight=3,
+        )
+
+        data = collector._convert_to_flamegraph_format()
+
+        self.assertEqual(data["thread_values"], {1: [2, 0], 2: [3, 0]})
+        children_by_line = {child["lineno"]: child for child in data["children"]}
+        self.assertEqual(children_by_line[10]["thread_values"], {1: [2, 2]})
+        self.assertEqual(children_by_line[20]["thread_values"], {2: [3, 3]})
+
     def test_flamegraph_collector_per_thread_gc_percentage(self):
         """Test that per-thread GC percentage uses total samples as denominator."""
         collector = FlamegraphCollector(sample_interval_usec=1000)

Misc/NEWS.d/next/Library/2026-07-19-13-00-00.gh-issue-154086.MJKrQV.rst

@@ -0,0 +1,2 @@
+Store per-thread sample counts in Tachyon flamegraphs so filtering a thread
+updates frame widths and totals.