Bug 1654591 - [taskgraph] Fallback all manifests if there was a timeout in the bugbug loader, r=marco

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

Author: ahal

taskcluster/taskgraph/util/chunking.py

@@ -6,8 +6,9 @@
 
 """Utility functions to handle test chunking."""
 
-import os
 import json
+import logging
+import os
 from abc import ABCMeta, abstractmethod
 
 import six
@@ -22,8 +23,9 @@
 
 from taskgraph import GECKO
 from taskgraph.util.backstop import is_backstop
-from taskgraph.util.bugbug import CT_LOW, push_schedules
+from taskgraph.util.bugbug import BugbugTimeoutException, CT_LOW, push_schedules
 
+logger = logging.getLogger(__name__)
 here = os.path.abspath(os.path.dirname(__file__))
 resolver = TestResolver.from_environment(cwd=here, loader_cls=TestManifestLoader)
 
@@ -210,15 +212,25 @@ class BugbugLoader(DefaultLoader):
     filter them based on a query to bugbug."""
     CONFIDENCE_THRESHOLD = CT_LOW
 
+    def __init__(self, *args, **kwargs):
+        super(BugbugLoader, self).__init__(*args, **kwargs)
+        self.timedout = False
+
     @memoize
     def get_manifests(self, suite, mozinfo):
         manifests = super(BugbugLoader, self).get_manifests(suite, mozinfo)
 
-        # Don't prune any manifests if we're on a backstop push.
-        if is_backstop(self.params):
+        # Don't prune any manifests if we're on a backstop push or there was a timeout.
+        if is_backstop(self.params) or self.timedout:
             return manifests
 
-        data = push_schedules(self.params['project'], self.params['head_rev'])
+        try:
+            data = push_schedules(self.params['project'], self.params['head_rev'])
+        except BugbugTimeoutException:
+            logger.warning("Timed out waiting for bugbug, loading all test manifests.")
+            self.timedout = True
+            return self.get_manifests(suite, mozinfo)
+
         bugbug_manifests = {m for m, c in data.get('groups', {}).items()
                             if c >= self.CONFIDENCE_THRESHOLD}