From 90feef0d0f67afa9deddc9f3ca31773a084c0d53 Mon Sep 17 00:00:00 2001 From: Julien Cristau Date: Thu, 8 Oct 2026 16:50:48 +0200 Subject: [PATCH 1/3] generator: check for dependency cycles Avoid surprising results (deadlocks or incomplete graphs) later. --- src/taskgraph/generator.py | 11 +++++++++++ test/test_generator.py | 25 +++++++++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/src/taskgraph/generator.py b/src/taskgraph/generator.py index 1368ef01c..dc86c98b9 100644 --- a/src/taskgraph/generator.py +++ b/src/taskgraph/generator.py @@ -446,6 +446,12 @@ def _run(self): for kind in kinds.values(): for dep in kind.config.get("kind-dependencies", []): edges.add((kind.name, dep, "kind-dependency")) + for _, dep, _ in edges: + if dep not in kinds: + message = f'Could not find the kind "{dep}"\nAvailable kinds:\n' + for k in sorted(kinds): + message += f' - "{k}"\n' + raise Exception(message) kind_graph = Graph(frozenset(kinds), frozenset(edges)) if target_kinds: @@ -453,6 +459,11 @@ def _run(self): set(target_kinds) | {"docker-image"} ) + # The task loaders wait for each kind's dependencies to be loaded + # before loading it, so a dependency loop would hang them or silently + # leave kinds out. This raises if there is one. + kind_graph.visit_postorder() + yield "kind_graph", kind_graph logger.info("Generating full task set") diff --git a/test/test_generator.py b/test/test_generator.py index 8a31c186a..95e88cfa8 100644 --- a/test/test_generator.py +++ b/test/test_generator.py @@ -437,6 +437,31 @@ def test_kind_graph(maketgg): } +def test_kind_graph_missing_kind(maketgg): + "A dependency on a kind that doesn't exist is an error" + tgg = maketgg( + kinds=[ + ("_fake2", {"kind-dependencies": ["_missing"]}), + ("_fake1", {"kind-dependencies": []}), + ] + ) + with pytest.raises(Exception, match='Could not find the kind "_missing"'): + tgg.kind_graph + + +def test_kind_graph_dependency_loop(maketgg): + "A dependency loop between kinds is an error" + tgg = maketgg( + kinds=[ + ("_fake3", {"kind-dependencies": ["_fake2"]}), + ("_fake2", {"kind-dependencies": ["_fake3"]}), + ("_fake1", {"kind-dependencies": []}), + ] + ) + with pytest.raises(Exception, match="Dependency loop detected"): + tgg.kind_graph + + def test_kind_graph_with_target_kinds(maketgg): "The kind_graph property respects target_kinds parameter" tgg = maketgg( From 1b49458cc3c09bd7dc5f7f0a22869954a255be93 Mon Sep 17 00:00:00 2001 From: Julien Cristau Date: Wed, 7 Oct 2026 23:35:13 +0200 Subject: [PATCH 2/3] perf(generator): fork a process per kind instead of using a process pool On Linux, kinds were loaded in a ProcessPoolExecutor whose workers are forked before any task exists, so the tasks of all the kinds a kind depends on were pickled and sent to it, and its own tasks pickled and sent back. In Firefox's taskgraph (~190 kinds, ~49,000 tasks) that meant 74,000 task pickles to the workers, all on the parent's threads under the GIL, with unrelated kinds waiting behind big transfers. Instead, fork a child for each kind once the kinds it depends on are loaded, up to the number of CPUs at once. The child shares every task loaded so far with the parent, so nothing is sent to it, and it sends its tasks back through a pipe of its own, read on the parent's main thread as children finish. Errors are reported as before, with the traceback from the child. A child exiting without sending its tasks, or tasks that can't be pickled, are reported as errors loading the kind, and the other children are killed on error. Kinds are still loaded by calling load_tasks on the Kind instances from _load_kinds, so subclasses keep working. TASKGRAPH_SERIAL, TASKGRAPH_USE_THREADS and other platforms are unchanged, except that the thread pool no longer needs Python 3.13's os.process_cpu_count. Part of https://bugzilla.mozilla.org/show_bug.cgi?id=2079635 --- src/taskgraph/generator.py | 187 ++++++++++++++++++++++++++++++++++--- test/test_generator.py | 160 +++++++++++++++++++++++++++++-- 2 files changed, 326 insertions(+), 21 deletions(-) diff --git a/src/taskgraph/generator.py b/src/taskgraph/generator.py index dc86c98b9..c90f3ac65 100644 --- a/src/taskgraph/generator.py +++ b/src/taskgraph/generator.py @@ -6,16 +6,20 @@ import inspect import logging import multiprocessing +import multiprocessing.connection import os +import pickle import platform +import signal +import sys +import traceback from concurrent.futures import ( FIRST_COMPLETED, - ProcessPoolExecutor, ThreadPoolExecutor, wait, ) from dataclasses import dataclass -from typing import Callable, Optional, Union +from typing import Callable, Optional, Union, cast from . import filter_tasks from .config import GraphConfig, load_graph_config @@ -40,6 +44,12 @@ class KindNotFound(Exception): """ +def _max_workers(): + # os.process_cpu_count is new in Python 3.13. + cpu_count = getattr(os, "process_cpu_count", os.cpu_count) + return cpu_count() or 1 + + @dataclass(frozen=True) class Kind: name: str @@ -124,6 +134,16 @@ def load(cls, root_dir, graph_config, kind_name): return cls(kind_name, path, config, graph_config) +class _RemoteTraceback(Exception): + """The traceback of an exception raised in a child process.""" + + def __init__(self, tb): + self.tb = tb + + def __str__(self): + return f'\n"""\n{self.tb}"""' + + class TaskGraphGenerator: """ The central controller for taskgraph. This handles all phases of graph @@ -399,6 +419,153 @@ def submit_ready_kinds(): return all_tasks + def _load_tasks_forked(self, kinds, kind_graph, parameters): + """Load each kind in a child process forked once the kinds it depends + on are loaded. + + The child shares every task loaded so far with the parent, so the + tasks of the kinds it depends on don't need to be sent to it, and it + sends its own tasks back through a pipe of its own. + """ + all_tasks = {} + tasks_by_kind = {} + # Pipe to each running child -> (kind name, pid) + running = {} + remaining = kinds.copy() + edges = set(kind_graph.edges) + max_workers = _max_workers() + + try: + while remaining or running: + kinds_with_deps = {edge[0] for edge in edges} + ready_kinds = sorted(set(remaining) - kinds_with_deps) + for name in ready_kinds[: max_workers - len(running)]: + kind = remaining.pop(name) + reader, pid = self._fork_kind( + kind, + parameters, + self._get_kind_dependencies_tasks(kind, tasks_by_kind), + ) + running[reader] = (name, pid) + + # Waiting on nothing would block forever. _run made sure there + # is no dependency loop, so some kind is always running here. + assert running + for ready in multiprocessing.connection.wait(list(running)): + # wait returns the objects it was given, but is typed as + # also returning sockets and file descriptors. + reader = cast(multiprocessing.connection.Connection, ready) + name, pid = running[reader] + try: + result = reader.recv() + except EOFError: + # The child exited without sending anything. + result = None + del running[reader] + reader.close() + _, status = os.waitpid(pid, 0) + new_tasks = self._kind_result(name, result, status) + + kind_tasks = tasks_by_kind.setdefault(name, {}) + for task in new_tasks: + if task.label in all_tasks: + raise Exception("duplicate tasks with label " + task.label) + all_tasks[task.label] = task + kind_tasks[task.label] = task + + edges = {e for e in edges if e[1] != name} + finally: + # Only left with running children on error. + for reader, (_, pid) in running.items(): + reader.close() + os.kill(pid, signal.SIGKILL) + os.waitpid(pid, 0) + + return all_tasks + + def _fork_kind(self, kind, parameters, kind_dependencies_tasks): + """Fork a child process loading the tasks of `kind`, and return the + pipe its result comes back through, and its pid.""" + reader, writer = multiprocessing.Pipe(duplex=False) + pid = os.fork() + if pid == 0: + # Never return into the caller's code in the child. + status = 1 + try: + reader.close() + status = self._load_kind_in_child( + kind, parameters, kind_dependencies_tasks, writer + ) + finally: + try: + sys.stdout.flush() + sys.stderr.flush() + finally: + os._exit(status) + writer.close() + return reader, pid + + def _load_kind_in_child(self, kind, parameters, kind_dependencies_tasks, writer): + """Load the tasks of `kind` and send them, or the exception loading + them raised, through `writer`. Return the child's exit status.""" + try: + tasks = kind.load_tasks( + parameters, kind_dependencies_tasks, self._write_artifacts + ) + except BaseException as e: + tb = traceback.format_exc() + try: + # Make sure the parent will be able to unpickle it. + pickle.loads(pickle.dumps(e)) + except Exception: + e = Exception(f"{type(e).__name__}: {e}") + writer.send((e, tb, None)) + return 1 + + try: + writer.send((None, None, tasks)) + except Exception as e: + writer.send( + ( + Exception( + f"Could not send the tasks of kind {kind.name} " + f"to the parent process: {type(e).__name__}: {e}" + ), + traceback.format_exc(), + None, + ) + ) + return 1 + return 0 + + @staticmethod + def _kind_result(name, result, status): + """Return the tasks the child loading kind `name` sent, given what it + sent and its exit status. Raise the exception loading them raised in + the child, if any.""" + if result is None: + code = os.waitstatus_to_exitcode(status) + if code < 0: + how = f"was killed by signal {signal.Signals(-code).name}" + else: + how = f"exited with status {code}" + raise Exception( + f"Process loading tasks for kind {name} {how} without sending them" + ) + + exc, tb, tasks = result + if exc is not None: + exc.__cause__ = _RemoteTraceback(tb) + if isinstance(exc, SchemaValidationError): + logger.error(f"Error loading tasks for kind {name}:\n{exc}") + else: + logger.error( + f"Error loading tasks for kind {name}:", + exc_info=(type(exc), exc, exc.__traceback__), + ) + raise exc + return tasks + def _run(self): logger.info("Loading graph configuration.") graph_config = load_graph_config(self.root_dir) @@ -471,7 +638,7 @@ def _run(self): # The next block deals with enabling parallel kind processing, which # currently has different support on different platforms. In summary: # * Parallel kind processing is supported and enabled by default on - # Linux. We use multiple processes by default, but experimental + # Linux. We fork a process per kind by default, but experimental # support for multiple threads can be enabled instead. # * On other platforms, we have experimental support for parallel # kind processing with multiple threads. @@ -491,21 +658,19 @@ def load_tasks(): if os.environ.get("TASKGRAPH_SERIAL"): return self._load_tasks_serial(kinds, kind_graph, parameters) elif os.environ.get("TASKGRAPH_USE_THREADS"): - executor = ThreadPoolExecutor(max_workers=os.process_cpu_count()) - else: - executor = ProcessPoolExecutor( - mp_context=multiprocessing.get_context("fork") + executor = ThreadPoolExecutor(max_workers=_max_workers()) + return self._load_tasks_parallel( + kinds, kind_graph, parameters, executor ) - return self._load_tasks_parallel( - kinds, kind_graph, parameters, executor - ) + else: + return self._load_tasks_forked(kinds, kind_graph, parameters) else: if os.environ.get("TASKGRAPH_SERIAL") or not os.environ.get( "TASKGRAPH_USE_THREADS" ): return self._load_tasks_serial(kinds, kind_graph, parameters) else: - executor = ThreadPoolExecutor(max_workers=os.process_cpu_count()) + executor = ThreadPoolExecutor(max_workers=_max_workers()) return self._load_tasks_parallel( kinds, kind_graph, parameters, executor ) diff --git a/test/test_generator.py b/test/test_generator.py index 95e88cfa8..ff547e651 100644 --- a/test/test_generator.py +++ b/test/test_generator.py @@ -5,10 +5,13 @@ import os import platform +import signal +import time from concurrent.futures import ProcessPoolExecutor import pytest from pytest_taskgraph import WithFakeKind, fake_load_graph_config +from pytest_taskgraph.fixtures.gen import FakeKind from taskgraph import generator, graph from taskgraph.generator import Kind, load_tasks_for_kind, load_tasks_for_kinds @@ -27,7 +30,7 @@ ) -class FakePPE(ProcessPoolExecutor): +class FakeTPE(ProcessPoolExecutor): loaded_kinds = [] def submit(self, kind_load_tasks, *args): @@ -35,18 +38,39 @@ def submit(self, kind_load_tasks, *args): return super().submit(kind_load_tasks, *args) -class FakeTPE(ProcessPoolExecutor): - loaded_kinds = [] +class RecordingKind(FakeKind): + """Records which process loaded it, and the kinds of the dependency + tasks it was given, in the file at `record_path`.""" - def submit(self, kind_load_tasks, *args): - self.loaded_kinds.append(kind_load_tasks.__self__.name) - return super().submit(kind_load_tasks, *args) + record_path = None + + def load_tasks(self, parameters, loaded_tasks, write_artifacts): + dep_kinds = sorted({t.kind for t in loaded_tasks.values()}) + line = f"{self.name} {os.getpid()} {','.join(dep_kinds)}\n" + fd = os.open(self.record_path, os.O_WRONLY | os.O_APPEND | os.O_CREAT) + try: + os.write(fd, line.encode()) + finally: + os.close(fd) + return super().load_tasks(parameters, loaded_tasks, write_artifacts) + + +def use_kind_class(tgg, cls): + """Make `tgg` load its kinds as instances of `cls`.""" + load_kinds = tgg._load_kinds + + def _load_kinds(graph_config, target_kinds=None): + for kind in load_kinds(graph_config, target_kinds): + yield cls(kind.name, kind.path, kind.config, kind.graph_config) + + tgg._load_kinds = _load_kinds @linuxonly -def test_kind_ordering_multiprocess(mocker, maketgg): - "When task kinds depend on each other, they are loaded in postorder" - mocked_ppe = mocker.patch.object(generator, "ProcessPoolExecutor", new=FakePPE) +def test_kind_ordering_forked(monkeypatch, tmp_path, maketgg): + """Each kind is loaded in a child process, with its subclass' load_tasks, + after the kinds it depends on, and is given their tasks.""" + monkeypatch.setattr(RecordingKind, "record_path", str(tmp_path / "record")) tgg = maketgg( kinds=[ ("_fake3", {"kind-dependencies": ["_fake2", "_fake1"]}), @@ -54,8 +78,124 @@ def test_kind_ordering_multiprocess(mocker, maketgg): ("_fake1", {"kind-dependencies": []}), ] ) + use_kind_class(tgg, RecordingKind) tgg._run_until("full_task_set") - assert mocked_ppe.loaded_kinds == ["_fake1", "_fake2", "_fake3"] + assert len(tgg.full_task_set.tasks) == 9 + + records = [ + line.split(" ") for line in (tmp_path / "record").read_text().splitlines() + ] + assert [(name, deps) for name, _, deps in records] == [ + ("_fake1", ""), + ("_fake2", "_fake1"), + ("_fake3", "_fake1,_fake2"), + ] + pids = {int(pid) for _, pid, _ in records} + assert len(pids) == 3 + assert os.getpid() not in pids + + +@linuxonly +def test_forked_max_workers(mocker, monkeypatch, tmp_path, maketgg): + "Independent kinds are all loaded when fewer can be loaded at once." + mocker.patch.object(generator, "_max_workers", return_value=1) + tgg = maketgg(kinds=[(f"_fake{i}", {}) for i in range(4)]) + assert len(tgg.full_task_set.tasks) == 12 + + +class UnpicklableError(Exception): + def __init__(self, a, b): + super().__init__(f"{a} {b}") + + +@linuxonly +def test_forked_exception_traceback(mocker, caplog, maketgg): + """An exception raised loading a kind in a child is raised in the parent, + with the traceback from the child.""" + + def fail_in_child(self, *args, **kwargs): + raise UnpicklableError("not", "picklable") + + mocker.patch.object(Kind, "load_tasks", fail_in_child) + tgg = maketgg() + + with caplog.at_level("ERROR"): + with pytest.raises(Exception, match="UnpicklableError: not picklable") as e: + tgg._run_until("full_task_set") + + assert "in fail_in_child" in str(e.value.__cause__) + records = [r for r in caplog.records if "Error loading tasks" in r.message] + assert len(records) == 1 + assert records[0].message == "Error loading tasks for kind _fake:" + assert "in fail_in_child" in caplog.text + + +@linuxonly +@pytest.mark.parametrize( + "exit,expected", + ( + pytest.param(lambda: os._exit(3), "exited with status 3", id="exit"), + pytest.param( + lambda: os.kill(os.getpid(), signal.SIGKILL), + "was killed by signal SIGKILL", + id="killed", + ), + ), +) +def test_forked_child_dies(mocker, maketgg, exit, expected): + "A child exiting without sending its tasks is reported as an error." + + def die(self, *args, **kwargs): + exit() + + mocker.patch.object(Kind, "load_tasks", die) + tgg = maketgg() + with pytest.raises( + Exception, match=f"Process loading tasks for kind _fake {expected}" + ): + tgg._run_until("full_task_set") + + +@linuxonly +def test_forked_unpicklable_result(mocker, maketgg): + "Tasks that can't be sent back are reported as an error loading the kind." + + def unpicklable(self, *args, **kwargs): + return [lambda: None] + + mocker.patch.object(Kind, "load_tasks", unpicklable) + tgg = maketgg() + with pytest.raises( + Exception, match="Could not send the tasks of kind _fake to the parent" + ): + tgg._run_until("full_task_set") + + +@linuxonly +def test_forked_children_killed_on_error(mocker, tmp_path, maketgg): + "When loading a kind fails, the children loading other kinds are killed." + pid_file = tmp_path / "pid" + load_tasks = Kind.load_tasks + + def fail_or_hang(self, *args, **kwargs): + if self.name == "_fail": + # Wait for the other child to be running. + while not pid_file.exists(): + time.sleep(0.01) + raise RuntimeError("failed") + if self.name == "_hang": + pid_file.write_text(str(os.getpid())) + time.sleep(60) + return load_tasks(self, *args, **kwargs) + + mocker.patch.object(Kind, "load_tasks", fail_or_hang) + tgg = maketgg(kinds=[("_fail", {}), ("_hang", {})]) + start = time.monotonic() + with pytest.raises(RuntimeError, match="failed"): + tgg._run_until("full_task_set") + assert time.monotonic() - start < 30 + with pytest.raises(ProcessLookupError): + os.kill(int(pid_file.read_text()), 0) @threadsonly From a8a2c832a07ce6c24a5dfbf8b253956607e65085 Mon Sep 17 00:00:00 2001 From: Julien Cristau Date: Wed, 7 Oct 2026 23:36:07 +0200 Subject: [PATCH 3/3] 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. Part of https://bugzilla.mozilla.org/show_bug.cgi?id=2079640 --- docs/reference/source/taskgraph.util.rst | 8 +++++ src/taskgraph/generator.py | 37 +++++++++++-------- 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, 148 insertions(+), 22 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 c90f3ac65..d61ac00e1 100644 --- a/src/taskgraph/generator.py +++ b/src/taskgraph/generator.py @@ -30,6 +30,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 @@ -487,21 +488,25 @@ def _fork_kind(self, kind, parameters, kind_dependencies_tasks): """Fork a child process loading the tasks of `kind`, and return the pipe its result comes back through, and its pid.""" reader, writer = multiprocessing.Pipe(duplex=False) - pid = os.fork() - if pid == 0: - # Never return into the caller's code in the child. - status = 1 - try: - reader.close() - status = self._load_kind_in_child( - kind, parameters, kind_dependencies_tasks, writer - ) - finally: + # The child inherits the disabled garbage collector, so it doesn't + # write to every page it shares with us traversing the objects in them. + with gc_disabled(): + pid = os.fork() + if pid == 0: + # Never leave this block, or return into the caller's code, in + # the child. + status = 1 try: - sys.stdout.flush() - sys.stderr.flush() + reader.close() + status = self._load_kind_in_child( + kind, parameters, kind_dependencies_tasks, writer + ) finally: - os._exit(status) + try: + sys.stdout.flush() + sys.stderr.flush() + finally: + os._exit(status) writer.close() return reader, pid @@ -780,7 +785,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 ff547e651..1fcd01cd2 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 import signal @@ -401,6 +402,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()