From c84c2fa00f6559f3624e5ffa7c5b78925f869c07 Mon Sep 17 00:00:00 2001 From: Julien Cristau Date: Wed, 7 Oct 2026 23:36:07 +0200 Subject: [PATCH] perf(generator): don't collect garbage while generating the task graph Generating a task graph creates millions of objects that stay alive until the end, so each collection traverses an ever growing heap without freeing anything. Disable the garbage collector while running each phase of the generator, restoring its previous state before yielding to the caller, and while loading a TaskGraph from JSON. Children loading kinds inherit the disabled collector, which also keeps them from writing to every page they share with the parent. Re-enabling the collector in children, or only freezing the heap after loading kinds and each phase (gc.freeze), measured slower in Firefox's taskgraph, with peak memory the same in all cases. --- docs/reference/source/taskgraph.util.rst | 8 +++++ src/taskgraph/generator.py | 7 +++- src/taskgraph/taskgraph.py | 18 +++++----- src/taskgraph/util/memory.py | 24 +++++++++++++ test/test_generator.py | 45 ++++++++++++++++++++++++ test/test_util_memory.py | 38 ++++++++++++++++++++ 6 files changed, 131 insertions(+), 9 deletions(-) create mode 100644 src/taskgraph/util/memory.py create mode 100644 test/test_util_memory.py diff --git a/docs/reference/source/taskgraph.util.rst b/docs/reference/source/taskgraph.util.rst index e9e363416..cd5d80920 100644 --- a/docs/reference/source/taskgraph.util.rst +++ b/docs/reference/source/taskgraph.util.rst @@ -60,6 +60,14 @@ taskgraph.util.keyed\_by module :undoc-members: :show-inheritance: +taskgraph.util.memory module +---------------------------- + +.. automodule:: taskgraph.util.memory + :members: + :undoc-members: + :show-inheritance: + taskgraph.util.parameterization module -------------------------------------- diff --git a/src/taskgraph/generator.py b/src/taskgraph/generator.py index 1368ef01c..81ff4f2a5 100644 --- a/src/taskgraph/generator.py +++ b/src/taskgraph/generator.py @@ -26,6 +26,7 @@ from .task import Task from .taskgraph import TaskGraph from .transforms.base import TransformConfig, TransformSequence +from .util.memory import gc_disabled from .util.python_path import find_object from .util.schema import SchemaValidationError from .util.verify import verifications @@ -604,7 +605,11 @@ def load_tasks(): def _run_until(self, name): while name not in self._run_results: try: - k, v = next(self._run) # type: ignore + # Generation creates millions of objects that stay alive, so + # collecting garbage meanwhile only traverses them for nothing. + # The callers' code runs with the collector in its usual state. + with gc_disabled(): + k, v = next(self._run) # type: ignore except StopIteration: raise AttributeError(f"No such run result {name}") self._run_results[k] = v diff --git a/src/taskgraph/taskgraph.py b/src/taskgraph/taskgraph.py index aee37a178..fdfabf1b2 100644 --- a/src/taskgraph/taskgraph.py +++ b/src/taskgraph/taskgraph.py @@ -6,6 +6,7 @@ from .graph import Graph from .task import Task +from .util.memory import gc_disabled @dataclass(frozen=True) @@ -58,13 +59,14 @@ def from_json(cls, tasks_dict): """ tasks = {} edges = set() - for key, value in tasks_dict.items(): - tasks[key] = Task.from_json(value) - if "task_id" in value: - tasks[key].task_id = value["task_id"] - for depname, dep in value["dependencies"].items(): - # Task filtering can cause dependencies to be removed from the graph. - if dep in tasks_dict: - edges.add((key, dep, depname)) + with gc_disabled(): + for key, value in tasks_dict.items(): + tasks[key] = Task.from_json(value) + if "task_id" in value: + tasks[key].task_id = value["task_id"] + for depname, dep in value["dependencies"].items(): + # Task filtering can cause dependencies to be removed from the graph. + if dep in tasks_dict: + edges.add((key, dep, depname)) task_graph = cls(tasks, Graph(frozenset(tasks), edges)) return tasks, task_graph diff --git a/src/taskgraph/util/memory.py b/src/taskgraph/util/memory.py new file mode 100644 index 000000000..3dc4117c0 --- /dev/null +++ b/src/taskgraph/util/memory.py @@ -0,0 +1,24 @@ +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, You can obtain one at http://mozilla.org/MPL/2.0/. + +import gc +from contextlib import contextmanager + + +@contextmanager +def gc_disabled(): + """Disable the garbage collector for the duration of the block, restoring + its previous state afterwards. + + Generating or loading a task graph creates millions of objects that stay + alive, so collections keep traversing a growing heap without freeing + anything. + """ + was_enabled = gc.isenabled() + gc.disable() + try: + yield + finally: + if was_enabled: + gc.enable() diff --git a/test/test_generator.py b/test/test_generator.py index 8a31c186a..ccbd8adb4 100644 --- a/test/test_generator.py +++ b/test/test_generator.py @@ -3,6 +3,7 @@ # file, You can obtain one at http://mozilla.org/MPL/2.0/. +import gc import os import platform from concurrent.futures import ProcessPoolExecutor @@ -261,6 +262,50 @@ def test_verifications(mocker, maketgg): m.assert_not_called() +@pytest.mark.parametrize("enabled", (True, False)) +def test_gc_disabled_during_generation(mocker, maketgg, enabled): + """The garbage collector is disabled while generating each phase, and + restored to its previous state between phases.""" + states = [] + mocker.patch.object( + generator, "verifications", side_effect=lambda *a: states.append(gc.isenabled()) + ) + tgg = maketgg(["_fake-t-2"]) + was_enabled = gc.isenabled() + if not enabled: + gc.disable() + try: + tgg.full_task_set + assert gc.isenabled() == enabled + tgg.morphed_task_graph + assert gc.isenabled() == enabled + finally: + if was_enabled: + gc.enable() + assert states and not any(states) + + +@linuxonly +def test_gc_disabled_in_forked_children(mocker, tmp_path, maketgg): + "Children loading kinds inherit the disabled garbage collector." + record = tmp_path / "gc" + load_tasks = Kind.load_tasks + + def record_gc(self, *args, **kwargs): + record.write_text(str(gc.isenabled())) + return load_tasks(self, *args, **kwargs) + + mocker.patch.object(Kind, "load_tasks", record_gc) + tgg = maketgg() + tgg._load_tasks_forked( + {k.name: k for k in tgg._load_kinds(tgg.graph_config)}, + tgg.kind_graph, + tgg.parameters, + ) + assert gc.isenabled() + assert record.read_text() == "False" + + def test_load_tasks_for_kind(monkeypatch): """ `load_tasks_for_kinds` will load the tasks for the provided kind diff --git a/test/test_util_memory.py b/test/test_util_memory.py new file mode 100644 index 000000000..8b5495d98 --- /dev/null +++ b/test/test_util_memory.py @@ -0,0 +1,38 @@ +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, You can obtain one at http://mozilla.org/MPL/2.0/. + +import gc + +import pytest + +from taskgraph.util.memory import gc_disabled + + +@pytest.fixture +def restore_gc(): + was_enabled = gc.isenabled() + yield + if was_enabled: + gc.enable() + else: + gc.disable() + + +def test_gc_disabled(restore_gc): + gc.enable() + with gc_disabled(): + assert not gc.isenabled() + assert gc.isenabled() + + with pytest.raises(RuntimeError): + with gc_disabled(): + raise RuntimeError() + assert gc.isenabled() + + +def test_gc_disabled_keeps_disabled(restore_gc): + gc.disable() + with gc_disabled(): + assert not gc.isenabled() + assert not gc.isenabled()