From 930c6423d512384690f6fac9b808107422edb2f1 Mon Sep 17 00:00:00 2001
From: Marco Castelluccio <mcastelluccio@mozilla.com>
Date: Thu, 24 Sep 2026 13:43:59 +0200
Subject: [PATCH] perf(generator): look up kind dependency tasks by kind

Before loading a kind, the generator filtered all the tasks loaded so
far to find the ones belonging to the kind's dependencies (copying them
first when loading kinds in parallel). This is proportional to the
number of kinds times the number of tasks, and in parallel mode it
happens on the main thread, delaying the submission of newly unblocked
kinds.

Now loaded tasks are also grouped by kind, so the tasks of each kind
dependency are looked up directly. They are still passed in the order in
which kinds were loaded.

With 150 kinds of 270 tasks each (40,500 tasks), each depending on three
other kinds, gathering the kind dependency tasks takes 0.01s in total
instead of 0.42s.
---
 src/taskgraph/generator.py | 31 ++++++++++++++++++++-----------
 1 file changed, 20 insertions(+), 11 deletions(-)

diff --git a/src/taskgraph/generator.py b/src/taskgraph/generator.py
index bc8f9c400..1368ef01c 100644
--- a/src/taskgraph/generator.py
+++ b/src/taskgraph/generator.py
@@ -282,8 +282,21 @@ def _load_kinds(self, graph_config, target_kinds=None):
                 except KindNotFound:
                     continue
 
+    @staticmethod
+    def _get_kind_dependencies_tasks(kind, tasks_by_kind):
+        """Return the tasks of the kinds `kind` depends on, in the order in
+        which they were loaded."""
+        kind_dependencies = kind.config.get("kind-dependencies", [])
+        return {
+            label: task
+            for kind_name, tasks in tasks_by_kind.items()
+            if kind_name in kind_dependencies
+            for label, task in tasks.items()
+        }
+
     def _load_tasks_serial(self, kinds, kind_graph, parameters):
         all_tasks = {}
+        tasks_by_kind = {}
         for kind_name in kind_graph.visit_postorder():
             logger.debug(f"Loading tasks for kind {kind_name}")
 
@@ -297,11 +310,7 @@ def _load_tasks_serial(self, kinds, kind_graph, parameters):
             try:
                 new_tasks = kind.load_tasks(
                     parameters,
-                    {
-                        k: t
-                        for k, t in all_tasks.items()
-                        if t.kind in kind.config.get("kind-dependencies", [])
-                    },
+                    self._get_kind_dependencies_tasks(kind, tasks_by_kind),
                     self._write_artifacts,
                 )
             except SchemaValidationError as exc:
@@ -310,15 +319,18 @@ def _load_tasks_serial(self, kinds, kind_graph, parameters):
             except Exception:
                 logger.exception(f"Error loading tasks for kind {kind_name}:")
                 raise
+            kind_tasks = tasks_by_kind.setdefault(kind_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
 
         return all_tasks
 
     def _load_tasks_parallel(self, kinds, kind_graph, parameters, executor):
         all_tasks = {}
+        tasks_by_kind = {}
         futures_to_kind = {}
         futures = set()
         edges = set(kind_graph.edges)
@@ -328,7 +340,6 @@ def _load_tasks_parallel(self, kinds, kind_graph, parameters, executor):
             def submit_ready_kinds():
                 """Create the next batch of tasks for kinds without dependencies."""
                 nonlocal kinds, edges, futures
-                loaded_tasks = all_tasks.copy()
                 kinds_with_deps = {edge[0] for edge in edges}
                 ready_kinds = (
                     set(kinds) - kinds_with_deps - set(futures_to_kind.values())
@@ -346,11 +357,7 @@ def submit_ready_kinds():
                     future = executor.submit(
                         kind.load_tasks,
                         dict(parameters),
-                        {
-                            k: t
-                            for k, t in loaded_tasks.items()
-                            if t.kind in kind.config.get("kind-dependencies", [])
-                        },
+                        self._get_kind_dependencies_tasks(kind, tasks_by_kind),
                         self._write_artifacts,
                     )
                     futures.add(future)
@@ -376,10 +383,12 @@ def submit_ready_kinds():
                     kind = futures_to_kind.pop(future)
                     futures.remove(future)
 
+                    kind_tasks = tasks_by_kind.setdefault(kind, {})
                     for task in future.result():
                         if task.label in all_tasks:
                             raise Exception("duplicate tasks with label " + task.label)
                         all_tasks[task.label] = task
+                        kind_tasks[task.label] = task
 
                     # Update state for next batch of futures.
                     del kinds[kind]