gh-distribute-toolchain

#!/usr/bin/env python3
# gh-distribute-toolchain - A script to distribute a toolchain to GitHub
# It does the following:
# 1. Download the toolchain artifacts from the given run
# 2. Create a git tag for the swiftwasm/swift repo with the patches
# 3. Push the tag to the swiftwasm/swift repo
# 4. Create a GitHub release for the tag with the toolchain artifacts
import os
import sys
import subprocess
import asyncio
import json
from dataclasses import dataclass
from typing import Optional, Tuple, List
from build.build_support.actions import derive_options_from_args, REPO_ROOT


USER_AGENT = "gh-distribute-toolchain by swiftwasm/swiftwasm-build"


@dataclass
class Secrets:
    GITHUB_TOKEN: str

    @staticmethod
    def derive():
        github_token = Secrets.env_value("GITHUB_TOKEN")
        if github_token is None:
            result = subprocess.run(
                ["gh", "auth", "token"], capture_output=True, check=False)
            if result.returncode != 0:
                raise Exception(("Failed to get GitHub token "
                                 "from GITHUB_TOKEN and gh auth token"))
            github_token = result.stdout.decode("utf-8").strip()

        return Secrets(
            GITHUB_TOKEN=github_token,
        )

    @staticmethod
    def env_value(key: str):
        return os.environ.get(key)


class GitHub:
    def __init__(self, token: str, repo: str = "swiftwasm/swiftwasm-build"):
        self.token = token
        self.repo = repo

    def list_artifacts(self, run_id: str):
        page = 1
        per_page = 30
        artifacts = []
        while True:
            url = ("https://api.github.com"
                   f"/repos/{self.repo}/actions/runs/{run_id}/artifacts"
                   f"?page={page}&per_page={per_page}")
            response = self.json_request("GET", url)
            if "artifacts" not in response:
                raise Exception(f"Unexpected response: {response}")
            artifacts += response["artifacts"]
            if len(response["artifacts"]) < per_page:
                break
            page += 1
        return artifacts

    def revision_at_run(self, run_id: str):
        url = ("https://api.github.com"
               f"/repos/{self.repo}/actions/runs/{run_id}")
        response = self.json_request("GET", url)
        return response["head_sha"]

    def get_release(self, tag_name: str):
        url = ("https://api.github.com"
               f"/repos/{self.repo}/releases/tags/{tag_name}")
        return self.json_request("GET", url)

    def get_release_by_id(self, id: str):
        url = (f"https://api.github.com/repos/{self.repo}/releases/{id}")
        return self.json_request("GET", url)

    def download_release_asset(self, url: str):
        return self.request("GET", url, headers={
            "Accept": "application/octet-stream"
        })

    def create_prerelease(self, tag_name: str, body: str):
        url = f"https://api.github.com/repos/{self.repo}/releases"
        return self.json_request("POST", url, body={
            "tag_name": tag_name,
            "name": tag_name,
            "prerelease": True,
            "body": body
        })

    def update_release_notes(self, release_id: str, body: str):
        url = (f"https://api.github.com/repos/{self.repo}/releases/{release_id}")
        return self.json_request("PATCH", url, body={"body": body})

    def workflow_runs(self, workflow_name: str, branch: str):
        url = ("https://api.github.com"
               f"/repos/{self.repo}/actions/workflows/{workflow_name}"
               f"/runs?branch={branch}")
        run = self.json_request("GET", url)
        workflow_runs = run["workflow_runs"]
        workflow_runs = sorted(
            workflow_runs, key=lambda run: run["run_number"], reverse=True)
        return workflow_runs

    async def download_artifact(self, artifact, path: str):
        curl_args = ["curl", "-L", "-s", "--show-error", "-o", path,
                     "--max-time", "600", "--retry", "5",
                     artifact["archive_download_url"],
                     "--header", f"Authorization: Bearer {self.token}"]
        proc = await asyncio.subprocess.create_subprocess_exec(
            curl_args[0], *curl_args[1:])
        returncode = await proc.wait()
        if returncode != 0:
            raise Exception(f"\"{' '.join(curl_args)}\" failed")

    def upload_release_asset(self, release_id: str, asset_path: str):
        content_len = os.path.getsize(asset_path)
        filename = os.path.basename(asset_path)
        url = (f"https://uploads.github.com/repos/{self.repo}/"
               f"releases/{release_id}/assets?name={filename}")
        curl_args = ["curl", "-XPOST",
                     "--max-time", "600", "--retry", "5",
                     "--header", f"Authorization: Bearer {self.token}",
                     "--header", f"Content-Length: {content_len}",
                     "--header", "Content-Type: application/x-gzip",
                     "--upload-file", asset_path,
                     url]
        subprocess.run(curl_args, check=True)

    def json_request(self, method: str, path: str,
                     body: Optional[dict] = None):
        import json
        headers = {"Accept": "application/vnd.github.v3+json"}
        data = None
        if body:
            data = json.dumps(body).encode("utf-8")
        resp = self.request(method, path, headers=headers, data=data)
        return json.loads(resp.read())

    def request(self, method: str, url: str, headers: dict,
                data: Optional[bytes] = None):
        headers = headers.copy()
        headers["Authorization"] = f"Bearer {self.token}"
        headers["User-Agent"] = USER_AGENT

        import urllib.request
        req = urllib.request.Request(url, headers=headers, method=method,
                                     data=data)
        resp = urllib.request.urlopen(req)
        if resp.status < 200 or 300 <= resp.status:
            raise Exception(
                ("GitHub API request failed: "
                 f"{resp.status} {resp.reason} {resp.read().decode('utf-8')}"))
        return resp


