1 files changed, 185 insertions, 0 deletions
diff --git a/taskcluster/gecko_taskgraph/actions/retrigger_custom.py b/taskcluster/gecko_taskgraph/actions/retrigger_custom.py
new file mode 100644
index 0000000000..677cb00a11
--- /dev/null
+++ b/taskcluster/gecko_taskgraph/actions/retrigger_custom.py
@@ -0,0 +1,185 @@
+# 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 json
+import logging
+
+from taskgraph.util.parameterization import resolve_task_references
+from taskgraph.util.taskcluster import get_task_definition
+
+from .registry import register_callback_action
+from .util import create_task_from_def, fetch_graph_and_labels
+
+logger = logging.getLogger(__name__)
+
+# Properties available for custom retrigger of any supported test suites
+basic_properties = {
+    "path": {
+        "type": "string",
+        "maxLength": 255,
+        "default": "",
+        "title": "Path name",
+        "description": "Path of test(s) to retrigger",
+    },
+    "logLevel": {
+        "type": "string",
+        "enum": ["debug", "info", "warning", "error", "critical"],
+        "default": "info",
+        "title": "Log level",
+        "description": "Log level for output (INFO is normal, DEBUG gives more detail)",
+    },
+    "environment": {
+        "type": "object",
+        "default": {"MOZ_LOG": ""},
+        "title": "Extra environment variables",
+        "description": "Extra environment variables to use for this run",
+        "additionalProperties": {"type": "string"},
+    },
+}
+
+# Additional properties available for custom retrigger of some additional test suites
+extended_properties = basic_properties.copy()
+extended_properties.update(
+    {
+        "runUntilFail": {
+            "type": "boolean",
+            "default": False,
+            "title": "Run until failure",
+            "description": (
+                "Runs the specified set of tests repeatedly "
+                "until failure (up to REPEAT times)"
+            ),
+        },
+        "repeat": {
+            "type": "integer",
+            "default": 0,
+            "minimum": 0,
+            "title": "Repeat test(s) N times",
+            "description": (
+                "Run test(s) repeatedly (usually used in "
+                "conjunction with runUntilFail)"
+            ),
+        },
+        "preferences": {
+            "type": "object",
+            "default": {"remote.log.level": "Info"},
+            "title": "Extra gecko (about:config) preferences",
+            "description": "Extra gecko (about:config) preferences to use for this run",
+            "additionalProperties": {"type": "string"},
+        },
+    }
+)
+
+
+@register_callback_action(
+    name="retrigger-custom",
+    title="Retrigger task with custom parameters",
+    symbol="rt",
+    description="Retriggers the specified task with custom environment and parameters",
+    context=[
+        {"test-type": "mochitest", "worker-implementation": "docker-worker"},
+        {"test-type": "reftest", "worker-implementation": "docker-worker"},
+        {"test-type": "geckoview-junit", "worker-implementation": "docker-worker"},
+    ],
+    order=10,
+    schema={
+        "type": "object",
+        "properties": extended_properties,
+        "additionalProperties": False,
+        "required": ["path"],
+    },
+)
+def extended_custom_retrigger_action(
+    parameters, graph_config, input, task_group_id, task_id
+):
+    handle_custom_retrigger(parameters, graph_config, input, task_group_id, task_id)
+
+
+@register_callback_action(
+    name="retrigger-custom (gtest)",
+    title="Retrigger gtest task with custom parameters",
+    symbol="rt",
+    description="Retriggers the specified task with custom environment and parameters",
+    context=[{"test-type": "gtest", "worker-implementation": "docker-worker"}],
+    order=10,
+    schema={
+        "type": "object",
+        "properties": basic_properties,
+        "additionalProperties": False,
+        "required": ["path"],
+    },
+)
+def basic_custom_retrigger_action_basic(
+    parameters, graph_config, input, task_group_id, task_id
+):
+    handle_custom_retrigger(parameters, graph_config, input, task_group_id, task_id)
+
+
+def handle_custom_retrigger(parameters, graph_config, input, task_group_id, task_id):
+    task = get_task_definition(task_id)
+    decision_task_id, full_task_graph, label_to_taskid = fetch_graph_and_labels(
+        parameters, graph_config
+    )
+
+    pre_task = full_task_graph.tasks[task["metadata"]["name"]]
+
+    # fix up the task's dependencies, similar to how optimization would
+    # have done in the decision
+    dependencies = {
+        name: label_to_taskid[label] for name, label in pre_task.dependencies.items()
+    }
+    new_task_definition = resolve_task_references(
+        pre_task.label, pre_task.task, task_id, decision_task_id, dependencies
+    )
+    new_task_definition.setdefault("dependencies", []).extend(dependencies.values())
+
+    # don't want to run mozharness tests, want a custom mach command instead
+    new_task_definition["payload"]["command"] += ["--no-run-tests"]
+
+    custom_mach_command = [task["tags"]["test-type"]]
+
+    # mochitests may specify a flavor
+    if new_task_definition["payload"]["env"].get("MOCHITEST_FLAVOR"):
+        custom_mach_command += [
+            "--keep-open=false",
+            "-f",
+            new_task_definition["payload"]["env"]["MOCHITEST_FLAVOR"],
+        ]
+
+    enable_e10s = json.loads(
+        new_task_definition["payload"]["env"].get("ENABLE_E10S", "true")
+    )
+    if not enable_e10s:
+        custom_mach_command += ["--disable-e10s"]
+
+    custom_mach_command += [
+        "--log-tbpl=-",
+        "--log-tbpl-level={}".format(input.get("logLevel", "debug")),
+    ]
+    if input.get("runUntilFail"):
+        custom_mach_command += ["--run-until-failure"]
+    if input.get("repeat"):
+        custom_mach_command += ["--repeat", str(input.get("repeat", 30))]
+
+    # add any custom gecko preferences
+    for (key, val) in input.get("preferences", {}).items():
+        custom_mach_command += ["--setpref", f"{key}={val}"]
+
+    custom_mach_command += [input["path"]]
+    new_task_definition["payload"]["env"]["CUSTOM_MACH_COMMAND"] = " ".join(
+        custom_mach_command
+    )
+
+    # update environment
+    new_task_definition["payload"]["env"].update(input.get("environment", {}))
+
+    # tweak the treeherder symbol
+    new_task_definition["extra"]["treeherder"]["symbol"] += "-custom"
+
+    logging.info("New task definition: %s", new_task_definition)
+
+    create_task_from_def(
+        new_task_definition, parameters["level"], action_tag="retrigger-custom-task"
+    )