ui: tp: add support for pushing SQL modules through RPC interface

This allows for out-of-tree modules to be loaded in the UI.

Change-Id: I4035227ac2a3f54bdf7535c4d0e767be6b50d333

protos/perfetto/trace_processor/trace_processor.proto

@@ -59,7 +59,8 @@ enum TraceProcessorApiVersion {
   // 11. Removal of experimental module from stdlib.
   // 12. Changed UI to be more aggresive about version matching.
   //     Added version_code.
-  TRACE_PROCESSOR_CURRENT_API_VERSION = 12;
+  // 13. Added TPM_REGISTER_SQL_MODULE method.
+  TRACE_PROCESSOR_CURRENT_API_VERSION = 13;
 }
 
 // At lowest level, the wire-format of the RPC protocol is a linear sequence of
@@ -101,6 +102,7 @@ message TraceProcessorRpc {
     TPM_DISABLE_AND_READ_METATRACE = 9;
     TPM_GET_STATUS = 10;
     TPM_RESET_TRACE_PROCESSOR = 11;
+    TPM_REGISTER_SQL_MODULE = 12;
   }
 
   oneof type {
@@ -132,6 +134,8 @@ message TraceProcessorRpc {
     EnableMetatraceArgs enable_metatrace_args = 106;
     // For TPM_RESET_TRACE_PROCESSOR.
     ResetTraceProcessorArgs reset_trace_processor_args = 107;
+    // For TPM_REGISTER_SQL_MODULE.
+    RegisterSqlModuleArgs register_sql_module_args = 108;
 
     // TraceProcessorMethod response args.
     // For TPM_APPEND_TRACE_DATA.
@@ -146,6 +150,8 @@ message TraceProcessorRpc {
     DisableAndReadMetatraceResult metatrace = 209;
     // For TPM_GET_STATUS.
     StatusResult status = 210;
+    // For TPM_REGISTER_SQL_MODULE.
+    RegisterSqlModuleResult register_sql_module_result = 211;
   }
 
   // Previously: RawQueryArgs for TPM_QUERY_RAW_DEPRECATED
@@ -328,3 +334,16 @@ message ResetTraceProcessorArgs {
   optional bool analyze_trace_proto_content = 3;
   optional bool ftrace_drop_until_all_cpus_valid = 4;
 }
+
+message RegisterSqlModuleArgs {
+  message Module {
+    optional string name = 1;
+    optional string sql = 2;
+  }
+  optional string top_level_package_name = 1;
+  repeated Module modules = 2;
+  optional bool allow_module_override = 3;
+}
+message RegisterSqlModuleResult {
+  optional string error = 1;
+}

src/trace_processor/rpc/rpc.cc

@@ -32,6 +32,7 @@
 #include "perfetto/ext/protozero/proto_ring_buffer.h"
 #include "perfetto/ext/trace_processor/rpc/query_result_serializer.h"
 #include "perfetto/protozero/field.h"
+#include "perfetto/protozero/proto_utils.h"
 #include "perfetto/protozero/scattered_heap_buffer.h"
 #include "perfetto/trace_processor/basic_types.h"
 #include "perfetto/trace_processor/metatrace_config.h"
@@ -323,6 +324,16 @@ void Rpc::ParseRpcRequest(const uint8_t* data, size_t len) {
       resp.Send(rpc_response_fn_);
       break;
     }
+    case RpcProto::TPM_REGISTER_SQL_MODULE: {
+      Response resp(tx_seq_id_++, req_type);
+      base::Status status = RegisterSqlModule(req.register_sql_module_args());
+      auto* res = resp->set_register_sql_module_result();
+      if (!status.ok()) {
+        res->set_error(status.message());
+      }
+      resp.Send(rpc_response_fn_);
+      break;
+    }
     default: {
       // This can legitimately happen if the client is newer. We reply with a
       // generic "unkown request" response, so the client can do feature
@@ -399,6 +410,18 @@ void Rpc::ResetTraceProcessor(const uint8_t* args, size_t len) {
   ResetTraceProcessorInternal(config);
 }
 
+base::Status Rpc::RegisterSqlModule(protozero::ConstBytes bytes) {
+  protos::pbzero::RegisterSqlModuleArgs::Decoder args(bytes);
+  SqlModule module;
+  module.name = args.top_level_package_name().ToStdString();
+  module.allow_module_override = args.allow_module_override();
+  for (auto it = args.modules(); it; ++it) {
+    protos::pbzero::RegisterSqlModuleArgs::Module::Decoder m(*it);
+    module.files.emplace_back(m.name().ToStdString(), m.sql().ToStdString());
+  }
+  return trace_processor_->RegisterSqlModule(module);
+}
+
 void Rpc::MaybePrintProgress() {
   if (eof_ || bytes_parsed_ - bytes_last_progress_ > kProgressUpdateBytes) {
     bytes_last_progress_ = bytes_parsed_;

src/trace_processor/rpc/rpc.h

@@ -27,6 +27,7 @@
 
 #include "perfetto/base/status.h"
 #include "perfetto/ext/protozero/proto_ring_buffer.h"
+#include "perfetto/protozero/field.h"
 #include "perfetto/trace_processor/basic_types.h"
 
 namespace perfetto {
@@ -104,7 +105,6 @@ class Rpc {
 
   base::Status Parse(const uint8_t*, size_t);
   base::Status NotifyEndOfFile();
-  void ResetTraceProcessor(const uint8_t*, size_t);
   std::string GetCurrentTraceName();
   std::vector<uint8_t> ComputeMetric(const uint8_t*, size_t);
   void EnableMetatrace(const uint8_t*, size_t);
@@ -130,6 +130,8 @@ class Rpc {
 
  private:
   void ParseRpcRequest(const uint8_t*, size_t);
+  void ResetTraceProcessor(const uint8_t*, size_t);
+  base::Status RegisterSqlModule(protozero::ConstBytes);
   void ResetTraceProcessorInternal(const Config&);
   void MaybePrintProgress();
   Iterator QueryInternal(const uint8_t*, size_t);

ui/src/controller/trace_controller.ts

@@ -430,6 +430,9 @@ export class TraceController extends Controller<States> {
         ftraceDropUntilAllCpusValid: FTRACE_DROP_UNTIL_FLAG.get(),
       });
     }
+    for (const p of globals.extraSqlPackages) {
+      await engine.registerSqlModules(p);
+    }
     this.engine = engine;
 
     if (isMetatracingEnabled()) {

ui/src/frontend/globals.ts

@@ -186,6 +186,16 @@ export const defaultTraceContext: TraceContext = {
   gpuCount: 0,
 };
 
+interface SqlModule {
+  readonly name: string;
+  readonly sql: string;
+}
+
+interface SqlPackage {
+  readonly name: string;
+  readonly modules: SqlModule[];
+}
+
 /**
  * Global accessors for state/dispatch in the frontend.
  */
@@ -236,6 +246,7 @@ class Globals implements AppContext {
   showPanningHint = false;
   permalinkHash?: string;
   showTraceErrorPopup = true;
+  extraSqlPackages: SqlPackage[] = [];
 
   traceContext = defaultTraceContext;
 

ui/src/protos/index.ts

@@ -71,6 +71,8 @@ import QueryServiceStateRequest = protos.perfetto.protos.QueryServiceStateReques
 import QueryServiceStateResponse = protos.perfetto.protos.QueryServiceStateResponse;
 import ReadBuffersRequest = protos.perfetto.protos.ReadBuffersRequest;
 import ReadBuffersResponse = protos.perfetto.protos.ReadBuffersResponse;
+import RegisterSqlModuleArgs = protos.perfetto.protos.RegisterSqlModuleArgs;
+import RegisterSqlModuleResult = protos.perfetto.protos.RegisterSqlModuleResult;
 import ResetTraceProcessorArgs = protos.perfetto.protos.ResetTraceProcessorArgs;
 import StatCounters = protos.perfetto.protos.SysStatsConfig.StatCounters;
 import StatusResult = protos.perfetto.protos.StatusResult;
@@ -138,6 +140,8 @@ export {
   QueryServiceStateResponse,
   ReadBuffersRequest,
   ReadBuffersResponse,
+  RegisterSqlModuleArgs,
+  RegisterSqlModuleResult,
   ResetTraceProcessorArgs,
   StatCounters,
   StatusResult,

ui/src/trace_processor/engine.ts

@@ -22,6 +22,7 @@ import {
   MetatraceCategories,
   QueryArgs,
   QueryResult as ProtoQueryResult,
+  RegisterSqlModuleArgs,
   ResetTraceProcessorArgs,
   TraceProcessorRpc,
   TraceProcessorRpcStream,
@@ -126,6 +127,7 @@ export abstract class EngineBase implements Engine {
   private pendingRestoreTables = new Array<Deferred<void>>();
   private pendingComputeMetrics = new Array<Deferred<string | Uint8Array>>();
   private pendingReadMetatrace?: Deferred<DisableAndReadMetatraceResult>;
+  private pendingRegisterSqlModule?: Deferred<void>;
   private _isMetatracingEnabled = false;
 
   constructor(tracker?: LoadingTracker) {
@@ -261,6 +263,15 @@ export abstract class EngineBase implements Engine {
         assertExists(this.pendingReadMetatrace).resolve(metatraceRes);
         this.pendingReadMetatrace = undefined;
         break;
+      case TPM.TPM_REGISTER_SQL_MODULE:
+        const registerResult = assertExists(rpc.registerSqlModuleResult);
+        const res = assertExists(this.pendingRegisterSqlModule);
+        if (registerResult.error && registerResult.error.length > 0) {
+          res.reject(registerResult.error);
+        } else {
+          res.resolve();
+        }
+        break;
       default:
         console.log(
           'Unexpected TraceProcessor response received: ',
@@ -461,6 +472,26 @@ export abstract class EngineBase implements Engine {
     return result;
   }
 
+  registerSqlModules(p: {
+    name: string;
+    modules: {name: string; sql: string}[];
+  }): Promise<void> {
+    if (this.pendingRegisterSqlModule) {
+      return Promise.reject(new Error('Already finalising a metatrace'));
+    }
+
+    const result = defer<void>();
+
+    const rpc = TraceProcessorRpc.create();
+    rpc.request = TPM.TPM_REGISTER_SQL_MODULE;
+    const args = (rpc.registerSqlModuleArgs = new RegisterSqlModuleArgs());
+    args.topLevelPackageName = p.name;
+    args.modules = p.modules;
+    this.pendingRegisterSqlModule = result;
+    this.rpcSendRequest(rpc);
+    return result;
+  }
+
   // Marshals the TraceProcessorRpc request arguments and sends the request
   // to the concrete Engine (Wasm or HTTP).
   private rpcSendRequest(rpc: TraceProcessorRpc) {