unlock unlimited code runs for $79, limited time →
NOSSLETTER.techLIVE
back to OSSpedia
python/cpython
mid8/10
data-structuresobservabilitytesting-strategy

gh-154086: Fix flamegraph thread sample counts

Project: python/cpython | Language: Python | Concepts: data-structures, observability, testing-strategy | Difficulty: mid
Link: https://github.com/python/cpython/pull/154104

The problem

The flamegraph visualizer for Python's sampling profiler did not accurately show per-thread sample counts. When filtering by thread, frame widths and totals were not updated to reflect only the selected thread's data, leading to misleading visualizations and incorrect statistics.

Why it's non-trivial

A naive fix might filter the tree nodes by thread but would miss updating the sample counts, self-times, and opcode stats for each thread. Without storing and propagating per-thread values at every node, the filtered view cannot correctly represent the thread's activity, and frame widths or totals will be wrong. The fix must ensure that all relevant data is tracked and the tree structure remains consistent and accurate after filtering.

The approach

1. Store per-thread sample counts and self-times in the collector's tree nodes

-        self._root = {"samples": 0, "children": {}, "threads": set()}
+        self._root = {
+            "samples": 0,
+            "children": {},
+            "threads": set(),
+            "thread_samples": collections.Counter(),
+            "thread_self": collections.Counter(),
+        }

This change adds two new fields to each node: thread_samples and thread_self, which track the total and self sample counts per thread. This allows accurate per-thread metrics to be stored as the profile is collected.

2. Update sample collection to increment per-thread counters

         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)
         ...
         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
         ...
         if current is not self._root:
             current["self"] += weight
+            current["thread_self"][thread_id] += weight

This ensures that each time a stack sample is processed, the relevant per-thread counters are updated for both total and self samples, as well as per-thread opcode counts.

3. Include per-thread values in the serialized flamegraph data

+                    "thread_values": {
+                        thread_id: [
+                            samples,
+                            node["thread_self"].get(thread_id, 0),
+                        ]
+                        for thread_id, samples in sorted(
+                            node["thread_samples"].items()
+                        )
+                    },
...
+                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()
+                        )
+                    }

Each node in the output tree now includes a thread_values dictionary mapping thread IDs to their sample and self counts, and a thread_opcodes dictionary for per-thread opcode stats.

4. Adjust pruning logic to preserve low-volume threads

-                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

Nodes with low global sample counts are pruned unless they are significant for any thread (at least 0.1% of that thread's samples or at least one sample). This ensures rare threads are not omitted.

5. Update the frontend to use per-thread values when filtering

-function selectFlamegraphData() {
-  const baseData = isShowingElided ? elidedFlamegraphData : normalData;
+function selectFlamegraphData(selectedThreadId = null) {
+  let baseData = isShowingElided ? elidedFlamegraphData : normalData;
+
+  if (selectedThreadId !== null) {
+    baseData = filterDataByThread(baseData, selectedThreadId);
+  }
...
-  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 selectedThreadId = currentThreadFilter !== 'all' ? parseInt(currentThreadFilter, 10) : null;
+  const selectedData = selectFlamegraphData(selectedThreadId);
+  const tooltip = createPythonTooltip(selectedData);
+  const chart = createFlamegraph(tooltip, selectedData.value, selectedData);
+  renderFlamegraph(chart, selectedData);
   populateThreadStats(selectedData, selectedThreadId);

The frontend now uses the per-thread filtered tree and displays the correct sample counts and frame widths for the selected thread.

6. Update tests to verify per-thread values and correct filtering

+    def test_flamegraph_nodes_include_per_thread_values(self):
+        ...
+        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):
+        ...
+        self.assertEqual(data["thread_values"], {1: [1999, 0], 2: [1, 0]})
+        ...
+    def test_flamegraph_nodes_include_per_thread_opcodes(self):
+        ...
+        self.assertEqual(
+            data["thread_opcodes"],
+            {1: {100: 2}, 2: {101: 3}},
+        )

Tests now check that per-thread values are present, correct, and that rare threads are not pruned away.

From the review thread

No notable review discussion on this PR.

Takeaways

  • When supporting filtering or aggregation by key (such as thread ID), store per-key values at every node to enable accurate filtered views.
  • Data pruning and tree transformations must consider all relevant keys to avoid omitting significant but low-volume data.
  • Tests should explicitly check new invariants and edge cases introduced by data model changes, such as correct per-key values and preservation of rare keys.