Convert all transforms to async

Author: ahal

src/taskgraph/generator.py

@@ -88,7 +88,7 @@ async def load_tasks(self, parameters, loaded_tasks, write_artifacts):
                 soft_dependencies=t.get("soft-dependencies"),
                 if_dependencies=t.get("if-dependencies"),
             )
-            for t in tasks
+            async for t in tasks
         ]
         logger.info(f"Generated {len(tasks)} tasks for kind {self.name}")
         return tasks

src/taskgraph/transforms/base.py

@@ -148,6 +148,9 @@ async def __call__(self, config, items):
                 items = xform(config, items)
             if items is None:
                 raise Exception(f"Transform {xform} is not a generator")
+
+        if not inspect.isasyncgen(items):
+            items = convert_async(items)
         return items
 
     def add(self, func):

src/taskgraph/transforms/cached_tasks.py

@@ -50,7 +50,7 @@ def format_task_digest(cached_task):
 
 
 @transforms.add
-def cache_task(config, tasks):
+async def cache_task(config, tasks):
     if taskgraph.fast:
         for task in tasks:
             yield task
@@ -61,7 +61,7 @@ def cache_task(config, tasks):
         if "cached_task" in task.attributes:
             digests[task.label] = format_task_digest(task.attributes["cached_task"])
 
-    for task in order_tasks(config, tasks):
+    for task in order_tasks(config, [t async for t in tasks]):
         cache = task.pop("cache", None)
         if cache is None:
             yield task

src/taskgraph/transforms/chunking.py

@@ -50,8 +50,8 @@
 
 
 @transforms.add
-def chunk_tasks(config, tasks):
-    for task in tasks:
+async def chunk_tasks(config, tasks):
+    async for task in tasks:
         chunk_config = task.pop("chunk", None)
         if not chunk_config:
             yield task

src/taskgraph/transforms/code_review.py

@@ -12,8 +12,8 @@
 
 
 @transforms.add
-def add_dependencies(config, jobs):
-    for job in jobs:
+async def add_dependencies(config, jobs):
+    async for job in jobs:
         job.setdefault("soft-dependencies", [])
         job["soft-dependencies"] += [
             dep_task.label

src/taskgraph/transforms/docker_image.py

@@ -65,7 +65,7 @@
 
 
 @transforms.add
-def fill_template(config, tasks):
+async def fill_template(config, tasks):
     available_packages = set()
     for task in config.kind_dependencies_tasks.values():
         if task.kind != "packages":
@@ -75,13 +75,11 @@ def fill_template(config, tasks):
 
     context_hashes = {}
 
-    tasks = list(tasks)
-
     if not taskgraph.fast and config.write_artifacts:
         if not os.path.isdir(CONTEXTS_DIR):
             os.makedirs(CONTEXTS_DIR)
 
-    for task in tasks:
+    async for task in tasks:
         image_name = task.pop("name")
         job_symbol = task.pop("symbol", None)
         args = task.pop("args", {})

src/taskgraph/transforms/fetch.py

@@ -78,9 +78,9 @@ def wrap(func):
 
 
 @transforms.add
-def process_fetch_job(config, jobs):
+async def process_fetch_job(config, jobs):
     # Converts fetch-url entries to the job schema.
-    for job in jobs:
+    async for job in jobs:
         typ = job["fetch"]["type"]
         name = job["name"]
         fetch = job.pop("fetch")
@@ -103,15 +103,15 @@ def configure_fetch(config, typ, name, fetch):
 
 
 @transforms.add
-def make_task(config, jobs):
+async def make_task(config, jobs):
     # Fetch tasks are idempotent and immutable. Have them live for
     # essentially forever.
     if config.params["level"] == "3":
         expires = "1000 years"
     else:
         expires = "28 days"
 
-    for job in jobs:
+    async for job in jobs:
         name = job["name"]
         artifact_prefix = job.get("artifact-prefix", "public")
         env = job.get("env", {})

src/taskgraph/transforms/from_deps.py

@@ -113,8 +113,8 @@
 
 
 @transforms.add
-def from_deps(config, tasks):
-    for task in tasks:
+async def from_deps(config, tasks):
+    async for task in tasks:
         # Setup and error handling.
         from_deps = task.pop("from-deps")
         kind_deps = config.config.get("kind-dependencies", [])

src/taskgraph/transforms/job/__init__.py

@@ -112,8 +112,8 @@
 
 
 @transforms.add
-def rewrite_when_to_optimization(config, jobs):
-    for job in jobs:
+async def rewrite_when_to_optimization(config, jobs):
+    async for job in jobs:
         when = job.pop("when", {})
         if not when:
             yield job
@@ -132,8 +132,8 @@ def rewrite_when_to_optimization(config, jobs):
 
 
 @transforms.add
-def set_implementation(config, jobs):
-    for job in jobs:
+async def set_implementation(config, jobs):
+    async for job in jobs:
         impl, os = worker_type_implementation(config.graph_config, job["worker-type"])
         if os:
             job.setdefault("tags", {})["os"] = os
@@ -148,8 +148,8 @@ def set_implementation(config, jobs):
 
 
 @transforms.add
-def set_label(config, jobs):
-    for job in jobs:
+async def set_label(config, jobs):
+    async for job in jobs:
         if "label" not in job:
             if "name" not in job:
                 raise Exception("job has neither a name nor a label")
@@ -160,8 +160,8 @@ def set_label(config, jobs):
 
 
 @transforms.add
-def add_resource_monitor(config, jobs):
-    for job in jobs:
+async def add_resource_monitor(config, jobs):
+    async for job in jobs:
         if job.get("attributes", {}).get("resource-monitor"):
             worker_implementation, worker_os = worker_type_implementation(
                 config.graph_config, job["worker-type"]
@@ -204,13 +204,13 @@ def get_attribute(dict, key, attributes, attribute_name):
 
 
 @transforms.add
-def use_fetches(config, jobs):
+async def use_fetches(config, jobs):
     artifact_names = {}
     aliases = {}
     extra_env = {}
 
+    jobs = [j async for j in jobs]
     if config.kind in ("toolchain", "fetch"):
-        jobs = list(jobs)
         for job in jobs:
             run = job.get("run", {})
             label = job["label"]
@@ -353,12 +353,12 @@ def cmp_artifacts(a):
 
 
 @transforms.add
-def make_task_description(config, jobs):
+async def make_task_description(config, jobs):
     """Given a build description, create a task description"""
     # import plugin modules first, before iterating over jobs
     import_sibling_modules(exceptions=("common.py",))
 
-    for job in jobs:
+    async for job in jobs:
         # always-optimized tasks never execute, so have no workdir
         if job["worker"]["implementation"] in ("docker-worker", "generic-worker"):
             job["run"].setdefault("workdir", "/builds/worker")

src/taskgraph/transforms/notify.py

@@ -140,8 +140,8 @@ def _convert_content(content):
 
 
 @transforms.add
-def add_notifications(config, tasks):
-    for task in tasks:
+async def add_notifications(config, tasks):
+    async for task in tasks:
         label = "{}-{}".format(config.kind, task["name"])
         if "notifications" in task:
             notify = _convert_legacy(config, task.pop("notifications"), label)