src/performance/lib/tp_shell/tp_utils.py - fuchsia - Git at Google

 # 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 hashlib
 import importlib.resources
 import logging
 import os
 import ssl
 import tempfile
 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 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,
         cache: bool = True,
     ) -> 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.
             cache: If True and trace_path is a URL, caches trace locally (default: True).
         """
         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)
         if not os.path.exists(self.tp_shell_path):
             raise FileNotFoundError(
                 f"Trace processor shell not found: {self.tp_shell_path}"
             )
         self.debug = debug

         delegate = FuchsiaPlatformDelegate(self.tp_shell_path)

         if trace_path.startswith("http://") or trace_path.startswith(
             "https://"
         ):
             if cache:
                 registry = delegate.default_resolver_registry()
                 self.trace_path = os.path.abspath(
                     self._resolve_and_cache_url(trace_path, registry)
                 )
             else:
                 self.trace_path = trace_path
         else:
             self.trace_path = os.path.abspath(trace_path)
             if not os.path.exists(self.trace_path):
                 raise FileNotFoundError(
                     f"Trace file not found: {self.trace_path}"
                 )

         # Override Perfetto's PlatformDelegate to point directly to our prebuilt shell binary
         perfetto.trace_processor.api.PLATFORM_DELEGATE = lambda: delegate

         # 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}"
         )
         try:
             self._tp = PerfettoTP(trace=self.trace_path, config=config)
         except Exception:
             if self._tp_shell_context is not None:
                 try:
                     self._tp_shell_context.__exit__(None, None, None)
                 except Exception:
                     pass
             raise

         # Register a finalizer to ensure the subprocess is cleaned up even if close() isn't called.
         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:
             _LOGGER.debug(
                 f"--- DEBUG SQL QUERY ---\n{query.strip()}\n-----------------------"
             )

         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 get_tables(self) -> set[str]:
         """Returns the set of table and view names available in the trace database."""
         rows = self.run_query(
             "SELECT name FROM sqlite_master WHERE type IN ('table', 'view')"
         )
         return {row["name"] for row in rows}

     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()

     @staticmethod
     def _resolve_and_cache_url(url: str, registry: ResolverRegistry) -> str:
         cache_key = hashlib.sha256(url.encode("utf-8")).hexdigest()
         cache_dir = os.path.join(tempfile.gettempdir(), "perf_analyze_cache")
         cached_file = os.path.join(cache_dir, f"{cache_key}.fxt")

         if os.path.exists(cached_file) and os.path.getsize(cached_file) > 0:
             _LOGGER.info("Using cached trace: %s", cached_file)
             return cached_file

         results = registry.resolve(url)
         if not results:
             raise ValueError(f"Could not resolve URL: {url}")

         os.makedirs(cache_dir, exist_ok=True)
         _LOGGER.info(
             "Downloading trace URL %s to cache %s...", url, cached_file
         )
         temp_file = f"{cached_file}.tmp.{os.getpid()}"
         try:
             with open(temp_file, "wb") as out:
                 for chunk in results[0].generator:
                     out.write(chunk)
             os.replace(temp_file, cached_file)
         except Exception:
             if os.path.exists(temp_file):
                 os.remove(temp_file)
             raise
         return cached_file
	# 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 hashlib
	import importlib.resources
	import logging
	import os
	import ssl
	import tempfile
	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 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,
	cache: bool = True,
	) -> 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.
	cache: If True and trace_path is a URL, caches trace locally (default: True).
	"""
	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)
	if not os.path.exists(self.tp_shell_path):
	raise FileNotFoundError(
	f"Trace processor shell not found: {self.tp_shell_path}"
	)
	self.debug = debug

	delegate = FuchsiaPlatformDelegate(self.tp_shell_path)

	if trace_path.startswith("http://") or trace_path.startswith(
	"https://"
	):
	if cache:
	registry = delegate.default_resolver_registry()
	self.trace_path = os.path.abspath(
	self._resolve_and_cache_url(trace_path, registry)
	)
	else:
	self.trace_path = trace_path
	else:
	self.trace_path = os.path.abspath(trace_path)
	if not os.path.exists(self.trace_path):
	raise FileNotFoundError(
	f"Trace file not found: {self.trace_path}"
	)

	# Override Perfetto's PlatformDelegate to point directly to our prebuilt shell binary
	perfetto.trace_processor.api.PLATFORM_DELEGATE = lambda: delegate

	# 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}"
	)
	try:
	self._tp = PerfettoTP(trace=self.trace_path, config=config)
	except Exception:
	if self._tp_shell_context is not None:
	try:
	self._tp_shell_context.__exit__(None, None, None)
	except Exception:
	pass
	raise

	# Register a finalizer to ensure the subprocess is cleaned up even if close() isn't called.
	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:
	_LOGGER.debug(
	f"--- DEBUG SQL QUERY ---\n{query.strip()}\n-----------------------"
	)

	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 get_tables(self) -> set[str]:
	"""Returns the set of table and view names available in the trace database."""
	rows = self.run_query(
	"SELECT name FROM sqlite_master WHERE type IN ('table', 'view')"
	)
	return {row["name"] for row in rows}

	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()

	@staticmethod
	def _resolve_and_cache_url(url: str, registry: ResolverRegistry) -> str:
	cache_key = hashlib.sha256(url.encode("utf-8")).hexdigest()
	cache_dir = os.path.join(tempfile.gettempdir(), "perf_analyze_cache")
	cached_file = os.path.join(cache_dir, f"{cache_key}.fxt")

	if os.path.exists(cached_file) and os.path.getsize(cached_file) > 0:
	_LOGGER.info("Using cached trace: %s", cached_file)
	return cached_file

	results = registry.resolve(url)
	if not results:
	raise ValueError(f"Could not resolve URL: {url}")

	os.makedirs(cache_dir, exist_ok=True)
	_LOGGER.info(
	"Downloading trace URL %s to cache %s...", url, cached_file
	)
	temp_file = f"{cached_file}.tmp.{os.getpid()}"
	try:
	with open(temp_file, "wb") as out:
	for chunk in results[0].generator:
	out.write(chunk)
	os.replace(temp_file, cached_file)
	except Exception:
	if os.path.exists(temp_file):
	os.remove(temp_file)
	raise
	return cached_file