benchmark_report.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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
import os
import platform
import re
import sys
import time
from collections import defaultdict
from pathlib import Path

import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from plot_style import (
    BAR_EDGE_COLOR,
    GROUP_BAR_WIDTH,
    GROUP_X,
    add_compact_legend,
    apply_benchmark_style,
    format_markdown_with_prettier,
    format_throughput_tick,
    save_benchmark_figure,
    serializer_offset,
    set_grouped_operation_axis,
    style_throughput_axis,
)

apply_benchmark_style(plt)

try:
    import psutil

    HAS_PSUTIL = True
except ImportError:
    HAS_PSUTIL = False

COLORS = {
    "fory": "#FF6f01",
    "protobuf": "#55BCC2",
    "msgpack": (0.55, 0.40, 0.45),
}
SERIALIZER_ORDER = ["fory", "protobuf", "msgpack"]
SERIALIZER_LABELS = {
    "fory": "fory",
    "protobuf": "protobuf",
    "msgpack": "msgpack",
}
DATATYPE_ORDER = [
    "struct",
    "sample",
    "mediacontent",
    "structlist",
    "samplelist",
    "mediacontentlist",
]
OPERATIONS = ["serialize", "deserialize"]
PAIR_CASE_PREFIXES = ("external_", "carrier_", "dynamic_")
PAIR_VARIANTS = ["self", "selected"]
UNIT_TO_NS = {
    "ps": 1e-3,
    "ns": 1.0,
    "us": 1e3,
    "µs": 1e3,
    "ms": 1e6,
    "s": 1e9,
}


def parse_args():
    parser = argparse.ArgumentParser(
        description="Generate Rust benchmark report from Criterion output"
    )
    parser.add_argument(
        "--log-file",
        default="results/cargo_bench.log",
        help="Criterion cargo bench log file",
    )
    parser.add_argument(
        "--size-file",
        default="results/serialized_sizes.txt",
        help="Serialized size table generated by fory_profiler",
    )
    parser.add_argument(
        "--output-dir",
        default="results",
        help="Output directory for report artifacts",
    )
    parser.add_argument(
        "--plot-prefix",
        default="",
        help="Image path prefix inside the markdown report",
    )
    return parser.parse_args()


def datatype_title(datatype):
    if datatype == "struct":
        return "NumericStruct"
    if datatype == "structlist":
        return "NumericStructList"
    if datatype == "mediacontent":
        return "MediaContent"
    if datatype == "mediacontentlist":
        return "MediaContentList"
    if datatype.endswith("list"):
        return f"{datatype[:-4].capitalize()}List"
    return datatype.capitalize()


def get_system_info(log_file):
    info = {
        "OS": f"{platform.system()} {platform.release()}",
        "Machine": platform.machine(),
        "Processor": platform.processor() or "Unknown",
    }
    if HAS_PSUTIL:
        info["CPU Cores (Physical)"] = psutil.cpu_count(logical=False)
        info["CPU Cores (Logical)"] = psutil.cpu_count(logical=True)
        info["Total RAM (GB)"] = round(psutil.virtual_memory().total / (1024**3), 2)
    if os.path.exists(log_file):
        info["Benchmark Date"] = time.strftime(
            "%Y-%m-%dT%H:%M:%S", time.localtime(os.path.getmtime(log_file))
        )
    return info


def parse_time_ns(measurement):
    tokens = measurement.replace("μ", "µ").split()
    if len(tokens) < 4:
        raise ValueError(f"unexpected criterion timing format: {measurement}")
    median_value = float(tokens[2])
    median_unit = tokens[3]
    return median_value * UNIT_TO_NS[median_unit]


def load_benchmark_results(log_file):
    pattern = re.compile(
        r"Benchmarking\s+([A-Za-z0-9_]+)/([A-Za-z0-9_]+).*?time:\s+\[([^\]]+)\]",
        re.DOTALL,
    )
    results = defaultdict(lambda: defaultdict(dict))
    pair_results = defaultdict(lambda: defaultdict(dict))

    with open(log_file, "r", encoding="utf-8") as file:
        content = file.read()

    for datatype, benchmark_name, measurement in pattern.findall(content):
        time_ns = parse_time_ns(measurement)
        if datatype.startswith(PAIR_CASE_PREFIXES):
            if "_" not in benchmark_name:
                continue
            variant, operation = benchmark_name.split("_", 1)
            if variant in PAIR_VARIANTS and operation in OPERATIONS:
                pair_results[datatype][operation][variant] = time_ns
            continue
        if datatype not in DATATYPE_ORDER:
            continue
        if "_" not in benchmark_name:
            continue
        serializer, operation = benchmark_name.split("_", 1)
        if serializer not in SERIALIZER_ORDER or operation not in OPERATIONS:
            continue
        results[datatype][operation][serializer] = time_ns
    return results, pair_results


