gh-154086: Fix flamegraph thread sample counts (GH-154104)

(cherry picked from commit e4b22ad)

Co-authored-by: Pablo Galindo Salgado <Pablogsal@gmail.com>

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);
 }
 
@@ -937,7 +946,9 @@ function formatDuration(seconds) {
 
 function populateProfileSummary(data) {
   const stats = data.stats || {};
-  const totalSamples = stats.total_samples || data.value || 0;
+  const totalSamples = currentThreadFilter !== 'all'
+    ? (data.value ?? 0)
+    : (stats.total_samples ?? data.value ?? 0);
   const duration = stats.duration_sec || 0;
   const sampleRate = stats.sample_rate || (duration > 0 ? totalSamples / duration : 0);
   const errorRate = stats.error_rate || 0;
@@ -1209,7 +1220,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 +1249,23 @@ 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 {
+      thread_values: _threadValues,
+      thread_opcodes: threadOpcodes,
+      ...sharedNode
+    } = node;
+    const filteredNode = {
+      ...sharedNode,
+      value: threadValues[0],
+      self: threadValues[1],
+      opcodes: threadOpcodes?.[threadId] ?? {},
+      children: []
+    };
 
     if (node.children && Array.isArray(node.children)) {
       filteredNode.children = node.children
@@ -1253,25 +1276,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 = {}
@@ -220,7 +226,18 @@ def convert_children(children, min_samples, path_info):
             out = []
             for func, node in children.items():
                 samples = node["samples"]
-                if samples < min_samples:
+                significant_for_thread = any(
+                    thread_samples >= max(
+                        1,
+                        int(
+                            self._root["thread_samples"][thread_id]
+                            * 0.001
+                        ),
+                    )
+                    for thread_id, thread_samples
+                    in node["thread_samples"].items()
+                )
+                if samples < min_samples and not significant_for_thread:
                     continue
 
                 # Intern all string components for maximum efficiency
@@ -243,6 +260,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)
@@ -255,6 +281,14 @@ def convert_children(children, min_samples, path_info):
                 opcodes = node.get("opcodes", {})
                 if opcodes:
                     child_entry["opcodes"] = dict(opcodes)
+                thread_opcodes = node.get("thread_opcodes")
+                if thread_opcodes:
+                    child_entry["thread_opcodes"] = {
+                        thread_id: dict(counts)
+                        for thread_id, counts in sorted(
+                            thread_opcodes.items()
+                        )
+                    }
 
                 # Recurse
                 child_entry["children"] = convert_children(
@@ -311,7 +345,25 @@ def convert_children(children, min_samples, path_info):
         opcode_mapping = get_opcode_mapping()
 
         # If we only have one root child, make it the root to avoid redundant level
-        if len(root_children) == 1:
+        root_thread_values = {
+            thread_id: [samples, 0]
+            for thread_id, samples in sorted(
+                self._root["thread_samples"].items()
+            )
+        }
+        sole_root_covers_profile = (
+            len(root_children) == 1
+            and root_children[0]["value"] == total_samples
+            and {
+                thread_id: values[0]
+                for thread_id, values
+                in root_children[0]["thread_values"].items()
+            } == {
+                thread_id: values[0]
+                for thread_id, values in root_thread_values.items()
+            }
+        )
+        if sole_root_covers_profile:
             main_child = root_children[0]
             # Update name and label to indicate it's the program root
             old_name = self._string_table.get_string(main_child["name"])
@@ -340,6 +392,7 @@ def convert_children(children, min_samples, path_info):
                 "per_thread_stats": per_thread_stats_with_pct
             },
             "threads": sorted(list(self._all_threads)),
+            "thread_values": root_thread_values,
             "strings": self._string_table.get_strings(),
             "opcode_mapping": opcode_mapping
         }
@@ -356,6 +409,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,18 +422,32 @@ 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:
                 node["opcodes"][opcode] += weight
+                thread_opcodes = node.setdefault("thread_opcodes", {})
+                thread_opcodes.setdefault(
+                    thread_id, collections.Counter()
+                )[opcode] += weight
 
             current = node
 
         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,87 @@ 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_pruning_preserves_low_volume_thread(self):
+        collector = FlamegraphCollector(sample_interval_usec=1000)
+        collector.process_frames(
+            [MockFrameInfo("app.py", 10, "busy")],
+            thread_id=1,
+            weight=1999,
+        )
+        collector.process_frames(
+            [MockFrameInfo("app.py", 20, "rare")],
+            thread_id=2,
+        )
+
+        data = collector._convert_to_flamegraph_format()
+
+        self.assertEqual(data["thread_values"], {1: [1999, 0], 2: [1, 0]})
+        children_by_line = {child["lineno"]: child for child in data["children"]}
+        self.assertEqual(children_by_line[10]["thread_values"], {1: [1999, 1999]})
+        self.assertEqual(children_by_line[20]["thread_values"], {2: [1, 1]})
+
+    def test_flamegraph_does_not_promote_incomplete_root(self):
+        collector = FlamegraphCollector(sample_interval_usec=1000)
+        collector.process_frames(
+            [MockFrameInfo("app.py", 1, "busy")],
+            thread_id=1,
+            weight=2000,
+        )
+        for line in range(2, 2002):
+            collector.process_frames(
+                [MockFrameInfo("app.py", line, f"fragment_{line}")],
+                thread_id=2,
+            )
+
+        data = collector._convert_to_flamegraph_format()
+
+        self.assertNotIn("filename", data)
+        self.assertEqual(data["thread_values"], {1: [2000, 0], 2: [2000, 0]})
+        self.assertEqual(len(data["children"]), 1)
+        self.assertEqual(data["children"][0]["thread_values"], {1: [2000, 2000]})
+
+    def test_flamegraph_nodes_include_per_thread_opcodes(self):
+        collector = FlamegraphCollector(sample_interval_usec=1000)
+        collector.process_frames(
+            [MockFrameInfo("app.py", 10, "worker", opcode=100)],
+            thread_id=1,
+            weight=2,
+        )
+        collector.process_frames(
+            [MockFrameInfo("app.py", 10, "worker", opcode=101)],
+            thread_id=2,
+            weight=3,
+        )
+
+        data = collector._convert_to_flamegraph_format()
+
+        self.assertEqual(data["opcodes"], {100: 2, 101: 3})
+        self.assertEqual(
+            data["thread_opcodes"],
+            {1: {100: 2}, 2: {101: 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.