Trace Redaction - Use Begin/End with CollectTimelineEvents

CollectTimlineEvents was one of the reason Begin/End was added to the
collect primitive. CollectTimelineEvents needs to be able to
create the timeline instance (smart pointer) once, events multiple
timelines, and then finalize it by sorting the events. The first
behaviours were possible using the old model, but the last option was
hard to support.

Having the timeline be search-able in build phase is required to
filter the frame timeline. However, previously the timeline was
converted from write-only to read-only, another build primitive
(optimize timeline). This would create a dependency between primitives
in the same phase - we want to avoid that.

Bug: 333404689

Change-Id: I3504e309d234e280345fce9771b0918b452e8071
diff --git a/src/trace_redaction/collect_timeline_events.cc b/src/trace_redaction/collect_timeline_events.cc
index c5a98a7..aac9752 100644
--- a/src/trace_redaction/collect_timeline_events.cc
+++ b/src/trace_redaction/collect_timeline_events.cc
@@ -108,15 +108,21 @@
 
 }  // namespace
 
-base::Status CollectTimelineEvents::Collect(const TracePacket::Decoder& packet,
-                                    Context* context) const {
-  // TODO(vaage): This should only be true on the first call. However, that
-  // means a branch is called N times when N-1 times it will be false. This may
-  // be common across Collect primitives. Having a "begin" and "end" end-points.
-  if (!context->timeline) {
-    context->timeline = std::make_unique<ProcessThreadTimeline>();
+base::Status CollectTimelineEvents::Begin(Context* context) const {
+  // This primitive is artifically limited to owning the timeline. In practice
+  // there is no reason why multiple primitives could contribute to the
+  // timeline.
+  if (context->timeline) {
+    return base::ErrStatus(
+        "CollectTimelineEvents: timeline was already initialized");
   }
 
+  context->timeline = std::make_unique<ProcessThreadTimeline>();
+  return base::OkStatus();
+}
+
+base::Status CollectTimelineEvents::Collect(const TracePacket::Decoder& packet,
+                                            Context* context) const {
   // Unlike ftrace events, process trees do not provide per-process or
   // per-thread timing information. The packet has timestamp and the process
   // tree has collection_end_timestamp (collection_end_timestamp > timestamp).
@@ -137,4 +143,11 @@
   return base::OkStatus();
 }
 
+base::Status CollectTimelineEvents::End(Context* context) const {
+  // Sort must be called in order to read from the timeline. If any more events
+  // are added after this, then sort will need to be called again.
+  context->timeline->Sort();
+  return base::OkStatus();
+}
+
 }  // namespace perfetto::trace_redaction
diff --git a/src/trace_redaction/collect_timeline_events.h b/src/trace_redaction/collect_timeline_events.h
index fb29062..27aa337 100644
--- a/src/trace_redaction/collect_timeline_events.h
+++ b/src/trace_redaction/collect_timeline_events.h
@@ -27,8 +27,12 @@
 // packets and stores them in a timeline.
 class CollectTimelineEvents : public CollectPrimitive {
  public:
+  base::Status Begin(Context* context) const override;
+
   base::Status Collect(const protos::pbzero::TracePacket::Decoder& packet,
                        Context* context) const override;
+
+  base::Status End(Context* context) const override;
 };
 
 }  // namespace perfetto::trace_redaction
diff --git a/src/trace_redaction/collect_timeline_events_unittest.cc b/src/trace_redaction/collect_timeline_events_unittest.cc
index 14aa982..ecf0905 100644
--- a/src/trace_redaction/collect_timeline_events_unittest.cc
+++ b/src/trace_redaction/collect_timeline_events_unittest.cc
@@ -87,10 +87,11 @@
 
 }  // namespace
 
