feat: cancel in-progress reviewbot pushes before new ones are started

bot/code_review_bot/workflow.py

@@ -2,6 +2,9 @@
 # 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/.
 
+import base64
+import json
+import re
 import time
 from datetime import datetime, timedelta
 from itertools import groupby
@@ -45,6 +48,13 @@
 TASKCLUSTER_NAMESPACE = "project.relman.{channel}.code-review.{name}"
 TASKCLUSTER_INDEX_TTL = 7  # in days
 
+DECISION_TASK_ROUTE = "gecko.v2.{repo}.revision.{revision}.taskgraph.decision"
+PUBLICATION_LOG_ARTIFACT = "public/logs/live.log"
+TREEHERDER_LINK_REGEX = re.compile(
+    rb"treeherder\.mozilla\.org/(?:#/)?jobs\?repo=(?P<repo>[\w-]+)"
+    rb"&revision=(?P<revision>[0-9a-f]{12,40})"
+)
+
 
 class Workflow:
     """
@@ -347,6 +357,11 @@ def start_analysis(self, revision):
         worker = MercurialWorker()
         output = worker.run(repository, build)
 
+        # Cancel any in-progress tasks from an earlier update
+        # This is done after pushing to try to avoid delaying runs of the
+        # new tasks.
+        self.cancel_previous(revision)
+
         # Update index when the patch has been pushed to try
         self.index(revision, state="pushed_to_try")
 
@@ -465,6 +480,165 @@ def publish(self, revision, issues, task_failures, notices, reviewers):
             else BuildState.Pass,
         )
 
+    def cancel_previous(self, revision):
+        """Cancel the try pushes triggered by earlier updates of a revision"""
+
+        # In order to cancel tasks from a previoush push we need the task group id
+        # of the previous push (which is the same as the decision task for that push).
+        # The only way to retrieve this through the Phabricator API is by a long series
+        # of requests:
+        # * Find all of the prior Build Targets (via builds, via buildables)
+        # * Find the log for each Build Target, which will contain a taskId of Code Review task
+        # * Pull the Code Review task log to fetch the treeherder link with the revision pushed to Try in it
+        # * Look up the decision task id in the task index via the revision
+        #
+        # (Despite the fact that the Phabricator UI shows the treeherder link in it, this is
+        #  not available through the API, so we have to take the long way to get here.)
+        try:
+            publication_tasks = self.list_previous_publication_tasks(revision)
+            logger.info("Found previous publication tasks", tasks=publication_tasks)
+        except Exception as e:
+            logger.warn(
+                "Failed to list previous publication tasks",
+                rev=str(revision),
+                error=str(e),
+            )
+            return
+
+        for task_id in publication_tasks:
+            task_group_id = self.find_try_decision_task(task_id)
+            if task_group_id is None:
+                continue
+
+            try:
+                self.queue_service.cancelTaskGroup(task_group_id)
+            except Exception as e:
+                logger.warn(
+                    "Failed to cancel a previous try push",
+                    task_group_id=task_group_id,
+                    error=str(e),
+                )
+                continue
+
+            logger.info("Cancelled a previous try push", task_group_id=task_group_id)
+
+    def list_previous_publication_tasks(self, revision):
+        logger.debug(
+            "Finding previous publication tasks", phid=revision.phabricator_phid
+        )
+        buildables = self.phabricator.request(
+            "harbormaster.buildable.search",
+            constraints={"containerPHIDs": [revision.phabricator_phid]},
+        )["data"]
+        if not buildables:
+            logger.debug("No buildables found", phid=revision.phabricator_phid)
+            return []
+
+        buildable_phids = [b["phid"] for b in buildables]
+
+        logger.debug(
+            "Found buildables",
+            phid=revision.phabricator_phid,
+            buildables=buildable_phids,
+        )
+        builds = self.phabricator.request(
+            "harbormaster.build.search",
+            constraints={"buildables": buildable_phids},
+        )["data"]
+        if not builds:
+            logger.debug("No builds found", buildables=buildable_phids)
+            return []
+
+        build_phids = [b["phid"] for b in builds]
+
+        logger.debug("Found builds", buildables=buildable_phids, builds=build_phids)
+        targets = [
+            target
+            for target in self.phabricator.request(
+                "harbormaster.target.search",
+                constraints={"buildPHIDs": build_phids},
+            )["data"]
+            # Ignore the current Build Target; there will be no code review
+            # task for it anyways, and we wouldn't want to cancel it even if
+            # there was one.
+            if target["phid"] != revision.build_target_phid
+        ]
+        if not targets:
+            return []
+
+        build_target_phids = [target["phid"] for target in targets]
+
+        logger.debug(
+            "Found build targets", builds=build_phids, build_targets=build_target_phids
+        )
+
+        logs = self.phabricator.request(
+            "harbormaster.log.search",
+            constraints={"buildTargetPHIDs": build_target_phids},
+        )["data"]
+
+        task_ids = []
+        for log in logs:
+            phid = log["fields"]["filePHID"]
+            logger.debug("Downloading log", file=phid)
+            blob = self.phabricator.request("file.download", phid=phid)
+            try:
+                payload = json.loads(base64.b64decode(blob).decode("utf-8", "replace"))
+            except (TypeError, ValueError):
+                logger.debug("Couldn't parse log", file=phid)
+                continue
+
+            if not isinstance(payload, dict):
+                logger.debug("Log is not an object", file=phid)
+                continue
+
+            task_id = payload.get("taskId")
+            if not task_id:
+                logger.debug("Couldn't find task id", file=phid)
+                continue
+
+            if task_id not in task_ids:
+                task_ids.append(task_id)
+
+        return task_ids
+
+    def find_try_decision_task(self, publication_task_id):
+        """
+        Find the decision task of the try push made by a publication task
+
+        The Treeherder link it published is only available in its own live log.
+        """
+        url = self.queue_service.buildUrl(
+            "getLatestArtifact", publication_task_id, PUBLICATION_LOG_ARTIFACT
+        )
+        # Allows HTTP_30x redirections retrieving the artifact
+        response = self.queue_service.session.get(
+            url, stream=True, allow_redirects=True
+        )
+        if not response.ok:
+            logger.warn(
+                "Failed to read the log of a publication task",
+                task=publication_task_id,
+                error=response.status_code,
+            )
+            return
+
+        match = TREEHERDER_LINK_REGEX.search(response.content)
+        if match is None:
+            logger.info(
+                "No try push found for a publication task", task=publication_task_id
+            )
+            return
+
+        route = DECISION_TASK_ROUTE.format(
+            repo=match.group("repo").decode("utf-8"),
+            revision=match.group("revision").decode("utf-8"),
+        )
+        try:
+            return self.index_service.findTask(route)["taskId"]
+        except Exception as e:
+            logger.warn("Failed to find a decision task", route=route, error=str(e))
+
     def index(self, revision, **kwargs):
         """
         Index current task on Taskcluster index