def load_serialized_sizes(size_file):
    if not os.path.exists(size_file):
        return {}

    pattern = re.compile(
        r"^\|\s*([^|]+?)\s*\|\s*(\d+|n/a)\s*\|\s*(\d+|n/a)\s*\|\s*(\d+|n/a)\s*\|$",
        re.IGNORECASE,
    )
    sizes = {}
    with open(size_file, "r", encoding="utf-8") as file:
        for line in file:
            match = pattern.match(line.strip())
            if not match:
                continue
            datatype, fory_size, protobuf_size, msgpack_size = match.groups()
            if datatype == "Datatype":
                continue
            sizes[datatype] = {
                "fory": parse_size_value(fory_size),
                "protobuf": parse_size_value(protobuf_size),
                "msgpack": parse_size_value(msgpack_size),
            }
    return sizes


def parse_size_value(value):
    return None if value.lower() == "n/a" else int(value)


def format_size_value(value):
    return str(value) if isinstance(value, int) and value >= 0 else "N/A"


def format_tps_tick(tps, _position):
    return format_throughput_tick(tps, _position)


def pair_case_title(case):
    return case.replace("_", " ").title()


def plot_throughput_grid_subplot(ax, results, datatype):
    if datatype not in results:
        ax.set_title(f"{datatype_title(datatype)}\nNo Data")
        ax.axis("off")
        return

    available = [
        serializer
        for serializer in SERIALIZER_ORDER
        if any(
            results[datatype][operation].get(serializer, 0) > 0
            for operation in OPERATIONS
        )
    ]
    if not available:
        ax.set_title(f"{datatype_title(datatype)}\nNo Data")
        ax.axis("off")
        return

    x = GROUP_X
    for index, serializer in enumerate(available):
        throughput = []
        for operation in OPERATIONS:
            time_ns = results[datatype][operation].get(serializer, 0)
            throughput.append(1e9 / time_ns if time_ns > 0 else 0)
        offset = serializer_offset(index, len(available))
        ax.bar(
            x + offset,
            throughput,
            GROUP_BAR_WIDTH,
            label=SERIALIZER_LABELS[serializer],
            color=COLORS.get(serializer, "#888888"),
            edgecolor=BAR_EDGE_COLOR,
            linewidth=0.8,
        )

    max_tps = max(
        1e9 / results[datatype][operation][serializer]
        for operation in OPERATIONS
        for serializer in available
        if results[datatype][operation].get(serializer, 0) > 0
    )
    ax.set_ylim(0, max_tps * 1.12)
    ax.set_title(datatype_title(datatype), pad=8)
    set_grouped_operation_axis(ax)
    style_throughput_axis(ax)
    ax.yaxis.set_major_formatter(FuncFormatter(format_tps_tick))
    add_compact_legend(ax)


def generate_plots(results, output_dir):
    os.makedirs(output_dir, exist_ok=True)

    fig, axes = plt.subplots(2, 3, figsize=(16.5, 9.0))
    for index, (ax, datatype) in enumerate(zip(axes.flat, DATATYPE_ORDER)):
        plot_throughput_grid_subplot(ax, results, datatype)
        if index % 3 == 0:
            ax.set_ylabel("Throughput (ops/sec)", labelpad=10)
    fig.suptitle(
        "Rust Serialization Throughput", fontsize=15, fontweight="normal", y=0.955
    )
    fig.tight_layout(rect=[0.02, 0.02, 0.995, 0.965], w_pad=1.2, h_pad=1.25)
    throughput_path = os.path.join(output_dir, "throughput.png")
    save_benchmark_figure(fig, throughput_path)
    plt.close(fig)

    return throughput_path


