perf(parameterization): speed up resolving task definition references

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
 

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