src/performance/lib/tp_shell/tp_utils.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.

"""Trace Processor wrapper utilities for Perfetto trace analysis."""

import importlib.resources
import logging
import os
import ssl
import urllib.request
import weakref
from types import TracebackType
from typing import Any

import perfetto.trace_processor.api
from perfetto.tools.download_trace import resolve_trace_url
from perfetto.trace_processor.api import TraceProcessor as PerfettoTP
from perfetto.trace_processor.api import TraceProcessorConfig
from perfetto.trace_processor.platform import PlatformDelegate
from perfetto.trace_uri_resolver import util as resolver_util
from perfetto.trace_uri_resolver.path import PathUriResolver
from perfetto.trace_uri_resolver.registry import ResolverRegistry
from perfetto.trace_uri_resolver.resolver import TraceUriResolver

_LOGGER = logging.getLogger(__name__)


class HttpUriResolver(TraceUriResolver):
    """URI Resolver that streams trace data from HTTP/HTTPS endpoints and permalinks."""

    PREFIX = "http"

    def __init__(self, uri: str) -> None:
        self.uri = uri

    @classmethod
    def from_trace_uri(cls, uri: str) -> "HttpUriResolver":
        return cls(uri)

    def resolve(self) -> list[TraceUriResolver.Result]:
        if "perfetto.dev" in self.uri:
            url = resolve_trace_url(self.uri)
        else:
            url = self.uri
        context = ssl._create_unverified_context()
        # Add a default timeout of 2 minutes (120 seconds) as a backstop to
        # avoid permanently blocking when loading the trace file.
        response = urllib.request.urlopen(url, context=context, timeout=120)
        return [
            TraceUriResolver.Result(
                trace=resolver_util.read_generator(response),
                metadata={"_url": url},
            )
        ]


class HttpsUriResolver(HttpUriResolver):
    PREFIX = "https"


class FuchsiaPlatformDelegate(PlatformDelegate):
    """PlatformDelegate that points directly to the host's prebuilt trace_processor_shell binary.

    This delegate is used to override the default PlatformDelegate in the
    third_party Perfetto Python SDK, ensuring that it uses our prebuilt, in-tree
    version of trace_processor_shell instead of trying to download or resolve it
    from the network or default paths.
    """

    host_tp_shell_path: str

    def __init__(self, host_tp_shell_path: str) -> None:
        super().__init__()
        self.host_tp_shell_path = host_tp_shell_path

    def get_shell_path(
        self, bin_path: str | None = None, fetch_latest: bool = False
    ) -> str:
        return self.host_tp_shell_path

    def get_resource(self, file: str) -> bytes:
        return (
            importlib.resources.files("perfetto.trace_processor")
            .joinpath(file)
            .read_bytes()
        )

    def default_resolver_registry(self) -> ResolverRegistry:
        return ResolverRegistry(
            resolvers=[PathUriResolver, HttpUriResolver, HttpsUriResolver]
        )


