gh-135953: Reduce memory usage of stack collectors (#138875)

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.

Co-authored-by: Pablo Galindo Salgado <pablogsal@gmail.com>
This commit is contained in:
László Kiss Kollár 2025-09-14 23:47:14 +01:00 committed by GitHub
parent efc08c5fbf
commit 3e06cfcaee
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 292 additions and 174 deletions

View file

@ -7,5 +7,6 @@
from .collector import Collector from .collector import Collector
from .pstats_collector import PstatsCollector from .pstats_collector import PstatsCollector
from .stack_collector import CollapsedStackCollector from .stack_collector import CollapsedStackCollector
from .string_table import StringTable
__all__ = ("Collector", "PstatsCollector", "CollapsedStackCollector") __all__ = ("Collector", "PstatsCollector", "CollapsedStackCollector", "StringTable")

View file

@ -1,5 +1,50 @@
const EMBEDDED_DATA = {{FLAMEGRAPH_DATA}}; const EMBEDDED_DATA = {{FLAMEGRAPH_DATA}};
// Global string table for resolving string indices
let stringTable = [];
// Function to resolve string indices to actual strings
function resolveString(index) {
if (typeof index === 'number' && index >= 0 && index < stringTable.length) {
return stringTable[index];
}
// Fallback for non-indexed strings or invalid indices
return String(index);
}
// Function to recursively resolve all string indices in flamegraph data
function resolveStringIndices(node) {
if (!node) return node;
// Create a copy to avoid mutating the original
const resolved = { ...node };
// Resolve string fields
if (typeof resolved.name === 'number') {
resolved.name = resolveString(resolved.name);
}
if (typeof resolved.filename === 'number') {
resolved.filename = resolveString(resolved.filename);
}
if (typeof resolved.funcname === 'number') {
resolved.funcname = resolveString(resolved.funcname);
}
// Resolve source lines if present
if (Array.isArray(resolved.source)) {
resolved.source = resolved.source.map(index =>
typeof index === 'number' ? resolveString(index) : index
);
}
// Recursively resolve children
if (Array.isArray(resolved.children)) {
resolved.children = resolved.children.map(child => resolveStringIndices(child));
}
return resolved;
}
// Python color palette - cold to hot // Python color palette - cold to hot
const pythonColors = [ const pythonColors = [
"#fff4bf", // Coldest - light yellow (<1%) "#fff4bf", // Coldest - light yellow (<1%)
@ -100,6 +145,10 @@ function createPythonTooltip(data) {
</div>`; </div>`;
} }
// Resolve strings for display
const funcname = resolveString(d.data.funcname) || resolveString(d.data.name);
const filename = resolveString(d.data.filename) || "";
const tooltipHTML = ` const tooltipHTML = `
<div> <div>
<div style="color: #3776ab; font-weight: 600; font-size: 16px; <div style="color: #3776ab; font-weight: 600; font-size: 16px;
@ -257,9 +306,9 @@ function updateSearchHighlight(searchTerm, searchInput) {
let matchCount = 0; let matchCount = 0;
d3.selectAll("#chart rect").each(function (d) { d3.selectAll("#chart rect").each(function (d) {
if (d && d.data) { if (d && d.data) {
const name = d.data.name || ""; const name = resolveString(d.data.name) || "";
const funcname = d.data.funcname || ""; const funcname = resolveString(d.data.funcname) || "";
const filename = d.data.filename || ""; const filename = resolveString(d.data.filename) || "";
const term = searchTerm.toLowerCase(); const term = searchTerm.toLowerCase();
const matches = const matches =
name.toLowerCase().includes(term) || name.toLowerCase().includes(term) ||
@ -317,12 +366,20 @@ function handleResize(chart, data) {
function initFlamegraph() { function initFlamegraph() {
ensureLibraryLoaded(); ensureLibraryLoaded();
const tooltip = createPythonTooltip(EMBEDDED_DATA);
const chart = createFlamegraph(tooltip, EMBEDDED_DATA.value); // Extract string table if present and resolve string indices
renderFlamegraph(chart, EMBEDDED_DATA); let processedData = EMBEDDED_DATA;
if (EMBEDDED_DATA.strings) {
stringTable = EMBEDDED_DATA.strings;
processedData = resolveStringIndices(EMBEDDED_DATA);
}
const tooltip = createPythonTooltip(processedData);
const chart = createFlamegraph(tooltip, processedData.value);
renderFlamegraph(chart, processedData);
attachPanelControls(); attachPanelControls();
initSearchHandlers(); initSearchHandlers();
handleResize(chart, EMBEDDED_DATA); handleResize(chart, processedData);
} }
if (document.readyState === "loading") { if (document.readyState === "loading") {
@ -338,7 +395,10 @@ function populateStats(data) {
const functionMap = new Map(); const functionMap = new Map();
function collectFunctions(node) { function collectFunctions(node) {
if (node.filename && node.funcname) { const filename = resolveString(node.filename);
const funcname = resolveString(node.funcname);
if (filename && funcname) {
// Calculate direct samples (this node's value minus children's values) // Calculate direct samples (this node's value minus children's values)
let childrenValue = 0; let childrenValue = 0;
if (node.children) { if (node.children) {
@ -347,7 +407,7 @@ function populateStats(data) {
const directSamples = Math.max(0, node.value - childrenValue); const directSamples = Math.max(0, node.value - childrenValue);
// Use file:line:funcname as key to ensure uniqueness // Use file:line:funcname as key to ensure uniqueness
const funcKey = `${node.filename}:${node.lineno || '?'}:${node.funcname}`; const funcKey = `${filename}:${node.lineno || '?'}:${funcname}`;
if (functionMap.has(funcKey)) { if (functionMap.has(funcKey)) {
const existing = functionMap.get(funcKey); const existing = functionMap.get(funcKey);
@ -355,15 +415,15 @@ function populateStats(data) {
existing.directPercent = (existing.directSamples / totalSamples) * 100; existing.directPercent = (existing.directSamples / totalSamples) * 100;
// Keep the most representative file/line (the one with more samples) // Keep the most representative file/line (the one with more samples)
if (directSamples > existing.maxSingleSamples) { if (directSamples > existing.maxSingleSamples) {
existing.filename = node.filename; existing.filename = filename;
existing.lineno = node.lineno || '?'; existing.lineno = node.lineno || '?';
existing.maxSingleSamples = directSamples; existing.maxSingleSamples = directSamples;
} }
} else { } else {
functionMap.set(funcKey, { functionMap.set(funcKey, {
filename: node.filename, filename: filename,
lineno: node.lineno || '?', lineno: node.lineno || '?',
funcname: node.funcname, funcname: funcname,
directSamples, directSamples,
directPercent: (directSamples / totalSamples) * 100, directPercent: (directSamples / totalSamples) * 100,
maxSingleSamples: directSamples maxSingleSamples: directSamples

View file

@ -7,51 +7,51 @@
import os import os
from .collector import Collector from .collector import Collector
from .string_table import StringTable
class StackTraceCollector(Collector): 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): def collect(self, stack_frames):
for frames in self._iter_all_frames(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): 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): def export(self, filename):
stack_counter = collections.Counter() lines = []
for call_tree in self.call_trees: for call_tree, count in self.stack_counter.items():
# Call tree is already in root->leaf order
stack_str = ";".join( stack_str = ";".join(
f"{os.path.basename(f[0])}:{f[2]}:{f[1]}" for f in call_tree 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: with open(filename, "w") as f:
for stack, count in stack_counter.items(): for stack, count in lines:
f.write(f"{stack} {count}\n") f.write(f"{stack} {count}\n")
print(f"Collapsed stack output written to {filename}") print(f"Collapsed stack output written to {filename}")
class FlamegraphCollector(StackTraceCollector): class FlamegraphCollector(StackTraceCollector):
def __init__(self): def __init__(self):
super().__init__()
self.stats = {} self.stats = {}
self._root = {"samples": 0, "children": {}}
self._total_samples = 0
self._func_intern = {}
self._string_table = StringTable()
def set_stats(self, sample_interval_usec, duration_sec, sample_rate, error_rate=None): def set_stats(self, sample_interval_usec, duration_sec, sample_rate, error_rate=None):
"""Set profiling statistics to include in flamegraph data.""" """Set profiling statistics to include in flamegraph data."""
@ -65,11 +65,13 @@ def set_stats(self, sample_interval_usec, duration_sec, sample_rate, error_rate=
def export(self, filename): def export(self, filename):
flamegraph_data = self._convert_to_flamegraph_format() flamegraph_data = self._convert_to_flamegraph_format()
# Debug output # Debug output with string table statistics
num_functions = len(flamegraph_data.get("children", [])) num_functions = len(flamegraph_data.get("children", []))
total_time = flamegraph_data.get("value", 0) total_time = flamegraph_data.get("value", 0)
string_count = len(self._string_table)
print( print(
f"Flamegraph data: {num_functions} root functions, total samples: {total_time}" f"Flamegraph data: {num_functions} root functions, total samples: {total_time}, "
f"{string_count} unique strings"
) )
if num_functions == 0: if num_functions == 0:
@ -98,105 +100,105 @@ def _format_function_name(func):
return f"{funcname} ({filename}:{lineno})" return f"{funcname} ({filename}:{lineno})"
def _convert_to_flamegraph_format(self): def _convert_to_flamegraph_format(self):
"""Convert call trees to d3-flamegraph format with optimized hierarchy building""" """Convert aggregated trie to d3-flamegraph format with string table optimization."""
if not self.call_trees: if self._total_samples == 0:
return {"name": "No Data", "value": 0, "children": []} return {
"name": self._string_table.intern("No Data"),
unique_functions = set() "value": 0,
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": [], "children": [],
"strings": self._string_table.get_strings()
} }
if "filename" in node: def convert_children(children, min_samples):
result.update( out = []
{ for func, node in children.items():
"filename": node["filename"], samples = node["samples"]
"lineno": node["lineno"], if samples < min_samples:
"funcname": node["funcname"], continue
}
# Intern all string components for maximum efficiency
filename_idx = self._string_table.intern(func[0])
funcname_idx = self._string_table.intern(func[2])
name_idx = self._string_table.intern(self._format_function_name(func))
child_entry = {
"name": name_idx,
"value": samples,
"children": [],
"filename": filename_idx,
"lineno": func[1],
"funcname": funcname_idx,
}
source = self._get_source_lines(func)
if source:
# Intern source lines for memory efficiency
source_indices = [self._string_table.intern(line) for line in source]
child_entry["source"] = source_indices
# Recurse
child_entry["children"] = convert_children(
node["children"], min_samples
) )
out.append(child_entry)
if source_code: # Sort by value (descending) then by name index for consistent ordering
result["source"] = source_code out.sort(key=lambda x: (-x["value"], x["name"]))
return out
# 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
# Filter out very small functions (less than 0.1% of total samples) # 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)) min_samples = max(1, int(total_samples * 0.001))
converted_root = convert_node(root, min_samples) root_children = convert_children(self._root["children"], min_samples)
if not root_children:
if not converted_root or not converted_root["children"]: return {
return {"name": "No significant data", "value": 0, "children": []} "name": self._string_table.intern("No significant data"),
"value": 0,
"children": [],
"strings": self._string_table.get_strings()
}
# If we only have one root child, make it the root to avoid redundant level # If we only have one root child, make it the root to avoid redundant level
if len(converted_root["children"]) == 1: if len(root_children) == 1:
main_child = converted_root["children"][0] main_child = root_children[0]
main_child["name"] = f"Program Root: {main_child['name']}" # Update the name to indicate it's the program root
old_name = self._string_table.get_string(main_child["name"])
new_name = f"Program Root: {old_name}"
main_child["name"] = self._string_table.intern(new_name)
main_child["stats"] = self.stats main_child["stats"] = self.stats
main_child["strings"] = self._string_table.get_strings()
return main_child return main_child
converted_root["name"] = "Program Root" return {
converted_root["stats"] = self.stats "name": self._string_table.intern("Program Root"),
return converted_root "value": total_samples,
"children": root_children,
"stats": self.stats,
"strings": self._string_table.get_strings()
}
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): def _get_source_lines(self, func):
filename, lineno, funcname = func filename, lineno, _ = func
try: try:
# Get several lines around the function definition
lines = [] lines = []
start_line = max(1, lineno - 2) start_line = max(1, lineno - 2)
end_line = lineno + 3 end_line = lineno + 3
@ -210,7 +212,6 @@ def _get_source_lines(self, func):
return lines if lines else None return lines if lines else None
except Exception: except Exception:
# If we can't get source code, return None
return None return None
def _create_flamegraph_html(self, data): def _create_flamegraph_html(self, data):

View file

@ -0,0 +1,53 @@
"""String table implementation for memory-efficient string storage in profiling data."""
class StringTable:
"""A string table for interning strings and reducing memory usage."""
def __init__(self):
self._strings = []
self._string_to_index = {}
def intern(self, string):
"""Intern a string and return its index.
Args:
string: The string to intern
Returns:
int: The index of the string in the table
"""
if not isinstance(string, str):
string = str(string)
if string in self._string_to_index:
return self._string_to_index[string]
index = len(self._strings)
self._strings.append(string)
self._string_to_index[string] = index
return index
def get_string(self, index):
"""Get a string by its index.
Args:
index: The index of the string
Returns:
str: The string at the given index, or empty string if invalid
"""
if 0 <= index < len(self._strings):
return self._strings[index]
return ""
def get_strings(self):
"""Get the list of all strings in the table.
Returns:
list: A copy of the strings list
"""
return self._strings.copy()
def __len__(self):
"""Return the number of strings in the table."""
return len(self._strings)

View file

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