CircleCI cancel redundant builds within workflows

circleci_2_0_cancel_redundant_builds.py

import json
import logging
import os
import sys
from os.path import dirname, join
from time import sleep

sys.path.insert(0, join(dirname(__file__), '.env'))

import arrow
import requests
from raven.contrib.awslambda import LambdaClient

logger = logging.getLogger()
logger.setLevel(logging.INFO)

client = LambdaClient(dsn=os.environ['SENTRY_DSN'])
circle_token = os.environ['CIRCLE_TOKEN']

queued_states = ['not_running', 'scheduled', 'queued']
running_states = ['running']
canceled_states = ['canceled']
failure_states = ['timedout', 'failed', 'not_run', 'retried', 'no_tests', '',
                  'infrastructure_fail'] + canceled_states
success_states = ['success', 'fixed']

active_states = queued_states + running_states
inactive_states = canceled_states + failure_states + success_states
all_states = active_states + inactive_states

base_url = 'https://circleci.com/api/v1.1/project/github/ORGANIZATION/PROJECT'

@client.capture_exceptions
def lambda_handler(event, context):
    """
    See example_payload.json in this directory for an example of what 'event' looks like.
    """
    print('Commit: {}'.format(json.dumps(event['head_commit'])))

    # Set to True if something went wrong and we should fetch again
    rerun = False
    for i in range(5):
        # Rerun at least once in case this lambda ran too fast after the push
        if i > 1:
            # Sleep a little bit in between tries to give things time to settle, such as new jobs to start
            sleep(1)

            # Rerun more times in case we requested that further down
            if not rerun:
                break

        print('i={}'.format(i))
        r = requests.get(
            base_url,
            params={
                'circle-token': circle_token,
                'limit': 100,
            })
        r.raise_for_status()
        build_data = r.json()

        running_builds = {}
        for build in build_data:
            assert build['status'] in all_states, build['status']
            if build['status'] in active_states:
                branch = build.get('branch')
                build_num = build['build_num']
                if not branch:
                    print('Build {} does not have a branch yet, going to rerun...'.format(build_num))
                    rerun = True
                    continue

                running_builds.setdefault(branch, {}).setdefault(build['committer_date'], []).append(build_num)

        print(json.dumps(running_builds))
        for branch, build_time_map in running_builds.items():
            all_timestamps = list(build_time_map.keys())
            cancel_timestamps = set(all_timestamps) - {max(all_timestamps, key=lambda ts: arrow.get(ts).datetime)}
            if cancel_timestamps:
                for ts in cancel_timestamps:
                    print('Branch {}, should cancel builds committed at {}: {}'.format(branch, ts, build_time_map[ts]))
                    for build_num in build_time_map[ts]:
                        cancel_build(build_num)
                        # Many builds are short-lived, so by the time we try to cancel them, they may already be done
                        # and have spawned follow-up tasks, so let's rerun to reap those too
                        rerun = True
            else:
                print('No builds to cancel on branch {}'.format(branch))


def cancel_build(build_num):
    print('Cancelling build {}'.format(build_num))
    r = requests.post(
        base_url + '/{}/cancel'.format(build_num),
        params={
            'circle-token': circle_token,
        })
    r.raise_for_status()
    print(r)