tp: implement async track set tracker

This CL implements the async track set tracker described in
go/perfetto-tp-scaling-async-slices.

Currently, only Android async slices have been migrated but in the
future, GPU slices will also use this tracker and Chrome *may* be
migrated over as well depending on the result of experiments.

Doc: go/perfetto-tp-scaling-async-slices
Bug: 169088987
Change-Id: Idf9fcc05aa55c3e2fb7ef4b6d80a51166da238f4
diff --git a/Android.bp b/Android.bp
index c1a92ef..e50aed9 100644
--- a/Android.bp
+++ b/Android.bp
@@ -7118,6 +7118,7 @@
     "src/trace_processor/importers/json/json_utils.cc",
     "src/trace_processor/importers/ninja/ninja_log_parser.cc",
     "src/trace_processor/importers/proto/args_table_utils.cc",
+    "src/trace_processor/importers/proto/async_track_set_tracker.cc",
     "src/trace_processor/importers/proto/heap_profile_tracker.cc",
     "src/trace_processor/importers/proto/metadata_tracker.cc",
     "src/trace_processor/importers/proto/packet_sequence_state.cc",
diff --git a/BUILD b/BUILD
index 0dbb850..241139e 100644
--- a/BUILD
+++ b/BUILD
@@ -1126,6 +1126,8 @@
         "src/trace_processor/importers/ninja/ninja_log_parser.h",
         "src/trace_processor/importers/proto/args_table_utils.cc",
         "src/trace_processor/importers/proto/args_table_utils.h",
+        "src/trace_processor/importers/proto/async_track_set_tracker.cc",
+        "src/trace_processor/importers/proto/async_track_set_tracker.h",
         "src/trace_processor/importers/proto/heap_profile_tracker.cc",
         "src/trace_processor/importers/proto/heap_profile_tracker.h",
         "src/trace_processor/importers/proto/metadata_tracker.cc",
diff --git a/src/trace_processor/BUILD.gn b/src/trace_processor/BUILD.gn
index d548cd8..0808519 100644
--- a/src/trace_processor/BUILD.gn
+++ b/src/trace_processor/BUILD.gn
@@ -88,6 +88,8 @@
     "importers/ninja/ninja_log_parser.h",
     "importers/proto/args_table_utils.cc",
     "importers/proto/args_table_utils.h",
+    "importers/proto/async_track_set_tracker.cc",
+    "importers/proto/async_track_set_tracker.h",
     "importers/proto/heap_profile_tracker.cc",
     "importers/proto/heap_profile_tracker.h",
     "importers/proto/metadata_tracker.cc",
diff --git a/src/trace_processor/export_json_unittest.cc b/src/trace_processor/export_json_unittest.cc
index 1efcb87..88b9451 100644
--- a/src/trace_processor/export_json_unittest.cc
+++ b/src/trace_processor/export_json_unittest.cc
@@ -227,9 +227,8 @@
 }
 
 TEST_F(ExportJsonTest, SystemEventsIgnored) {
-  constexpr int64_t kCookie = 22;
-  TrackId track = context_.track_tracker->InternAndroidAsyncTrack(
-      /*name=*/kNullStringId, /*upid=*/0, kCookie);
+  TrackId track = context_.track_tracker->CreateAndroidAsyncTrack(
+      /*name=*/kNullStringId, /*upid=*/0);
   context_.args_tracker->Flush();  // Flush track args.
 
   // System events have no category.
diff --git a/src/trace_processor/importers/common/slice_tracker.cc b/src/trace_processor/importers/common/slice_tracker.cc
index a1082d2..0a5ef04 100644
--- a/src/trace_processor/importers/common/slice_tracker.cc
+++ b/src/trace_processor/importers/common/slice_tracker.cc
@@ -53,7 +53,8 @@
   });
 }
 
