Cancel bitrise tasks on taskcluster cancellation

bitrisescript/src/bitrisescript/bitrise.py

@@ -236,14 +236,19 @@ async def wait_and_download_workflow_log(artifacts_dir: str, build_slug: str) ->
         artifacts_dir (str): Directory to download artifacts to.
         build_slug (str): Identifier of workflow to run.
     """
+    should_retrieve_log = True
     try:
         await wait_for_build_finish(build_slug)
         log.info(f"Build '{build_slug}' is successful. Retrieving artifacts...")
         await download_artifacts(build_slug, artifacts_dir)
+    except asyncio.CancelledError:
+        should_retrieve_log = False
+        raise
     finally:
-        log.info(f"Retrieving bitrise log for '{build_slug}'...")
-        await download_log(build_slug, artifacts_dir)
-        await dump_perfherder_data(artifacts_dir)
+        if should_retrieve_log:
+            log.info(f"Retrieving bitrise log for '{build_slug}'...")
+            await download_log(build_slug, artifacts_dir)
+            await dump_perfherder_data(artifacts_dir)
 
 
 async def run_build(artifacts_dir: str, workflow_id: str, **build_params: Any) -> None:
@@ -268,15 +273,20 @@ async def run_build(artifacts_dir: str, workflow_id: str, **build_params: Any) -
         "build_params": build_params,
     }
 
-    response = await client.request("/builds", method="post", json=data)
-    if response.get("status", "") != "ok":
-        raise Exception(f"Bitrise status for '{workflow_id}' is not ok. Got: {response}")
+    build_slug = None
+    try:
+        response = await client.request("/builds", method="post", json=data)
+        if response.get("status", "") != "ok":
+            raise Exception(f"Bitrise status for '{workflow_id}' is not ok. Got: {response}")
 
-    build_slug = response["build_slug"]
+        build_slug = response["build_slug"]
 
-    log.info(f"Created new job for '{workflow_id}'. Slug: {build_slug}")
+        log.info(f"Created new job for '{workflow_id}'. Slug: {build_slug}")
 
-    await wait_and_download_workflow_log(artifacts_dir, build_slug)
+        await wait_and_download_workflow_log(artifacts_dir, build_slug)
+    except asyncio.CancelledError:
+        if build_slug is not None:
+            await abort_build(build_slug, "Build cancelled")
 
 
 async def get_running_builds(workflow_id: str, **kwargs) -> Optional[str]:
@@ -313,3 +323,21 @@ def find_running_build(running_builds: list[dict], build_params: Any):
             return build["slug"]
     # Nothing found
     return None
+
+
+async def abort_build(build_slug: str, reason: str) -> None:
+    """Abort a specific build
+
+    Args:
+        build_slug (str): The Bitrise build to abort
+        reason (str): The reason for the build cancellation
+    """
+    log.info(f'Aborting build {build_slug} with reason "{reason}"')
+    client = BitriseClient()
+    build_abort_params = {
+        "abort_reason": reason,
+        "abort_with_success": False,
+        "skip_notifications": False,
+    }
+
+    await client.request(f"/builds/{build_slug}/abort", method="post", json=build_abort_params)

bitrisescript/src/bitrisescript/script.py

@@ -4,6 +4,7 @@
 import asyncio
 import logging
 import os
+import signal
 
 from bitrisescript.bitrise import BitriseClient, find_running_build, get_running_builds, run_build, wait_and_download_workflow_log
 from bitrisescript.task import get_artifact_dir, get_bitrise_app, get_bitrise_workflows, get_build_params
@@ -12,6 +13,11 @@
 log = logging.getLogger(__name__)
 
 
+def handle_sigterm(futures):
+    log.info("SIGTERM received, cancelling futures")
+    futures.cancel()
+
+
 async def async_main(config, task):
     app = get_bitrise_app(config, task)
     log.info(f"Bitrise app: '{app}'")
@@ -36,7 +42,10 @@ async def async_main(config, task):
                     futures.append(run_build(artifact_dir, **build_params))
 
         await client.set_app_prefix(app)
-        await asyncio.gather(*futures)
+        future_group = asyncio.gather(*futures)
+        loop = asyncio.get_event_loop()
+        loop.add_signal_handler(signal.SIGTERM, handle_sigterm, future_group)
+        await future_group
     finally:
         if client:
             await client.close()

bitrisescript/tests/test_bitrise.py

@@ -1,3 +1,4 @@
+import asyncio
 import inspect
 import logging
 from asyncio import Future
@@ -352,3 +353,47 @@ async def test_get_running_builds(responses):
 async def test_find_running_build(responses, running_builds, build_params, expected):
     result = bitrise.find_running_build(running_builds, build_params)
     assert result == expected
+
+
+@pytest.mark.asyncio
+async def test_abort_build(mocker, client):
+    build_slug = "123"
+
+    m_request = mocker.patch.object(client, "request", return_value=mocker.AsyncMock())
+
+    await bitrise.abort_build(build_slug, "out of baguettes")
+
+    m_request.assert_called_once_with(
+        f"/builds/{build_slug}/abort", method="post", json={"abort_reason": "out of baguettes", "abort_with_success": False, "skip_notifications": False}
+    )
+
+
+@pytest.mark.asyncio
+async def test_cancel_running_build(mocker, client, tmp_path):
+    build_slug = "abc"
+    build_response = {"status": "ok", "build_slug": build_slug}
+    artifacts_dir = tmp_path / "artifacts"
+    wf_id = "test"
+
+    build_event = asyncio.Event()
+
+    async def wait_for_build_finish(*args):
+        build_event.set()
+        await asyncio.sleep(1)
+        assert False, "The task should have gotten cancelled"
+
+    m_request = mocker.patch.object(client, "request", return_value=mocker.AsyncMock())
+    m_request.return_value = build_response
+
+    m_wait = mocker.patch.object(bitrise, "wait_for_build_finish", wait_for_build_finish)
+    m_dl_log = mocker.patch.object(bitrise, "download_log", return_value=mocker.AsyncMock())
+    m_abort_build = mocker.patch.object(bitrise, "abort_build", return_value=mocker.AsyncMock())
+
+    loop = asyncio.get_event_loop()
+
+    task = loop.create_task(bitrise.run_build(artifacts_dir, wf_id, foo="bar"))
+    await build_event.wait()
+    task.cancel()
+    await task
+
+    m_abort_build.assert_called_once_with(build_slug, "Build cancelled")