split me up: collection of fixes that will hopefully fix integration …

…tests

taskcluster/fxci_config_taskgraph/transforms/firefoxci_artifact.py

@@ -4,6 +4,7 @@
 
 import json
 import os
+import re
 from collections import defaultdict
 from copy import deepcopy
 
@@ -26,6 +27,7 @@ def make_firefoxci_artifact_tasks(config, tasks):
         include_attrs = task.pop("include-attrs", {})
         exclude_attrs = task.pop("exclude-attrs", {})
         include_deps = task.pop("include-deps", [])
+        mirror_public_artifacts = [re.compile(r) for r in task.pop("mirror-public-artifacts", [])]
         for decision_index_path in task.pop("decision-index-paths"):
             for _, task_def in find_tasks(
                 decision_index_path,
@@ -36,14 +38,10 @@ def make_firefoxci_artifact_tasks(config, tasks):
                 # Add docker images
                 if "image" in task_def["payload"]:
                     image = task_def["payload"]["image"]
-                    if not isinstance(image, dict) or "taskId" not in image:
-                        continue
-
-                    task_id = image["taskId"]
-                    if task_id in tasks_to_create:
-                        continue
-
-                    tasks_to_create[task_id] = [image["path"]]
+                    if isinstance(image, dict) and "taskId" in image:
+                        task_id = image["taskId"]
+                        if task_id not in tasks_to_create:
+                            tasks_to_create[task_id] = [image["path"]]
 
                 # Add private artifacts
                 if "MOZ_FETCHES" in task_def["payload"].get("env", {}):
@@ -52,7 +50,8 @@ def make_firefoxci_artifact_tasks(config, tasks):
                     )
                     for fetch in fetches:
                         if fetch["artifact"].startswith("public"):
-                            continue
+                            if not any([pat.match(task_def["metadata"]["name"]) for pat in mirror_public_artifacts]):
+                                continue
 
                         task_id = fetch["task"]
                         tasks_to_create[task_id].append(fetch["artifact"])

taskcluster/fxci_config_taskgraph/transforms/integration_test.py

@@ -210,14 +210,17 @@ def rewrite_mirrored_dependencies(
     prefix: str,
     dependencies: dict[str, str],
     tasks: dict[str, Any],
-    patterns: list[re.Pattern],
+    include_deps: list[str],
+    artifact_tasks,
 ):
     """Re-write dependencies and fetches of tasks that are being re-run in the
     staging instance that are also being re-run in the staging instance. Without
     this, the downstream tasks will attempt to refer to firefoxci task ids that
     do not exist in the staging cluster, and task submission will fail.
     """
+    patterns = [re.compile(p) for p in include_deps]
     modified_deps = set()
+    mirrored_deps = set()
     # First, update any dependencies that are also being run as part of this integration test
     for upstream_task_id in dependencies:
         if upstream_task_id in tasks:
