perf: optionally use 'orjson' if available

ahal · ahal · commit a2df1a0c8223 · 2025-05-27T11:17:21.000-04:00
This creates a "json" shim such that using it guarantees compatibility
between orjson and the stdlib json.

The interface of the shim is compatible with the stdlib json, though the
output will be formatted slightly differently.
diff --git a/pyproject.toml b/pyproject.toml
@@ -34,9 +34,8 @@ dependencies = [
 dynamic = ["version"]  # set in src/taskgraph/__init__.py
 
 [project.optional-dependencies]
-load-image = [
-  "zstandard"
-]
+load-image = ["zstandard"]
+orjson = ["orjson"]
 
 [project.scripts]
 taskgraph = "taskgraph.main:main"
diff --git a/src/taskgraph/actions/registry.py b/src/taskgraph/actions/registry.py
@@ -4,7 +4,6 @@
 
 
 import functools
-import json
 from collections import namedtuple
 from types import FunctionType
 
@@ -13,7 +12,7 @@
 from taskgraph import create
 from taskgraph.config import load_graph_config
 from taskgraph.parameters import Parameters
-from taskgraph.util import hash, taskcluster, yaml
+from taskgraph.util import hash, json, taskcluster, yaml
 from taskgraph.util.python_path import import_sibling_modules
 
 actions = []
diff --git a/src/taskgraph/create.py b/src/taskgraph/create.py
@@ -3,14 +3,14 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 
-import json
 import logging
 import os
 import sys
 from concurrent import futures
 
 from slugid import nice as slugid
 
+from taskgraph.util import json
 from taskgraph.util.parameterization import resolve_timestamps
 from taskgraph.util.taskcluster import CONCURRENCY, get_session
 from taskgraph.util.time import current_json_time
@@ -116,8 +116,7 @@ def create_task(session, task_id, label, task_def):
             [task_id, task_def],
             sys.stdout,
             sort_keys=True,
-            indent=4,
-            separators=(",", ": "),
+            indent=2,
         )
         # add a newline
         print("")
diff --git a/src/taskgraph/decision.py b/src/taskgraph/decision.py
@@ -3,7 +3,6 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 
-import json
 import logging
 import os
 import pathlib
@@ -19,6 +18,7 @@
 from taskgraph.generator import TaskGraphGenerator
 from taskgraph.parameters import Parameters, get_version
 from taskgraph.taskgraph import TaskGraph
+from taskgraph.util import json
 from taskgraph.util.python_path import find_object
 from taskgraph.util.schema import Schema, validate_schema
 from taskgraph.util.vcs import Repository, get_repository
@@ -374,7 +374,7 @@ def write_artifact(filename, data):
             yaml.safe_dump(data, f, allow_unicode=True, default_flow_style=False)
     elif filename.endswith(".json"):
         with open(path, "w") as f:
-            json.dump(data, f, sort_keys=True, indent=2, separators=(",", ": "))
+            json.dump(data, f, sort_keys=True, indent=2)
     elif filename.endswith(".gz"):
         import gzip
 
diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py
@@ -3,7 +3,6 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 
-import json
 import os
 import shlex
 import subprocess
@@ -18,7 +17,7 @@
 except ImportError as e:
     zstd = e
 
