src/performance/lib/storage/block_io.py

# Copyright 2026 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Extracts and publishes cumulative block/storage I/O rate metrics from Inspect."""

import json
import logging
import os
import pathlib
import typing
import zipfile
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Protocol, Self

from perf_publish import publish
from reporting import metrics
from trace_processing import trace_metrics, trace_model

_LOGGER: logging.Logger = logging.getLogger(__name__)

# Maps block driver Inspect tree name (metadata.name) to root child node name.
_BLOCK_DRIVER_NODES: dict[str, str] = {
    "ftl": "ftl",
    "sdmmc": "sdmmc_core",
    "ufs": "ufs",
    "virtio-block": "virtio_block",
}

# (op_key, op_label) tuples for Inspect block operation counters and labels.
_BLOCK_OPERATIONS: tuple[tuple[str, str], ...] = (
    ("read", "Read"),
    ("write", "Write"),
    ("flush", "Flush"),
    ("trim", "Trim"),
)

# Default block operations published as headline structured metrics
# (fuchsiaperf.json). Headline metrics track write and flush operation rates
# (ops/sec), which represent the primary indicators of flash wear, filesystem
# journaling pressure, and sync latency. All operations (read, write, flush,
# trim) across both ops/sec and bytes/sec continue to be recorded as freeform
# metrics in storage_io.freeform.json.
_DEFAULT_HEADLINE_OPERATIONS: tuple[tuple[str, str], ...] = (
    ("write", "Write"),
    ("flush", "Flush"),
)

STORAGE_INSPECT_SELECTORS: list[str] = [
    "bootstrap/*-drivers*:[name=ftl]root",
    "bootstrap/*-drivers*:[name=sdmmc]root",
    "bootstrap/*-drivers*:[name=ufs]root",
    "bootstrap/*-drivers*:[name=virtio-block]root",
]


class InspectDataProvider(Protocol):
    """Protocol for a Honeydew device object that can query Inspect data."""

    def get_inspect_data(
        self,
        *,
        selectors: list[str] | None = None,
    ) -> Any:
        """Returns an InspectDataCollection for the requested selectors."""


@dataclass(frozen=True)
class InspectEntry:
    """Normalized Inspect entry from a snapshot zip or Honeydew query."""

    moniker: str | None
    name: str | None
    timestamp_ns: int
    root: Mapping[str, Any]

    @classmethod
    def from_dict(cls, entry: Mapping[str, Any]) -> "InspectEntry | None":
        """Parses a raw JSON Inspect entry dictionary into an `InspectEntry`."""
        metadata = entry.get("metadata")
        if not isinstance(metadata, Mapping):
            return None
        payload = entry.get("payload")
        if not isinstance(payload, Mapping):
            return None
        root = payload.get("root")
        if not isinstance(root, Mapping):
            return None
        timestamp = metadata.get("timestamp", 0)
        name = metadata.get("name")
        moniker = entry.get("moniker")
        return cls(
            moniker=str(moniker) if moniker is not None else None,
            name=str(name) if name is not None else None,
            timestamp_ns=int(timestamp or 0),
            root=root,
        )


