trace_processor: Migrate track_event_module to the new base class

Bug: 141459049
Change-Id: Ie94a89b04d70c49d2f22854bf26cb9f13e3c4f1e
diff --git a/Android.bp b/Android.bp
index bf2d0bb..9a586c5 100644
--- a/Android.bp
+++ b/Android.bp
@@ -4801,6 +4801,7 @@
     "src/trace_processor/importers/proto/proto_trace_parser.cc",
     "src/trace_processor/importers/proto/proto_trace_tokenizer.cc",
     "src/trace_processor/importers/proto/system_probes_parser.cc",
+    "src/trace_processor/importers/proto/track_event_module.cc",
     "src/trace_processor/importers/proto/track_event_parser.cc",
     "src/trace_processor/importers/proto/track_event_tokenizer.cc",
     "src/trace_processor/importers/systrace/systrace_parser.cc",
diff --git a/BUILD b/BUILD
index ed835ea..ed26485 100644
--- a/BUILD
+++ b/BUILD
@@ -842,6 +842,7 @@
         "src/trace_processor/importers/proto/system_probes_module.h",
         "src/trace_processor/importers/proto/system_probes_parser.cc",
         "src/trace_processor/importers/proto/system_probes_parser.h",
+        "src/trace_processor/importers/proto/track_event_module.cc",
         "src/trace_processor/importers/proto/track_event_module.h",
         "src/trace_processor/importers/proto/track_event_parser.cc",
         "src/trace_processor/importers/proto/track_event_parser.h",
diff --git a/src/trace_processor/BUILD.gn b/src/trace_processor/BUILD.gn
index fe0a7ea..6e215f3 100644
--- a/src/trace_processor/BUILD.gn
+++ b/src/trace_processor/BUILD.gn
@@ -102,6 +102,7 @@
     "importers/proto/proto_trace_tokenizer.h",
     "importers/proto/system_probes_module.h",
     "importers/proto/system_probes_parser.h",
+    "importers/proto/track_event_module.cc",
     "importers/proto/track_event_module.h",
     "importers/proto/track_event_parser.cc",
     "importers/proto/track_event_parser.h",
diff --git a/src/trace_processor/importers/proto/proto_trace_parser.cc b/src/trace_processor/importers/proto/proto_trace_parser.cc
index b4bbaf9..ea35587 100644
--- a/src/trace_processor/importers/proto/proto_trace_parser.cc
+++ b/src/trace_processor/importers/proto/proto_trace_parser.cc
@@ -37,7 +37,6 @@
 #include "src/trace_processor/importers/proto/heap_graph_module.h"
 #include "src/trace_processor/importers/proto/packet_sequence_state.h"
 #include "src/trace_processor/importers/proto/system_probes_module.h"
-#include "src/trace_processor/importers/proto/track_event_module.h"
 #include "src/trace_processor/metadata.h"
 #include "src/trace_processor/process_tracker.h"
 #include "src/trace_processor/slice_tracker.h"
@@ -206,9 +205,6 @@
   if (!context_->ftrace_module->ParsePacket(packet, ttp).ignored())
     return;
 
-  if (!context_->track_event_module->ParsePacket(packet, ttp).ignored())
-    return;
-
   if (!context_->system_probes_module->ParsePacket(packet, ttp).ignored())
     return;
 
diff --git a/src/trace_processor/importers/proto/proto_trace_parser_unittest.cc b/src/trace_processor/importers/proto/proto_trace_parser_unittest.cc
index d19edae..edb53a7 100644
--- a/src/trace_processor/importers/proto/proto_trace_parser_unittest.cc
+++ b/src/trace_processor/importers/proto/proto_trace_parser_unittest.cc
@@ -263,8 +263,6 @@
 #endif  // PERFETTO_BUILDFLAG(PERFETTO_TP_GRAPHICS)
     context_.ftrace_module.reset(
         new ProtoImporterModule<FtraceModule>(&context_));
