xref: /aosp_15_r20/external/perfetto/python/perfetto/trace_processor/api.py (revision 6dbdd20afdafa5e3ca9b8809fa73465d530080dc)
1# Copyright (C) 2020 The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#      http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import dataclasses as dc
16from urllib.parse import urlparse
17from typing import List, Optional
18
19from perfetto.common.exceptions import PerfettoException
20from perfetto.common.query_result_iterator import QueryResultIterator
21from perfetto.trace_processor.http import TraceProcessorHttp
22from perfetto.trace_processor.platform import PlatformDelegate
23from perfetto.trace_processor.protos import ProtoFactory
24from perfetto.trace_processor.shell import load_shell
25from perfetto.trace_uri_resolver import registry
26from perfetto.trace_uri_resolver.registry import ResolverRegistry
27
28# Defining this field as a module variable means this can be changed by
29# implementations at startup and used for all TraceProcessor objects
30# without having to specify on each one.
31# In Google3, this field is rewritten using Copybara to a implementation
32# which can integrates with internal infra.
33PLATFORM_DELEGATE = PlatformDelegate
34
35TraceReference = registry.TraceReference
36
37# Custom exception raised if any trace_processor functions return a
38# response with an error defined
39TraceProcessorException = PerfettoException
40
41@dc.dataclass
42class TraceProcessorConfig:
43  bin_path: Optional[str]
44  unique_port: bool
45  verbose: bool
46  ingest_ftrace_in_raw: bool
47  enable_dev_features: bool
48  resolver_registry: Optional[ResolverRegistry]
49  load_timeout: int
50  extra_flags: Optional[List[str]]
51
52  def __init__(self,
53               bin_path: Optional[str] = None,
54               unique_port: bool = True,
55               verbose: bool = False,
56               ingest_ftrace_in_raw: bool = False,
57               enable_dev_features=False,
58               resolver_registry: Optional[ResolverRegistry] = None,
59               load_timeout: int = 2,
60               extra_flags: Optional[List[str]] = None):
61    self.bin_path = bin_path
62    self.unique_port = unique_port
63    self.verbose = verbose
64    self.ingest_ftrace_in_raw = ingest_ftrace_in_raw
65    self.enable_dev_features = enable_dev_features
66    self.resolver_registry = resolver_registry
67    self.load_timeout = load_timeout
68    self.extra_flags = extra_flags
69
70
71class TraceProcessor:
72  QueryResultIterator = QueryResultIterator
73  Row = QueryResultIterator.Row
74
75  def __init__(self,
76               trace: Optional[TraceReference] = None,
77               addr: Optional[str] = None,
78               config: TraceProcessorConfig = TraceProcessorConfig(),
79               file_path: Optional[str] = None):
80    """Create a trace processor instance.
81
82    Args:
83      trace: reference to a trace to be loaded into the trace
84        processor instance.
85
86        One of several types is supported:
87        1) path to a trace file to open and read
88        2) a file like object (file, io.BytesIO or similar) to read
89        3) a generator yielding bytes
90        4) a trace URI which resolves to one of the above types
91        5) a trace URI resolver; this is a subclass of
92        resolver.TraceUriResolver which generates a reference to a
93        trace when the |resolve| method is called on it.
94
95        An URI is similar to a connection string (e.g. for a web
96        address or SQL database) which specifies where to lookup traces
97        and which traces to pick from this data source. The format of a
98        string should be as follows:
99        resolver_name:key_1=list,of,values;key_2=value
100
101        Custom resolvers can be provided to handle URIs via
102        |config.resolver_registry|.
103      addr: address of a running trace processor instance. Useful to query an
104        already loaded trace.
105      config: configuration options which customize functionality of trace
106        processor and the Python binding.
107      file_path (deprecated): path to a trace file to load. Use
108        |trace| instead of this field: specifying both will cause
109        an exception to be thrown.
110    """
111
112    if trace and file_path:
113      raise TraceProcessorException(
114          "trace and file_path cannot both be specified.")
115
116    self.config = config
117    self.platform_delegate = PLATFORM_DELEGATE()
118    self.protos = ProtoFactory(self.platform_delegate)
119    self.resolver_registry = config.resolver_registry or \
120      self.platform_delegate.default_resolver_registry()
121    self.http = self._create_tp_http(addr)
122
123    if trace or file_path:
124      try:
125        self._parse_trace(trace if trace else file_path)
126      except TraceProcessorException as ex:
127        self.close()
128        raise ex
129
130  def query(self, sql: str):
131    """Executes passed in SQL query using class defined HTTP API, and returns
132    the response as a QueryResultIterator. Raises TraceProcessorException if
133    the response returns with an error.
134
135    Args:
136      sql: SQL query written as a String
137
138    Returns:
139      A class which can iterate through each row of the results table. This
140      can also be converted to a pandas dataframe by calling the
141      as_pandas_dataframe() function after calling query.
142    """
143    response = self.http.execute_query(sql)
144    if response.error:
145      raise TraceProcessorException(response.error)
146
147    return TraceProcessor.QueryResultIterator(response.column_names,
148                                              response.batch)
149
150  def metric(self, metrics: List[str]):
151    """Returns the metrics data corresponding to the passed in trace metric.
152    Raises TraceProcessorException if the response returns with an error.
153
154    Args:
155      metrics: A list of valid metrics as defined in TraceMetrics
156
157    Returns:
158      The metrics data as a proto message
159    """
160    response = self.http.compute_metric(metrics)
161    if response.error:
162      raise TraceProcessorException(response.error)
163
164    metrics = self.protos.TraceMetrics()
165    metrics.ParseFromString(response.metrics)
166    return metrics
167
168  def enable_metatrace(self):
169    """Enable metatrace for the currently running trace_processor.
170    """
171    return self.http.enable_metatrace()
172
173  def disable_and_read_metatrace(self):
174    """Disable and return the metatrace formed from the currently running
175    trace_processor. This must be enabled before attempting to disable. This
176    returns the serialized bytes of the metatrace data directly. Raises
177    TraceProcessorException if the response returns with an error.
178    """
179    response = self.http.disable_and_read_metatrace()
180    if response.error:
181      raise TraceProcessorException(response.error)
182
183    return response.metatrace
184
185  def _create_tp_http(self, addr: str) -> TraceProcessorHttp:
186    if addr:
187      p = urlparse(addr)
188      parsed = p.netloc if p.netloc else p.path
189      return TraceProcessorHttp(parsed, protos=self.protos)
190
191    url, self.subprocess = load_shell(self.config.bin_path,
192                                      self.config.unique_port,
193                                      self.config.verbose,
194                                      self.config.ingest_ftrace_in_raw,
195                                      self.config.enable_dev_features,
196                                      self.platform_delegate,
197                                      self.config.load_timeout,
198                                      self.config.extra_flags)
199    return TraceProcessorHttp(url, protos=self.protos)
200
201  def _parse_trace(self, trace: TraceReference):
202    resolved_lst = self.resolver_registry.resolve(trace)
203    if not resolved_lst:
204      raise TraceProcessorException(
205          'trace argument did not resolve to a trace.')
206
207    if len(resolved_lst) > 1:
208      raise TraceProcessorException(
209          'trace argument resolved to more than one trace. Trace processor '
210          'only supports loading a single trace; please use '
211          'BatchTraceProcessor to operate on multiple traces.')
212
213    resolved = resolved_lst[0]
214    for chunk in resolved.generator:
215      result = self.http.parse(chunk)
216      if result.error:
217        raise TraceProcessorException(
218            f'Failed while parsing trace. Error message: {result.error}')
219    self.http.notify_eof()
220
221  def __enter__(self):
222    return self
223
224  def __exit__(self, a, b, c):
225    del a, b, c  # Unused.
226    self.close()
227    return False
228
229  def close(self):
230    if hasattr(self, 'subprocess'):
231      self.subprocess.kill()
232      self.subprocess.wait()
233
234    if hasattr(self, 'http'):
235      self.http.conn.close()
236
237  def __del__(self):
238    self.close()
239