1 files changed, 78 insertions, 0 deletions
diff --git a/taskcluster/android_taskgraph/transforms/chunk.py b/taskcluster/android_taskgraph/transforms/chunk.py
new file mode 100644
index 0000000000..bb17a99f4e
--- /dev/null
+++ b/taskcluster/android_taskgraph/transforms/chunk.py
@@ -0,0 +1,78 @@
+# 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 http://mozilla.org/MPL/2.0/.
+
+
+from copy import deepcopy
+
+from taskgraph import MAX_DEPENDENCIES
+from taskgraph.transforms.base import TransformSequence
+from taskgraph.util.treeherder import add_suffix
+
+# XXX Docker images may be added after this transform, so we allow one more dep to be added
+MAX_NUMBER_OF_DEPS = MAX_DEPENDENCIES - 1
+
+transforms = TransformSequence()
+
+
+def build_task_definition(orig_task, deps, soft_deps, count):
+    task = deepcopy(orig_task)
+    task["dependencies"] = {label: label for label in deps}
+    task["soft-dependencies"] = list(soft_deps)
+    task["name"] = "{}-{}".format(orig_task["name"], count)
+    if "treeherder" in task:
+        task["treeherder"]["symbol"] = add_suffix(
+            task["treeherder"]["symbol"], f"-{count}"
+        )
+
+    task["attributes"]["is_final_chunked_task"] = False
+    return task
+
+
+def get_chunked_label(config, chunked_task):
+    return "{}-{}".format(config.kind, chunked_task["name"])
+
+
+@transforms.add
+def add_dependencies(config, tasks):
+    for task in tasks:
+        count = 1
+        soft_deps = set()
+        regular_deps = set()
+        chunked_labels = set()
+
+        soft_dep_labels = list(task.pop("soft-dependencies", []))
+        regular_dep_labels = list(task.get("dependencies", {}).keys())
+        # sort for deterministic chunking
+        all_dep_labels = sorted(set(soft_dep_labels + regular_dep_labels))
+
+        for dep_label in all_dep_labels:
+            if dep_label in regular_dep_labels:
+                regular_deps.add(dep_label)
+            else:
+                soft_deps.add(dep_label)
+
+            if len(regular_deps) + len(soft_deps) == MAX_NUMBER_OF_DEPS:
+                chunked_task = build_task_definition(
+                    task, regular_deps, soft_deps, count
+                )
+                chunked_label = get_chunked_label(config, chunked_task)
+                chunked_labels.add(chunked_label)
+                yield chunked_task
+                soft_deps.clear()
+                regular_deps.clear()
+                count += 1
+
+        if regular_deps or soft_deps:
+            chunked_task = build_task_definition(task, regular_deps, soft_deps, count)
+            chunked_label = get_chunked_label(config, chunked_task)
+            chunked_labels.add(chunked_label)
+            yield chunked_task
+
+        task["dependencies"] = {label: label for label in chunked_labels}
+        # Chunk yields a last task that doesn't have a number appended to it.
+        # It helps configuring Github which waits on a single label.
+        # Setting this attribute also enables multi_dep to select the right
+        # task to depend on.
+        task["attributes"]["is_final_chunked_task"] = True
+        yield task