Bug 1526028 - Support opting into or out of specific caches in run transforms

packages/pytest-taskgraph/src/pytest_taskgraph/fixtures/gen.py

@@ -85,6 +85,12 @@ def fake_load_graph_config(root_dir):
                         "os": "linux",
                         "worker-type": "linux",
                     },
+                    "t-win": {
+                        "provisioner": "taskgraph-t",
+                        "implementation": "generic-worker",
+                        "os": "windows",
+                        "worker-type": "win",
+                    },
                 }
             },
             "task-priority": "low",

src/taskgraph/transforms/run/common.py

@@ -9,6 +9,7 @@
 
 import hashlib
 import json
+import re
 
 from taskgraph.util.taskcluster import get_artifact_prefix
 
@@ -31,7 +32,11 @@ def add_cache(task, taskdesc, name, mount_point, skip_untrusted=False):
         skip_untrusted (bool): Whether cache is used in untrusted environments
             (default: False). Only applies to docker-worker.
     """
-    if not task["run"].get("use-caches", True):
+    use_caches = task["run"].get("use-caches", True)
+    if isinstance(use_caches, list):
+        use_caches = any(re.match(pattern, name) for pattern in use_caches)
+
+    if not use_caches:
         return
 
     worker = task["worker"]

src/taskgraph/transforms/run/run_task.py

@@ -28,8 +28,10 @@
         # tend to hide their caches.  This cache is never added for level-1 tasks.
         # TODO Once bug 1526028 is fixed, this and 'use-caches' should be merged.
         Required("cache-dotcache"): bool,
-        # Whether or not to use caches.
-        Optional("use-caches"): bool,
+        # Whether or not to use caches. If a boolean, all caches will be
+        # enabled or disabled. If a list, entries are regexes. Only cache names
+        # that match one of the listed regexes are enabled.
+        Optional("use-caches"): Any(bool, [str]),
         # if true (the default), perform a checkout on the worker
         Required("checkout"): Any(bool, {str: dict}),
         Optional(

test/test_transforms_run.py

@@ -47,8 +47,8 @@ def transform(monkeypatch, run_transform):
     monkeypatch.setenv("TASK_ID", "fakeid")
 
     def inner(task_input):
-        task = deepcopy(TASK_DEFAULTS)
-        task.update(task_input)
+        defaults = deepcopy(TASK_DEFAULTS)
+        task = merge(defaults, task_input)
 
         with patch("taskgraph.transforms.run.configure_taskdesc_for_run") as m:
             # This forces the generator to be evaluated
@@ -58,32 +58,74 @@ def inner(task_input):
     return inner
 
 
-@pytest.mark.parametrize(
-    "task",
-    [
-        {"worker-type": "t-linux"},
-        pytest.param(
-            {"worker-type": "releng-hardware/gecko-t-win10-64-hw"},
-            marks=pytest.mark.xfail,
-        ),
-    ],
-    ids=["docker-worker", "generic-worker"],
-)
-def test_worker_caches(task, transform):
-    config, task, taskdesc, impl = transform(task)
-    add_cache(task, taskdesc, "cache1", "/cache1")
-    add_cache(task, taskdesc, "cache2", "/cache2", skip_untrusted=True)
+@pytest.fixture
+def run_caches(transform):
+    def inner(task):
+        config, task, taskdesc, impl = transform(task)
+        add_cache(task, taskdesc, "cache1", "/cache1")
+        add_cache(task, taskdesc, "cache2", "/cache2", skip_untrusted=True)
+
+        if impl not in ("docker-worker", "generic-worker"):
+            pytest.xfail(f"caches not implemented for '{impl}'")
+
+        key = "caches" if impl == "docker-worker" else "mounts"
+
+        result = taskdesc["worker"].get(key)
+
+        if result:
+            print("Dumping for copy/paste:")
+            pprint(result, indent=2)
+
+            # Create a new schema object with just the part relevant to caches.
+            partial_schema = Schema(payload_builders[impl].schema.schema[key])
+            validate_schema(partial_schema, taskdesc["worker"][key], "validation error")
+
+        return result
+
+    return inner
+
 
-    if impl not in ("docker-worker", "generic-worker"):
-        pytest.xfail(f"caches not implemented for '{impl}'")
+def test_caches_docker_worker(run_caches):
+    assert run_caches({"worker-type": "t-linux"}) == [
+        {
+            "mount-point": "/cache1",
+            "name": "cache1",
+            "skip-untrusted": False,
+            "type": "persistent",
+        },
+        {
+            "mount-point": "/cache2",
+            "name": "cache2",
+            "skip-untrusted": True,
+            "type": "persistent",
+        },
+    ]
+
+
+def test_caches_generic_worker(run_caches):
+    assert run_caches({"worker-type": "t-win"}) == [
+        {"cache-name": "cache1", "directory": "/cache1"},
+        {"cache-name": "cache2", "directory": "/cache2"},
+    ]
 
-    key = "caches" if impl == "docker-worker" else "mounts"
-    assert key in taskdesc["worker"]
-    assert len(taskdesc["worker"][key]) == 2
 
-    # Create a new schema object with just the part relevant to caches.
-    partial_schema = Schema(payload_builders[impl].schema.schema[key])
-    validate_schema(partial_schema, taskdesc["worker"][key], "validation error")
+def test_caches_regex(run_caches):
+    task = {
+        "run": {
+            "use-caches": ["cache[^1]"],
+        },
+        "worker-type": "t-win",
+    }
+    caches = run_caches(task)
+    assert len(caches) == 1
+    assert caches[0]["cache-name"] == "cache2"
+
+
+@pytest.mark.parametrize("worker_type", ("t-linux", "t-win"))
+def test_caches_disabled(run_caches, worker_type):
+    assert (
+        run_caches({"run": {"use-caches": False}, "worker-type": worker_type}) is None
+    )
 
 
 def test_rewrite_when_to_optimization(run_transform, make_transform_config):