gh-154086: Make thread filtering exact

Author: pablogsal

Lib/profiling/sampling/_flamegraph_assets/flamegraph.js

@@ -946,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;
@@ -1252,10 +1254,16 @@ function filterDataByThread(data, threadId) {
       return null;
     }
 
+    const {
+      thread_values: _threadValues,
+      thread_opcodes: threadOpcodes,
+      ...sharedNode
+    } = node;
     const filteredNode = {
-      ...node,
+      ...sharedNode,
       value: threadValues[0],
       self: threadValues[1],
+      opcodes: threadOpcodes?.[threadId] ?? {},
       children: []
     };
 

Lib/profiling/sampling/stack_collector.py

@@ -226,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
@@ -270,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(
@@ -326,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"])
@@ -355,12 +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": {
-                thread_id: [samples, 0]
-                for thread_id, samples in sorted(
-                    self._root["thread_samples"].items()
-                )
-            },
+            "thread_values": root_thread_values,
             "strings": self._string_table.get_strings(),
             "opcode_mapping": opcode_mapping
         }
@@ -406,6 +438,10 @@ def process_frames(self, frames, thread_id, weight=1):
 
             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
 

Lib/test/test_profiling/test_sampling_profiler/test_collectors.py

@@ -1357,6 +1357,66 @@ def test_flamegraph_nodes_include_per_thread_values(self):
         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)