class PerfettoTraceProcessor:
    """A wrapper around Perfetto's official Python TraceProcessor API.

    This class provides a way to interact with Perfetto's TraceProcessor API,
    handling the setup and teardown of the trace_processor_shell backend.
    The trace file is parsed only once, and backend process is torn down
    when the processor is closed.

    Attributes:
        trace_path: Path to the trace file to ingest.
        tp_shell_path: Path to the trace_processor_shell binary.
        debug: If True, prints SQL queries.
    """

    trace_path: str
    tp_shell_path: str
    debug: bool
    _tp: PerfettoTP
    _finalizer: weakref.finalize
    _tp_shell_context: Any

    def __init__(
        self,
        trace_path: str,
        tp_shell_path: str | None = None,
        debug: bool = False,
    ) -> None:
        """Initializes PerfettoTraceProcessor.

        Args:
            trace_path: Path to the trace file to ingest.
            tp_shell_path: Optional path to the trace_processor_shell binary.
            debug: If True, prints SQL queries.
        """
        if trace_path.startswith("http://") or trace_path.startswith(
            "https://"
        ):
            self.trace_path = trace_path
        else:
            self.trace_path = os.path.abspath(trace_path)
        self._tp_shell_context = None

        if tp_shell_path is None:
            try:
                # The logic here involving _tp_shell_context and its __enter__/__exit__
                # methods is necessary to correctly manage the lifecycle of the extracted
                # trace_processor_shell binary when loaded from package resources.
                # importlib.resources.as_file returns a context manager that extracts
                # the resource to a temporary location. We must call __enter__() to get
                # the path and ensure __exit__() is called for cleanup.
                resource = importlib.resources.files("tp_shell.bin").joinpath(
                    "trace_processor_shell"
                )
                self._tp_shell_context = importlib.resources.as_file(resource)
                # Enter the extraction context to obtain a real filesystem path
                extracted_path = self._tp_shell_context.__enter__()
                tp_shell_path = str(extracted_path)
                # Ensure the extracted file is executable (necessary in ZIP contexts)
                os.chmod(tp_shell_path, 0o755)
            except Exception as e:
                if self._tp_shell_context is not None:
                    try:
                        self._tp_shell_context.__exit__(None, None, None)
                    except Exception:
                        pass
                raise FileNotFoundError(
                    "trace_processor_shell was not found in the packaged resources. "
                    "The binary must either be explicitly specified or packaged as a data source dependency."
                ) from e

        self.tp_shell_path = os.path.abspath(tp_shell_path)
        self.debug = debug

        if not (
            self.trace_path.startswith("http://")
            or self.trace_path.startswith("https://")
        ) and not os.path.exists(self.trace_path):
            raise FileNotFoundError(f"Trace file not found: {self.trace_path}")
        if not os.path.exists(self.tp_shell_path):
            raise FileNotFoundError(
                f"Trace processor shell not found: {self.tp_shell_path}"
            )

        # Override Perfetto's PlatformDelegate to point directly to our prebuilt shell binary
        perfetto.trace_processor.api.PLATFORM_DELEGATE = (
            lambda: FuchsiaPlatformDelegate(self.tp_shell_path)
        )

        # Configure TraceProcessor to use a unique port to avoid conflicts
        config = TraceProcessorConfig(unique_port=True)

        _LOGGER.info(
            f"Initializing Perfetto TraceProcessor for trace: {self.trace_path}"
        )
        self._tp = PerfettoTP(trace=self.trace_path, config=config)
        # Register a finalizer to ensure the subprocess is cleaned up even if close() isn't called.

        # This cleanup is handled either in the except block on failure or by the weakref
        # finalizer via _cleanup.
        self._finalizer = weakref.finalize(
            self, self._cleanup, self._tp, self._tp_shell_context
        )

    @staticmethod
    def _cleanup(tp: PerfettoTP, context: Any = None) -> None:
        """Safely tears down the Perfetto TraceProcessor shell process and deletes temporary extraction files."""
        _LOGGER.info("Tearing down Perfetto TraceProcessor shell process...")
        try:
            tp.close()
        except Exception as e:
            _LOGGER.error(f"Error closing trace processor: {e}")

        if context is not None:
            _LOGGER.info(
                "Cleaning up temporary trace_processor_shell extraction..."
            )
            try:
                context.__exit__(None, None, None)
            except Exception as e:
                _LOGGER.error(f"Error cleaning up extraction context: {e}")

    def run_query(self, query: str) -> list[dict[str, Any]]:
        """Runs a SQL query against the trace and returns the result as a list of dicts."""
        if not hasattr(self, "_finalizer") or not self._finalizer.alive:
            raise RuntimeError("Trace processor is closed.")

        if self.debug:
            print("--- DEBUG SQL QUERY ---")
            print(query.strip())
            print("-----------------------")

        try:
            result_iterator = self._tp.query(query)
            # Row.__dict__ on instance contains exactly the dynamic attributes (columns)
            return [row.__dict__ for row in result_iterator]
        except Exception as e:
            _LOGGER.error(f"Error running query: {e}")
            raise

    def __enter__(self) -> "PerfettoTraceProcessor":
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        self.close()

    def close(self) -> None:
        """Tears down the trace processor shell process."""
        if hasattr(self, "_finalizer") and self._finalizer.alive:
            self._finalizer()