-void SliceTracker::BeginLegacyUnnestable(tables::SliceTable::Row row) {
+void SliceTracker::BeginLegacyUnnestable(tables::SliceTable::Row row,
+                                         SetArgsCallback args_callback) {
   // Ensure that the duration is pending for this row.
   // TODO(lalitm): change this to eventually use null instead of -1.
   row.dur = kPendingDuration;
@@ -68,7 +69,7 @@
   // Ensure that StartSlice knows that this track is unnestable.
   stacks_[row.track_id].is_legacy_unnestable = true;
 
-  StartSlice(row.ts, row.track_id, SetArgsCallback{}, [this, &row]() {
+  StartSlice(row.ts, row.track_id, args_callback, [this, &row]() {
     return context_->storage->mutable_slice_table()->Insert(row).id;
   });
 }
diff --git a/src/trace_processor/importers/common/slice_tracker.h b/src/trace_processor/importers/common/slice_tracker.h
index c54575e..197204a 100644
--- a/src/trace_processor/importers/common/slice_tracker.h
+++ b/src/trace_processor/importers/common/slice_tracker.h
@@ -50,7 +50,8 @@
   // as the latest time we saw a begin event. For legacy Android use only. See
   // the comment in SystraceParser::ParseSystracePoint for information on why
   // this method exists.
-  void BeginLegacyUnnestable(tables::SliceTable::Row row);
+  void BeginLegacyUnnestable(tables::SliceTable::Row row,
+                             SetArgsCallback args_callback);
 
   void BeginGpu(tables::GpuSliceTable::Row row,
                 SetArgsCallback args_callback = SetArgsCallback());
diff --git a/src/trace_processor/importers/common/track_tracker.cc b/src/trace_processor/importers/common/track_tracker.cc
index d4b3868..21eabba 100644
--- a/src/trace_processor/importers/common/track_tracker.cc
+++ b/src/trace_processor/importers/common/track_tracker.cc
@@ -150,24 +150,12 @@
   return id;
 }
 
-TrackId TrackTracker::InternAndroidAsyncTrack(StringId name,
-                                              UniquePid upid,
-                                              int64_t cookie) {
-  AndroidAsyncTrackTuple tuple{upid, cookie, name};
-
-  auto it = android_async_tracks_.find(tuple);
-  if (it != android_async_tracks_.end())
-    return it->second;
-
+TrackId TrackTracker::CreateAndroidAsyncTrack(StringId name, UniquePid upid) {
   tables::ProcessTrackTable::Row row(name);
   row.upid = upid;
   auto id = context_->storage->mutable_process_track_table()->Insert(row).id;
-  android_async_tracks_[tuple] = id;
-
-  context_->args_tracker->AddArgsTo(id)
-      .AddArg(source_key_, Variadic::String(android_source_))
-      .AddArg(source_id_key_, Variadic::Integer(cookie));
-
+  context_->args_tracker->AddArgsTo(id).AddArg(
+      source_key_, Variadic::String(android_source_));
   return id;
 }
 
diff --git a/src/trace_processor/importers/common/track_tracker.h b/src/trace_processor/importers/common/track_tracker.h
index b525325..6329937 100644
--- a/src/trace_processor/importers/common/track_tracker.h
+++ b/src/trace_processor/importers/common/track_tracker.h
@@ -50,10 +50,8 @@
                                        bool source_id_is_process_scoped,
                                        StringId source_scope);
 
-  // Interns a Android async track into the storage.
-  TrackId InternAndroidAsyncTrack(StringId name,
-                                  UniquePid upid,
-                                  int64_t cookie);
+  // Creates and inserts a Android async track into the storage.
+  TrackId CreateAndroidAsyncTrack(StringId name, UniquePid upid);
 
   // Interns a track for legacy Chrome process-scoped instant events into the
   // storage.
@@ -214,17 +212,6 @@
              std::tie(r.source_id, r.upid, r.source_scope);
     }
   };