@@ -227,6 +230,15 @@ def rewrite_mirrored_dependencies(
                 upstream_task_label = f"{prefix}-{name}"
                 taskdesc["dependencies"][upstream_task_label] = upstream_task_label
 
+        artifact_task_label = f"firefoxci-artifact-{prefix}-{upstream_task_id}"
+        # TODO: shouldn't add docker images here; they get added already elsewhere
+        if artifact_task_label in artifact_tasks:
+            fetched_artifacts = artifact_tasks[artifact_task_label].task["payload"]["env"]["FETCH_FIREFOXCI_ARTIFACTS"]
+            if "image.tar.zst" in fetched_artifacts:
+                continue
+            mirrored_deps.add(upstream_task_id)
+            taskdesc["dependencies"][artifact_task_label] = artifact_task_label
+
     # Second, update any fetches that point to dependencies that are also being run as part
     # of this integration test
     updated_fetches = []
@@ -241,6 +253,9 @@ def rewrite_mirrored_dependencies(
                 fetch_task_label = tasks[fetch_task_id]["metadata"]["name"]
                 fetch["task"] = f"<{prefix}-{fetch_task_label}>"
 
+            if fetch_task_id in mirrored_deps:
+                fetch["task"] = f"<firefoxci-artifact-{prefix}-{fetch_task_id}>"
+
             updated_fetches.append(fetch)
 
         taskdesc["task"]["payload"]["env"]["MOZ_FETCHES"] = {
@@ -253,6 +268,8 @@ def make_integration_test_description(
     name_prefix: str,
     tasks: dict[str, Any],
     include_deps: list[str],
+    should_patch_root_url: bool,
+    artifact_tasks,
 ):
     """Schedule a task on the staging Taskcluster instance.
 
@@ -274,30 +291,6 @@ def make_integration_test_description(
     if "treeherder" in task_def["extra"]:
         del task_def["extra"]["treeherder"]
 
-    include_patterns = [re.compile(p) for p in include_deps]
-    # Tasks may only have 1 root url set, which is primarily used to decide
-    # where to find `MOZ_FETCHES`. When we're including dependencies, we peek
-    # at upstream tasks to figure out what this should be set to. If any
-    # upstream task is present that doesn't match an include pattern, its
-    # dependencies must be pulled from the firefox ci cluster, and we must
-    # patch the root. If a task matches an include pattern, its dependencies
-    # will be in the staging cluster, and we must not patch the root.
-    # If we have a mix of tasks, we cannot proceed.
-    matches = set()
-    for upstream_task_id in orig_dependencies:
-        if upstream_task_id in tasks:
-            name = tasks[upstream_task_id]["metadata"]["name"]
-            matches.add(any([pat.match(name) for pat in include_patterns]))
-
-    should_patch_root_url = True
-    if len(matches) == 2:
-        raise Exception(
-            f"task '{task_def['metadata']['name']}' has both mirrored and unmirrored dependencies; it will not work correctly"
-        )
-    elif len(matches) == 1:
-        if matches.pop() is True:
-            should_patch_root_url = False
-
     if should_patch_root_url:
         patch_root_url(task_def)
     rewrite_mounts(task_def)
@@ -328,7 +321,7 @@ def make_integration_test_description(
     rewrite_docker_image(taskdesc)
     rewrite_private_fetches(taskdesc)
     rewrite_mirrored_dependencies(
-        taskdesc, name_prefix, orig_dependencies, tasks, include_patterns
+        taskdesc, name_prefix, orig_dependencies, tasks, include_deps, artifact_tasks,
     )
     return taskdesc
 
@@ -341,10 +334,12 @@ def schedule_tasks_at_index(config, tasks):
     if os.environ["TASKCLUSTER_ROOT_URL"] != STAGING_ROOT_URL:
         return
 
+    artifact_tasks = {k: v for k, v in config.kind_dependencies_tasks.items() if k.startswith("firefoxci-artifact")}
     for task in tasks:
         include_attrs = task.pop("include-attrs", {})
         exclude_attrs = task.pop("exclude-attrs", {})
         include_deps = task.pop("include-deps", [])
+        patch_root_url = task.pop("patch-root-url", True)
         for decision_index_path in task.pop("decision-index-paths"):
             found_tasks = find_tasks(
                 decision_index_path,
@@ -356,5 +351,5 @@ def schedule_tasks_at_index(config, tasks):
                 # task_def is copied to avoid modifying the version in `tasks`, which
                 # may be used to modify parts of the new task description
                 yield make_integration_test_description(
-                    copy.deepcopy(task_def), task["name"], found_tasks, include_deps
+                    copy.deepcopy(task_def), task["name"], found_tasks, include_deps, patch_root_url, artifact_tasks
                 )

taskcluster/kinds/firefoxci-artifact/kind.yml

@@ -60,3 +60,5 @@ tasks:
       # expire, which causes us to try to mirror expired docker image tasks
       # into the staging cluster
       - ^(?!(Decision|Action|PR action|build-docker-image|fetch)).*
+    mirror-public-artifacts:
+      - ^toolchain.*

taskcluster/kinds/integration-test/kind.yml

@@ -42,3 +42,4 @@ tasks:
       # expire, which causes us to try to mirror expired docker image tasks
       # into the staging cluster
       - ^(?!(Decision|Action|PR action|build-docker-image|fetch)).*
+    patch-root-url: false

taskcluster/test/test_transforms_integration_test.py

@@ -379,7 +379,7 @@ def run_include_deps_test(run_test, *args, **kwargs):
                     "taskId": "ghi",
                     "path": "public/image.tar.zst",
                 },
-                "command": ["run-task", "build-a-thing"],
+                "command": ["run-task", "build-a-thing", "foo"],
             },
             "tags": {},
         },
@@ -394,7 +394,7 @@ def run_include_deps_test(run_test, *args, **kwargs):
                     "taskId": "jkl",
                     "path": "public/image.tar.zst",
                 },
-                "command": ["run-task", "test-a-thing"],
+                "command": ["run-task", "test-a-thing", "foo"],
                 "env": {
                     "MOZ_FETCHES": json.dumps(
                         [
@@ -420,7 +420,7 @@ def run_include_deps_test(run_test, *args, **kwargs):
                     "taskId": "mno",
                     "path": "public/image.tar.zst",
                 },
-                "command": ["run-task", "sign-a-thing"],
+                "command": ["run-task", "sign-a-thing", "foo"],
                 "env": {
                     "MOZ_FETCHES": json.dumps(
                         [
@@ -502,11 +502,11 @@ def test_include_some_deps(run_test):
                         "taskId": task_id,
                         "path": artifact,
                     },
-                    "command": ["run-task", "foo"],
+                    "command": ["run-task", "foo", "bar"],
                 },
             },
         },
-        include_deps=["^sign", "^test"],
+        include_deps=["^(?!(build))"],
         name="gecko",
     )
     expected = set(["gecko-foo", "gecko-test-thing", "gecko-sign-thing"])
@@ -521,6 +521,8 @@ def test_include_some_deps(run_test):
         in sign_task["task"]["payload"]["env"]["MOZ_FETCHES"]["task-reference"]
     )
     assert "TASKCLUSTER_ROOT_URL" in test_task["task"]["payload"]["command"][2]
+    assert "TASKCLUSTER_ROOT_URL" not in sign_task["task"]["payload"]["command"][2]
+    assert "TASKCLUSTER_ROOT_URL" not in foo_task["task"]["payload"]["command"][2]
 
 
 def test_no_deps(run_test):