Add v2 metrics support to trace_processor

Add support for perfetto v2 metric definitions and outputting trace
summaries to the trace_processor probe.

Bug: 413156992

Change-Id: I3fe5dc93de2e47a80cee638161117b553128a4bd
Reviewed-on: https://chromium-review.googlesource.com/c/crossbench/+/6567346
Reviewed-by: Camillo Bruni <[email protected]>
Commit-Queue: Kameron Lutes <[email protected]>

Author: kalutes

config/doc/probe/metric_def.textproto

@@ -0,0 +1,12 @@
+metric_spec: {
+  id: "file_metric_def_textproto"
+  value: "metric_column"
+  query {
+    id: "query_id"
+    table {
+      table_name: "metric_table"
+      module_name: "module_name"
+      column_names: "metric_column"
+    }
+  }
+}

config/doc/probe/trace_processor.config.hjson

@@ -2,69 +2,97 @@
 {
   probes: {
     trace_processor: {
+      // If necessary, a custom path for trace_processor_shell can be passed
+      // trace_processor_bin: "/optional/path/to/trace_processor_shell"
       queries: [
         // Queries can be specified as a reference to a file in
         // crossbench/probes/perfetto/trace_processor/queries:
-        "speedometer_cpu_time",
+        speedometer_cpu_time
         // Or specified inline:
         {
-          name: "my_query",
-          sql: "select dur from slice where slice.name = 'my_slice'",
-        },
-      ],
-      metrics: ["trace_stats"],
+          name: my_query
+          sql: select dur from slice where slice.name = 'my_slice'
+        }
+      ]
+      metrics: [
+        trace_stats
+      ]
       // Additional directories containing trace processor modules
       // can be specified:
       module_paths: [
-        "/my_project/modules/ext",
+        /my_project/modules/ext
       ]
-      // If necessary, a custom path for trace_processor_shell
-      // can be passed
-      // trace_processor_bin: "/optional/path/to/trace_processor_shell"
-    },
+      // Perfetto v2 metric definitions can be specified either inline
+      // or as a path:
+      metric_definitions: [
+        ./metric_def.textproto
+        '''
+          metric_spec: {
+            id: "inline_textproto_metric"
+            value: "metric_column"
+            query {
+              id: "query_id"
+              table {
+                table_name: "metric_table"
+                module_name: "module_name"
+                column_names: "metric_column"
+              }
+            }
+          }
+        '''
+      ]
+      // Metrics to include only in the trace summary can be specified.
+      // summary_metrics will include everything in the 'metrics' list as well
+      // as additional metrics (such as v2 metrics) specified here:
+      summary_metrics: [
+        file_metric_def_textproto
+        inline_textproto_metric
+      ]
+    }
     perfetto: {
-      textproto: '''
-        buffers {
-          size_kb: 300000
-          fill_policy: DISCARD
-        }
-        data_sources {
-          config {
-            name: "org.chromium.trace_metadata"
+      textproto:
+        '''
+          buffers {
+            size_kb: 300000
+            fill_policy: DISCARD
           }
-        }
-        data_sources {
-          config {
-            name: "track_event"
-            track_event_config {
-              disabled_categories: "*"
-              enabled_categories: "blink.user_timing"
-              enabled_categories: "toplevel"
-              enabled_categories: "__metadata"
-              timestamp_unit_multiplier: 1000
-              enable_thread_time_sampling: true
-              filter_debug_annotations: false
-              filter_dynamic_event_names: false
+          data_sources {
+            config {
+              name: "org.chromium.trace_metadata"
             }
           }
-        }
-        data_sources: {
-          config {
-              name: "linux.ftrace"
-            ftrace_config {
-              ftrace_events: "sched/sched_switch"
-              ftrace_events: "power/suspend_resume"
-              ftrace_events: "sched/sched_wakeup"
-              ftrace_events: "sched/sched_wakeup_new"
-              ftrace_events: "sched/sched_waking"
-              ftrace_events: "sched/sched_process_exit"
-              ftrace_events: "sched/sched_process_free"
-              ftrace_events: "task/task_newtask"
-              ftrace_events: "task/task_rename"
+          data_sources {
+            config {
+              name: "track_event"
+              track_event_config {
+                disabled_categories: "*"
+                enabled_categories: "blink.user_timing"
+                enabled_categories: "toplevel"
+                enabled_categories: "__metadata"
+                timestamp_unit_multiplier: 1000
+                enable_thread_time_sampling: true
+                filter_debug_annotations: false
+                filter_dynamic_event_names: false
+              }
             }
           }
-        }
-      '''
+          data_sources: {
+            config {
+                name: "linux.ftrace"
+              ftrace_config {
+                ftrace_events: "sched/sched_switch"
+                ftrace_events: "power/suspend_resume"
+                ftrace_events: "sched/sched_wakeup"
+                ftrace_events: "sched/sched_wakeup_new"
+                ftrace_events: "sched/sched_waking"
+                ftrace_events: "sched/sched_process_exit"
+                ftrace_events: "sched/sched_process_free"
+                ftrace_events: "task/task_newtask"
+                ftrace_events: "task/task_rename"
+              }
+            }
+          }
+        '''
     }
   }
 }

crossbench/probes/perfetto/trace_processor/trace_processor.py

@@ -26,7 +26,7 @@
 from crossbench.parse import ObjectParser, PathParser
 from crossbench.probes.metric import MetricsMerger
 from crossbench.probes.probe import Probe, ProbeConfigParser, ProbeContext
-from crossbench.probes.results import LocalProbeResult, ProbeResult
+from crossbench.probes.results import EmptyProbeResult, LocalProbeResult, ProbeResult
 from crossbench.replacements import Replacements
 
 if TYPE_CHECKING:
@@ -135,6 +135,21 @@ def config_parser(cls) -> ProbeConfigParser[Self]:
         is_list=True,
         default=tuple(),
         help="Name of metric to be run (can be any metric from Perfetto)")
