soong_trace.py

#!/usr/bin/env python3
#
# Copyright (C) 2022 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
from dataclasses import dataclass
import gzip
import json
from pathlib import Path
from typing import TextIO

from paths import TOOLCHAIN_MODULE_LIST_PATH
from utils import ResolvedPath, ms_to_hms

SUFFIXES_RUST = [
    ".rs",
    ".rlib",
    ".rlib.clippy",
    ".dylib.so",
]

#
# Classes
#

@dataclass
class TraceInfo:
    targets_total:     int = 0
    targets_rust:      int = 0
    duration_total_ms: int = 0
    duration_rust_ms:  int = 0

#
# Helper functions
#

def open_trace(trace_path: Path) -> TextIO:
    if not trace_path.exists():
        print(f"Trace file does not exist: {trace_path.as_posix()}")
        exit(-1)

    if trace_path.suffix == ".gz":
       return gzip.open(trace_path, mode="rt")
    else:
        return open(trace_path, mode="r")

def is_rust_target(target_name: str, module_prefixes: list[str]) -> bool:
    target_path = Path(target_name)

    if target_path.suffix != "":
        for suffix in SUFFIXES_RUST:
            if target_path.name.endswith(suffix):
                return True

        return False
    else:
        for prefix in module_prefixes:
            if target_name.startswith(prefix):
                return True

        return False


def target_in_known_rust_module(target_name: str, module_prefixes: list[str]) -> bool:
    for prefix in module_prefixes:
        if target_name.startswith(prefix):
            return True
    return False

#
# Program logic
#

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Produce a summary of Rust-related information from a Soong build trace")

    parser.add_argument("trace", metavar="TRACE", type=ResolvedPath, help="Soong trace file to process")

    return parser.parse_args()


# The module list is generated using the Bazel query command:
#
# $ m queryview
# $ bazel query --config=queryview 'kind("rust_(?!default).*", //...)' --output label_kind
#
# `rust_defaults` targets are excluded as they are abstract and used
# to avoid duplicating common properties among targets.
def load_target_prefixes() -> list[str]:
    rust_target_prefixes = []
    TRACE_PATH_PREFIX = "out/soong/.intermediates"
    if TOOLCHAIN_MODULE_LIST_PATH.exists():
        with open(TOOLCHAIN_MODULE_LIST_PATH, "r") as ml:
            # Example Line:
            # rust_binary_host rule //external/rust/crates/grpcio-compiler:grpc_rust_plugin--linux_glibc_x86_64
            #
            # comps    = ["rust_binary_host", "rule", "//external/rust/crates/grpcio-compiler:grpc_rust_plugin--linux_glibc_x86_64"]
            # mod_path = "//external/rust/crates/grpcio-compiler"
            # rest     = "grpc_rust_plugin--linux_glibc_x86_64"
            # mod_name = "grpc_rust_plugin"
            # variant  = "linux_glibc_x86_64"
            #
            # target_prefix = out/soong/.intermediates/external/rust/crates/grpcio-compiler/grpc_rust_plugin/linux_glibc_x86_64
            for line in ml.readlines():
                comps             = line.strip().split(" ", 3)
                mod_path, rest    = comps[2].split(":")
                mod_name, variant = rest.split("--")

                target_prefix = f"{TRACE_PATH_PREFIX}/{mod_path[2:]}/{mod_name}/{variant}"
                rust_target_prefixes.append(target_prefix)

    return rust_target_prefixes


def process_trace(trace_path: Path, rust_target_prefixes: list[str]) -> TraceInfo:
    info = TraceInfo()
    with open_trace(trace_path) as fd:
        for item in json.load(fd):
            if "dur" in item:
                info.targets_total += 1
                info.duration_total_ms += item["dur"]

                if is_rust_target(item["name"], rust_target_prefixes):
                    info.targets_rust += 1
                    info.duration_rust_ms += item["dur"]

    return info


def main() -> None:
    args = parse_args()

    # Load list of known Rust targets if it exists.
    rust_target_prefixes = load_target_prefixes()
    trace_info = process_trace(args.trace, rust_target_prefixes)

    rust_dur_parts  = ms_to_hms(trace_info.duration_rust_ms)
    total_dur_parts = ms_to_hms(trace_info.duration_total_ms)

    print(f"Rust targets: {trace_info.targets_rust}")
    print(f"Rust duration (ms): {trace_info.duration_rust_ms}")
    print(f"Rust duration (dd hh:mm:ss.ms): {rust_dur_parts[0]} {rust_dur_parts[1]}:{rust_dur_parts[2]}:{rust_dur_parts[3]:02.2f}")
    print(f"Total targets: {trace_info.targets_total}")
    print(f"Total duration (ms): {trace_info.duration_total_ms}")
    print(f"Total duration (dd hh:mm:ss.ms): {total_dur_parts[0]} {total_dur_parts[1]}:{total_dur_parts[2]}:{total_dur_parts[3]:02.2f}")
    print(f"Rust percent: {(float(trace_info.duration_rust_ms) / float(trace_info.duration_total_ms) * 100):02.2f}%")
    print()


if __name__ == '__main__':
    main()