bot/tests/conftest.py

@@ -468,6 +468,14 @@ def raise_for_status(self):
         else:
             return json.dumps(self.body)
 
+    @property
+    def ok(self):
+        return self.code < 300
+
+    @property
+    def status_code(self):
+        return self.code
+
     @property
     def content(self):
         return self.body.encode()
@@ -508,6 +516,7 @@ class MockQueue:
     def __init__(self):
         self._artifacts = {}
         self.session = SessionMock()
+        self.cancelled_groups = []
 
     def configure(self, relations):
         # Reset the session mock
@@ -584,10 +593,19 @@ def listArtifacts(self, task_id, run_id):
     def listLatestArtifacts(self, task_id):
         return self._artifacts.get(task_id, {})
 
-    def buildUrl(self, route_name, task, run, name):
+    def buildUrl(self, route_name, *args):
+        if route_name == "getLatestArtifact":
+            task, name = args
+            return f"http://tc.test/{task}/artifacts/{name}"
+
         assert route_name == "getArtifact"
+        task, run, name = args
         return f"http://tc.test/{task}/{run}/artifacts/{name}"
 
+    def cancelTaskGroup(self, group_id):
+        self.cancelled_groups.append(group_id)
+        return {"taskGroupId": group_id, "taskIds": []}
+
     def createArtifact(self, task_id, run_id, name, payload):
         if task_id not in self._artifacts:
             self._artifacts[task_id] = {"artifacts": []}