-from taskgraph.util import docker
+from taskgraph.util import docker, json
 from taskgraph.util.taskcluster import (
     get_artifact_url,
     get_root_url,
diff --git a/src/taskgraph/main.py b/src/taskgraph/main.py
@@ -4,7 +4,6 @@
 
 import argparse
 import atexit
-import json
 import logging
 import os
 import re
@@ -56,9 +55,9 @@ def format_taskgraph_labels(taskgraph):
 
 
 def format_taskgraph_json(taskgraph):
-    return json.dumps(
-        taskgraph.to_json(), sort_keys=True, indent=2, separators=(",", ": ")
-    )
+    from taskgraph.util import json
+
+    return json.dumps(taskgraph.to_json(), sort_keys=True, indent=2)
 
 
 def format_taskgraph_yaml(taskgraph):
@@ -757,6 +756,7 @@ def actions(args):
     from taskgraph.actions import render_actions_json
     from taskgraph.generator import TaskGraphGenerator
     from taskgraph.parameters import parameters_loader
+    from taskgraph.util import json
 
     if args.pop("verbose", False):
         logging.root.setLevel(logging.DEBUG)
@@ -766,7 +766,7 @@ def actions(args):
         tgg = TaskGraphGenerator(root_dir=args.get("root"), parameters=parameters)
 
         actions = render_actions_json(tgg.parameters, tgg.graph_config, "DECISION-TASK")
-        print(json.dumps(actions, sort_keys=True, indent=2, separators=(",", ": ")))
+        print(json.dumps(actions, sort_keys=True, indent=2))
     except Exception:
         traceback.print_exc()
         sys.exit(1)
@@ -784,6 +784,7 @@ def actions(args):
 def action_callback(options):
     from taskgraph.actions import trigger_action_callback
     from taskgraph.actions.util import get_parameters
+    from taskgraph.util import json
 
     try:
         # the target task for this action (or null if it's a group action)
@@ -833,7 +834,7 @@ def test_action_callback(options):
     import taskgraph.actions
     import taskgraph.parameters
     from taskgraph.config import load_graph_config
-    from taskgraph.util import yaml
+    from taskgraph.util import json, yaml
 
     def load_data(filename):
         with open(filename) as f:
diff --git a/src/taskgraph/parameters.py b/src/taskgraph/parameters.py
@@ -4,7 +4,6 @@
 
 import gzip
 import hashlib
-import json
 import os
 import time
 from datetime import datetime
@@ -18,7 +17,7 @@
 import mozilla_repo_urls
 from voluptuous import ALLOW_EXTRA, Any, Optional, Required, Schema
 
-from taskgraph.util import yaml
+from taskgraph.util import json, yaml
 from taskgraph.util.readonlydict import ReadOnlyDict
 from taskgraph.util.schema import validate_schema
 from taskgraph.util.taskcluster import find_task_id, get_artifact_url
diff --git a/src/taskgraph/transforms/docker_image.py b/src/taskgraph/transforms/docker_image.py
@@ -2,8 +2,6 @@
 # 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 json
 import logging
 import os
 import re
@@ -12,6 +10,7 @@
 
 import taskgraph
 from taskgraph.transforms.base import TransformSequence
+from taskgraph.util import json
 from taskgraph.util.docker import create_context_tar, generate_context_hash
 from taskgraph.util.schema import Schema
 
diff --git a/src/taskgraph/transforms/run/__init__.py b/src/taskgraph/transforms/run/__init__.py
@@ -10,14 +10,14 @@
 """
 
 import copy
-import json
 import logging
 
 from voluptuous import Exclusive, Extra, Optional, Required
 
 from taskgraph.transforms.base import TransformSequence
 from taskgraph.transforms.cached_tasks import order_tasks
 from taskgraph.transforms.task import task_description_schema
+from taskgraph.util import json
 from taskgraph.util import path as mozpath
 from taskgraph.util.python_path import import_sibling_modules
 from taskgraph.util.schema import Schema, validate_schema
diff --git a/src/taskgraph/transforms/run/common.py b/src/taskgraph/transforms/run/common.py
@@ -7,11 +7,10 @@
 consistency.
 """
 
-import json
 from typing import Any, Dict, List, Union
 
 from taskgraph.transforms.base import TransformConfig
-from taskgraph.util import path
+from taskgraph.util import json, path
 from taskgraph.util.caches import CACHES, get_checkout_dir
 from taskgraph.util.taskcluster import get_artifact_prefix
 
diff --git a/src/taskgraph/util/json.py b/src/taskgraph/util/json.py
@@ -0,0 +1,73 @@
+# 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/.
+
+"""
+This module is a compatibility shim between stdlib json and orjson.
+It uses orjson if available, and falls back to stdlib json.
+
+It only supports features and parameters that both modules support,
+and attempts to produce identical output regardless of which module
+is used.
+
+The interface is a subset of what's supported by stdlib json.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any, Callable
+
+if TYPE_CHECKING:
+    from _typeshed import SupportsRead, SupportsWrite
+
+try:
+    import orjson
+except ImportError:
+    orjson = None
+
+
+def loads(s: str | bytes | bytearray) -> Any:
+    if orjson:
+        return orjson.loads(s)
+
+    return json.loads(s)
+
+
+def load(fh: SupportsRead[str | bytes]) -> Any:
+    if orjson:
+        return loads(fh.read())
+    return json.load(fh)
+
+
+def dumps(
+    obj: Any,
+    default: Callable[[Any], Any] | None = None,
+    indent: int | None = None,
+    sort_keys: bool = False,
+) -> str:
+    if indent and indent != 2:
+        raise ValueError("An indent other than 2 is not supported!")
+
+    if orjson:
+        option = 0
+        if indent:
+            option |= orjson.OPT_INDENT_2
+
+        if sort_keys:
+            option |= orjson.OPT_SORT_KEYS
+
+        return orjson.dumps(obj, default=default, option=option).decode("utf-8")
+
+    separators = (",", ": ") if indent else (",", ":")
+    return json.dumps(
+        obj,
+        default=default,
+        indent=indent,
+        sort_keys=sort_keys,
+        separators=separators,
+    )
+
+
+def dump(obj: Any, fh: SupportsWrite[str], **kwargs) -> None:
+    fh.write(dumps(obj, **kwargs))
diff --git a/taskcluster/docker/decision/system-setup.sh b/taskcluster/docker/decision/system-setup.sh
@@ -9,7 +9,7 @@ apt-get install -y --force-yes --no-install-recommends \
     python3-pip
 
 pushd /setup/taskgraph
-uv export --no-cache --no-dev > /setup/requirements.txt
+uv export --no-cache --extra orjson --no-dev > /setup/requirements.txt
 uv pip install --no-cache --system --break-system-packages -r /setup/requirements.txt
 uv pip install --no-cache --system --break-system-packages --no-deps .
 popd
diff --git a/taskcluster/kinds/check/kind.yml b/taskcluster/kinds/check/kind.yml
@@ -38,4 +38,4 @@ tasks:
             symbol: check(type)
         run:
             command: >-
-                uv run pyright
+                uv run --all-extras pyright
diff --git a/test/test_transforms_run_run_task.py b/test/test_transforms_run_run_task.py
@@ -80,7 +80,7 @@ def assert_docker_worker(task):
                 "CI_REPOSITORY_TYPE": "hg",
                 "HG_STORE_PATH": "/builds/worker/checkouts/hg-store",
                 "MOZ_SCM_LEVEL": "1",
-                "REPOSITORIES": '{"ci": "Taskgraph"}',
+                "REPOSITORIES": '{"ci":"Taskgraph"}',
                 "VCS_PATH": "/builds/worker/checkouts/vcs",
             },
             "implementation": "docker-worker",
@@ -116,7 +116,7 @@ def assert_generic_worker(task):
                 "CI_REPOSITORY_TYPE": "hg",
                 "HG_STORE_PATH": "y:/hg-shared",
                 "MOZ_SCM_LEVEL": "1",
-                "REPOSITORIES": '{"ci": "Taskgraph"}',
+                "REPOSITORIES": '{"ci":"Taskgraph"}',
                 "VCS_PATH": "build/src",
             },
             "implementation": "generic-worker",
diff --git a/test/test_util_json.py b/test/test_util_json.py
@@ -0,0 +1,76 @@
+from itertools import count
+from textwrap import dedent
+from unittest.mock import Mock
+
+import pytest
+
+from taskgraph.util.json import dump, dumps, load, loads
+
+
+@pytest.fixture(autouse=True, params=["json", "orjson"])
+def with_module(request, mocker):
+    module = request.param
+    if module == "orjson":
+        pytest.importorskip("orjson")
+
+    if module == "json":
+        mocker.patch("taskgraph.util.json.orjson", None)
+
+
+@pytest.mark.parametrize(
+    "input_data, expected",
+    [
+        ('{"key": "value"}', {"key": "value"}),
+    ],
+    ids=count(),
+)
+def test_loads(input_data, expected):
+    assert loads(input_data) == expected
+
+
+@pytest.mark.parametrize(
+    "input_data, expected",
+    [
+        ('{"key": "value"}', {"key": "value"}),
+    ],
+    ids=count(),
+)
+def test_load(input_data, expected):
+    mock_file = Mock(read=Mock(return_value=input_data))
+    assert load(mock_file) == expected
+
+
+@pytest.mark.parametrize(
+    "input_data, json_args, expected",
+    [
+        ({"b": 1, "a": 2}, {}, '{"b":1,"a":2}'),
+        ({"b": 1, "a": 2}, {"sort_keys": True}, '{"a":2,"b":1}'),
+        (
+            {"b": 1, "a": 2},
+            {"indent": 2},
+            dedent("""
+            {
+              "b": 1,
+              "a": 2
+            }""").lstrip(),
+        ),
+    ],
+    ids=count(),
+)
+def test_dumps(input_data, json_args, expected):
+    assert dumps(input_data, **json_args) == expected
+
+
+@pytest.mark.parametrize(
+    "input_data, json_args, expected",
+    [
+        ({"key": "value"}, {}, '{"key":"value"}'),
+    ],
+    ids=count(),
+)
+def test_dump(input_data, json_args, expected):
+    mock_file = Mock()
+
+    dump(input_data, mock_file, **json_args)
+
+    mock_file.write.assert_called_once_with(expected)
diff --git a/uv.lock b/uv.lock
