Bug 1752677 - Add a task config to set the 'existing_tasks' parameter, r=releng-reviewers,bhearsum,jcristau

This adds the `-E/--use-existing-tasks` flag to the `fuzzy` or `chooser`
selectors. For example:

  ./mach try fuzzy -E

This queries the Treeherder API to find your latest try push, then any tasks
that exist on that push, but aren't explicitly selected this time, will be
replaced with those older tasks.

You can also run:

  ./mach try fuzzy -E task-id=<task id>

Where `<task id>` is the id of the Decision task from the push you would like
to use as the base.

Differential Revision: https://phabricator.services.mozilla.com/D200380

python/sites/python-test.txt

@@ -1,4 +1,5 @@
 pypi:pytest==7.0.1
+pypi:pytest-mock==3.12.0
 pypi:Flask==2.1.3
 # (indirect) avoids dependency on markupsafe >= 2.1.0, which is currently incompatible with glean-parser
 pypi:MarkupSafe==2.0.1

tools/tryselect/push.py

@@ -97,8 +97,17 @@ def generate_try_task_config(method, labels, params=None, routes=None):
     # True). Their dependencies can be optimized though.
     params.setdefault("optimize_target_tasks", False)
 
+    # Remove selected labels from 'existing_tasks' parameter if present
+    if "existing_tasks" in params:
+        params["existing_tasks"] = {
+            label: tid
+            for label, tid in params["existing_tasks"].items()
+            if label not in labels
+        }
+
     try_config = params.setdefault("try_task_config", {})
     try_config.setdefault("env", {})["TRY_SELECTOR"] = method
+
     try_config["tasks"] = sorted(labels)
 
     if routes:

tools/tryselect/selectors/chooser/__init__.py

@@ -31,6 +31,7 @@ class ChooserParser(BaseTryParser):
         "chemspill-prio",
         "disable-pgo",
         "env",
+        "existing-tasks",
         "gecko-profile",
         "path",
         "pernosco",

tools/tryselect/selectors/fuzzy.py

@@ -114,6 +114,7 @@ class FuzzyParser(BaseTryParser):
         "chemspill-prio",
         "disable-pgo",
         "env",
+        "existing-tasks",
         "gecko-profile",
         "path",
         "pernosco",

tools/tryselect/task_config.py

@@ -19,14 +19,17 @@ from contextlib import contextmanager
 from textwrap import dedent
 
 import mozpack.path as mozpath
+import requests
 import six
 from mozbuild.base import BuildEnvironmentNotFoundException, MozbuildObject
 from mozversioncontrol import Repository
+from taskgraph.util import taskcluster
 
 from .tasks import resolve_tests_by_suite
+from .util.ssh import get_ssh_user
 
-here = os.path.abspath(os.path.dirname(__file__))
-build = MozbuildObject.from_environment(cwd=here)
+here = pathlib.Path(__file__).parent
+build = MozbuildObject.from_environment(cwd=str(here))
 
 
 @contextmanager
@@ -155,13 +158,7 @@ class Pernosco(TryConfig):
                 # log in. Prevent people with non-Mozilla addresses from using this
                 # flag so they don't end up consuming time and resources only to
                 # realize they can't actually log in and see the reports.