-  struct AndroidAsyncTrackTuple {
-    UniquePid upid;
-    int64_t cookie;
-    StringId name;
-
-    friend bool operator<(const AndroidAsyncTrackTuple& l,
-                          const AndroidAsyncTrackTuple& r) {
-      return std::tie(l.upid, l.cookie, l.name) <
-             std::tie(r.upid, r.cookie, r.name);
-    }
-  };
   struct DescriptorTrackReservation {
     uint64_t parent_uuid = 0;
     base::Optional<uint32_t> pid;
@@ -269,7 +256,6 @@
 
   std::map<GpuTrackTuple, TrackId> gpu_tracks_;
   std::map<ChromeTrackTuple, TrackId> chrome_tracks_;
-  std::map<AndroidAsyncTrackTuple, TrackId> android_async_tracks_;
   std::map<UniquePid, TrackId> chrome_process_instant_tracks_;
   base::Optional<TrackId> chrome_global_instant_track_id_;
   std::map<uint64_t /* uuid */, DescriptorTrackReservation>
diff --git a/src/trace_processor/importers/proto/async_track_set_tracker.cc b/src/trace_processor/importers/proto/async_track_set_tracker.cc
new file mode 100644
index 0000000..3b2067d
--- /dev/null
+++ b/src/trace_processor/importers/proto/async_track_set_tracker.cc
@@ -0,0 +1,110 @@
+/*
+ * Copyright (C) 2020 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/async_track_set_tracker.h"
+
+#include "src/trace_processor/importers/common/track_tracker.h"
+#include "src/trace_processor/types/trace_processor_context.h"
+
+namespace perfetto {
+namespace trace_processor {
+
+AsyncTrackSetTracker::AsyncTrackSetTracker(TraceProcessorContext* context)
+    : context_(context) {}
+
+AsyncTrackSetTracker::TrackSetId AsyncTrackSetTracker::InternAndroidSet(
+    UniquePid upid,
+    StringId name) {
+  AndroidTuple tuple{upid, name};
+
+  auto it = android_track_set_ids_.find(tuple);
+  if (it != android_track_set_ids_.end())
+    return it->second;
+
+  uint32_t id = static_cast<uint32_t>(track_sets_.size());
+
+  TrackSet set;
+  set.android_tuple = tuple;
+  set.type = TrackSetType::kAndroid;
+  track_sets_.emplace_back(set);
+  android_track_set_ids_[tuple] = id;
+  return id;
+}
+
+TrackId AsyncTrackSetTracker::Begin(TrackSetId id,
+                                    int64_t cookie,
+                                    NestingBehaviour nest) {
+  PERFETTO_DCHECK(id < track_sets_.size());
+
+  TrackSet& set = track_sets_[id];
+  auto it = std::find_if(
+      set.tracks.begin(), set.tracks.end(),
+      [cookie](const TrackState& state) { return state.cookie == cookie; });
+  if (it == set.tracks.end()) {
+    TrackState& state = set.tracks[GetOrCreateEmptyTrack(set, cookie)];
+    PERFETTO_DCHECK(state.nest_count == 0);
+    state.nest_count = 1;
+    return state.id;
+  }
+
+  switch (nest) {
+    case NestingBehaviour::kLegacySaturatingUnnestable:
+      PERFETTO_DCHECK(it->nest_count <= 1);
+      break;
+    case NestingBehaviour::kUnnestable:
+      PERFETTO_DCHECK(it->nest_count == 0);
+      break;
+  }
+  return it->id;
+}
+
+TrackId AsyncTrackSetTracker::End(TrackSetId id, int64_t cookie) {
+  PERFETTO_DCHECK(id < track_sets_.size());
+
+  TrackSet& set = track_sets_[id];
+  auto it = std::find_if(
+      set.tracks.begin(), set.tracks.end(),
+      [cookie](const TrackState& state) { return state.cookie == cookie; });
+  if (it == set.tracks.end())
+    return set.tracks[GetOrCreateEmptyTrack(set, cookie)].id;
+  return it->id;
+}
+
+uint32_t AsyncTrackSetTracker::GetOrCreateEmptyTrack(TrackSet& set,
+                                                     int64_t cookie) {
+  auto it = std::find_if(
+      set.tracks.begin(), set.tracks.end(),
+      [](const TrackState& state) { return state.nest_count == 0; });
+  if (it != set.tracks.end())
+    return static_cast<uint32_t>(std::distance(set.tracks.begin(), it));
+
+  TrackState state;
+  state.cookie = cookie;
+  state.nest_count = 0;
+  switch (set.type) {
+    case TrackSetType::kAndroid:
+      state.id = context_->track_tracker->CreateAndroidAsyncTrack(
+          set.android_tuple.name, set.android_tuple.upid);
+      break;
+  }
+
+  uint32_t idx = static_cast<uint32_t>(set.tracks.size());
+  set.tracks.emplace_back(state);
+  return idx;
+}
+
+}  // namespace trace_processor
+}  // namespace perfetto
diff --git a/src/trace_processor/importers/proto/async_track_set_tracker.h b/src/trace_processor/importers/proto/async_track_set_tracker.h
new file mode 100644
index 0000000..92c6311
--- /dev/null
+++ b/src/trace_processor/importers/proto/async_track_set_tracker.h
@@ -0,0 +1,135 @@
+/*
+ * Copyright (C) 2020 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.
+ */
+
+#ifndef SRC_TRACE_PROCESSOR_IMPORTERS_PROTO_ASYNC_TRACK_SET_TRACKER_H_
+#define SRC_TRACE_PROCESSOR_IMPORTERS_PROTO_ASYNC_TRACK_SET_TRACKER_H_
+
+#include <unordered_map>
+
+#include "src/trace_processor/storage/trace_storage.h"
+
+namespace perfetto {
+namespace trace_processor {
+
+class TraceProcessorContext;
+
+// Tracker used to reduce the number of trace processor tracks corresponding
+// to a single "UI track".
+//
+// UIs using trace processor want to display all slices in the same context
+// (e.g. same upid) and same name into a single track. However, because trace
+// processor does not allow parallel slices on a single track (because it breaks
+// things like span join, self time computation etc.), at the trace processor
+// level these parallel slices are put on different tracks.
+//
+// Creating a new track for every event, however, leads to an explosion of
+// tracks which is undesirable. This class exists to multiplex slices so that
+// n events correspond to a single track in a way which minimises the number of
+// tracks which needs to be merged by the UI.
+//
+// The intended usage of this class is for callers to first call one of the
+// Intern* methods to obtain a TrackSetId followed by Begin/End just before
+// calling into SliceTracker's Begin/End respectively. For example:
+//  TrackSetId set_id = track_set_tracker->InternAndroidSet(upid, name);
+//  if (event.begin) {
+//    TrackId id = track_set_tracker->Begin(set_id, cookie);
+//    slice_tracker->Begin(ts, id, ...)
+//  } else {
+//    ... (same thing with end)
+//  }
+class AsyncTrackSetTracker {
+ public:
+  using TrackSetId = uint32_t;
+
+  // Indicates the nesting behaviour of slices associated to a single
+  // cookie.
+  enum class NestingBehaviour {
+    // Indicates that slices are unnestable; that is, it is an error
+    // to call Begin -> Begin with a single cookie without End inbetween.
+    // This pattern should be the default behaviour that most async slices
+    // should use.
+    kUnnestable,
+
+    // Indicates that slices are unnestable but also saturating; that is
+    // calling Begin -> Begin only causes a single Begin to be recorded.
+    // This is only really useful for Android async slices which have this
+    // behaviour for legacy reasons. See the comment in
+    // SystraceParser::ParseSystracePoint for information on why
+    // this behaviour exists.
+    kLegacySaturatingUnnestable,
+  };
+
+  explicit AsyncTrackSetTracker(TraceProcessorContext* context);
+  ~AsyncTrackSetTracker() = default;
+
+  // Interns a set of Android async slice tracks assocaited with the given
+  // upid and name.
+  TrackSetId InternAndroidSet(UniquePid, StringId name);
+
+  // Starts a new slice on the given async track set which has the given cookie
+  // and nesting behaviour.
+  TrackId Begin(TrackSetId id,
+                int64_t cookie,
+                NestingBehaviour = NestingBehaviour::kUnnestable);
+
+  // Ends a new slice on the given async track set which has the given cookie
+  // and nesting behaviour.
+  TrackId End(TrackSetId id, int64_t cookie);
+
+ private:
+  struct AndroidTuple {
+    UniquePid upid;
+    StringId name;
+
+    friend bool operator<(const AndroidTuple& l, const AndroidTuple& r) {
+      return std::tie(l.upid, l.name) < std::tie(r.upid, r.name);
+    }
+  };
+
+  enum class TrackSetType {
+    kAndroid,
+  };
+
+  struct TrackState {
+    TrackId id;
+    int64_t cookie;
+    uint32_t nest_count;
+  };
+
+  struct TrackSet {
+    TrackSetType type;
+    union {
+      // Only set when |type| == |TrackSetType::kAndroid|.
+      AndroidTuple android_tuple;
+    };
+    std::vector<TrackState> tracks;
+  };
+
+  // Gets an empty track (i.e. a track with no event started on it) or creates a
+  // new track in the set if none exists.
+  // Returns the index to the associated TrackState in the |tracks| vector.
+  uint32_t GetOrCreateEmptyTrack(TrackSet& set, int64_t cookie);
+
+  std::map<AndroidTuple, TrackSetId> android_track_set_ids_;
+  std::vector<TrackSet> track_sets_;
+
+  TraceProcessorContext* const context_;
+};
+
+}  // namespace trace_processor
+}  // namespace perfetto
+
+#endif  // SRC_TRACE_PROCESSOR_IMPORTERS_PROTO_ASYNC_TRACK_SET_TRACKER_H_
diff --git a/src/trace_processor/importers/proto/async_track_set_tracker_unittest.cc b/src/trace_processor/importers/proto/async_track_set_tracker_unittest.cc
new file mode 100644
index 0000000..669a7d0
--- /dev/null
+++ b/src/trace_processor/importers/proto/async_track_set_tracker_unittest.cc
@@ -0,0 +1,87 @@
+/*
+ * Copyright (C) 2020 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/async_track_set_tracker.h"
+
+#include "src/trace_processor/importers/common/track_tracker.h"
+#include "src/trace_processor/types/trace_processor_context.h"
+#include "test/gtest_and_gmock.h"
+
+namespace perfetto {
+namespace trace_processor {
+namespace {
+
+class AsyncTrackSetTrackerUnittest {
+ public:
+  AsyncTrackSetTrackerUnittest() {
+    context_.storage.reset(new TraceStorage());
+    context_.track_tracker.reset(new TrackTracker(&context_));
+    context_.async_track_set_tracker.reset(new AsyncTrackSetTracker(&context_));
+
+    storage_ = context_.storage.get();
+    tracker_ = context_.async_track_set_tracker.get();
+  }
+
+ protected:
+  TraceStorage* storage_ = nullptr;
+  AsyncTrackSetTracker* tracker_ = nullptr;
+
+ private:
+  TraceProcessorContext context_;
+};
+
+TEST_F(AsyncTrackSetTrackerUnittest, Smoke) {
+  auto set_id = tracker_->InternAndroidSet(1, storage_->InternString("test"));
+
+  auto begin = tracker_->Begin(
+      set_id, 1,
+      AsyncTrackSetTracker::NestingType::kLegacySaturatingUnnestable);
+  auto end = tracker_->End(set_id, 1);
+
+  ASSERT_EQ(begin, end);
+
+  uint32_t row = *storage_->process_track_table().id().IndexOf(begin);
+  ASSERT_EQ(storage_->process_track_table().upid()[row], 1);
+  ASSERT_EQ(storage_->process_track_table().name()[row],
+            storage_->InternString("test"));
+}
+
+TEST_F(AsyncTrackSetTrackerUnittest, EndFirst) {
+  auto set_id = tracker_->InternAndroidSet(1, storage_->InternString("test"));
+  auto end = tracker_->End(set_id, 1);
+
+  uint32_t row = *storage_->process_track_table().id().IndexOf(end);
+  ASSERT_EQ(storage_->process_track_table().upid()[row], 1);
+  ASSERT_EQ(storage_->process_track_table().name()[row],
+            storage_->InternString("test"));
+}
+
+TEST_F(AsyncTrackSetTrackerUnittest, LegacySaturating) {
+  auto set_id = tracker_->InternAndroidSet(1, storage_->InternString("test"));
+
+  auto begin = tracker_->Begin(
+      set_id, 1,
+      AsyncTrackSetTracker::NestingType::kLegacySaturatingUnnestable);
+  auto begin_2 = tracker_->Begin(
+      set_id, 1,
+      AsyncTrackSetTracker::NestingType::kLegacySaturatingUnnestable);
+
+  ASSERT_EQ(begin, begin_2);
+}
+
+}  // namespace
+}  // namespace trace_processor
+}  // namespace perfetto
diff --git a/src/trace_processor/importers/systrace/systrace_parser.cc b/src/trace_processor/importers/systrace/systrace_parser.cc
index 6e8ac7c..f2370d3 100644
--- a/src/trace_processor/importers/systrace/systrace_parser.cc
+++ b/src/trace_processor/importers/systrace/systrace_parser.cc
@@ -21,6 +21,7 @@
 #include "src/trace_processor/importers/common/process_tracker.h"
 #include "src/trace_processor/importers/common/slice_tracker.h"
 #include "src/trace_processor/importers/common/track_tracker.h"
+#include "src/trace_processor/importers/proto/async_track_set_tracker.h"
 
 namespace perfetto {
 namespace trace_processor {
@@ -28,7 +29,8 @@
 SystraceParser::SystraceParser(TraceProcessorContext* ctx)
     : context_(ctx),
       lmk_id_(ctx->storage->InternString("mem.lmk")),
-      screen_state_id_(ctx->storage->InternString("ScreenState")) {}
+      screen_state_id_(ctx->storage->InternString("ScreenState")),
+      cookie_id_(ctx->storage->InternString("cookie")) {}
 
 SystraceParser::~SystraceParser() = default;
 
@@ -158,8 +160,9 @@
       UniquePid upid =
           context_->process_tracker->GetOrCreateProcess(point.tgid);
 
-      TrackId track_id = context_->track_tracker->InternAndroidAsyncTrack(
-          name_id, upid, cookie);
+      auto track_set_id =
+          context_->async_track_set_tracker->InternAndroidSet(upid, name_id);
+
       if (point.phase == 'S') {
         // Historically, async slices on Android did not support nesting async
         // slices (i.e. you could not have a stack of async slices). If clients
@@ -177,10 +180,18 @@
         // issues. No other code should ever use this method.
         tables::SliceTable::Row row;
         row.ts = ts;
-        row.track_id = track_id;
+        row.track_id = context_->async_track_set_tracker->Begin(
+            track_set_id, cookie,
+            AsyncTrackSetTracker::NestingBehaviour::
+                kLegacySaturatingUnnestable);
         row.name = name_id;
-        context_->slice_tracker->BeginLegacyUnnestable(row);
+        context_->slice_tracker->BeginLegacyUnnestable(
+            row, [this, cookie](ArgsTracker::BoundInserter* inserter) {
+              inserter->AddArg(cookie_id_, Variadic::Integer(cookie));
+            });
       } else {
+        TrackId track_id =
+            context_->async_track_set_tracker->End(track_set_id, cookie);
         context_->slice_tracker->End(ts, track_id);
       }
       break;
diff --git a/src/trace_processor/importers/systrace/systrace_parser.h b/src/trace_processor/importers/systrace/systrace_parser.h
index bc098fd..e55a58e 100644
--- a/src/trace_processor/importers/systrace/systrace_parser.h
+++ b/src/trace_processor/importers/systrace/systrace_parser.h
@@ -254,6 +254,7 @@
   TraceProcessorContext* const context_;
   const StringId lmk_id_;
   const StringId screen_state_id_;
+  const StringId cookie_id_;
 };
 
 }  // namespace trace_processor
diff --git a/src/trace_processor/trace_processor_context.cc b/src/trace_processor/trace_processor_context.cc
index 4187959..e2df18f 100644
--- a/src/trace_processor/trace_processor_context.cc
+++ b/src/trace_processor/trace_processor_context.cc
@@ -27,6 +27,7 @@
 #include "src/trace_processor/importers/common/slice_tracker.h"
 #include "src/trace_processor/importers/common/track_tracker.h"
 #include "src/trace_processor/importers/ftrace/ftrace_module.h"
+#include "src/trace_processor/importers/proto/async_track_set_tracker.h"
 #include "src/trace_processor/importers/proto/heap_profile_tracker.h"
 #include "src/trace_processor/importers/proto/metadata_tracker.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
diff --git a/src/trace_processor/trace_processor_storage_impl.cc b/src/trace_processor/trace_processor_storage_impl.cc
index 9eea8df..ae1f9cc 100644
--- a/src/trace_processor/trace_processor_storage_impl.cc
+++ b/src/trace_processor/trace_processor_storage_impl.cc
@@ -27,6 +27,7 @@
 #include "src/trace_processor/importers/common/track_tracker.h"
 #include "src/trace_processor/importers/default_modules.h"
 #include "src/trace_processor/importers/proto/args_table_utils.h"
+#include "src/trace_processor/importers/proto/async_track_set_tracker.h"
 #include "src/trace_processor/importers/proto/heap_profile_tracker.h"
 #include "src/trace_processor/importers/proto/metadata_tracker.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
@@ -42,6 +43,7 @@
   context_.config = cfg;
   context_.storage.reset(new TraceStorage(context_.config));
   context_.track_tracker.reset(new TrackTracker(&context_));
+  context_.async_track_set_tracker.reset(new AsyncTrackSetTracker(&context_));
   context_.args_tracker.reset(new ArgsTracker(&context_));
   context_.slice_tracker.reset(new SliceTracker(&context_));
   context_.flow_tracker.reset(new FlowTracker(&context_));
diff --git a/src/trace_processor/types/trace_processor_context.h b/src/trace_processor/types/trace_processor_context.h
index 294a804..f2a74b1 100644
--- a/src/trace_processor/types/trace_processor_context.h
+++ b/src/trace_processor/types/trace_processor_context.h
@@ -27,6 +27,7 @@
 namespace trace_processor {
 
 class ArgsTracker;
+class AsyncTrackSetTracker;
 class AndroidProbesTracker;
 class ChunkedTraceReader;
 class ClockTracker;
@@ -68,6 +69,7 @@
   std::unique_ptr<ArgsTracker> args_tracker;
 
   std::unique_ptr<TrackTracker> track_tracker;
+  std::unique_ptr<AsyncTrackSetTracker> async_track_set_tracker;
   std::unique_ptr<SliceTracker> slice_tracker;
   std::unique_ptr<FlowTracker> flow_tracker;
   std::unique_ptr<ProcessTracker> process_tracker;