perf(generator): don't collect garbage while generating the task graph

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

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

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

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

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

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