def load_inspect_entries(
    dut: InspectDataProvider | None = None,
    snapshot_dirs: Sequence[str | os.PathLike[str]] = (),
) -> list[InspectEntry]:
    """Loads Inspect entries from the newest `Snapshot_*.zip` in `snapshot_dirs`,
    falling back to `dut.get_inspect_data()` if no snapshot zip is found.

    Checking multiple directories (e.g. `[self._teardown_class_artifacts, self.log_path]`)
    supports base test classes that save teardown snapshots in either
    `<log_path>/teardown_class` or `<log_path>`.

    Logs a warning and returns `[]` on any snapshot read or Inspect query error
    rather than raising.
    """
    for directory in snapshot_dirs:
        if not directory:
            continue
        dir_path = pathlib.Path(directory)
        if not dir_path.exists():
            continue
        snapshot_zips = sorted(dir_path.glob("Snapshot_*.zip"))
        for snapshot_zip in reversed(snapshot_zips):
            try:
                with zipfile.ZipFile(snapshot_zip, "r") as zf:
                    if "inspect.json" in zf.namelist():
                        data = json.loads(zf.read("inspect.json"))
                        if isinstance(data, list):
                            entries = [
                                parsed
                                for item in data
                                if isinstance(item, Mapping)
                                and (parsed := InspectEntry.from_dict(item))
                                is not None
                            ]
                            if entries:
                                return entries
            except Exception as e:
                _LOGGER.warning(
                    "Failed to read inspect.json from %s: %s", snapshot_zip, e
                )

    if dut is None:
        return []

    try:
        collection = dut.get_inspect_data(selectors=STORAGE_INSPECT_SELECTORS)
        entries = []
        for item in collection.data:
            payload = item.payload
            root = payload.get("root") if isinstance(payload, Mapping) else None
            if not isinstance(root, Mapping):
                continue
            name = getattr(item.metadata, "name", None)
            if name is None:
                for driver_name, node_name in _BLOCK_DRIVER_NODES.items():
                    if isinstance(root.get(node_name), Mapping):
                        name = driver_name
                        break
            entries.append(
                InspectEntry(
                    moniker=getattr(item, "moniker", None),
                    name=name,
                    timestamp_ns=int(item.metadata.timestamp.nanoseconds()),
                    root=root,
                )
            )
        return entries
    except Exception as e:
        _LOGGER.warning("Failed to query fallback storage inspect data: %s", e)
        return []


def _resolve_driver_and_node(
    entry: InspectEntry,
) -> tuple[str, str, Mapping[str, Any]] | None:
    """Resolves `(driver_name, node_name, node_dict)` from an `InspectEntry`.

    Entries loaded from `Snapshot_*.zip` have `entry.name` set to the Inspect
    tree name (e.g. `"ftl"`), whereas entries constructed from `fuchsia_inspect`
    without `metadata.name` may have `entry.name=None`. In both cases,
    `entry.root` contains the driver's unique child node (`"ftl"`,
    `"sdmmc_core"`, `"ufs"`, or `"virtio_block"`).
    """
    if entry.timestamp_ns <= 0:
        return None
    if entry.name in _BLOCK_DRIVER_NODES:
        node_name = _BLOCK_DRIVER_NODES[entry.name]
        node = entry.root.get(node_name)
        if isinstance(node, Mapping):
            return entry.name, node_name, node
    for driver_name, node_name in _BLOCK_DRIVER_NODES.items():
        node = entry.root.get(node_name)
        if isinstance(node, Mapping):
            return driver_name, node_name, node
    return None


def _has_block_op_counters(node: Mapping[str, Any]) -> bool:
    """Returns True if `node` has `.kib` and `.count` counters for all block ops."""
    return all(
        f"block.{op_key}.kib" in node and f"block.{op_key}.count" in node
        for op_key, _ in _BLOCK_OPERATIONS
    )


def _extract_hardware_block_rates(
    entry: InspectEntry,
) -> dict[str, tuple[float, float]] | None:
    """Attempts to extract `{op_label: (bytes_per_sec, ops_per_sec)}` for Read,
    Write, Flush, and Trim from a hardware block driver `InspectEntry` (`ftl`,
    `sdmmc_core`, `ufs`, or `virtio_block`).
    """
    resolved = _resolve_driver_and_node(entry)
    if resolved is None:
        return None
    _, _, node = resolved
    if not _has_block_op_counters(node):
        return None
    uptime_sec = float(entry.timestamp_ns) / 1e9
    return {
        op_label: (
            (int(node[f"block.{op_key}.kib"]) * 1024) / uptime_sec,
            int(node[f"block.{op_key}.count"]) / uptime_sec,
        )
        for op_key, op_label in _BLOCK_OPERATIONS
    }


