Merge "Add ability to store arbitrary user data on the TrackEventTlsS…

…tate" into main

include/perfetto/tracing/event_context.h

@@ -119,6 +119,17 @@ class PERFETTO_EXPORT_COMPONENT EventContext {
                          std::forward<T>(value));
   }
 
+  // Read arbitrary user data that is associated with the thread-local per
+  // instance state of the track event. `key` must be non-null and unique
+  // per TrackEventTlsStateUserData subclass.
+  TrackEventTlsStateUserData* GetTlsUserData(const void* key);
+
+  // Set arbitrary user data that is associated with the thread-local per
+  // instance state of the track event. `key` must be non-null and unique
+  // per TrackEventTlsStateUserData subclass.
+  void SetTlsUserData(const void* key,
+                      std::unique_ptr<TrackEventTlsStateUserData> data);
+
  private:
   template <typename, size_t, typename, typename>
   friend class TrackEventInternedDataIndex;
@@ -129,7 +140,7 @@ class PERFETTO_EXPORT_COMPONENT EventContext {
 
   EventContext(TracePacketHandle,
                internal::TrackEventIncrementalState*,
-               const internal::TrackEventTlsState*);
+               internal::TrackEventTlsState*);
   EventContext(const EventContext&) = delete;
 
   protos::pbzero::DebugAnnotation* AddDebugAnnotation(const char* name);
@@ -142,7 +153,7 @@ class PERFETTO_EXPORT_COMPONENT EventContext {
   // TODO(mohitms): Make it const-reference instead of pointer, once we
   // are certain that it cannot be nullptr. Once we switch to client library in
   // chrome, we can make that happen.
-  const internal::TrackEventTlsState* tls_state_ = nullptr;
+  internal::TrackEventTlsState* tls_state_ = nullptr;
   // TODO(kraskevich): Come up with a more precise name once we have more than
   // one usecase.
   // TODO(kraskevich): Remove once Chromium has fully switched to client lib.

include/perfetto/tracing/internal/track_event_data_source.h

@@ -882,7 +882,7 @@ class TrackEventDataSource
     const Category* static_category =
         CatTraits::GetStaticCategory(Registry, category);
 
-    const TrackEventTlsState& tls_state = *ctx.GetCustomTlsState();
+    TrackEventTlsState& tls_state = *ctx.GetCustomTlsState();
     TraceWriterBase* trace_writer = ctx.tls_inst_->trace_writer.get();
     // Make sure incremental state is valid.
     TrackEventIncrementalState* incr_state = ctx.GetIncrementalState();

include/perfetto/tracing/internal/track_event_internal.h

@@ -83,6 +83,18 @@ class PERFETTO_EXPORT_COMPONENT TrackEventSessionObserver {
       const DataSourceBase::ClearIncrementalStateArgs&);
 };
 