class SDKArtifactFormat:
    def swift_sdk_artifact_id(self, tag_name: str):
        raise NotImplementedError()

    def swift_sdk_artifact_filename(self, tag_name: str):
        raise NotImplementedError()


class SDKArtifactFormatV3(SDKArtifactFormat):
    def __init__(self, target_triple: str):
        self.target_triple = target_triple

    def swift_sdk_artifact_id(self, tag_name: str):
        return tag_name.removeprefix("swift-wasm-") + "-" + self.target_triple

    def swift_sdk_artifact_filename(self, tag_name: str):
        return f"{tag_name}-{self.target_triple}.artifactbundle"


class Distribution:
    def __init__(self, github: GitHub, swift_github: GitHub,
                 run_id: str, secrets: Secrets, dry_run: bool = False, verbose: bool = False):
        self.github = github
        self.swift_github = swift_github
        self.run_id = run_id
        self.secrets = secrets
        self.dry_run = dry_run
        self.verbose = verbose
        self.distribution_dir = os.path.join(
            os.path.dirname(REPO_ROOT), "build", "Distribution")
        self.artifacts_dir = os.path.join(
            self.distribution_dir, "artifacts", self.run_id)
        self.checkout_dir = os.path.join(
            self.distribution_dir, "checkout")

    async def run(self, options):
        downloads = []
        for artifact in self.toolchain_artifacts(options):
            if not artifact["name"].endswith("-artifactbundle"):
                print(f"Skipping {artifact['name']} because it's not an"
                      " artifactbundle")
                continue

            try:
                scheme, _ = derive_platform_suffix_and_scheme(
                    artifact["name"], options.scheme)
            except Exception:
                print(f"Skipping {artifact['name']} because it's not a valid artifactbundle for {options.scheme}")
                continue

            if options.scheme != scheme:
                print(f"Skipping {artifact['name']} because it's not scheme {options.scheme}")
                # Skip unrelated artifact
                continue

            async def download_work(artifact):
                downloaded_path = await self.download_artifact(artifact)
                return [artifact, downloaded_path]

            print((f"Downloading {artifact['name']} from "
                   f"{artifact['archive_download_url']}..."))
            downloads.append(
                asyncio.create_task(download_work(artifact)))

        try:
            downloaded_paths = await asyncio.gather(*downloads)
        except Exception as e:
            print(f"Download task failed: {e}, cancelling all tasks")
            for task in downloads:
                task.cancel()
            raise e

        if options.verbose:
            print("Downloaded paths:")
            for artifact, artifact_path in downloaded_paths:
                print(f"  {artifact['name']} -> {artifact_path}")

        release = None
        tag_name = options.override_name

        # Create tag and release on GitHub
        if not tag_name:
            for artifact, artifact_path in downloaded_paths:
                if not artifact["name"].endswith("-artifactbundle"):
                    continue
                _, target_triple = derive_platform_suffix_and_scheme(
                    artifact["name"], options.scheme)
                bundle_path = await self.unpack_artifactbundle(artifact_path)
                dirents = os.listdir(bundle_path)
                dirents.remove("info.json")
                sdk_artifact_id = dirents[0]
                tag_name = "swift-wasm-" + sdk_artifact_id.removesuffix("-" + target_triple)
                break

        if not tag_name:
            raise Exception("Could not determine tag name")

        release, swift_version, swiftwasm_build_version = self.create_tag_and_prerelease(tag_name, options)

        # Move to artifacts directory because "zip" does not have
        # --chdir option unlike "tar"
        os.chdir(self.artifacts_dir)

        packaging_tasks = []
        # Package and upload Swift SDK artifactbundles
        for downloaded in downloaded_paths:
            artifact, _ = downloaded
            if not artifact["name"].endswith("-artifactbundle"):
                continue

            async def package_sdk_work(downloaded):
                artifact, artifact_path = downloaded
                scheme, target_triple = derive_platform_suffix_and_scheme(
                    artifact["name"], options.scheme)
                artifact_format = SDKArtifactFormatV3(target_triple)
                bundle_path = await self.unpack_artifactbundle(artifact_path)
                package = await self.package_artifactbundle(
                    bundle_path, tag_name, artifact_format)
                if options.dry_run:
                    print(f"Skip uploading actual artifact \"{package}\"")
                    return
                self.upload_to_release(package, release)

                # Upload sha256 checksum
                import hashlib
                checksum_path = package + ".sha256"
                with open(package, "rb") as f:
                    checksum = hashlib.sha256(f.read()).hexdigest()
                with open(checksum_path, "w") as f:
                    f.write(checksum)
                self.upload_to_release(checksum_path, release)

            print(f"Packaging {artifact['name']}...")
            packaging_tasks.append(
                asyncio.create_task(package_sdk_work(downloaded)))

        try:
            await asyncio.gather(*packaging_tasks)
        except Exception as e:
            print(f"Some of packaging tasks failed: {e}, cancelling all tasks")
            for task in packaging_tasks:
                task.cancel()
            raise e

        # Update release notes with checksums
        self.update_release_notes_with_checksums(
            release, swift_version, swiftwasm_build_version)

    async def package_artifactbundle(
        self, artifact_path: str, tag_name: str,
        artifact_format: SDKArtifactFormat
    ):
        import shutil
        print(f"Packaging {artifact_path}")
        dirents = os.listdir(artifact_path)
        if len(dirents) != 2:
            raise Exception((
                f"Unexpected number of files in {artifact_path}:"
                f" {len(dirents)} (expected 2)"))
        info_json = os.path.join(artifact_path, "info.json")
        dirents.remove(os.path.basename(info_json))
        sdk_artifact_id = dirents[0]
        sdk_path = os.path.join(artifact_path, sdk_artifact_id)

        expected_artifact_id = artifact_format.swift_sdk_artifact_id(tag_name)
        if sdk_artifact_id != expected_artifact_id:
            print((f"Re-packaging artifactbundle: "
                  f"{sdk_artifact_id} -> {expected_artifact_id}"))
            new_sdk_path = os.path.join(artifact_path, expected_artifact_id)
            shutil.move(sdk_path, new_sdk_path)

            # Replace artifact id in info.json
            self.rename_artifact_id(
                info_json, sdk_artifact_id, expected_artifact_id)
            sdk_path = new_sdk_path
            sdk_artifact_id = expected_artifact_id

        # Rename .artifactbundle file name
        artifact_filename = artifact_format.swift_sdk_artifact_filename(tag_name)
        tagged_artifact_path = os.path.join(
            os.path.dirname(artifact_path), artifact_filename)
        bundlezip_path = tagged_artifact_path + ".zip"

        if not os.path.exists(tagged_artifact_path):
            shutil.move(artifact_path, tagged_artifact_path)

        if os.path.exists(bundlezip_path):
            return bundlezip_path
        # Re-zip artifactbundle
        print((f"Re-zipping artifactbundle {tagged_artifact_path}"
               f" to {bundlezip_path}"))

        # Due to lack of "--chdir", we chdir to artifacts dir before
        # packaging works.
        assert os.getcwd() == self.artifacts_dir
        relative_tagged_artifact_path = os.path.relpath(tagged_artifact_path)
        assert relative_tagged_artifact_path == os.path.basename(tagged_artifact_path)

        zip_proc = await asyncio.subprocess.create_subprocess_exec(
            "zip", "-r", "-y", "-q",
            bundlezip_path, relative_tagged_artifact_path)
        if await zip_proc.wait() != 0:
            raise Exception(f"Failed to zip {tagged_artifact_path}")
        return bundlezip_path

    def rename_artifact_id(self, info_json_path, old_id, new_id):
        print(f"Renaming artifact id: {old_id} -> {new_id}")
        with open(info_json_path, "r") as f:
            contents = f.read()
        contents = contents.replace(old_id, new_id)
        with open(info_json_path, "w") as f:
            f.write(contents)

    def unzip_artifact(self, artifact_zip):
        print(f"Unpacking {artifact_zip}...")
        # Get the list of files in the .zip file
        contents = subprocess.check_output(["zipinfo", "-1", artifact_zip])
        contents = contents.decode("utf-8").splitlines()
        if len(contents) != 1:
            raise Exception(
                (f"Unexpected number of files in {artifact_zip}:"
                 f" {len(contents)} (expected 1)"))

        content_name = contents[0]
        content_path = os.path.join(
            os.path.dirname(artifact_zip), content_name)

        if not os.path.exists(content_path):
            self.effect_run(["unzip", artifact_zip, "-d",
                            os.path.dirname(content_path)], check=True)
        return [content_name, content_path]

    async def unpack_artifactbundle(self, zip_bundlezip):
        artifact_name, bundlezip_path = self.unzip_artifact(zip_bundlezip)
        bundle_path, _ = os.path.splitext(bundlezip_path)
        if os.path.exists(bundle_path):
            print(f"Already unzipped: {bundle_path}")
            return bundle_path
        print(f"Unzipping .artifactbundle.zip: {bundlezip_path}")
        await self.check_async_subprocess(
            "unzip", "-q", bundlezip_path, "-d", os.path.dirname(bundle_path))
        return bundle_path

    async def download_artifact(self, artifact):
        if not os.path.exists(self.artifacts_dir):
            os.makedirs(self.artifacts_dir, exist_ok=True)
        path = os.path.join(self.artifacts_dir, artifact["name"])
        if os.path.exists(path):
            print(f"Artifact already exists at {path}")
            return path
        await self.github.download_artifact(artifact, path)
        print(f"Downloaded {artifact['name']}")
        return path

    def create_tag_and_prerelease(self, tag_name, options):
        print(f"Creating tag for {tag_name}...")
        swift_version, swiftwasm_build_version = \
            self.git_push_swift_source(options, tag_name)
        try:
            # If release already exists, use it
            release = self.swift_github.get_release(tag_name)
            print(f"Use existing prerelease for {tag_name}...")
            return [release, swift_version, swiftwasm_build_version]
        except Exception:
            pass
        print(f"Creating prerelease for {tag_name}...")
        body = self.make_release_note(swift_version, swiftwasm_build_version)
        print(f"Release note:\n{body}")
        return [self.swift_github.create_prerelease(tag_name, body), swift_version, swiftwasm_build_version]

    def make_release_note(self, swift_version, swiftwasm_build_version,
                          checksums: Optional[List[tuple]]=None):
        projects = [
            ["apple/swift",
                f"releases/tag/{swift_version}"],
            ["swiftwasm/swiftwasm-build",
                f"commit/{swiftwasm_build_version}"]
        ]
        body = "| Project | Version |\n"
        body += "|:--|:--|\n"
        for repo, subpath in projects:
            body += f"| `{repo}` | https://github.com/{repo}/{subpath} |\n"
        if checksums:
            body += "\n"
            body += """\
### Installation

You can install Swift SDKs for WebAssembly using the following commands:

"""
            for name, download_url, checksum in checksums:
                body += f"**{name}**\n\n"
                body += "```console\n"
                body += f"$ swift sdk install {download_url} --checksum {checksum}\n"
                body += "```\n"
        return body

    def update_release_notes_with_checksums(self, release, swift_version, swiftwasm_build_version):
        checksums = []
        release = self.swift_github.get_release_by_id(release["id"])
        for artifact in release["assets"]:
            if not artifact["name"].endswith(".artifactbundle.zip"):
                continue
            name = artifact["name"]
            artifact_download_url = artifact["browser_download_url"]
            checksum_download_url = artifact_download_url + ".sha256"
            checksum = self.github.download_release_asset(checksum_download_url)
            checksum = checksum.read().decode("utf-8").strip()
            checksums.append([name, artifact_download_url, checksum])

        body = self.make_release_note(swift_version, swiftwasm_build_version, checksums)
        self.swift_github.update_release_notes(release["id"], body)

    def git_push_swift_source(self, options, tag_name):
        build_repo = "https://github.com/swiftwasm/swiftwasm-build.git"
        build_rev = self.github.revision_at_run(self.run_id)

        os.makedirs(self.checkout_dir, exist_ok=True)
        build_repo_dir = os.path.join(self.checkout_dir, "swiftwasm-build")
        if not os.path.exists(build_repo_dir):
            self.subprocess_run(["git", "clone", build_repo,
                                build_repo_dir], check=True)

        self.subprocess_run(["git", "-C", build_repo_dir,
                            "fetch", "origin"], check=True)
        self.subprocess_run(["git", "-C", build_repo_dir,
                             "checkout", build_rev], check=True)

        # Checkout swift repository and apply patches at the build revision
        git_swift_workspace = os.path.join(
            build_repo_dir, "tools", "git-swift-workspace")
        self.subprocess_run([git_swift_workspace, "--scheme",
                            options.scheme], check=True)

        swift_repo_dir = os.path.join(self.checkout_dir, "swift")

        fork_repo = 'git@github.com:swiftwasm/swift.git'
        status = self.subprocess_run(['git', '-C', swift_repo_dir,
                                      'remote', 'get-url', 'swiftwasm']).returncode
        if status != 0:
            self.effect_run(['git', '-C', swift_repo_dir, 'remote',
                            'add', 'swiftwasm', fork_repo], check=True)

        # Fetch the tag from the fork. This can fail if the tag doesn't exist
        self.subprocess_run(['git', '-C', swift_repo_dir,
                            'fetch', 'swiftwasm', 'tag', tag_name, '--no-tags'])

        status = subprocess.run([
            'git', '-C', swift_repo_dir,
            'tag', '--list', '--contains', tag_name]).returncode
        if status != 0:
            self.effect_run(["git", "-C", swift_repo_dir, "tag", tag_name], check=True)
            self.effect_run(["git", "-C", swift_repo_dir,
                            "push", "swiftwasm", tag_name], check=True)
        else:
            print(f"Tag {tag_name} already exists")

        manifest = os.path.join(
            build_repo_dir, "schemes", options.scheme, "manifest.json")
        with open(manifest, "r") as f:
            manifest_json = json.load(f)
            base_tag = manifest_json["base-tag"]

        return [base_tag, build_rev]

    def upload_to_release(self, artifact_path, release):
        name = os.path.basename(artifact_path)
        if "assets" in release:
            for asset in release["assets"]:
                asset_name = asset["name"]
                if os.path.basename(artifact_path) == asset_name:
                    print((f"{name} is already uploaded"
                           f" to release {release['name']}"))
                    return
        print(f"Uploading {name} to release {release['name']}...")
        self.swift_github.upload_release_asset(release['id'], artifact_path)

    def toolchain_artifacts(self, options):
        artifacts = self.github.list_artifacts(self.run_id)
        for artifact in artifacts:
            artifact_name = artifact["name"]
            should_yield = artifact_name.endswith("-artifactbundle")
            if should_yield:
                yield artifact

    def subprocess_run(self, args, **kwargs):
        """
        Run a non-mutating subprocess and print the command if
        verbose or dry_run is True.
        """
        if self.verbose or self.dry_run:
            print(" ".join(args))
        return subprocess.run(args, **kwargs)

    def effect_run(self, args, **kwargs):
        """
        Run a mutating subprocess and print the command if
        verbose or dry_run is True.
        """
        if self.verbose or self.dry_run:
            print(" ".join(args))
        if self.dry_run:
            return
        return subprocess.run(args, **kwargs)

    async def check_async_subprocess(self, program, *args):
        if self.verbose or self.dry_run:
            print(f"[async] {program} {' '.join(args)}")
        if self.dry_run:
            return
        proc = await asyncio.subprocess.create_subprocess_exec(program, *args)
        retcode = await proc.wait()
        if retcode != 0:
            raise Exception(f"Failed to execute: {' '.join(args)}")