def _normalize_entries(
    inspect_entries: Sequence[InspectEntry | Mapping[str, Any]],
) -> list[InspectEntry]:
    """Converts a sequence of `InspectEntry` or raw dict entries into a
    deterministically ordered list of `InspectEntry` objects.
    """
    normalized: list[InspectEntry] = []
    for entry in inspect_entries:
        parsed: InspectEntry | None = None
        if isinstance(entry, InspectEntry):
            parsed = entry
        elif isinstance(entry, Mapping):
            parsed = InspectEntry.from_dict(entry)
        if parsed is None:
            continue
        if parsed.name is None:
            resolved = _resolve_driver_and_node(parsed)
            if resolved is not None:
                parsed = InspectEntry(
                    moniker=parsed.moniker,
                    name=resolved[0],
                    timestamp_ns=parsed.timestamp_ns,
                    root=parsed.root,
                )
        normalized.append(parsed)
    return sorted(
        normalized,
        key=lambda e: (e.name or "", e.moniker or "", e.timestamp_ns),
    )


def diff_inspect_entries(
    start_entries: Sequence[InspectEntry | Mapping[str, Any]],
    end_entries: Sequence[InspectEntry | Mapping[str, Any]],
) -> list[InspectEntry]:
    """Computes per-interval delta `InspectEntry` items between `start_entries` and `end_entries`.

    For each hardware block driver entry in `end_entries`, subtracts the baseline
    cumulative bytes/ops from `start_entries` and sets `timestamp_ns` to
    `end.timestamp_ns - start.timestamp_ns`. Prefers active storage drivers
    (non-zero cumulative operations since boot) over idle secondary block
    interfaces (such as SDIO Wi-Fi controllers) so zero-I/O test windows still
    produce `0.0` rate metrics for the active storage device.
    """
    valid_end: list[tuple[InspectEntry, str, Mapping[str, Any], int]] = []
    for entry in _normalize_entries(end_entries):
        resolved = _resolve_driver_and_node(entry)
        if resolved is None:
            continue
        driver_name, node_name, node = resolved
        if not _has_block_op_counters(node):
            continue
        total_ops = sum(
            int(node[f"block.{op_key}.count"])
            for op_key, _ in _BLOCK_OPERATIONS
        )
        normalized_end = InspectEntry(
            moniker=entry.moniker,
            name=driver_name,
            timestamp_ns=entry.timestamp_ns,
            root=entry.root,
        )
        valid_end.append((normalized_end, node_name, node, total_ops))

    if any(total_ops > 0 for _, _, _, total_ops in valid_end):
        valid_end = [item for item in valid_end if item[3] > 0]

    start_by_key: dict[tuple[str | None, str | None], InspectEntry] = {}
    start_by_name: dict[str | None, InspectEntry] = {}
    for s_entry in _normalize_entries(start_entries):
        resolved = _resolve_driver_and_node(s_entry)
        if resolved is None:
            continue
        s_driver_name, _, s_node = resolved
        if not _has_block_op_counters(s_node):
            continue
        normalized_start = InspectEntry(
            moniker=s_entry.moniker,
            name=s_driver_name,
            timestamp_ns=s_entry.timestamp_ns,
            root=s_entry.root,
        )
        start_by_key[(s_driver_name, s_entry.moniker)] = normalized_start
        s_total_ops = sum(
            int(s_node[f"block.{op_key}.count"])
            for op_key, _ in _BLOCK_OPERATIONS
        )
        if s_driver_name not in start_by_name or s_total_ops > 0:
            start_by_name[s_driver_name] = normalized_start

    delta_entries: list[InspectEntry] = []
    for end_entry, node_name, end_node, _ in valid_end:
        start_entry = start_by_key.get(
            (end_entry.name, end_entry.moniker)
        ) or start_by_name.get(end_entry.name)
        if (
            start_entry is None
            or end_entry.timestamp_ns <= start_entry.timestamp_ns
        ):
            delta_entries.append(end_entry)
            continue
        start_node = start_entry.root.get(node_name)
        if not isinstance(start_node, Mapping):
            delta_entries.append(end_entry)
            continue
        if any(
            int(end_node[f"block.{op_key}.kib"])
            < int(start_node[f"block.{op_key}.kib"])
            or int(end_node[f"block.{op_key}.count"])
            < int(start_node[f"block.{op_key}.count"])
            for op_key, _ in _BLOCK_OPERATIONS
        ):
            delta_entries.append(end_entry)
            continue

        delta_node = dict(end_node)
        for op_key, _ in _BLOCK_OPERATIONS:
            delta_node[f"block.{op_key}.kib"] = int(
                end_node[f"block.{op_key}.kib"]
            ) - int(start_node[f"block.{op_key}.kib"])
            delta_node[f"block.{op_key}.count"] = int(
                end_node[f"block.{op_key}.count"]
            ) - int(start_node[f"block.{op_key}.count"])
        delta_entries.append(
            InspectEntry(
                moniker=end_entry.moniker,
                name=end_entry.name,
                timestamp_ns=end_entry.timestamp_ns - start_entry.timestamp_ns,
                root={**end_entry.root, node_name: delta_node},
            )
        )
    return delta_entries


