trace_processor: add support for proto and sql metric from shell

This CL introduces support for metrics and metric protos to be added to
the trace processor from the shell. Only single files of each are
supported - this restriction may be relaxed in the future.

Bug: 129747127
Change-Id: Ia80c378bc0ed7289e85f8886c9efbe08bd749fbe
diff --git a/BUILD b/BUILD
index c5b611d..c9f3970 100644
--- a/BUILD
+++ b/BUILD
@@ -579,6 +579,8 @@
         "//third_party/perfetto/protos:trace_sys_stats_zero_cc_proto",
         "//third_party/perfetto/protos:trace_track_event_zero_cc_proto",
         "//third_party/perfetto/protos:trace_zero_cc_proto",
+        "//third_party/protobuf",
+        "//third_party/protobuf:libprotoc",
         "//third_party/sqlite",
         "//third_party/sqlite:sqlite_ext_percentile",
     ],
diff --git a/BUILD.gn b/BUILD.gn
index 4a2c960..846d132 100644
--- a/BUILD.gn
+++ b/BUILD.gn
@@ -108,6 +108,7 @@
     deps = [
       "src/trace_processor:trace_processor_shell",
     ]
+    testonly = true  # We need this for proto full.
   }
 }
 
diff --git a/src/trace_processor/BUILD.gn b/src/trace_processor/BUILD.gn
index 5fe14a2..806f16c 100644
--- a/src/trace_processor/BUILD.gn
+++ b/src/trace_processor/BUILD.gn
@@ -199,9 +199,11 @@
 
 if (current_toolchain == host_toolchain) {
   executable("trace_processor_shell_host") {
+    testonly = true  # We need this for proto full.
     deps = [
       ":lib",
       "../../gn:default_deps",
+      "../../gn:protobuf_full_deps",
       "../base",
     ]
     if (perfetto_build_standalone) {
@@ -224,6 +226,7 @@
 }
 
 copy("trace_processor_shell") {
+  testonly = true  # We need this for proto full.
   host_out_dir_ = get_label_info(":trace_processor_shell_host($host_toolchain)",
                                  "root_out_dir")
   deps = [
diff --git a/src/trace_processor/metrics/descriptors.cc b/src/trace_processor/metrics/descriptors.cc
index 35104e1..6b1d393 100644
--- a/src/trace_processor/metrics/descriptors.cc
+++ b/src/trace_processor/metrics/descriptors.cc
@@ -145,6 +145,9 @@
   using FieldDescriptorProto = protos::pbzero::FieldDescriptorProto;
   for (auto& descriptor : descriptors_) {
     for (auto& field : *descriptor.mutable_fields()) {
+      if (!field.resolved_type_name().empty())
+        continue;
+
       if (field.type() == FieldDescriptorProto::TYPE_MESSAGE ||
           field.type() == FieldDescriptorProto::TYPE_ENUM) {
         auto opt_desc =
diff --git a/src/trace_processor/trace_processor_impl.cc b/src/trace_processor/trace_processor_impl.cc
index 844f464..2b9022e 100644
--- a/src/trace_processor/trace_processor_impl.cc
+++ b/src/trace_processor/trace_processor_impl.cc
@@ -22,6 +22,7 @@
 #include <functional>
 
 #include "perfetto/base/logging.h"
+#include "perfetto/base/string_splitter.h"
 #include "perfetto/base/string_utils.h"
 #include "perfetto/base/time.h"
 #include "perfetto/protozero/scattered_heap_buffer.h"
@@ -445,11 +446,19 @@
   }
   auto no_ext_name = basename.substr(0, sql_idx);
 
+  std::string stripped_sql;
+  for (base::StringSplitter sp(sql, '\n'); sp.Next();) {
+    if (strncmp(sp.cur_token(), "--", 2) != 0) {
+      stripped_sql.append(sp.cur_token());
+      stripped_sql.push_back('\n');
+    }
+  }
+
   metrics::SqlMetricFile metric;
   metric.path = path;
   metric.proto_field_name = no_ext_name;
   metric.output_table_name = no_ext_name + "_output";
-  metric.sql = sql;
+  metric.sql = stripped_sql;
   sql_metrics_.emplace_back(metric);
   return util::OkStatus();
 }
diff --git a/src/trace_processor/trace_processor_shell.cc b/src/trace_processor/trace_processor_shell.cc
index 8edba60..14fa56d 100644
--- a/src/trace_processor/trace_processor_shell.cc
+++ b/src/trace_processor/trace_processor_shell.cc
@@ -26,6 +26,9 @@
 #include <iostream>
 #include <vector>
 
+#include <google/protobuf/compiler/parser.h>
+#include <google/protobuf/io/zero_copy_stream_impl.h>
+
 #include "perfetto/base/build_config.h"
 #include "perfetto/base/file_utils.h"
 #include "perfetto/base/logging.h"
@@ -254,6 +257,50 @@
   return 0;
 }
 
+class ErrorPrinter : public google::protobuf::io::ErrorCollector {
+  void AddError(int line, int col, const std::string& msg) override {
+    PERFETTO_ELOG("%d:%d: %s", line, col, msg.c_str());
+  }
+
+  void AddWarning(int line, int col, const std::string& msg) override {
+    PERFETTO_ILOG("%d:%d: %s", line, col, msg.c_str());
+  }
+};
+
+util::Status RegisterMetric(const std::string& register_metric) {
+  std::string sql;
+  base::ReadFile(register_metric, &sql);
+
+  std::string path = "shell/";
+  auto slash_idx = register_metric.rfind('/');
+  path += slash_idx == std::string::npos
+              ? register_metric
+              : register_metric.substr(slash_idx + 1);
+
+  return g_tp->RegisterMetric(path, sql);
+}
+
+util::Status ExtendMetricsProto(const std::string& extend_metrics_proto) {
+  google::protobuf::FileDescriptorSet desc_set;
+
+  base::ScopedFile file(base::OpenFile(extend_metrics_proto, O_RDONLY));
+
+  google::protobuf::io::FileInputStream stream(file.get());
+  ErrorPrinter printer;
+  google::protobuf::io::Tokenizer tokenizer(&stream, &printer);
+
+  auto* proto = desc_set.add_file();
+  google::protobuf::compiler::Parser parser;
+  parser.Parse(&tokenizer, proto);
+
+  std::vector<uint8_t> metric_proto;
+  metric_proto.resize(static_cast<size_t>(desc_set.ByteSize()));
+  desc_set.SerializeToArray(metric_proto.data(),
+                            static_cast<int>(metric_proto.size()));
+
+  return g_tp->ExtendMetricsProto(metric_proto.data(), metric_proto.size());
+}
+
 int RunMetrics(const std::vector<std::string>& metric_names) {
   std::vector<uint8_t> metric_result;
   util::Status status = g_tp->ComputeMetric(metric_names, &metric_result);
@@ -511,25 +558,34 @@
       "Interactive trace processor shell.\n"
       "Usage: %s [OPTIONS] trace_file.pb\n\n"
       "Options:\n"
-      " -h|--help            Prints this usage.\n"
-      " -v|--version         Prints the version of trace processor.\n"
-      " -d|--debug           Enable virtual table debugging.\n"
-      " -p|--perf-file FILE  Writes the time taken to ingest the trace and"
-      "execute the queries to the given file. Only valid with -q or "
+      " -h|--help                          Prints this usage.\n"
+      " -v|--version                       Prints the version of trace "
+      "processor.\n"
+      " -d|--debug                         Enable virtual table debugging.\n"
+      " -p|--perf-file FILE                Writes the time taken to ingest the "
+      "trace and execute the queries to the given file. Only valid with -q or "
       "--run-metrics and the file will only be written if the execution is "
       "successful\n"
-      " -q|--query-file FILE Read and execute an SQL query from a file.\n"
-      " -i|--interactive     Starts interactive mode even after a query file "
-      "is specified with -q.\n"
-      " -e|--export FILE     Export the trace into a SQLite database.\n"
-      " --run-metrics x,y,z  Runs a comma separated list of metrics and "
-      "prints the result as a TraceMetrics proto to stdout.\n",
+      " -q|--query-file FILE               Read and execute an SQL query from "
+      "a file.\n"
+      " -i|--interactive                   Starts interactive mode even after "
+      "a query file is specified with -q.\n"
+      " -e|--export FILE                   Export the trace into a SQLite "
+      "database.\n"
+      " --run-metrics x,y,z                Runs a comma separated list of "
+      "metrics and prints the result as a TraceMetrics proto to stdout.\n"
+      " --register-metric FILE             Registers the given SQL file with "
+      "trace processor to allow this file to be run as a metric.\n"
+      " --extend-metrics-proto FILE        Extends the TraceMetrics proto "
+      "using the extensions in the given proto file.\n",
       argv[0]);
 }
 
 int TraceProcessorMain(int argc, char** argv) {
   enum LongOption {
     OPT_RUN_METRICS = 1000,
+    OPT_REGISTER_METRIC,
+    OPT_EXTEND_METRICS_PROTO,
   };
 
   static const struct option long_options[] = {
@@ -541,12 +597,17 @@
       {"query-file", required_argument, nullptr, 'q'},
       {"export", required_argument, nullptr, 'e'},
       {"run-metrics", required_argument, nullptr, OPT_RUN_METRICS},
+      {"register-metric", required_argument, nullptr, OPT_REGISTER_METRIC},
+      {"extend-metrics-proto", required_argument, nullptr,
+       OPT_EXTEND_METRICS_PROTO},
       {nullptr, 0, nullptr, 0}};
 
   std::string perf_file_path;
   std::string query_file_path;
   std::string sqlite_file_path;
   std::string metric_names;
+  std::string register_metric;
+  std::string extend_metrics_proto;
   bool explicit_interactive = false;
   int option_index = 0;
   for (;;) {
@@ -591,6 +652,16 @@
       continue;
     }
 
+    if (option == OPT_REGISTER_METRIC) {
+      register_metric = optarg;
+      continue;
+    }
+
+    if (option == OPT_EXTEND_METRICS_PROTO) {
+      extend_metrics_proto = optarg;
+      continue;
+    }
+
     PrintUsage(argv);
     return option == 'h' ? 0 : 1;
   }
@@ -699,8 +770,22 @@
 
   auto t_run_start = base::GetWallTimeNs();
 
-  // First, see if we have some metrics to run. If we do, just run them and
-  // return.
+  if (!extend_metrics_proto.empty()) {
+    util::Status status = ExtendMetricsProto(extend_metrics_proto);
+    if (!status.ok()) {
+      PERFETTO_ELOG("Error when extending proto: %s", status.c_message());
+      return 1;
+    }
+  }
+
+  if (!register_metric.empty()) {
+    util::Status status = RegisterMetric(register_metric);
+    if (!status.ok()) {
+      PERFETTO_ELOG("Error when registering metric: %s", status.c_message());
+      return 1;
+    }
+  }
+
   if (!metric_names.empty()) {
     std::vector<std::string> metrics;
     for (base::StringSplitter ss(metric_names, ','); ss.Next();) {