def derive_platform_suffix_and_scheme(
    artifact_name: str, target_scheme: str
) -> Optional[Tuple[str, str]]:
    """
    Returns platform suffix, scheme, and target triple derived from the
    artifact name.
    e.g.
        "main-wasm32-unknown-wasi-artifactbundle", scheme="main"
          -> ["main", "wasm32-unknown-wasi"]
    """

    # v3 artifact name: <scheme>-<target-triple>-artifactbundle
    # Note that <scheme> can contain '-'.
    name: str = artifact_name
    artifact_suffixes = ("-artifactbundle")
    if not name.endswith(artifact_suffixes):
        raise Exception((f"Unexpected artifact name {name}"
                         f", expected to have one of \"{artifact_suffixes}\""
                         " suffix"))

    if name.startswith(target_scheme):
        # v3 artifact name
        rest = name[len(target_scheme) + 1:]
        if not rest.endswith("-artifactbundle"):
            raise Exception(
                f"Unexpected artifact name {name} with format v3 should"
                " end with -artifactbundle")
        target_triple = rest[:-len("-artifactbundle")]
        return [target_scheme, target_triple]
    raise Exception(f"Unexpected artifact name {name}")


def latest_success_run_id(gh: GitHub, workflow_name: str, branch: str, scheme: str):
    """
    Find the latest successful run ID for the given workflow name and branch,
    also containing a toolchain artifact built for the given scheme
    """
    for run in gh.workflow_runs(workflow_name, branch):
        if "head_branch" not in run or run["head_branch"] != branch:
            continue
        if "conclusion" not in run or run["conclusion"] != "success":
            continue
        artifacts = gh.list_artifacts(run["id"])
        for artifact in artifacts:
            artifact_name = artifact["name"]
            try:
                artifact_scheme, _ = derive_platform_suffix_and_scheme(
                    artifact_name, scheme)
                if artifact_scheme == scheme:
                    return run["id"]
            except Exception:
                continue
    print(f"Could not find a successful run for {workflow_name} on {branch} branch with scheme {scheme}")
    return None


async def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('run_id', type=str)
    parser.add_argument(
        '--override-name', type=str,
        default=os.environ.get("GH_DISTRIBUTE_TOOLCHAIN_OVERRIDE_NAME"))
    parser.add_argument(
        '--only-swift-sdk', action='store_true')
    options = derive_options_from_args(sys.argv[1:], parser)

    secrets = Secrets.derive()
    gh = GitHub(secrets.GITHUB_TOKEN, "swiftwasm/swiftwasm-build")
    swift_gh = GitHub(secrets.GITHUB_TOKEN, "swiftwasm/swift")
    swiftwasm_build_dir = os.path.join(os.path.dirname(__file__), "..")
    os.chdir(swiftwasm_build_dir)
    run_id = options.run_id
    if run_id == "latest":
        run_id = latest_success_run_id(
            gh, workflow_name="build-toolchain.yml", branch="main", scheme=options.scheme)
        if run_id is None:
            # No successful run found. This scheme might not be built recently.
            return
        run_id = str(run_id)
        print(f"Automatically determined run_id: {run_id}")
    await Distribution(gh, swift_gh, run_id, secrets,
                       options.dry_run, options.verbose).run(options)


if __name__ == '__main__':
    asyncio.run(main())