+    parser.add_argument(
+        "metric_definitions",
+        type=ObjectParser.str_or_file_contents,
+        is_list=True,
+        default=tuple(),
+        help=("Textproto for perfetto metrics v2 definition files. "
+              "Can be inline textproto or a path to a .textproto file."))
+    parser.add_argument(
+        "summary_metrics",
+        type=str,
+        is_list=True,
+        default=tuple(),
+        help=("Additional metrics to only include in the trace summary. "
+              "Includes all of <metrics>. These can be v2 metrics if the "
+              "corresponding metric definition is supplied."))
     parser.add_argument(
         "queries",
         type=TraceProcessorQueryConfig,
@@ -156,13 +171,18 @@ def config_parser(cls) -> ProbeConfigParser[Self]:
 
   def __init__(self,
                batch: bool,
+               metric_definitions: Iterable[str],
+               summary_metrics: Iterable[str],
                metrics: Iterable[str],
                queries: Iterable[TraceProcessorQueryConfig],
                module_paths: Iterable[pth.LocalPath],
                trace_processor_bin: Optional[pth.LocalPath] = None) -> None:
     super().__init__()
     self._batch = batch
     self._metrics = tuple(metrics)
+    self._metric_definitions: tuple[str, ...] = tuple(metric_definitions)
+    self._summary_metrics: tuple[str,
+                                 ...] = tuple(metrics) + tuple(summary_metrics)
     ObjectParser.unique_sequence([query.name for query in queries],
                                  name="query names")
     self._queries = tuple(queries)
@@ -184,13 +204,22 @@ def metrics(self) -> tuple[str, ...]:
   def queries(self) -> tuple[TraceProcessorQueryConfig, ...]:
     return self._queries
 
+  @property
+  def metric_definitions(self) -> tuple[str, ...]:
+    return self._metric_definitions
+
+  @property
+  def summary_metrics(self) -> tuple[str, ...]:
+    return self._summary_metrics
+
   @property
   def module_paths(self) -> tuple[pth.LocalPath, ...]:
     return self._module_paths
 
   @property
   def has_work(self) -> bool:
-    return len(self._queries) != 0 or len(self._metrics) != 0
+    return len(self._queries) != 0 or len(self._metrics) != 0 or len(
+        self._summary_metrics) != 0 or len(self._metric_definitions) != 0
 
   @property
   def needs_tp_run(self) -> bool:
@@ -209,7 +238,7 @@ def tp_config(self) -> TraceProcessorConfig:
     extra_flags = []
 
     for module_path in self.module_paths:
-      extra_flags.append("--override-sql-module")
+      extra_flags.append("--add-sql-module")
       extra_flags.append(str(module_path))
 
     return TraceProcessorConfig(
@@ -240,6 +269,13 @@ def _check_sql(self) -> None:
       for query in self.queries:
         tp.query(query.sql)
 
+      metric_ids: Optional[list[str]] = None
+      if len(self.summary_metrics):
+        metric_ids = list(self.summary_metrics)
+
+      tp.trace_summary(
+          specs=list(self.metric_definitions), metric_ids=metric_ids)
+
   def _add_cb_columns(self, df: pd.DataFrame, run: Run) -> pd.DataFrame:
     df["cb_browser"] = run.browser.unique_name
     df["cb_story"] = run.story.name
@@ -373,12 +409,13 @@ def _merge_trace_files(self) -> LocalProbeResult:
 
   def _maybe_run_tp(self) -> ProbeResult:
     if not self.probe.needs_tp_run:
-      return LocalProbeResult()
+      return EmptyProbeResult()
 
     with TraceProcessor(
         trace=CrossbenchTraceUriResolver(self),
         config=self.probe.tp_config) as tp:
-      return self._run_queries(tp).merge(self._run_metrics(tp))
+      return self._run_queries(tp).merge(self._run_metrics(tp)).merge(
+          self._summarize_trace(tp))
 
   def _run_queries(self, tp: TraceProcessor) -> LocalProbeResult:
 
@@ -406,6 +443,32 @@ def run_metric(metric: str):
       files = tuple(map(run_metric, self.probe.metrics))
       return LocalProbeResult(json=files)
 
+  def _summarize_trace(self, tp: TraceProcessor) -> ProbeResult:
+    if not self.probe.summary_metrics and not self.probe.metric_definitions:
+      return EmptyProbeResult()
+
+    with self.run.actions(
+        "TRACE_PROCESSOR: Running trace summary", verbose=True):
+
+      # Trace processor interprets an empty list as 'emit no metrics' and
+      # 'None' as emit all metrics specified in the metric definitions.
+      # When no metric IDs are explicitly given, default to the more
+      # sensible option of emitting every metric.
+      metric_ids: Optional[list[str]] = None
+      if len(self.probe.summary_metrics):
+        metric_ids = list(self.probe.summary_metrics)
+
+      proto_result = tp.trace_summary(
+          specs=list(self.probe.metric_definitions), metric_ids=metric_ids)
+
+      proto_file = self.local_result_path / "v2_metrics.proto"
+
+      with proto_file.open("wb") as f:
+        f.write(proto_result.SerializeToString())
+
+      return LocalProbeResult(file=[proto_file])
+
+
   @property
   def merged_trace_path(self) -> pth.LocalPath:
     return self.local_result_path / "merged_trace.zip"

tests/crossbench/probes/test_trace_processor.py

@@ -55,6 +55,17 @@ def test_parse_example_config(self):
         r".*\/crossbench\/probes\/perfetto\/trace_processor\/modules\/ext")
     self.assertEqual(str(probe.module_paths[1]), "/my_project/modules/ext")
 
+    metric_definitions = probe.metric_definitions
+    self.assertEqual(len(metric_definitions), 2)
+    self.assertTrue("file_textproto_metric" in metric_definitions[0])
+    self.assertTrue("inline_textproto_metric" in metric_definitions[1])
+
+    summary_metrics = probe.summary_metrics
+    # Should contain everything in 'metrics' plus 'summary_metrics'
+    self.assertListEqual(
+        list(summary_metrics),
+        ["trace_stats", "file_textproto_metric", "inline_textproto_metric"])
+
   def test_query_config_duplicate_name_raises(self):
     with self.assertRaisesRegex(ArgumentTypeError,
                                 "Unexpected duplicates in query names"):