Reduce memory usage of stack collectors

The stack collector base class keeps all frames until export() is
called, which causes significant unnecessary memory usage. Instead, we
can process the frames on the fly in the collect call by dispatching the
aggregation logic to the subclass through the process_frames method.

Author: lkollar

Lib/profiling/sampling/stack_collector.py

@@ -10,48 +10,46 @@
 
 
 class StackTraceCollector(Collector):
-    def __init__(self):
-        self.call_trees = []
-        self.function_samples = collections.defaultdict(int)
-
-    def _process_frames(self, frames):
-        """Process a single thread's frame stack."""
-        if not frames:
-            return
-
-        # Store the complete call stack (reverse order - root first)
-        call_tree = list(reversed(frames))
-        self.call_trees.append(call_tree)
-
-        # Count samples per function
-        for frame in frames:
-            self.function_samples[frame] += 1
-
     def collect(self, stack_frames):
         for frames in self._iter_all_frames(stack_frames):
-            self._process_frames(frames)
+            if not frames:
+                continue
+            self.process_frames(frames)
+
+    def process_frames(self, frames):
+        pass
 
 
 class CollapsedStackCollector(StackTraceCollector):
+    def __init__(self):
+        self.stack_counter = collections.Counter()
+
+    def process_frames(self, frames):
+        call_tree = tuple(reversed(frames))
+        self.stack_counter[call_tree] += 1
+
     def export(self, filename):
-        stack_counter = collections.Counter()
-        for call_tree in self.call_trees:
-            # Call tree is already in root->leaf order
+        lines = []
+        for call_tree, count in self.stack_counter.items():
             stack_str = ";".join(
                 f"{os.path.basename(f[0])}:{f[2]}:{f[1]}" for f in call_tree
             )
-            stack_counter[stack_str] += 1
+            lines.append((stack_str, count))
+
+        lines.sort(key=lambda x: (-x[1], x[0]))
 
         with open(filename, "w") as f:
-            for stack, count in stack_counter.items():
+            for stack, count in lines:
                 f.write(f"{stack} {count}\n")
         print(f"Collapsed stack output written to {filename}")
 
 
 class FlamegraphCollector(StackTraceCollector):
     def __init__(self):
-        super().__init__()
         self.stats = {}
+        self._root = {"samples": 0, "children": {}}
+        self._total_samples = 0
+        self._func_intern = {}
 
     def set_stats(self, sample_interval_usec, duration_sec, sample_rate, error_rate=None):
         """Set profiling statistics to include in flamegraph data."""
@@ -98,105 +96,78 @@ def _format_function_name(func):
         return f"{funcname} ({filename}:{lineno})"
 
     def _convert_to_flamegraph_format(self):
-        """Convert call trees to d3-flamegraph format with optimized hierarchy building"""
-        if not self.call_trees:
+        """Convert aggregated trie to d3-flamegraph format."""
+        if self._total_samples == 0:
             return {"name": "No Data", "value": 0, "children": []}
 