+// A class that the embedder can store arbitrary data user data per thread.
+class PERFETTO_EXPORT_COMPONENT TrackEventTlsStateUserData {
+ public:
+  TrackEventTlsStateUserData() = default;
+  // Not clonable.
+  TrackEventTlsStateUserData(const TrackEventTlsStateUserData&) = delete;
+  TrackEventTlsStateUserData& operator=(const TrackEventTlsStateUserData&) =
+      delete;
+
+  virtual ~TrackEventTlsStateUserData();
+};
+
 namespace internal {
 class TrackEventCategoryRegistry;
 
@@ -104,6 +116,7 @@ struct TrackEventTlsState {
   bool filter_dynamic_event_names = false;
   uint64_t timestamp_unit_multiplier = 1;
   uint32_t default_clock;
+  std::map<const void*, std::unique_ptr<TrackEventTlsStateUserData>> user_data;
 };
 
 struct TrackEventIncrementalState {
@@ -206,7 +219,7 @@ class PERFETTO_EXPORT_COMPONENT TrackEventInternal {
   static perfetto::EventContext WriteEvent(
       TraceWriterBase*,
       TrackEventIncrementalState*,
-      const TrackEventTlsState& tls_state,
+      TrackEventTlsState& tls_state,
       const Category* category,
       perfetto::protos::pbzero::TrackEvent::Type,
       const TraceTimestamp& timestamp,

src/tracing/event_context.cc

@@ -25,7 +25,7 @@ namespace perfetto {
 EventContext::EventContext(
     EventContext::TracePacketHandle trace_packet,
     internal::TrackEventIncrementalState* incremental_state,
-    const internal::TrackEventTlsState* tls_state)
+    internal::TrackEventTlsState* tls_state)
     : trace_packet_(std::move(trace_packet)),
       event_(trace_packet_->set_track_event()),
       incremental_state_(incremental_state),
@@ -68,4 +68,22 @@ protos::pbzero::DebugAnnotation* EventContext::AddDebugAnnotation(
   return annotation;
 }
 
+TrackEventTlsStateUserData* EventContext::GetTlsUserData(const void* key) {
+  PERFETTO_CHECK(tls_state_);
+  PERFETTO_CHECK(key);
+  auto it = tls_state_->user_data.find(key);
+  if (it != tls_state_->user_data.end()) {
+    return it->second.get();
+  }
+  return nullptr;
+}
+
+void EventContext::SetTlsUserData(
+    const void* key,
+    std::unique_ptr<TrackEventTlsStateUserData> data) {
+  PERFETTO_CHECK(tls_state_);
+  PERFETTO_CHECK(key);
+  tls_state_->user_data[key] = std::move(data);
+}
+
 }  // namespace perfetto

src/tracing/internal/track_event_internal.cc

@@ -42,6 +42,8 @@ void TrackEventSessionObserver::OnStop(const DataSourceBase::StopArgs&) {}
 void TrackEventSessionObserver::WillClearIncrementalState(
     const DataSourceBase::ClearIncrementalStateArgs&) {}
 
+TrackEventTlsStateUserData::~TrackEventTlsStateUserData() = default;
+
 namespace internal {
 
 BaseTrackEventInternedDataIndex::~BaseTrackEventInternedDataIndex() = default;
@@ -522,7 +524,7 @@ void TrackEventInternal::WriteEventName(perfetto::DynamicString event_name,
 EventContext TrackEventInternal::WriteEvent(
     TraceWriterBase* trace_writer,
     TrackEventIncrementalState* incr_state,
-    const TrackEventTlsState& tls_state,
+    TrackEventTlsState& tls_state,
     const Category* category,
     perfetto::protos::pbzero::TrackEvent::Type type,
     const TraceTimestamp& timestamp,

src/tracing/test/api_integrationtest.cc

@@ -4208,6 +4208,58 @@ TEST_P(PerfettoApiTest, CustomIncrementalState) {
       ClearDataSourceTlsStateOnReset<TestIncrementalDataSource>();
 }
 
+const void* const kKey1 = &kKey1;
+const void* const kKey2 = &kKey2;
+
+TEST_P(PerfettoApiTest, TrackEventUserData) {
+  // Create a new trace session.
+  auto* tracing_session = NewTraceWithCategories({"foo"});
+  tracing_session->get()->StartBlocking();
+  perfetto::TrackEventTlsStateUserData* data_1_ptr = nullptr;
+  perfetto::TrackEventTlsStateUserData* data_2_ptr = nullptr;
+
+  TRACE_EVENT_BEGIN(
+      "foo", "E", [&data_1_ptr, &data_2_ptr](perfetto::EventContext& ctx) {
+        EXPECT_EQ(nullptr, ctx.GetTlsUserData(kKey1));
+        EXPECT_EQ(nullptr, ctx.GetTlsUserData(kKey2));
+        std::unique_ptr<perfetto::TrackEventTlsStateUserData> data_1 =
+            std::make_unique<perfetto::TrackEventTlsStateUserData>();
+        data_1_ptr = data_1.get();
+        std::unique_ptr<perfetto::TrackEventTlsStateUserData> data_2 =
+            std::make_unique<perfetto::TrackEventTlsStateUserData>();
+        data_2_ptr = data_2.get();
+        ctx.SetTlsUserData(kKey1, std::move(data_1));
+        ctx.SetTlsUserData(kKey2, std::move(data_2));
+        EXPECT_EQ(data_1_ptr, ctx.GetTlsUserData(kKey1));
+        EXPECT_EQ(data_2_ptr, ctx.GetTlsUserData(kKey2));
+      });
+  TRACE_EVENT_END("foo");
+  TRACE_EVENT_BEGIN("foo", "F",
+                    [&data_1_ptr, &data_2_ptr](perfetto::EventContext& ctx) {
+                      EXPECT_EQ(data_1_ptr, ctx.GetTlsUserData(kKey1));
+                      EXPECT_EQ(data_2_ptr, ctx.GetTlsUserData(kKey2));
+                    });
+  TRACE_EVENT_END("foo");
+
+  std::vector<char> raw_trace = StopSessionAndReturnBytes(tracing_session);
+
+  EXPECT_THAT(ReadSlicesFromTrace(raw_trace),
+              ElementsAre("B:foo.E", "E", "B:foo.F", "E"));
+
+  // Expect that the TLS User Data is cleared between tracing sessions.
+  tracing_session = NewTraceWithCategories({"foo"});
+  tracing_session->get()->StartBlocking();
+
+  TRACE_EVENT_BEGIN("foo", "E", [](perfetto::EventContext& ctx) {
+    EXPECT_EQ(nullptr, ctx.GetTlsUserData(kKey1));
+    EXPECT_EQ(nullptr, ctx.GetTlsUserData(kKey2));
+  });
+  TRACE_EVENT_END("foo");
+
+  raw_trace = StopSessionAndReturnBytes(tracing_session);
+  EXPECT_THAT(ReadSlicesFromTrace(raw_trace), ElementsAre("B:foo.E", "E"));
+}
+
 TEST_P(PerfettoApiTest, OnFlush) {
   auto* data_source = &data_sources_["my_data_source"];