feat: allow for both batched and single queries in task replacement

This patch removes the need for migrating all call sites, especially
in `gecko`.

This is the result of more discussion following the two previous
patches¹. Turns out that while duck typing in interfaces is not great,
leaking implementation details to the user is probably worse.

[1] 41a10aa
[1] 5207917

Author: Alphare

src/taskgraph/docker.py

@@ -17,10 +17,8 @@
 
 from taskgraph.util import docker
 from taskgraph.util.taskcluster import (
-    find_task_id_batched,
     get_artifact_url,
     get_session,
-    status_task_batched,
 )
 
 DEPLOY_WARNING = """
@@ -66,15 +64,7 @@ def load_image_by_name(image_name, tag=None):
     task = tasks[f"build-docker-image-{image_name}"]
 
     indexes = task.optimization.get("index-search", [])
-    index_to_taskid = {}
-    task_id = False
-    if indexes:
-        indexes = list(indexes)
-        index_to_taskid = find_task_id_batched(indexes)
-        taskid_to_status = status_task_batched(list(index_to_taskid.values()))
-
-        arg = (indexes, index_to_taskid, taskid_to_status)
-        task_id = IndexSearch().should_replace_task(task, {}, None, arg)
+    task_id = IndexSearch().should_replace_task(task, {}, None, indexes)
 
     if task_id in (True, False):
         print(

src/taskgraph/optimize/strategies.py

@@ -3,6 +3,7 @@
 
 from taskgraph.optimize.base import OptimizationStrategy, register_strategy
 from taskgraph.util.path import match as match_path
+from taskgraph.util.taskcluster import find_task_id, status_task
 
 logger = logging.getLogger(__name__)
 
@@ -23,12 +24,28 @@ class IndexSearch(OptimizationStrategy):
 
     def should_replace_task(self, task, params, deadline, arg):
         "Look for a task with one of the given index paths"
-        index_paths, label_to_taskid, taskid_to_status = arg
+        batched = False
+        # Appease static checker that doesn't understand that this is not needed
+        label_to_taskid = {}
+        taskid_to_status = {}
+
+        if isinstance(arg, tuple) and len(arg) == 3:
+            # allow for a batched call optimization instead of two queries
+            # per index path
+            index_paths, label_to_taskid, taskid_to_status = arg
+            batched = True
+        else:
+            index_paths = arg
 
         for index_path in index_paths:
             try:
-                task_id = label_to_taskid[index_path]
-                status = taskid_to_status[task_id]
+                if batched:
+                    task_id = label_to_taskid[index_path]
+                    status = taskid_to_status[task_id]
+                else:
+                    # 404 is raised as `KeyError` also end up here
+                    task_id = find_task_id(index_path)
+                    status = status_task(task_id)
                 # status can be `None` if we're in `testing` mode
                 # (e.g. test-action-callback)
                 if not status or status.get("state") in ("exception", "failed"):

test/test_optimize_strategies.py

@@ -1,6 +1,7 @@
 # Any copyright is dedicated to the public domain.
 # http://creativecommons.org/publicdomain/zero/1.0/
 
+import os
 from datetime import datetime
 from test.fixtures.gen import make_task
 from time import mktime
@@ -43,9 +44,29 @@ def params():
         ),
     ),
 )
-def test_index_search(state, expires, expected):
+def test_index_search(responses, params, state, expires, expected):
     taskid = "abc"
     index_path = "foo.bar.latest"
+
+    responses.add(
+        responses.GET,
+        f"{os.environ['TASKCLUSTER_ROOT_URL']}/api/index/v1/task/{index_path}",
+        json={"taskId": taskid},
+        status=200,
+    )
+
+    responses.add(
+        responses.GET,
+        f"{os.environ['TASKCLUSTER_ROOT_URL']}/api/queue/v1/task/{taskid}/status",
+        json={
+            "status": {
+                "state": state,
+                "expires": expires,
+            }
+        },
+        status=200,
+    )
+
     label_to_taskid = {index_path: taskid}
     taskid_to_status = {
         taskid: {
@@ -58,10 +79,12 @@ def test_index_search(state, expires, expected):
     deadline = "2021-06-07T19:03:20.482Z"
     assert (
         opt.should_replace_task(
-            {}, params, deadline, ((index_path,), label_to_taskid, taskid_to_status)
+            {}, params, deadline, ([index_path], label_to_taskid, taskid_to_status)
         )
         == expected
     )
+    # test the non-batched variant as well
+    assert opt.should_replace_task({}, params, deadline, [index_path]) == expected
 
 
 @pytest.mark.parametrize(