def _extract_block_io_metrics(
    inspect_entries: Sequence[InspectEntry | Mapping[str, Any]],
    *,
    baseline_entries: Sequence[InspectEntry | Mapping[str, Any]] | None = None,
) -> dict[str, tuple[float, float]] | None:
    """Selects the primary hardware block driver and extracts
    `{op_label: (bytes_per_sec, ops_per_sec)}`.
    """
    if baseline_entries is not None:
        inspect_entries = diff_inspect_entries(
            baseline_entries, inspect_entries
        )

    rates: dict[str, tuple[float, float]] | None = None

    # Target devices only have one primary block storage driver, but entries in
    # `inspect.json` are not guaranteed to be ordered. Sort entries first, and
    # prefer an active driver (non-zero total I/O ops) over an idle one.
    for entry in _normalize_entries(inspect_entries):
        candidate = _extract_hardware_block_rates(entry)
        if candidate is None:
            continue
        if rates is None:
            rates = candidate
        if sum(ops for _, ops in candidate.values()) > 0:
            rates = candidate
            break

    return rates


class StorageIoMetricsProcessor(trace_metrics.MetricsProcessor):
    """Computes block/storage I/O rate metrics from Inspect data.

    Publishes headline block write and flush operation rates
    (`BlockWriteOpsPerSec` and `BlockFlushOpsPerSec`) as structured metrics for
    regression detection, and all block operation rates (`read`, `write`,
    `flush`, `trim` bytes/sec and ops/sec) as freeform metrics.
    """

    FREEFORM_METRICS_FILENAME = "storage_io"

    def __init__(
        self,
        dut: InspectDataProvider | None = None,
        *,
        label_prefix: str | None = "Storage/System",
        start_entries: Sequence[InspectEntry | Mapping[str, Any]] | None = None,
        end_entries: Sequence[InspectEntry | Mapping[str, Any]] | None = None,
    ) -> None:
        self._dut = dut
        self._label_prefix = label_prefix
        self._start_entries: Sequence[
            InspectEntry | Mapping[str, Any]
        ] | None = start_entries
        self._end_entries: Sequence[
            InspectEntry | Mapping[str, Any]
        ] | None = end_entries

    @property
    @typing.override
    def event_patterns(self) -> set[str]:
        return set()

    def capture_start(
        self, dut: InspectDataProvider | None = None
    ) -> list[InspectEntry]:
        """Captures baseline block driver Inspect entries before a test workload."""
        target_dut = dut if dut is not None else self._dut
        entries = load_inspect_entries(dut=target_dut)
        self._start_entries = entries
        self._end_entries = None
        return entries

    def capture_end(
        self, dut: InspectDataProvider | None = None
    ) -> list[InspectEntry]:
        """Captures final block driver Inspect entries after a test workload."""
        target_dut = dut if dut is not None else self._dut
        entries = load_inspect_entries(dut=target_dut)
        self._end_entries = entries
        return entries

    def __enter__(self) -> Self:
        self.capture_start()
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        self.capture_end()

    async def __aenter__(self) -> Self:
        self.capture_start()
        return self

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        self.capture_end()

    def _compute_rates(
        self,
    ) -> tuple[dict[str, tuple[float, float]], bool] | None:
        if self._end_entries is None and self._dut is not None:
            self.capture_end()
        entries: Sequence[InspectEntry | Mapping[str, Any]] | None
        if self._end_entries is not None:
            entries = self._end_entries
            baseline = self._start_entries
        else:
            entries = self._start_entries
            baseline = None
        if not entries:
            return None
        rates = _extract_block_io_metrics(entries, baseline_entries=baseline)
        if rates is None:
            return None
        return rates, baseline is not None

    @typing.override
    def process_metrics(
        self, model: trace_model.Model | None = None
    ) -> Sequence[metrics.TestCaseResult]:
        computed = self._compute_rates()
        if computed is None:
            return []
        rates, has_baseline = computed

        prefix = f"{self._label_prefix}/" if self._label_prefix else ""
        period_str = "interval" if has_baseline else "device uptime"
        results: list[metrics.TestCaseResult] = []
        for op_key, op_label in _DEFAULT_HEADLINE_OPERATIONS:
            _, ops_per_sec = rates[op_label]
            results.append(
                metrics.TestCaseResult(
                    label=f"{prefix}Block{op_label}OpsPerSec",
                    unit=metrics.Unit.count,
                    values=[ops_per_sec],
                    doc=f"Block {op_key} operations per second of {period_str}.",
                    direction=metrics.Direction.smallerIsBetter,
                )
            )
        return results

    @typing.override
    def process_freeform_metrics(
        self, model: trace_model.Model | None = None
    ) -> tuple[str, metrics.JSON]:
        computed = self._compute_rates()
        if computed is None:
            return self.FREEFORM_METRICS_FILENAME, {}
        rates, _ = computed

        breakdown: dict[str, metrics.JSON] = {
            op_key: {
                "bytes_per_sec": rates[op_label][0],
                "ops_per_sec": rates[op_label][1],
            }
            for op_key, op_label in _BLOCK_OPERATIONS
        }
        return self.FREEFORM_METRICS_FILENAME, breakdown


