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()