-class CollectTimelineEventsTest : public testing::Test,
-                          public testing::WithParamInterface<TestParams> {
+class CollectTimelineEventsTest
+    : public testing::Test,
+      public testing::WithParamInterface<TestParams> {
  protected:
-  base::Status PushProcessTreePacket(uint64_t timestamp) {
+  std::string CreateProcessTreePacket(uint64_t timestamp) {
     protos::gen::TracePacket packet;
     packet.set_trusted_uid(9999);
     packet.set_timestamp(timestamp);
@@ -118,12 +119,10 @@
 
     process_tree->set_collection_end_timestamp(timestamp);
 
-    std::string packet_str = packet.SerializeAsString();
-    return build_.Collect(protos::pbzero::TracePacket::Decoder(packet_str),
-                          &context_);
+    return packet.SerializeAsString();
   }
 
-  base::Status PushSchedProcessFreePacket(uint64_t timestamp) {
+  std::string CreateSchedProcessFreePacket(uint64_t timestamp) {
     protos::gen::TracePacket packet;
 
     packet.set_trusted_uid(9999);
@@ -142,26 +141,32 @@
     process_free->set_pid(kUnityTid);
     process_free->set_prio(120);
 
-    std::string packet_str = packet.SerializeAsString();
-    return build_.Collect(protos::pbzero::TracePacket::Decoder(packet_str),
-                          &context_);
+    return packet.SerializeAsString();
   }
-
-  CollectTimelineEvents build_;
-  Context context_;
 };
 
-class CollectTimelineEventsWithProcessTree : public CollectTimelineEventsTest {};
+class CollectTimelineEventsWithProcessTree : public CollectTimelineEventsTest {
+};
 
 TEST_P(CollectTimelineEventsWithProcessTree, FindsOpenSpans) {
   auto params = GetParam();
 
-  auto result = PushProcessTreePacket(kProcessTreeTimestamp);
-  ASSERT_OK(result) << result.message();
+  auto packet_str = CreateProcessTreePacket(kProcessTreeTimestamp);
 
-  context_.timeline->Sort();
+  protos::pbzero::TracePacket::Decoder packet(packet_str);
 
-  auto slice = context_.timeline->Search(params.ts(), params.pid());
+  Context context;
+  CollectTimelineEvents collector;
+  auto begin_status = collector.Begin(&context);
+  ASSERT_OK(begin_status) << begin_status.message();
+
+  auto packet_status = collector.Collect(packet, &context);
+  ASSERT_OK(packet_status) << packet_status.message();
+
+  auto end_status = collector.End(&context);
+  ASSERT_OK(end_status) << end_status.message();
+
+  auto slice = context.timeline->Search(params.ts(), params.pid());
   ASSERT_EQ(slice.pid, params.pid());
   ASSERT_EQ(slice.uid, params.uid());
 }
@@ -186,20 +191,33 @@
         TestParams(kProcessTreeTimestamp + 1, kUnityTid, kUnityPackage)));
 
 // Assumes all CollectTimelineEventsWithProcessTree tests pass.
-class CollectTimelineEventsWithFreeProcess : public CollectTimelineEventsTest {};
+class CollectTimelineEventsWithFreeProcess : public CollectTimelineEventsTest {
+};
 
 TEST_P(CollectTimelineEventsWithFreeProcess, FindsClosedSpans) {
   auto params = GetParam();
 
-  auto result = PushProcessTreePacket(kProcessTreeTimestamp);
-  ASSERT_OK(result) << result.message();
+  auto packet_1_str = CreateProcessTreePacket(kProcessTreeTimestamp);
+  auto packet_2_str = CreateSchedProcessFreePacket(kThreadFreeTimestamp);
 
-  result = PushSchedProcessFreePacket(kThreadFreeTimestamp);
-  ASSERT_OK(result) << result.message();
+  protos::pbzero::TracePacket::Decoder packet_1(packet_1_str);
+  protos::pbzero::TracePacket::Decoder packet_2(packet_2_str);
 
-  context_.timeline->Sort();
+  Context context;
+  CollectTimelineEvents collector;
+  auto begin_status = collector.Begin(&context);
+  ASSERT_OK(begin_status) << begin_status.message();
 
-  auto slice = context_.timeline->Search(params.ts(), params.pid());
+  auto packet_1_status = collector.Collect(packet_1, &context);
+  ASSERT_OK(packet_1_status) << packet_1_status.message();
+
+  auto packet_2_status = collector.Collect(packet_2, &context);
+  ASSERT_OK(packet_2_status) << packet_2_status.message();
+
+  auto end_status = collector.End(&context);
+  ASSERT_OK(end_status) << end_status.message();
+
+  auto slice = context.timeline->Search(params.ts(), params.pid());
   ASSERT_EQ(slice.pid, params.pid());
   ASSERT_EQ(slice.uid, params.uid());
 }