def extract_storage_io_metrics(
    inspect_entries: Sequence[InspectEntry | Mapping[str, Any]],
    label_prefix: str | None = None,
    *,
    baseline_entries: Sequence[InspectEntry | Mapping[str, Any]] | None = None,
) -> list[metrics.TestCaseResult]:
    """Extracts block/storage I/O rate metrics (`Block{Read,Write,Flush,Trim}{Bytes,Ops}PerSec`)
    from hardware block driver Inspect entries (`ftl`, `sdmmc_core`, `ufs`,
    `virtio_block`).

    When `baseline_entries` is provided, computes per-interval delta rates
    between `baseline_entries` and `inspect_entries` via `diff_inspect_entries`.
    Otherwise, computes cumulative rates over device uptime.

    Target devices have a single primary block storage driver. To guarantee
    deterministic behavior independent of the input order in `inspect_entries`,
    entries are sorted by `(name, moniker, timestamp_ns)` and any active block
    driver with non-zero I/O operations takes precedence over idle instances.
    """
    rates = _extract_block_io_metrics(
        inspect_entries, baseline_entries=baseline_entries
    )
    if rates is None:
        return []

    prefix = f"{label_prefix}/" if label_prefix else ""
    results: list[metrics.TestCaseResult] = []
    for op_key, op_label in _BLOCK_OPERATIONS:
        bytes_per_sec, ops_per_sec = rates[op_label]
        results.extend(
            [
                metrics.TestCaseResult(
                    label=f"{prefix}Block{op_label}BytesPerSec",
                    unit=metrics.Unit.bytesPerSecond,
                    values=[bytes_per_sec],
                    doc=f"Cumulative block {op_key} bytes per second of device uptime.",
                    direction=metrics.Direction.smallerIsBetter,
                ),
                metrics.TestCaseResult(
                    label=f"{prefix}Block{op_label}OpsPerSec",
                    unit=metrics.Unit.count,
                    values=[ops_per_sec],
                    doc=f"Cumulative block {op_key} operations per second of device uptime.",
                    direction=metrics.Direction.smallerIsBetter,
                ),
            ]
        )
    return results


