Merge changes I50e1fcd7,I32f4eadb

* changes:
  btp: augment query exception with metadata about trace
  btp: add support for incrementing stats on query/execute as well

Author: Treehugger Robot

python/perfetto/batch_trace_processor/api.py

@@ -38,44 +38,55 @@
 PLATFORM_DELEGATE = PlatformDelegate
 
 TraceListReference = registry.TraceListReference
+Metadata = Dict[str, str]
 
 
-# Enum encoding how errors while loading traces in BatchTraceProcessor should
-# be handled.
-class LoadFailureHandling(Enum):
-  # If any trace fails to load, raises an exception causing the entire batch
-  # trace processor to fail.
+# Enum encoding how errors while loading/querying traces in BatchTraceProcessor
+# should be handled.
+class FailureHandling(Enum):
+  # If any trace fails to load or be queried, raises an exception causing the
+  # entire batch trace processor to fail.
   # This is the default behaviour and the method which should be preferred for
   # any interactive use of BatchTraceProcessor.
   RAISE_EXCEPTION = 0
 
-  # If a trace fails to load, the trace processor for that trace is dropped but
-  # loading of other traces is unaffected. |load_failures| is incremented in the
-  # Stats class for the batch trace processor instance.
+  # If a trace fails to load or be queried, the trace processor for that trace
+  # is dropped but loading of other traces is unaffected. A failure integer is
+  # incremented in the Stats class for the batch trace processor instance.
   INCREMENT_STAT = 1
 
 
 @dc.dataclass
 class BatchTraceProcessorConfig:
   tp_config: TraceProcessorConfig
-  load_failure_handling: LoadFailureHandling
-
-  def __init__(self,
-               tp_config: TraceProcessorConfig = TraceProcessorConfig(),
-               load_failure_handling: LoadFailureHandling = LoadFailureHandling
-               .RAISE_EXCEPTION):
+  load_failure_handling: FailureHandling
+  query_failure_handling: FailureHandling
+
+  def __init__(
+      self,
+      tp_config: TraceProcessorConfig = TraceProcessorConfig(),
+      load_failure_handling: FailureHandling = FailureHandling.RAISE_EXCEPTION,
+      execute_failure_handling: FailureHandling = FailureHandling
+      .RAISE_EXCEPTION,
+  ):
     self.tp_config = tp_config
     self.load_failure_handling = load_failure_handling
+    self.execute_failure_handling = execute_failure_handling
 
 
 # Contains stats about the events which happened during the use of
 # BatchTraceProcessor.
 @dc.dataclass
 class Stats:
   # The number of traces which failed to load; only non-zero if
-  # LoadFailureHandling.INCREMENT_STAT is chosen as the handling type.
+  # FailureHanding.INCREMENT_STAT is chosen as the load failure handling type.
   load_failures: int = 0
 
+  # The number of traces which failed while executing (query, metric or
+  # arbitary function); only non-zero if FailureHanding.INCREMENT_STAT is
+  # chosen as the execute failure handling type.
+  execute_failures: int = 0
+
 
 class BatchTraceProcessor:
   """Run ad-hoc SQL queries across many Perfetto traces.
@@ -121,7 +132,7 @@ def __init__(self,
         trace processor and underlying trace processors.
     """
 
-    self.tps = None
+    self.tps_and_metadata = None
     self.closed = False
     self._stats = Stats()
 
@@ -148,8 +159,7 @@ def __init__(self,
         len(resolved)) or query_executor
 
     self.query_executor = query_executor
-    self.metadata = [t.metadata for t in resolved]
-    self.tps = [
+    self.tps_and_metadata = [
         x for x in load_exectuor.map(self._create_tp, resolved) if x is not None
     ]
 
@@ -266,7 +276,12 @@ def execute(self, fn: Callable[[TraceProcessor], Any]) -> List[Any]:
       A list of values with the result of executing the fucntion (one per
       trace).
     """
-    return list(self.query_executor.map(fn, self.tps))
+
+    def wrapped(pair: Tuple[TraceProcessor, Metadata]):
+      (tp, metadata) = pair
+      return self._execute_handling_failure(fn, tp, metadata)
+
+    return list(self.query_executor.map(wrapped, self.tps_and_metadata))
 
   def execute_and_flatten(self, fn: Callable[[TraceProcessor], pd.DataFrame]
                          ) -> pd.DataFrame:
@@ -285,15 +300,15 @@ def execute_and_flatten(self, fn: Callable[[TraceProcessor], pd.DataFrame]
       be added to the dataframe (see |query_and_flatten| for details).
     """
 
-    def wrapped(pair: Tuple[TraceProcessor, Dict[str, str]]):
+    def wrapped(pair: Tuple[TraceProcessor, Metadata]):
       (tp, metadata) = pair
-      df = fn(tp)
+      df = self._execute_handling_failure(fn, tp, metadata)
       for key, value in metadata.items():
         df[key] = value
       return df
 
     df = pd.concat(
-        list(self.query_executor.map(wrapped, zip(self.tps, self.metadata))))
+        list(self.query_executor.map(wrapped, self.tps_and_metadata)))
     return df.reset_index(drop=True)
 
   def close(self):
@@ -309,8 +324,8 @@ def close(self):
       return
     self.closed = True
 
