crosvm/tools/infra/binary_size

#!/usr/bin/env python3
# Copyright 2023 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import argparse
import json
import tempfile
import os
import sys
import time
import pathlib

sys.path.append(os.path.dirname(sys.path[0]))

from impl.common import cwd, CROSVM_ROOT, cmd

BOT_MASTERS = "crosvm.ci"
METRIC_NAME = "file_size"
TEST_SUITE_NAME = "crosvm.binary_output"
CHROMEPERF_URL = "https://chromeperf.appspot.com"


# TODO(zihanchen): upload binary size to db for backup
# TODO(zihanchen): Add usage


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--base-dir",
        help="Base dir of cargo target, used to convert target path to shorter relative path, must "
        "be absolute",
    )
    parser.add_argument(
        "--target-name",
        help="Target name of binary",
    )
    parser.add_argument(
        "--target-path",
        help="Path to binary, must be absolute if base dir is provided",
        required=True,
    )
    parser.add_argument("--upload", action="store_true")
    parser.add_argument(
        "--catapult-tool-path",
        help="Path to catapult dashboard upload tool",
        default="/tools/catapult",
    )
    parser.add_argument(
        "--builder-name",
        help="Name of builder generating binaries",
    )
    parser.add_argument("--log-url", help="Url to build execution")
    parser.add_argument("--build-version", help="Version of crosvm corresponding to the results")

    args = parser.parse_args()
    if args.base_dir:
        if not os.path.isabs(args.base_dir):
            raise AssertionError("Base dir mush be and absolute path")
        if not os.path.isabs(args.target_path):
            raise AssertionError("When base dir is provided, target path must be an absolute path")
        if not os.path.commonpath([args.base_dir, args.target_path]) == os.path.normpath(
            args.base_dir
        ):
            raise AssertionError("Base dir must be a parent directory of target path")
    output_path = (
        os.path.relpath(args.target_path, start=args.base_dir)
        if args.base_dir
        else os.path.normpath(args.target_path)
    )

    # output path can contain duplicated target names, dedup before printing
    if len(output_path.split("/")) > 2 and output_path.split("/")[0] == output_path.split("/")[1]:
        output_path = "/".join(output_path.split("/")[1:])

    file_size = os.path.getsize(args.target_path)
    size_dict_for_gerrit = {output_path: file_size}
    print(json.dumps(size_dict_for_gerrit))

    if not args.upload:
        return

    fuchiaperf_fd, fuchiaperf_file_path = tempfile.mkstemp(suffix=".fuchisaperf.json", text=True)

    fuchiaperf_file = os.fdopen(fuchiaperf_fd, "w")

    fuchsiaperf_data = [
        {
            "test_name": output_path,
            "test_suite": TEST_SUITE_NAME,
            "unit": "bytes",
            "values": [file_size],
            "metric": METRIC_NAME,
        }
    ]

    json.dump(fuchsiaperf_data, fuchiaperf_file)
    fuchiaperf_file.close()

    with open(fuchiaperf_file_path, "r") as f:
        print(f.read(), file=sys.stderr)

    _, catapult_file_path = tempfile.mkstemp(suffix=".json", text=True)

    with cwd(CROSVM_ROOT / "tools" / "impl" / "catapult_converter"):
        cmd(
            "cargo",
            "run",
            "--",
            "--input",
            fuchiaperf_file_path,
            "--execution-timestamp-ms",
            str(int(time.time() * 1000)),
            "--masters",
            BOT_MASTERS,
            "--bots",
            args.builder_name,
            "--log-url",
            args.log_url,
            *(
                (
                    "--product-versions",
                    args.build_version,
                )
                if args.build_version
                else ()
            )
        ).write_to(pathlib.Path(catapult_file_path))

    # print catapult file to stderr so we retain a copy
    with open(catapult_file_path, "r") as f:
        print(f.read(), file=sys.stderr)

    print(
        cmd(
            args.catapult_tool_path,
            "upload",
            '--service-account-json=":gce"',
            "--url",
            CHROMEPERF_URL,
            catapult_file_path,
        ).stdout(),
        file=sys.stderr,
    )


if __name__ == "__main__":
    main()