tp: fix nest count handling for async slices

This CL fixes the quite broken state of nest counts for async slices. We
didn't notice this until now as we only ever used legacy unnestable
slices but as we're now using unnestable slices moving forward, we need
to fix this.

Change-Id: If559f8fcfb3d64078c2ff88301dd2e9950aa36d9
diff --git a/Android.bp b/Android.bp
index eb4a59c..9516297 100644
--- a/Android.bp
+++ b/Android.bp
@@ -7480,6 +7480,7 @@
     "src/trace_processor/importers/memory_tracker/graph_unittest.cc",
     "src/trace_processor/importers/memory_tracker/raw_process_memory_node_unittest.cc",
     "src/trace_processor/importers/proto/args_table_utils_unittest.cc",
+    "src/trace_processor/importers/proto/async_track_set_tracker_unittest.cc",
     "src/trace_processor/importers/proto/heap_graph_tracker_unittest.cc",
     "src/trace_processor/importers/proto/heap_profile_tracker_unittest.cc",
     "src/trace_processor/importers/proto/proto_trace_parser_unittest.cc",
diff --git a/src/trace_processor/BUILD.gn b/src/trace_processor/BUILD.gn
index 2d18f6d..234f159 100644
--- a/src/trace_processor/BUILD.gn
+++ b/src/trace_processor/BUILD.gn
@@ -378,6 +378,7 @@
     "importers/memory_tracker/graph_unittest.cc",
     "importers/memory_tracker/raw_process_memory_node_unittest.cc",
     "importers/proto/args_table_utils_unittest.cc",
+    "importers/proto/async_track_set_tracker_unittest.cc",
     "importers/proto/heap_graph_tracker_unittest.cc",
     "importers/proto/heap_profile_tracker_unittest.cc",
     "importers/proto/proto_trace_parser_unittest.cc",
diff --git a/src/trace_processor/importers/proto/async_track_set_tracker.cc b/src/trace_processor/importers/proto/async_track_set_tracker.cc
index 3b2067d..dc0493b 100644
--- a/src/trace_processor/importers/proto/async_track_set_tracker.cc
+++ b/src/trace_processor/importers/proto/async_track_set_tracker.cc
@@ -50,25 +50,26 @@
   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;
-  }
+  TrackState& state = it == set.tracks.end()
+                          ? set.tracks[GetOrCreateEmptyTrack(set, cookie)]
+                          : *it;
+  PERFETTO_DCHECK(it != set.tracks.end() || state.nest_count == 0);
 
   switch (nest) {
     case NestingBehaviour::kLegacySaturatingUnnestable:
-      PERFETTO_DCHECK(it->nest_count <= 1);
+      PERFETTO_DCHECK(state.nest_count <= 1);
+      state.nest_count = 1;
       break;
     case NestingBehaviour::kUnnestable:
-      PERFETTO_DCHECK(it->nest_count == 0);
+      PERFETTO_DCHECK(state.nest_count == 0);
+      state.nest_count++;
       break;
   }
-  return it->id;
+  return state.id;
 }
 
 TrackId AsyncTrackSetTracker::End(TrackSetId id, int64_t cookie) {
@@ -80,6 +81,16 @@
       [cookie](const TrackState& state) { return state.cookie == cookie; });
   if (it == set.tracks.end())
     return set.tracks[GetOrCreateEmptyTrack(set, cookie)].id;
+
+  // It's possible to have a nest count of 0 even when we know about the track.
+  // Suppose the following sequence of events for some |id| and |cookie|:
+  //   Begin
+  //   (trace starts)
+  //   Begin
+  //   End
+  //   End <- nest count == 0 here even though we have a record of this track.
+  if (it->nest_count > 0)
+    it->nest_count--;
   return it->id;
 }
 
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
index 669a7d0..9056859 100644
--- a/src/trace_processor/importers/proto/async_track_set_tracker_unittest.cc
+++ b/src/trace_processor/importers/proto/async_track_set_tracker_unittest.cc
@@ -16,6 +16,8 @@
 
 #include "src/trace_processor/importers/proto/async_track_set_tracker.h"
 
+#include "src/trace_processor/importers/common/args_tracker.h"
+#include "src/trace_processor/importers/common/global_args_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"
@@ -24,10 +26,12 @@
 namespace trace_processor {
 namespace {
 
-class AsyncTrackSetTrackerUnittest {
+class AsyncTrackSetTrackerUnittest : public testing::Test {
  public:
   AsyncTrackSetTrackerUnittest() {
     context_.storage.reset(new TraceStorage());
+    context_.global_args_tracker.reset(new GlobalArgsTracker(&context_));
+    context_.args_tracker.reset(new ArgsTracker(&context_));
     context_.track_tracker.reset(new TrackTracker(&context_));
     context_.async_track_set_tracker.reset(new AsyncTrackSetTracker(&context_));
 
@@ -48,13 +52,13 @@
 
   auto begin = tracker_->Begin(
       set_id, 1,
-      AsyncTrackSetTracker::NestingType::kLegacySaturatingUnnestable);
+      AsyncTrackSetTracker::NestingBehaviour::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().upid()[row], 1u);
   ASSERT_EQ(storage_->process_track_table().name()[row],
             storage_->InternString("test"));
 }
@@ -64,7 +68,7 @@
   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().upid()[row], 1u);
   ASSERT_EQ(storage_->process_track_table().name()[row],
             storage_->InternString("test"));
 }
@@ -74,14 +78,39 @@
 
   auto begin = tracker_->Begin(
       set_id, 1,
-      AsyncTrackSetTracker::NestingType::kLegacySaturatingUnnestable);
+      AsyncTrackSetTracker::NestingBehaviour::kLegacySaturatingUnnestable);
   auto begin_2 = tracker_->Begin(
       set_id, 1,
-      AsyncTrackSetTracker::NestingType::kLegacySaturatingUnnestable);
+      AsyncTrackSetTracker::NestingBehaviour::kLegacySaturatingUnnestable);
 
   ASSERT_EQ(begin, begin_2);
 }
 
+TEST_F(AsyncTrackSetTrackerUnittest, Unnestable) {
+  auto set_id = tracker_->InternAndroidSet(1, storage_->InternString("test"));
+
+  auto begin = tracker_->Begin(
+      set_id, 1, AsyncTrackSetTracker::NestingBehaviour::kUnnestable);
+  auto end = tracker_->End(set_id, 1);
+  auto begin_2 = tracker_->Begin(
+      set_id, 1, AsyncTrackSetTracker::NestingBehaviour::kUnnestable);
+
+  ASSERT_EQ(begin, end);
+  ASSERT_EQ(begin, begin_2);
+}
+
+TEST_F(AsyncTrackSetTrackerUnittest, UnnestableMultipleEndAfterBegin) {
+  auto set_id = tracker_->InternAndroidSet(1, storage_->InternString("test"));
+
+  auto begin = tracker_->Begin(
+      set_id, 1, AsyncTrackSetTracker::NestingBehaviour::kUnnestable);
+  auto end = tracker_->End(set_id, 1);
+  auto end_2 = tracker_->End(set_id, 1);
+
+  ASSERT_EQ(begin, end);
+  ASSERT_EQ(end, end_2);
+}
+
 }  // namespace
 }  // namespace trace_processor
 }  // namespace perfetto