-        unique_functions = set()
-        for call_tree in self.call_trees:
-            unique_functions.update(call_tree)
-
-        func_to_name = {
-            func: self._format_function_name(func) for func in unique_functions
-        }
-
-        root = {"name": "root", "children": {}, "samples": 0}
-
-        for call_tree in self.call_trees:
-            current_node = root
-            current_node["samples"] += 1
-
-            for func in call_tree:
-                func_name = func_to_name[func]  # Use pre-computed name
-
-                if func_name not in current_node["children"]:
-                    current_node["children"][func_name] = {
-                        "name": func_name,
-                        "func": func,
-                        "children": {},
-                        "samples": 0,
-                        "filename": func[0],
-                        "lineno": func[1],
-                        "funcname": func[2],
-                    }
-
-                current_node = current_node["children"][func_name]
-                current_node["samples"] += 1
-
-        def convert_node(node, min_samples=1):
-            if node["samples"] < min_samples:
-                return None
-
-            source_code = None
-            if "func" in node:
-                source_code = self._get_source_lines(node["func"])
-
-            result = {
-                "name": node["name"],
-                "value": node["samples"],
-                "children": [],
-            }
-
-            if "filename" in node:
-                result.update(
-                    {
-                        "filename": node["filename"],
-                        "lineno": node["lineno"],
-                        "funcname": node["funcname"],
-                    }
+        def convert_children(children, min_samples):
+            out = []
+            for func, node in children.items():
+                samples = node["samples"]
+                if samples < min_samples:
+                    continue
+
+                name = self._format_function_name(func)
+                child_entry = {
+                    "name": name,
+                    "value": samples,
+                    "children": [],
+                    "filename": func[0],
+                    "lineno": func[1],
+                    "funcname": func[2],
+                }
+
+                source = self._get_source_lines(func)
+                if source:
+                    child_entry["source"] = source
+
+                # Recurse
+                child_entry["children"] = convert_children(
+                    node["children"], min_samples
                 )
+                out.append(child_entry)
 
-            if source_code:
-                result["source"] = source_code
-
-            # Recursively convert children
-            child_nodes = []
-            for child_name, child_node in node["children"].items():
-                child_result = convert_node(child_node, min_samples)
-                if child_result:
-                    child_nodes.append(child_result)
-
-            # Sort children by sample count (descending)
-            child_nodes.sort(key=lambda x: x["value"], reverse=True)
-            result["children"] = child_nodes
-
-            return result
+            out.sort(key=lambda x: (-x["value"], x["name"]))
+            return out
 
         # Filter out very small functions (less than 0.1% of total samples)
-        total_samples = len(self.call_trees)
+        total_samples = self._total_samples
         min_samples = max(1, int(total_samples * 0.001))
 
-        converted_root = convert_node(root, min_samples)
-
-        if not converted_root or not converted_root["children"]:
+        root_children = convert_children(self._root["children"], min_samples)
+        if not root_children:
             return {"name": "No significant data", "value": 0, "children": []}
 
         # If we only have one root child, make it the root to avoid redundant level
-        if len(converted_root["children"]) == 1:
-            main_child = converted_root["children"][0]
+        if len(root_children) == 1:
+            main_child = root_children[0]
             main_child["name"] = f"Program Root: {main_child['name']}"
             main_child["stats"] = self.stats
             return main_child
 
-        converted_root["name"] = "Program Root"
-        converted_root["stats"] = self.stats
-        return converted_root
+        return {"name": "Program Root", "value": total_samples, "children": root_children, "stats": self.stats}
+
+    def process_frames(self, frames):
+        # Reverse to root->leaf
+        call_tree = reversed(frames)
+        self._root["samples"] += 1
+        self._total_samples += 1
+
+        current = self._root
+        for func in call_tree:
+            func = self._func_intern.setdefault(func, func)
+            children = current["children"]
+            node = children.get(func)
+            if node is None:
+                node = {"samples": 0, "children": {}}
+                children[func] = node
+            node["samples"] += 1
+            current = node
 
     def _get_source_lines(self, func):
-        filename, lineno, funcname = func
+        filename, lineno, _ = func
 
         try:
-            # Get several lines around the function definition
             lines = []
             start_line = max(1, lineno - 2)
             end_line = lineno + 3
@@ -210,7 +181,6 @@ def _get_source_lines(self, func):
             return lines if lines else None
 
         except Exception:
-            # If we can't get source code, return None
             return None
 
     def _create_flamegraph_html(self, data):

Lib/test/test_profiling/test_sampling_profiler.py

@@ -272,25 +272,26 @@ def test_collapsed_stack_collector_with_empty_and_deep_stacks(self):
 
         # Test with empty frames
         collector.collect([])
-        self.assertEqual(len(collector.call_trees), 0)
+        self.assertEqual(len(collector.stack_counter), 0)
 
         # Test with single frame stack
         test_frames = [MockInterpreterInfo(0, [MockThreadInfo(1, [("file.py", 10, "func")])])]
         collector.collect(test_frames)
-        self.assertEqual(len(collector.call_trees), 1)
-        self.assertEqual(collector.call_trees[0], [("file.py", 10, "func")])
+        self.assertEqual(len(collector.stack_counter), 1)
+        ((path,), count), = collector.stack_counter.items()
+        self.assertEqual(path, ("file.py", 10, "func"))
+        self.assertEqual(count, 1)
 
         # Test with very deep stack
         deep_stack = [(f"file{i}.py", i, f"func{i}") for i in range(100)]
         test_frames = [MockInterpreterInfo(0, [MockThreadInfo(1, deep_stack)])]
         collector = CollapsedStackCollector()
         collector.collect(test_frames)
-        self.assertEqual(len(collector.call_trees[0]), 100)
-        # Check it's properly reversed
-        self.assertEqual(
-            collector.call_trees[0][0], ("file99.py", 99, "func99")
-        )
-        self.assertEqual(collector.call_trees[0][-1], ("file0.py", 0, "func0"))
+        # One aggregated path with 100 frames (reversed)
+        (path_tuple,), = (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"))
 
     def test_pstats_collector_basic(self):
         """Test basic PstatsCollector functionality."""
@@ -382,27 +383,20 @@ def test_collapsed_stack_collector_basic(self):
         collector = CollapsedStackCollector()
 
         # Test empty state
-        self.assertEqual(len(collector.call_trees), 0)
-        self.assertEqual(len(collector.function_samples), 0)
+        self.assertEqual(len(collector.stack_counter), 0)
 
         # Test collecting sample data
         test_frames = [
             MockInterpreterInfo(0, [MockThreadInfo(1, [("file.py", 10, "func1"), ("file.py", 20, "func2")])])
         ]
         collector.collect(test_frames)
 
-        # Should store call tree (reversed)
-        self.assertEqual(len(collector.call_trees), 1)
-        expected_tree = [("file.py", 20, "func2"), ("file.py", 10, "func1")]
-        self.assertEqual(collector.call_trees[0], expected_tree)
-
-        # Should count function samples
-        self.assertEqual(
-            collector.function_samples[("file.py", 10, "func1")], 1
-        )
-        self.assertEqual(
-            collector.function_samples[("file.py", 20, "func2")], 1
-        )
+        # Should store one reversed path
+        self.assertEqual(len(collector.stack_counter), 1)
+        (path, count), = collector.stack_counter.items()
+        expected_tree = (("file.py", 20, "func2"), ("file.py", 10, "func1"))
+        self.assertEqual(path, expected_tree)
+        self.assertEqual(count, 1)
 
     def test_collapsed_stack_collector_export(self):
         collapsed_out = tempfile.NamedTemporaryFile(delete=False)
@@ -441,9 +435,9 @@ def test_flamegraph_collector_basic(self):
         """Test basic FlamegraphCollector functionality."""
         collector = FlamegraphCollector()
 
-        # Test empty state (inherits from StackTraceCollector)
-        self.assertEqual(len(collector.call_trees), 0)
-        self.assertEqual(len(collector.function_samples), 0)
+        # Empty collector should produce 'No Data'
+        data = collector._convert_to_flamegraph_format()
+        self.assertIn(data["name"], ("No Data", "No significant data"))
 
         # Test collecting sample data
         test_frames = [
@@ -454,18 +448,18 @@ def test_flamegraph_collector_basic(self):
         ]
         collector.collect(test_frames)
 
-        # Should store call tree (reversed)
-        self.assertEqual(len(collector.call_trees), 1)
-        expected_tree = [("file.py", 20, "func2"), ("file.py", 10, "func1")]
-        self.assertEqual(collector.call_trees[0], expected_tree)
-
-        # Should count function samples
-        self.assertEqual(
-            collector.function_samples[("file.py", 10, "func1")], 1
-        )
-        self.assertEqual(
-            collector.function_samples[("file.py", 20, "func2")], 1
-        )
+        # Convert and verify structure: func2 -> func1 with counts = 1
+        data = collector._convert_to_flamegraph_format()
+        # Expect promotion: root is the single child (func2), with func1 as its only child
+        name = data.get("name", "")
+        self.assertIsInstance(name, str)
+        self.assertTrue(name.startswith("Program Root: "))
+        self.assertIn("func2 (file.py:20)", name)  # formatted name
+        children = data.get("children", [])
+        self.assertEqual(len(children), 1)
+        child = children[0]
+        self.assertIn("func1 (file.py:10)", child["name"])  # formatted name
+        self.assertEqual(child["value"], 1)
 
     def test_flamegraph_collector_export(self):
         """Test flamegraph HTML export functionality."""
@@ -1508,28 +1502,29 @@ def test_collapsed_stack_with_recursion(self):
         for frames in recursive_frames:
             collector.collect([frames])
 
-        # Should capture both call trees
-        self.assertEqual(len(collector.call_trees), 2)
-
-        # First tree should be longer (deeper recursion)
-        tree1 = collector.call_trees[0]
-        tree2 = collector.call_trees[1]
+        # Should capture both call paths
+        self.assertEqual(len(collector.stack_counter), 2)
 
-        # Trees should be different lengths due to different recursion depths
-        self.assertNotEqual(len(tree1), len(tree2))
+        # First path should be longer (deeper recursion) than the second
+        paths = list(collector.stack_counter.keys())
+        lengths = [len(p) for p in paths]
+        self.assertNotEqual(lengths[0], lengths[1])
 
         # Both should contain factorial calls
-        self.assertTrue(any("factorial" in str(frame) for frame in tree1))
-        self.assertTrue(any("factorial" in str(frame) for frame in tree2))
+        self.assertTrue(any(any(f[2] == "factorial" for f in p) for p in paths))
 
-        # Function samples should count all occurrences
+        # Verify total occurrences via aggregation
         factorial_key = ("factorial.py", 10, "factorial")
         main_key = ("main.py", 5, "main")
 
-        # factorial appears 5 times total (3 + 2)
-        self.assertEqual(collector.function_samples[factorial_key], 5)
-        # main appears 2 times total
-        self.assertEqual(collector.function_samples[main_key], 2)
+        def total_occurrences(func):
+            total = 0
+            for path, count in collector.stack_counter.items():
+                total += sum(1 for f in path if f == func) * count
+            return total
+
+        self.assertEqual(total_occurrences(factorial_key), 5)
+        self.assertEqual(total_occurrences(main_key), 2)
 
 
 @requires_subprocess()