Bug 1661493 - Support extending OptimizationSchema in taskgraph proje…

…cts. r=taskgraph-reviewers,ahal

In order for Thunderbird to effectively cut down on excess builds, there needs
to be a way to define taskgraph optimization strategies beyond what is made
available in the mozilla repository.
taskgraph.optimize.register_strategy gets half of the job done. The other piece
is task schema validation which checks optimizations against OptimizationSchema.

OptimizationSchema gets moved to taskgraph.optimize.schema so that it is not set
too early by taskgraph.util.schema. Projects then call set_optimization_schema()
to replace the default schema object. set_optimization_schema() needs to be called
in the project's taskgraph:register function before any transform code is loaded.

Differential Revision: https://phabricator.services.mozilla.com/D88429

taskcluster/taskgraph/optimize/schema.py

@@ -0,0 +1,61 @@
+#  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 __future__ import absolute_import, print_function, unicode_literals
+
+import logging
+import voluptuous
+from six import text_type
+
+from mozbuild import schedules
+
+logger = logging.getLogger(__name__)
+
+
+default_optimizations = (
+    # always run this task (default)
+    None,
+    # always optimize this task
+    {'always': None},
+    # optimize strategy aliases for build kind
+    {'build': list(schedules.ALL_COMPONENTS)},
+    {'build-fuzzing': None},
+    # search the index for the given index namespaces, and replace this task if found
+    # the search occurs in order, with the first match winning
+    {'index-search': [text_type]},
+    # never optimize this task
+    {'never': None},
+    # skip the task except for every Nth push
+    {'skip-unless-expanded': None},
+    {'skip-unless-backstop': None},
+    # skip this task if none of the given file patterns match
+    {'skip-unless-changed': [text_type]},
+    # skip this task if unless the change files' SCHEDULES contains any of these components
+    {'skip-unless-schedules': list(schedules.ALL_COMPONENTS)},
+    # optimize strategy aliases for the test kind
+    {'test': list(schedules.ALL_COMPONENTS)},
+    {'test-inclusive': list(schedules.ALL_COMPONENTS)},
+    # optimize strategy alias for test-verify tasks
+    {'test-verify': list(schedules.ALL_COMPONENTS)},
+    # optimize strategy alias for upload-symbols tasks
+    {'upload-symbols': None},
+)
+
+OptimizationSchema = voluptuous.Any(*default_optimizations)
+
+
+def set_optimization_schema(schema_tuple):
+    """Sets OptimizationSchema so it can be imported by the task transform.
+    This function is called by projects that extend Firefox's taskgraph.
+    It should be called by the project's taskgraph:register function before
+    any transport or job runner code is imported.
+
+    :param tuple schema_tuple: Tuple of possible optimization strategies
+    """
+    global OptimizationSchema
+    if OptimizationSchema.validators == default_optimizations:
+        logger.info("OptimizationSchema updated.")
+        OptimizationSchema = voluptuous.Any(*schema_tuple)
+    else:
+        raise Exception('Can only call set_optimization_schema once.')

taskcluster/taskgraph/transforms/task.py

@@ -31,9 +31,9 @@
     Schema,
     optionally_keyed_by,
     resolve_keyed_by,
-    OptimizationSchema,
     taskref_or_string,
 )
+from taskgraph.optimize.schema import OptimizationSchema
 from taskgraph.util.partners import get_partners_to_be_published
 from taskgraph.util.scriptworker import (
     BALROG_ACTIONS,

taskcluster/taskgraph/transforms/tests.py

@@ -36,14 +36,15 @@
 from taskgraph.transforms.base import TransformSequence
 from taskgraph.util.attributes import match_run_on_projects, keymatch
 from taskgraph.util.keyed_by import evaluate_keyed_by
-from taskgraph.util.schema import resolve_keyed_by, OptimizationSchema
 from taskgraph.util.templates import merge
 from taskgraph.util.treeherder import split_symbol, join_symbol
 from taskgraph.util.platforms import platform_family
 from taskgraph.util.schema import (
+    resolve_keyed_by,
     optionally_keyed_by,
     Schema,
 )
+from taskgraph.optimize.schema import OptimizationSchema
 from taskgraph.util.chunking import (
     chunk_manifests,
     get_runtimes,

taskcluster/taskgraph/util/schema.py

@@ -13,8 +13,6 @@
 
 import taskgraph
 
-from mozbuild import schedules
-
 from .keyed_by import evaluate_keyed_by
 
 
@@ -215,35 +213,6 @@ def __getitem__(self, item):
         return self.schema[item]
 
 
-OptimizationSchema = voluptuous.Any(
-    # always run this task (default)
-    None,
-    # always optimize this task
-    {'always': None},
-    # optimize strategy aliases for build kind
-    {'build': list(schedules.ALL_COMPONENTS)},
-    {'build-fuzzing': None},
-    # search the index for the given index namespaces, and replace this task if found
-    # the search occurs in order, with the first match winning
-    {'index-search': [text_type]},
-    # never optimize this task
-    {'never': None},
-    # skip the task except for every Nth push
-    {'skip-unless-expanded': None},
-    {'skip-unless-backstop': None},
-    # skip this task if none of the given file patterns match
-    {'skip-unless-changed': [text_type]},
-    # skip this task if unless the change files' SCHEDULES contains any of these components
-    {'skip-unless-schedules': list(schedules.ALL_COMPONENTS)},
-    # optimize strategy aliases for the test kind
-    {'test': list(schedules.ALL_COMPONENTS)},
-    {'test-inclusive': list(schedules.ALL_COMPONENTS)},
-    # optimize strategy alias for test-verify tasks
-    {'test-verify': list(schedules.ALL_COMPONENTS)},
-    # optimize strategy alias for upload-symbols tasks
-    {'upload-symbols': None},
-)
-
 # shortcut for a string where task references are allowed
 taskref_or_string = voluptuous.Any(
     text_type,