-                cmd = ["ssh", "-G", "hg.mozilla.org"]
-                output = subprocess.check_output(
-                    cmd, universal_newlines=True
-                ).splitlines()
-                address = [
-                    l.rsplit(" ", 1)[-1] for l in output if l.startswith("user")
-                ][0]
+                address = get_ssh_user()
                 if not address.endswith("@mozilla.com"):
                     print(
                         dedent(
@@ -260,6 +257,70 @@ class Environment(TryConfig):
         }
 
 
+class ExistingTasks(ParameterConfig):
+    TREEHERDER_PUSH_ENDPOINT = (
+        "https://treeherder.mozilla.org/api/project/try/push/?count=1&author={user}"
+    )
+    TREEHERDER_PUSH_URL = (
+        "https://treeherder.mozilla.org/jobs?repo={branch}&revision={revision}"
+    )
+
+    arguments = [
+        [
+            ["-E", "--use-existing-tasks"],
+            {
+                "const": "last_try_push",
+                "default": None,
+                "nargs": "?",
+                "help": """
+                    Use existing tasks from a previous push. Without args this
+                uses your most recent try push. You may also specify
+                `rev=<revision>` where <revision> is the head revision of the
+                try push or `task-id=<task id>` where <task id> is the Decision
+                task id of the push. This last method even works for non-try
+                branches.
+                """,
+            },
+        ]
+    ]
+
+    def find_decision_task(self, use_existing_tasks):
+        branch = "try"
+        if use_existing_tasks == "last_try_push":
+            # Use existing tasks from user's previous try push.
+            user = get_ssh_user()
+            url = self.TREEHERDER_PUSH_ENDPOINT.format(user=user)
+            res = requests.get(url, headers={"User-Agent": "gecko-mach-try/1.0"})
+            res.raise_for_status()
+            data = res.json()
+            if data["meta"]["count"] == 0:
+                raise Exception(f"Could not find a try push for '{user}'!")
+            revision = data["results"][0]["revision"]
+
+        elif use_existing_tasks.startswith("rev="):
+            revision = use_existing_tasks[len("rev=") :]
+
+        else:
+            raise Exception("Unable to parse '{use_existing_tasks}'!")
+
+        url = self.TREEHERDER_PUSH_URL.format(branch=branch, revision=revision)
+        print(f"Using existing tasks from: {url}")
+        index_path = f"gecko.v2.{branch}.revision.{revision}.taskgraph.decision"
+        return taskcluster.find_task_id(index_path)
+
+    def get_parameters(self, use_existing_tasks, **kwargs):
+        if not use_existing_tasks:
+            return
+
+        if use_existing_tasks.startswith("task-id="):
+            tid = use_existing_tasks[len("task-id=") :]
+        else:
+            tid = self.find_decision_task(use_existing_tasks)
+
+        label_to_task_id = taskcluster.get_artifact(tid, "public/label-to-taskid.json")
+        return {"existing_tasks": label_to_task_id}
+
+
 class RangeAction(Action):
     def __init__(self, min, max, *args, **kwargs):
         self.min = min
@@ -418,7 +479,7 @@ class GeckoProfile(TryConfig):
         gecko_profile_entries,
         gecko_profile_features,
         gecko_profile_threads,
-        **kwargs
+        **kwargs,
     ):
         if profile or not all(
             s is None for s in (gecko_profile_features, gecko_profile_threads)
@@ -548,6 +609,7 @@ all_task_configs = {
     "chemspill-prio": ChemspillPrio,
     "disable-pgo": DisablePgo,
     "env": Environment,
+    "existing-tasks": ExistingTasks,
     "gecko-profile": GeckoProfile,
     "path": Path,
     "pernosco": Pernosco,

tools/tryselect/test/conftest.py

@@ -8,12 +8,19 @@ from unittest.mock import MagicMock
 import pytest
 import yaml
 from moztest.resolve import TestResolver
+from responses import RequestsMock
 from taskgraph.graph import Graph
 from taskgraph.task import Task
 from taskgraph.taskgraph import TaskGraph
 from tryselect import push
 
 
+@pytest.fixture
+def responses():
+    with RequestsMock() as rsps:
+        yield rsps
+
+
 @pytest.fixture
 def tg(request):
     if not hasattr(request.module, "TASKS"):

tools/tryselect/test/python.toml

@@ -20,6 +20,8 @@ subsuite = "try"
 # shouldn't be run in parallel with those other tests.
 sequential = true
 
+["test_push.py"]
+
 ["test_release.py"]
 
 ["test_scriptworker.py"]

tools/tryselect/test/test_push.py

@@ -0,0 +1,54 @@
+import mozunit
+import pytest
+from tryselect import push
+
+
+@pytest.mark.parametrize(
+    "method,labels,params,routes,expected",
+    (
+        pytest.param(
+            "fuzzy",
+            ["task-foo", "task-bar"],
+            None,
+            None,
+            {
+                "parameters": {
+                    "optimize_target_tasks": False,
+                    "try_task_config": {
+                        "env": {"TRY_SELECTOR": "fuzzy"},
+                        "tasks": ["task-bar", "task-foo"],
+                    },
+                },
+                "version": 2,
+            },
+            id="basic",
+        ),
+        pytest.param(
+            "fuzzy",
+            ["task-foo"],
+            {"existing_tasks": {"task-foo": "123", "task-bar": "abc"}},
+            None,
+            {
+                "parameters": {
+                    "existing_tasks": {"task-bar": "abc"},
+                    "optimize_target_tasks": False,
+                    "try_task_config": {
+                        "env": {"TRY_SELECTOR": "fuzzy"},
+                        "tasks": ["task-foo"],
+                    },
+                },
+                "version": 2,
+            },
+            id="existing_tasks",
+        ),
+    ),
+)
+def test_generate_try_task_config(method, labels, params, routes, expected):
+    assert (
+        push.generate_try_task_config(method, labels, params=params, routes=routes)
+        == expected
+    )
+
+
+if __name__ == "__main__":
+    mozunit.main()

tools/tryselect/test/test_task_configs.py

@@ -3,7 +3,6 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 import inspect
-import subprocess
 from argparse import ArgumentParser
 from textwrap import dedent
 
@@ -11,6 +10,9 @@ import mozunit
 import pytest
 from tryselect.task_config import Pernosco, all_task_configs
 
+TC_URL = "https://firefox-ci-tc.services.mozilla.com"
+TH_URL = "https://treeherder.mozilla.org"
+
 # task configs have a list of tests of the form (input, expected)
 TASK_CONFIG_TESTS = {
     "artifact": [
@@ -130,26 +132,25 @@ def test_task_configs(config_patch_resolver, task_config, args, expected):
 
 
 @pytest.fixture
-def patch_pernosco_email_check(monkeypatch):
-    def inner(val):
-        def fake_check_output(*args, **kwargs):
-            return val
-
-        monkeypatch.setattr(subprocess, "check_output", fake_check_output)
+def patch_ssh_user(mocker):
+    def inner(user):
+        mock_stdout = mocker.Mock()
+        mock_stdout.stdout = dedent(
+            f"""
+            key1 foo
+            user {user}
+            key2 bar
+            """
+        )
+        return mocker.patch(
+            "tryselect.util.ssh.subprocess.run", return_value=mock_stdout
+        )
 
     return inner
 
 
-def test_pernosco(patch_pernosco_email_check):
-    patch_pernosco_email_check(
-        dedent(
-            """
-        user foobar@mozilla.com
-        hostname hg.mozilla.com
-    """
-        )
-    )
-
+def test_pernosco(patch_ssh_user):
+    patch_ssh_user("user@mozilla.com")
     parser = ArgumentParser()
 
     cfg = Pernosco()
@@ -159,5 +160,94 @@ def test_pernosco(patch_pernosco_email_check):
     assert params == {"try_task_config": {"env": {"PERNOSCO": "1"}}}
 
 
+def test_exisiting_tasks(responses, patch_ssh_user):
+    parser = ArgumentParser()
+    cfg = all_task_configs["existing-tasks"]()
+    cfg.add_arguments(parser)
+
+    user = "user@example.com"
+    rev = "a" * 40
+    task_id = "abc"
+    label_to_taskid = {"task-foo": "123", "task-bar": "456"}
+
+    args = ["--use-existing-tasks"]
+    args = parser.parse_args(args)
+
+    responses.add(
+        responses.GET,
+        f"{TH_URL}/api/project/try/push/?count=1&author={user}",
+        json={"meta": {"count": 1}, "results": [{"revision": rev}]},
+    )
+
+    responses.add(
+        responses.GET,
+        f"{TC_URL}/api/index/v1/task/gecko.v2.try.revision.{rev}.taskgraph.decision",
+        json={"taskId": task_id},
+    )
+
+    responses.add(
+        responses.GET,
+        f"{TC_URL}/api/queue/v1/task/{task_id}/artifacts/public/label-to-taskid.json",
+        json=label_to_taskid,
+    )
+
+    m = patch_ssh_user(user)
+    params = cfg.get_parameters(**vars(args))
+    assert params == {"existing_tasks": label_to_taskid}
+
+    m.assert_called_once_with(
+        ["ssh", "-G", "hg.mozilla.org"], text=True, check=True, capture_output=True
+    )
+
+
+def test_exisiting_tasks_task_id(responses):
+    parser = ArgumentParser()
+    cfg = all_task_configs["existing-tasks"]()
+    cfg.add_arguments(parser)
+
+    task_id = "abc"
+    label_to_taskid = {"task-foo": "123", "task-bar": "456"}
+
+    args = ["--use-existing-tasks", f"task-id={task_id}"]
+    args = parser.parse_args(args)
+
+    responses.add(
+        responses.GET,
+        f"{TC_URL}/api/queue/v1/task/{task_id}/artifacts/public/label-to-taskid.json",
+        json=label_to_taskid,
+    )
+
+    params = cfg.get_parameters(**vars(args))
+    assert params == {"existing_tasks": label_to_taskid}
+
+
+def test_exisiting_tasks_rev(responses):
+    parser = ArgumentParser()
+    cfg = all_task_configs["existing-tasks"]()
+    cfg.add_arguments(parser)
+
+    rev = "aaaaaa"
+    task_id = "abc"
+    label_to_taskid = {"task-foo": "123", "task-bar": "456"}
+
+    args = ["--use-existing-tasks", f"rev={rev}"]
+    args = parser.parse_args(args)
+
+    responses.add(
+        responses.GET,
+        f"{TC_URL}/api/index/v1/task/gecko.v2.try.revision.{rev}.taskgraph.decision",
+        json={"taskId": task_id},
+    )
+
+    responses.add(
+        responses.GET,
+        f"{TC_URL}/api/queue/v1/task/{task_id}/artifacts/public/label-to-taskid.json",
+        json=label_to_taskid,
+    )
+
+    params = cfg.get_parameters(**vars(args))
+    assert params == {"existing_tasks": label_to_taskid}
+
+
 if __name__ == "__main__":
     mozunit.main()

tools/tryselect/util/ssh.py

@@ -0,0 +1,24 @@
+# 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 https://mozilla.org/MPL/2.0/.
+
+import subprocess
+
+
+def get_ssh_user(host="hg.mozilla.org"):
+    ssh_config = subprocess.run(
+        ["ssh", "-G", host],
+        text=True,
+        check=True,
+        capture_output=True,
+    ).stdout
+
+    lines = [l.strip() for l in ssh_config.splitlines()]
+    for line in lines:
+        if not line:
+            continue
+        key, value = line.split(" ", 1)
+        if key.lower() == "user":
+            return value
+
+    raise Exception(f"Could not detect ssh user for '{host}'!")