-    if self.tps:
-      for tp in self.tps:
+    if self.tps_and_metadata:
+      for tp, _ in self.tps_and_metadata:
         tp.close()
 
   def stats(self):
@@ -319,17 +334,28 @@ def stats(self):
     See |Stats| class definition for the list of the statistics available."""
     return self._stats
 
-  def _create_tp(self,
-                 trace: ResolverRegistry.Result) -> Optional[TraceProcessor]:
+  def _create_tp(self, trace: ResolverRegistry.Result
+                ) -> Optional[Tuple[TraceProcessor, Metadata]]:
     try:
-      return TraceProcessor(trace=trace.generator, config=self.config.tp_config)
+      return TraceProcessor(
+          trace=trace.generator, config=self.config.tp_config), trace.metadata
     except TraceProcessorException as ex:
-      if self.config.load_failure_handling == \
-        LoadFailureHandling.RAISE_EXCEPTION:
+      if self.config.load_failure_handling == FailureHandling.RAISE_EXCEPTION:
         raise ex
       self._stats.load_failures += 1
       return None
 
+  def _execute_handling_failure(self, fn: Callable[[TraceProcessor], Any],
+                                tp: TraceProcessor, metadata: Metadata):
+    try:
+      return fn(tp)
+    except TraceProcessorException as ex:
+      if self.config.execute_failure_handling == \
+          FailureHandling.RAISE_EXCEPTION:
+        raise TraceProcessorException(f'{metadata} {ex}') from None
+      self._stats.execute_failures += 1
+      return pd.DataFrame()
+
   def __enter__(self):
     return self
 

python/test/api_integrationtest.py

@@ -20,8 +20,8 @@
 import pandas as pd
 
 from perfetto.batch_trace_processor.api import BatchTraceProcessor
-from perfetto.batch_trace_processor.api import LoadFailureHandling
 from perfetto.batch_trace_processor.api import BatchTraceProcessorConfig
+from perfetto.batch_trace_processor.api import FailureHandling
 from perfetto.batch_trace_processor.api import TraceListReference
 from perfetto.trace_processor.api import PLATFORM_DELEGATE
 from perfetto.trace_processor.api import TraceProcessor
@@ -90,17 +90,18 @@ def resolve(self):
 
 def create_batch_tp(
     traces: TraceListReference,
-    failure_handling: LoadFailureHandling = LoadFailureHandling.RAISE_EXCEPTION
+    load_failure_handling: FailureHandling = FailureHandling.RAISE_EXCEPTION,
+    execute_failure_handling: FailureHandling = FailureHandling.RAISE_EXCEPTION,
 ):
-  default = PLATFORM_DELEGATE().default_resolver_registry()
-  default.register(SimpleResolver)
-  default.register(RecursiveResolver)
-  return BatchTraceProcessor(
-      traces=traces,
-      config=BatchTraceProcessorConfig(
-          load_failure_handling=failure_handling,
-          tp_config=TraceProcessorConfig(
-              bin_path=os.environ["SHELL_PATH"], resolver_registry=default)))
+  registry = PLATFORM_DELEGATE().default_resolver_registry()
+  registry.register(SimpleResolver)
+  registry.register(RecursiveResolver)
+  config = BatchTraceProcessorConfig(
+      load_failure_handling=load_failure_handling,
+      execute_failure_handling=execute_failure_handling,
+      tp_config=TraceProcessorConfig(
+          bin_path=os.environ["SHELL_PATH"], resolver_registry=registry))
+  return BatchTraceProcessor(traces=traces, config=config)
 
 
 def create_tp(trace: TraceReference):
@@ -229,13 +230,32 @@ def test_recursive_resolver(self):
       df = btp.query_and_flatten('select dur from slice limit 1')
       pd.testing.assert_frame_equal(df, expected, check_dtype=False)
 
-  def test_btp_failure(self):
+  def test_btp_load_failure(self):
     f = io.BytesIO(b'<foo></foo>')
     with self.assertRaises(TraceProcessorException):
       _ = create_batch_tp(traces=f)
 
-  def test_btp_failure_increment_stat(self):
+  def test_btp_load_failure_increment_stat(self):
     f = io.BytesIO(b'<foo></foo>')
     btp = create_batch_tp(
-        traces=f, failure_handling=LoadFailureHandling.INCREMENT_STAT)
+        traces=f, load_failure_handling=FailureHandling.INCREMENT_STAT)
     self.assertEqual(btp.stats().load_failures, 1)
+
+  def test_btp_query_failure(self):
+    btp = create_batch_tp(traces=example_android_trace_path())
+    with self.assertRaises(TraceProcessorException):
+      _ = btp.query('select * from sl')
+
+  def test_btp_query_failure_increment_stat(self):
+    btp = create_batch_tp(
+        traces=example_android_trace_path(),
+        execute_failure_handling=FailureHandling.INCREMENT_STAT)
+    _ = btp.query('select * from sl')
+    self.assertEqual(btp.stats().execute_failures, 1)
+
+  def test_btp_query_failure_message(self):
+    btp = create_batch_tp(
+        traces='simple:path={}'.format(example_android_trace_path()))
+    with self.assertRaisesRegex(
+        TraceProcessorException, expected_regex='.*source.*generator.*'):
+      _ = btp.query('select * from sl')