def write_report(system_info, results, pair_results, sizes, output_dir, plot_prefix):
    pair_rows = []
    for case in pair_results:
        for operation in OPERATIONS:
            self_time = pair_results[case][operation].get("self", 0)
            selected_time = pair_results[case][operation].get("selected", 0)
            if self_time > 0 and selected_time > 0:
                pair_rows.append((case, operation, self_time, selected_time))

    report = [
        "# Rust Benchmark Performance Report\n\n",
        f"_Generated on {time.strftime('%Y-%m-%d %H:%M:%S')}_\n\n",
        "## How to Generate This Report\n\n",
        "```bash\n",
        "cd benchmarks/rust\n",
        "cargo bench --manifest-path xlang/Cargo.toml --bench serialization_bench 2>&1 | tee results/cargo_bench.log\n",
        "cargo bench --manifest-path local/Cargo.toml --bench external_type_bench 2>&1 | tee -a results/cargo_bench.log\n",
        "cargo run --release --manifest-path xlang/Cargo.toml --bin fory_profiler -- --print-all-serialized-sizes | tee results/serialized_sizes.txt\n",
        "python benchmark_report.py --log-file results/cargo_bench.log --size-file results/serialized_sizes.txt --output-dir results\n",
        "```\n\n",
        "## Benchmark Plot\n\n",
        "The plot shows throughput (ops/sec); higher is better.\n\n",
        f"![Throughput]({plot_prefix}throughput.png)\n\n",
        "## Hardware & OS Info\n\n",
        "| Key | Value |\n",
        "|-----|-------|\n",
    ]

    for key, value in system_info.items():
        report.append(f"| {key} | {value} |\n")

    report.append("\n## Benchmark Results\n\n")
    report.append("### Timing Results (nanoseconds)\n\n")
    report.append(
        "| Datatype | Operation | fory (ns) | protobuf (ns) | msgpack (ns) | Fastest |\n"
    )
    report.append(
        "|----------|-----------|-----------|---------------|--------------|---------|\n"
    )

    for datatype in DATATYPE_ORDER:
        if datatype not in results:
            continue
        for operation in OPERATIONS:
            times = {
                serializer: results[datatype][operation].get(serializer, 0)
                for serializer in SERIALIZER_ORDER
            }
            positive = {name: value for name, value in times.items() if value > 0}
            fastest = min(positive, key=positive.get) if positive else "N/A"
            report.append(
                "| "
                + f"{datatype_title(datatype)} | {operation.capitalize()} | "
                + " | ".join(
                    f"{times[serializer]:.1f}" if times[serializer] > 0 else "N/A"
                    for serializer in SERIALIZER_ORDER
                )
                + f" | {fastest} |\n"
            )

    report.append("\n### Throughput Results (ops/sec)\n\n")
    report.append(
        "| Datatype | Operation | fory TPS | protobuf TPS | msgpack TPS | Fastest |\n"
    )
    report.append(
        "|----------|-----------|----------|--------------|-------------|---------|\n"
    )

    for datatype in DATATYPE_ORDER:
        if datatype not in results:
            continue
        for operation in OPERATIONS:
            throughput = {}
            for serializer in SERIALIZER_ORDER:
                time_ns = results[datatype][operation].get(serializer, 0)
                throughput[serializer] = 1e9 / time_ns if time_ns > 0 else 0
            positive = {name: value for name, value in throughput.items() if value > 0}
            fastest = max(positive, key=positive.get) if positive else "N/A"
            report.append(
                "| "
                + f"{datatype_title(datatype)} | {operation.capitalize()} | "
                + " | ".join(
                    f"{throughput[serializer]:,.0f}"
                    if throughput[serializer] > 0
                    else "N/A"
                    for serializer in SERIALIZER_ORDER
                )
                + f" | {fastest} |\n"
            )

    if pair_rows:
        report.append("\n### External-Type and Carrier Comparisons\n\n")
        report.append(
            "These comparisons use equivalent schemas, wire work, and benchmark configuration. "
            "The selected lane uses an external structural serializer, custom serializer, carrier "
            "serializer, or registered external target as named by the case. The self lane uses "
            "the equivalent self-provided target. Lower time is better.\n\n"
        )
        report.append(
            "| Case | Operation | Self-provided (ns) | Selected (ns) | Selected delta |\n"
        )
        report.append(
            "|------|-----------|--------------------|---------------|----------------|\n"
        )
        for case, operation, self_time, selected_time in pair_rows:
            delta = (selected_time / self_time - 1) * 100
            report.append(
                f"| {pair_case_title(case)} | {operation.capitalize()} | "
                f"{self_time:.1f} | {selected_time:.1f} | {delta:+.2f}% |\n"
            )

    if sizes:
        report.append("\n### Serialized Data Sizes (bytes)\n\n")
        report.append("| Datatype | fory | protobuf | msgpack |\n")
        report.append("|----------|------|----------|---------|\n")
        for datatype in DATATYPE_ORDER:
            title = datatype_title(datatype)
            if title not in sizes:
                continue
            entry = sizes[title]
            report.append(
                f"| {title} | {format_size_value(entry['fory'])} | "
                + f"{format_size_value(entry['protobuf'])} | "
                + f"{format_size_value(entry['msgpack'])} |\n"
            )

    report_path = os.path.join(output_dir, "README.md")
    with open(report_path, "w", encoding="utf-8") as file:
        file.writelines(report)

    format_markdown_with_prettier(report_path)

    return report_path


def main():
    args = parse_args()
    results, pair_results = load_benchmark_results(args.log_file)
    sizes = load_serialized_sizes(args.size_file)
    system_info = get_system_info(args.log_file)
    generate_plots(results, args.output_dir)
    report_path = write_report(
        system_info,
        results,
        pair_results,
        sizes,
        args.output_dir,
        args.plot_prefix,
    )
    print(f"✅ Plots saved in: {args.output_dir}")
    print(f"📄 Markdown report generated at: {report_path}")


if __name__ == "__main__":
    main()