-    context_.track_event_module.reset(
-        new ProtoImporterModule<TrackEventModule>(&context_));
     context_.system_probes_module.reset(
         new ProtoImporterModule<SystemProbesModule>(&context_));
     context_.android_probes_module.reset(
@@ -272,6 +270,7 @@
     context_.heap_graph_module.reset(
         new ProtoImporterModule<HeapGraphModule>(&context_));
 
+    context_.modules.emplace_back(new TrackEventModule(&context_));
 #if PERFETTO_BUILDFLAG(PERFETTO_TP_GRAPHICS)
     context_.modules.emplace_back(new GraphicsEventModule(&context_));
 #endif  // PERFETTO_BUILDFLAG(PERFETTO_TP_GRAPHICS)
diff --git a/src/trace_processor/importers/proto/proto_trace_tokenizer.cc b/src/trace_processor/importers/proto/proto_trace_tokenizer.cc
index b75ccf4..e6ba1a0 100644
--- a/src/trace_processor/importers/proto/proto_trace_tokenizer.cc
+++ b/src/trace_processor/importers/proto/proto_trace_tokenizer.cc
@@ -32,7 +32,6 @@
 #include "src/trace_processor/importers/ftrace/ftrace_module.h"
 #include "src/trace_processor/importers/proto/packet_sequence_state.h"
 #include "src/trace_processor/importers/proto/proto_incremental_state.h"
-#include "src/trace_processor/importers/proto/track_event_module.h"
 #include "src/trace_processor/stats.h"
 #include "src/trace_processor/trace_sorter.h"
 #include "src/trace_processor/trace_storage.h"
@@ -289,16 +288,11 @@
   if (!res.ignored())
     return res.ToStatus();
 
-  res = context_->track_event_module->TokenizePacket(decoder, &packet,
-                                                     timestamp, state);
-  if (!res.ignored())
-    return res.ToStatus();
-
   auto& modules = context_->modules_by_field;
   for (uint32_t field_id = 1; field_id < modules.size(); ++field_id) {
     if (modules[field_id] && decoder.Get(field_id).valid()) {
-      modules[field_id]->TokenizePacket(decoder, &packet, timestamp, state,
-                                        field_id);
+      res = modules[field_id]->TokenizePacket(decoder, &packet, timestamp,
+                                              state, field_id);
       if (!res.ignored())
         return res.ToStatus();
     }
diff --git a/src/trace_processor/importers/proto/track_event_module.cc b/src/trace_processor/importers/proto/track_event_module.cc
new file mode 100644
index 0000000..8e80a80
--- /dev/null
+++ b/src/trace_processor/importers/proto/track_event_module.cc
@@ -0,0 +1,76 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "src/trace_processor/importers/proto/track_event_module.h"
+#include "perfetto/base/build_config.h"
+#include "src/trace_processor/timestamped_trace_piece.h"
+
+#include "protos/perfetto/config/trace_config.pbzero.h"
+#include "protos/perfetto/trace/trace_packet.pbzero.h"
+
+namespace perfetto {
+namespace trace_processor {
+
+using perfetto::protos::pbzero::TracePacket;
+
+TrackEventModule::TrackEventModule(TraceProcessorContext* context)
+    : tokenizer_(context), parser_(context) {
+  RegisterForField(TracePacket::kTrackEventFieldNumber, context);
+  RegisterForField(TracePacket::kTrackDescriptorFieldNumber, context);
+  RegisterForField(TracePacket::kThreadDescriptorFieldNumber, context);
+  RegisterForField(TracePacket::kProcessDescriptorFieldNumber, context);
+}
+
+TrackEventModule::~TrackEventModule() = default;
+
+ModuleResult TrackEventModule::TokenizePacket(
+    const TracePacket::Decoder& decoder,
+    TraceBlobView* packet,
+    int64_t packet_timestamp,
+    PacketSequenceState* state,
+    uint32_t field_id) {
+  switch (field_id) {
+    case TracePacket::kTrackDescriptorFieldNumber:
+      tokenizer_.TokenizeTrackDescriptorPacket(decoder);
+      return ModuleResult::Handled();
+    case TracePacket::kTrackEventFieldNumber:
+      tokenizer_.TokenizeTrackEventPacket(state, decoder, packet,
+                                          packet_timestamp);
+      return ModuleResult::Handled();
+    case TracePacket::kThreadDescriptorFieldNumber:
+      // TODO(eseckler): Remove these once Chrome has switched fully over to
+      // TrackDescriptors.
+      tokenizer_.TokenizeThreadDescriptorPacket(state, decoder);
+      return ModuleResult::Handled();
+    case TracePacket::kProcessDescriptorFieldNumber:
+      tokenizer_.TokenizeProcessDescriptorPacket(decoder);
+      return ModuleResult::Handled();
+  }
+  return ModuleResult::Ignored();
+}
+
+void TrackEventModule::ParsePacket(const TracePacket::Decoder& decoder,
+                                   const TimestampedTracePiece& ttp,
+                                   uint32_t field_id) {
+  if (field_id == TracePacket::kTrackEventFieldNumber) {
+    parser_.ParseTrackEvent(
+        ttp.timestamp, ttp.thread_timestamp, ttp.thread_instruction_count,
+        ttp.packet_sequence_state, ttp.packet_sequence_state_generation,
+        decoder.track_event());
+  }
+}
+
+}  // namespace trace_processor
+}  // namespace perfetto
diff --git a/src/trace_processor/importers/proto/track_event_module.h b/src/trace_processor/importers/proto/track_event_module.h
index b58b8bc..05dbb80 100644
--- a/src/trace_processor/importers/proto/track_event_module.h
+++ b/src/trace_processor/importers/proto/track_event_module.h
@@ -17,66 +17,31 @@
 #ifndef SRC_TRACE_PROCESSOR_IMPORTERS_PROTO_TRACK_EVENT_MODULE_H_
 #define SRC_TRACE_PROCESSOR_IMPORTERS_PROTO_TRACK_EVENT_MODULE_H_
 
-#include "perfetto/base/build_config.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 #include "src/trace_processor/importers/proto/track_event_parser.h"
 #include "src/trace_processor/importers/proto/track_event_tokenizer.h"
-#include "src/trace_processor/timestamped_trace_piece.h"
 
-#include "protos/perfetto/config/trace_config.pbzero.h"
 #include "protos/perfetto/trace/trace_packet.pbzero.h"
 
 namespace perfetto {
 namespace trace_processor {
 
-class TrackEventModule : public ProtoImporterModuleBase</*IsEnabled=*/1> {
+class TrackEventModule : public NewProtoImporterModule {
  public:
-  explicit TrackEventModule(TraceProcessorContext* context)
-      : ProtoImporterModuleBase(context),
-        tokenizer_(context),
-        parser_(context) {}
+  explicit TrackEventModule(TraceProcessorContext* context);
+
+  ~TrackEventModule() override;
 
   ModuleResult TokenizePacket(
       const protos::pbzero::TracePacket::Decoder& decoder,
       TraceBlobView* packet,
       int64_t packet_timestamp,
-      PacketSequenceState* state) {
-    if (decoder.has_track_descriptor()) {
-      tokenizer_.TokenizeTrackDescriptorPacket(decoder);
-      return ModuleResult::Handled();
-    }
+      PacketSequenceState* state,
+      uint32_t field_id) override;
 
-    if (decoder.has_track_event()) {
-      tokenizer_.TokenizeTrackEventPacket(state, decoder, packet,
-                                          packet_timestamp);
-      return ModuleResult::Handled();
-    }
-
-    // TODO(eseckler): Remove these once Chrome has switched fully over to
-    // TrackDescriptors.
-    if (decoder.has_thread_descriptor()) {
-      tokenizer_.TokenizeThreadDescriptorPacket(state, decoder);
-      return ModuleResult::Handled();
-    }
-    if (decoder.has_process_descriptor()) {
-      tokenizer_.TokenizeProcessDescriptorPacket(decoder);
-      return ModuleResult::Handled();
-    }
-
-    return ModuleResult::Ignored();
-  }
-
-  ModuleResult ParsePacket(const protos::pbzero::TracePacket::Decoder& decoder,
-                           const TimestampedTracePiece& ttp) {
-    if (decoder.has_track_event()) {
-      parser_.ParseTrackEvent(
-          ttp.timestamp, ttp.thread_timestamp, ttp.thread_instruction_count,
-          ttp.packet_sequence_state, ttp.packet_sequence_state_generation,
-          decoder.track_event());
-      return ModuleResult::Handled();
-    }
-    return ModuleResult::Ignored();
-  }
+  void ParsePacket(const protos::pbzero::TracePacket::Decoder& decoder,
+                   const TimestampedTracePiece& ttp,
+                   uint32_t field_id) override;
 
  private:
   TrackEventTokenizer tokenizer_;
diff --git a/src/trace_processor/trace_processor_context.h b/src/trace_processor/trace_processor_context.h
index 784b130..6b6ed88 100644
--- a/src/trace_processor/trace_processor_context.h
+++ b/src/trace_processor/trace_processor_context.h
@@ -76,7 +76,6 @@
   std::unique_ptr<BinderTracker> binder_tracker;
 
   std::unique_ptr<ProtoImporterModule<FtraceModule>> ftrace_module;
-  std::unique_ptr<ProtoImporterModule<TrackEventModule>> track_event_module;
   std::unique_ptr<ProtoImporterModule<SystemProbesModule>> system_probes_module;
   std::unique_ptr<ProtoImporterModule<AndroidProbesModule>>
       android_probes_module;
diff --git a/src/trace_processor/trace_processor_storage_impl.cc b/src/trace_processor/trace_processor_storage_impl.cc
index da5521b..0b3e3bd 100644
--- a/src/trace_processor/trace_processor_storage_impl.cc
+++ b/src/trace_processor/trace_processor_storage_impl.cc
@@ -70,8 +70,6 @@
 #endif  // PERFETTO_BUILDFLAG(PERFETTO_TP_GRAPHICS)
   context_.ftrace_module.reset(
       new ProtoImporterModule<FtraceModule>(&context_));
-  context_.track_event_module.reset(
-      new ProtoImporterModule<TrackEventModule>(&context_));
   context_.system_probes_module.reset(
       new ProtoImporterModule<SystemProbesModule>(&context_));
   context_.android_probes_module.reset(
@@ -79,6 +77,7 @@
   context_.heap_graph_module.reset(
       new ProtoImporterModule<HeapGraphModule>(&context_));
 
+  context_.modules.emplace_back(new TrackEventModule(&context_));
 #if PERFETTO_BUILDFLAG(PERFETTO_TP_GRAPHICS)
   context_.modules.emplace_back(new GraphicsEventModule(&context_));
 #endif  // PERFETTO_BUILDFLAG(PERFETTO_TP_GRAPHICS)