From 40c338530e41ab94ffd2a842bf3519b2afff8b5a Mon Sep 17 00:00:00 2001
From: Marco Castelluccio <mcastelluccio@mozilla.com>
Date: Thu, 24 Sep 2026 14:36:27 +0200
Subject: [PATCH] perf(parameterization): speed up resolving task definition
 references

_recurse walks every task definition when resolving task references
during optimization, and again when resolving timestamps during task
creation. For each single-key dictionary it iterated over all the
parameter functions, building a set of the dictionary keys to compare
against each of them. Now the key is looked up directly in the
parameter functions. Dictionaries are also checked before lists, as
they are much more common in task definitions.

On a typical task definition, resolve_task_references takes 14.5us
instead of 18.3us, and resolve_timestamps 20.5us instead of 23.1us.
---
 src/taskgraph/util/parameterization.py | 22 ++++++------
 test/test_graph_perf.py                | 47 ++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 12 deletions(-)

diff --git a/src/taskgraph/util/parameterization.py b/src/taskgraph/util/parameterization.py
index 48e8f8966..ab9df23e3 100644
--- a/src/taskgraph/util/parameterization.py
+++ b/src/taskgraph/util/parameterization.py
@@ -15,20 +15,18 @@
 
 def _recurse(val, param_fns):
     def recurse(val):
-        if isinstance(val, list):
-            return [recurse(v) for v in val]
-        elif isinstance(val, dict):
+        if isinstance(val, dict):
             if len(val) == 1:
-                for param_key, param_fn in param_fns.items():
-                    if set(val.keys()) == {param_key}:
-                        if isinstance(val[param_key], dict):
-                            # handle `{"task-reference": {"<foo>": "bar"}}`
-                            return {
-                                param_fn(key): recurse(v)
-                                for key, v in val[param_key].items()
-                            }
-                        return param_fn(val[param_key])
+                ((key, value),) = val.items()
+                param_fn = param_fns.get(key)
+                if param_fn is not None:
+                    if isinstance(value, dict):
+                        # handle `{"task-reference": {"<foo>": "bar"}}`
+                        return {param_fn(k): recurse(v) for k, v in value.items()}
+                    return param_fn(value)
             return {k: recurse(v) for k, v in val.items()}
+        elif isinstance(val, list):
+            return [recurse(v) for v in val]
         else:
             return val
 
diff --git a/test/test_graph_perf.py b/test/test_graph_perf.py
index 83d400f1b..3559f2d7b 100644
--- a/test/test_graph_perf.py
+++ b/test/test_graph_perf.py
@@ -12,6 +12,11 @@
 from taskgraph.task import Task
 from taskgraph.taskgraph import TaskGraph
 from taskgraph.transforms.base import TransformSequence
+from taskgraph.util.parameterization import (
+    resolve_task_references,
+    resolve_timestamps,
+)
+from taskgraph.util.time import current_json_time
 
 # ---------------------------------------------------------------------------
 # Graph builders – each returns (tasks_dict, Graph, TaskGraph) for 1000 nodes
@@ -179,6 +184,48 @@ def test_taskgraph_to_json(geometry):
     assert len(data) == N
 
 
+# ---------------------------------------------------------------------------
+# Benchmarks – task definition parameterization
+# ---------------------------------------------------------------------------
+
+
+def _make_task_def(i):
+    return {
+        "created": {"relative-datestamp": "0 seconds"},
+        "deadline": {"relative-datestamp": "1 day"},
+        "expires": {"relative-datestamp": "28 days"},
+        "metadata": {"name": f"task-{i}", "description": "d", "owner": "o"},
+        "routes": [f"index.domain.v2.project.task-{i}"],
+        "payload": {
+            "command": ["run-task", "--", "bash", "-c", f"echo {i}"],
+            "env": {
+                "BUILD": {"task-reference": "<build>"},
+                "ARTIFACT": {"artifact-reference": "<build/public/target.zip>"},
+                **{f"VAR{j}": f"value-{j}" for j in range(20)},
+            },
+            "artifacts": [{"name": f"public/a{j}", "path": "/x"} for j in range(5)],
+        },
+    }
+
+
+TASK_DEFS = [_make_task_def(i) for i in range(N)]
+
+
+@pytest.mark.benchmark
+def test_resolve_task_references():
+    for task_def in TASK_DEFS:
+        resolve_task_references(
+            "label", task_def, "task-id", "decision-id", {"build": "build-id"}
+        )
+
+
+@pytest.mark.benchmark
+def test_resolve_timestamps():
+    now = current_json_time(datetime_format=True)
+    for task_def in TASK_DEFS:
+        resolve_timestamps(now, task_def)
+
+
 # ---------------------------------------------------------------------------
 # Benchmarks – TransformSequence with a simple transform
 # ---------------------------------------------------------------------------