From 3a0e47aa37df57fe5575f9dc23aa245ebf67e90e Mon Sep 17 00:00:00 2001
From: Zimuzo Ezeozue <zezeozue@google.com>
Date: Sun, 8 Dec 2024 18:29:44 +0000
Subject: [PATCH] Add utid args to async atrace slices

- For begin and instant events, added a utid arg.
- For end events, added an end_utid arg.

Change-Id: I41a28a00115cf31364d9b18cf42028f06c2ba394
Test: tools/diff_test_trace_processor.py out/android/trace_processor_shell
---
 .../importers/systrace/systrace_parser.cc     | 32 +++++++++++++++----
 .../importers/systrace/systrace_parser.h      |  2 ++
 .../diff_tests/tables/tests.py                | 13 ++++++++
 3 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/src/trace_processor/importers/systrace/systrace_parser.cc b/src/trace_processor/importers/systrace/systrace_parser.cc
index e7a2dab698..477ef7158f 100644
--- a/src/trace_processor/importers/systrace/systrace_parser.cc
+++ b/src/trace_processor/importers/systrace/systrace_parser.cc
@@ -41,7 +41,9 @@ namespace perfetto::trace_processor {
 SystraceParser::SystraceParser(TraceProcessorContext* ctx)
     : context_(ctx),
       lmk_id_(ctx->storage->InternString("mem.lmk")),
-      cookie_id_(ctx->storage->InternString("cookie")) {}
+      cookie_id_(ctx->storage->InternString("cookie")),
+      utid_id_(ctx->storage->InternString("utid")),
+      end_utid_id_(ctx->storage->InternString("end_utid")) {}
 
 SystraceParser::~SystraceParser() = default;
 
@@ -248,16 +250,33 @@ void SystraceParser::ParseSystracePoint(
       if (point.phase == 'N') {
         TrackId track_id =
             context_->async_track_set_tracker->Scoped(track_set_id, ts, 0);
-        context_->slice_tracker->Scoped(ts, track_id, kNullStringId, name_id,
-                                        0);
+        auto utid = context_->process_tracker->GetOrCreateThread(pid);
+        context_->slice_tracker->Scoped(
+            ts, track_id, kNullStringId, name_id, 0,
+            [this, utid](ArgsTracker::BoundInserter* inserter) {
+              inserter->AddArg(utid_id_, Variadic::UnsignedInteger(utid),
+                               ArgsTracker::UpdatePolicy::kSkipIfExists);
+            });
       } else if (point.phase == 'G') {
         TrackId track_id = context_->async_track_set_tracker->Begin(
             track_set_id, point.int_value);
-        context_->slice_tracker->Begin(ts, track_id, kNullStringId, name_id);
+        auto utid = context_->process_tracker->GetOrCreateThread(pid);
+        context_->slice_tracker->Begin(
+            ts, track_id, kNullStringId, name_id,
+            [this, utid](ArgsTracker::BoundInserter* inserter) {
+              inserter->AddArg(utid_id_, Variadic::UnsignedInteger(utid),
+                               ArgsTracker::UpdatePolicy::kSkipIfExists);
+            });
       } else if (point.phase == 'H') {
         TrackId track_id = context_->async_track_set_tracker->End(
             track_set_id, point.int_value);
-        context_->slice_tracker->End(ts, track_id);
+        auto utid = context_->process_tracker->GetOrCreateThread(pid);
+        context_->slice_tracker->End(
+            ts, track_id, {}, {},
+            [this, utid](ArgsTracker::BoundInserter* inserter) {
+              inserter->AddArg(end_utid_id_, Variadic::UnsignedInteger(utid),
+                               ArgsTracker::UpdatePolicy::kSkipIfExists);
+            });
       }
       break;
     }
@@ -313,8 +332,7 @@ void SystraceParser::ParseSystracePoint(
           ts, static_cast<double>(point.int_value), track_id,
           [this, opt_utid](ArgsTracker::BoundInserter* inserter) {
             if (opt_utid) {
-              inserter->AddArg(context_->storage->InternString("utid"),
-                               Variadic::UnsignedInteger(*opt_utid),
+              inserter->AddArg(utid_id_, Variadic::UnsignedInteger(*opt_utid),
                                ArgsTracker::UpdatePolicy::kSkipIfExists);
             }
           });
diff --git a/src/trace_processor/importers/systrace/systrace_parser.h b/src/trace_processor/importers/systrace/systrace_parser.h
index 380abb16a3..72cbaa35ed 100644
--- a/src/trace_processor/importers/systrace/systrace_parser.h
+++ b/src/trace_processor/importers/systrace/systrace_parser.h
@@ -322,6 +322,8 @@ class SystraceParser : public Destructible {
   TraceProcessorContext* const context_;
   const StringId lmk_id_;
   const StringId cookie_id_;
+  const StringId utid_id_;
+  const StringId end_utid_id_;
 };
 
 }  // namespace trace_processor
diff --git a/test/trace_processor/diff_tests/tables/tests.py b/test/trace_processor/diff_tests/tables/tests.py
index e0007fef26..2f6fb9ec2f 100644
--- a/test/trace_processor/diff_tests/tables/tests.py
+++ b/test/trace_processor/diff_tests/tables/tests.py
@@ -487,3 +487,16 @@ def test_cpu_track_table_machine_id(self):
         "__intrinsic_cpu_track",4096,0,1
         "__intrinsic_cpu_track",4097,1,1
         """))
+
+  def test_async_slice_utid_arg_set_id(self):
+    return DiffTestBlueprint(
+        trace=DataPath('android_monitor_contention_trace.atr'),
+        query="""
+        SELECT COUNT(DISTINCT extract_arg(arg_set_id, 'utid')) AS utid_count,
+        COUNT(DISTINCT extract_arg(arg_set_id, 'end_utid')) AS end_utid_count
+        FROM counter
+        """,
+        out=Csv("""
+        "utid_count","end_utid_count"
+        89,0
+        """))