def capture(
    dut: InspectDataProvider | None = None,
    snapshot_dirs: Sequence[str | os.PathLike[str]] = (),
    label_prefix: str | None = None,
    inspect_entries: Sequence[InspectEntry | Mapping[str, Any]] | None = None,
    baseline_entries: Sequence[InspectEntry | Mapping[str, Any]] | None = None,
) -> metrics.Report:
    """Captures block/storage I/O rate metrics as a `metrics.Report`.

    Can be combined with trace or memory metrics via `report.extend(storage.capture(dut=self.dut))`.
    """
    entries = (
        inspect_entries
        if inspect_entries is not None
        else load_inspect_entries(dut=dut, snapshot_dirs=snapshot_dirs)
    )
    structured = extract_storage_io_metrics(
        entries,
        label_prefix=label_prefix,
        baseline_entries=baseline_entries,
    )
    description = metrics.describe_callable(capture)
    return metrics.Report(
        structured,
        {},
        [
            metrics.MetricsProcessorDescription(
                doc=description["doc"],
                code_path=description["code_path"],
                line_no=description["line_no"],
                metrics=[tcr.describe() for tcr in structured],
            )
        ],
    )


def publish_teardown_storage_io_metrics(
    test_instance: Any,
    test_suite: str,
    expected_metric_names_filename: str | None = None,
    label_prefix: str | None = None,
    runtime_deps_dir: str | os.PathLike[str] | None = None,
    test_data_module: Any | None = None,
    baseline_entries: Sequence[InspectEntry | Mapping[str, Any]] | None = None,
) -> list[metrics.TestCaseResult]:
    """One-line helper for `FuchsiaBaseTest` subclasses' `teardown_class()` that
    loads the teardown snapshot, extracts storage I/O rates, writes
    `<teardown_class_artifacts>/<test_suite>-benchmark.fuchsiaperf.json`, and
    publishes it via `publish.publish_fuchsiaperf()`.
    """
    try:
        snapshot_dirs: list[str] = []
        for attr in ("_teardown_class_artifacts", "log_path"):
            val = getattr(test_instance, attr, None)
            if val:
                snapshot_dirs.append(str(val))

        inspect_entries = load_inspect_entries(
            dut=getattr(test_instance, "dut", None),
            snapshot_dirs=snapshot_dirs,
        )
        resolved_prefix = (
            test_instance.__class__.__name__
            if label_prefix is None
            else label_prefix
        )
        storage_metrics = extract_storage_io_metrics(
            inspect_entries,
            label_prefix=resolved_prefix,
            baseline_entries=baseline_entries,
        )
        if not storage_metrics:
            _LOGGER.warning(
                "No block/storage I/O inspect metrics found in teardown snapshot."
            )
            return []

        output_dir = getattr(
            test_instance, "_teardown_class_artifacts", None
        ) or getattr(test_instance, "log_path")
        os.makedirs(output_dir, exist_ok=True)
        fuchsiaperf_json_path = (
            pathlib.Path(output_dir)
            / f"{test_suite}-benchmark.fuchsiaperf.json"
        )
        metrics.TestCaseResult.write_fuchsiaperf_json(
            results=storage_metrics,
            test_suite=test_suite,
            output_path=fuchsiaperf_json_path,
        )

        publish_kwargs: dict[str, Any] = {
            "fuchsia_perf_file_paths": [fuchsiaperf_json_path],
            "expected_metric_names_filename": (
                expected_metric_names_filename or f"{test_suite}.txt"
            ),
        }
        if runtime_deps_dir is not None:
            publish_kwargs["runtime_deps_dir"] = runtime_deps_dir
        if test_data_module is not None:
            publish_kwargs["test_data_module"] = test_data_module

        publish.publish_fuchsiaperf(**publish_kwargs)
        return storage_metrics
    except Exception as e:
        _LOGGER.warning("Failed to publish storage I/O metrics: %s", e)
        return []