Add TRACE_EVENT_BEGIN/TRACE_EVENT_END functions that allow typed arguments.

These macros allow users to write small lambda functions that get direct access
to the protozero TrackEvent. They can then access and fill in the values as
needed.

This seems to have no noticeable performance impact:
https://docs.google.com/spreadsheets/d/1YBB7V2hinQ-lcf9WBFrdMzxakosJHou0B3U10Gzpme0/edit?usp=sharing

And binary size is slightly negative (probably more zero like).

Change-Id: I2e9266daad902ac64bdaa1a1f1bf5678dc48ac9f
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1921984
Commit-Queue: Stephen Nusko <nuskos@chromium.org>
Reviewed-by: Eric Seckler <eseckler@chromium.org>
Reviewed-by: Sami Kyöstilä <skyostil@chromium.org>
Auto-Submit: Stephen Nusko <nuskos@chromium.org>
Cr-Commit-Position: refs/heads/master@{#718914}

base/trace_event/trace_log.cc

@@ -1081,6 +1081,69 @@ void TraceLog::UseNextTraceBuffer() {
   thread_shared_chunk_index_ = 0;
 }
 
+bool TraceLog::ShouldAddAfterUpdatingState(
+    char phase,
+    const unsigned char* category_group_enabled,
+    const char* name,
+    unsigned long long id,
+    int thread_id,
+    TraceArguments* args) {
+  if (!*category_group_enabled)
+    return false;
+
+  // Avoid re-entrance of AddTraceEvent. This may happen in GPU process when
+  // ECHO_TO_CONSOLE is enabled: AddTraceEvent -> LOG(ERROR) ->
+  // GpuProcessLogMessageHandler -> PostPendingTask -> TRACE_EVENT ...
+  if (thread_is_in_trace_event_.Get())
+    return false;
+
+  DCHECK(name);
+
+  // Check and update the current thread name only if the event is for the
+  // current thread to avoid locks in most cases.
+  if (thread_id == static_cast<int>(PlatformThread::CurrentId())) {
+    const char* new_name =
+        ThreadIdNameManager::GetInstance()->GetNameForCurrentThread();
+    // Check if the thread name has been set or changed since the previous
+    // call (if any), but don't bother if the new name is empty. Note this will
+    // not detect a thread name change within the same char* buffer address: we
+    // favor common case performance over corner case correctness.
+    static auto* current_thread_name = new ThreadLocalPointer<const char>();
+    if (new_name != current_thread_name->Get() && new_name && *new_name) {
+      current_thread_name->Set(new_name);
+
+      AutoLock thread_info_lock(thread_info_lock_);
+
+      auto existing_name = thread_names_.find(thread_id);
+      if (existing_name == thread_names_.end()) {
+        // This is a new thread id, and a new name.
+        thread_names_[thread_id] = new_name;
+      } else {
+        // This is a thread id that we've seen before, but potentially with a
+        // new name.
+        std::vector<StringPiece> existing_names = base::SplitStringPiece(
+            existing_name->second, ",", base::KEEP_WHITESPACE,
+            base::SPLIT_WANT_NONEMPTY);
+        if (!Contains(existing_names, new_name)) {
+          if (!existing_names.empty())
+            existing_name->second.push_back(',');
+          existing_name->second.append(new_name);
+        }
+      }
+    }
+  }
+
+#if defined(OS_WIN)
+  // This is done sooner rather than later, to avoid creating the event and
+  // acquiring the lock, which is not needed for ETW as it's already threadsafe.
+  if (*category_group_enabled & TraceCategory::ENABLED_FOR_ETW_EXPORT) {
+    TraceEventETWExport::AddEvent(phase, category_group_enabled, name, id,
+                                  args);
+  }
+#endif  // OS_WIN
+  return true;
+}
+
 TraceEventHandle TraceLog::AddTraceEvent(
     char phase,
     const unsigned char* category_group_enabled,
@@ -1159,20 +1222,14 @@ TraceEventHandle TraceLog::AddTraceEventWithThreadIdAndTimestamp(
     TraceArguments* args,
     unsigned int flags) {
   TraceEventHandle handle = {0, 0, 0};
-  if (!*category_group_enabled)
-    return handle;
-
-  // Avoid re-entrance of AddTraceEvent. This may happen in GPU process when
-  // ECHO_TO_CONSOLE is enabled: AddTraceEvent -> LOG(ERROR) ->
-  // GpuProcessLogMessageHandler -> PostPendingTask -> TRACE_EVENT ...
-  if (thread_is_in_trace_event_.Get())
+  if (!ShouldAddAfterUpdatingState(phase, category_group_enabled, name, id,
+                                   thread_id, args)) {
     return handle;
+  }
+  DCHECK(!timestamp.is_null());
 
   AutoThreadLocalBoolean thread_is_in_trace_event(&thread_is_in_trace_event_);
 
-  DCHECK(name);
-  DCHECK(!timestamp.is_null());
-
   if (flags & TRACE_EVENT_FLAG_MANGLE_ID) {
     if ((flags & TRACE_EVENT_FLAG_FLOW_IN) ||
         (flags & TRACE_EVENT_FLAG_FLOW_OUT))
@@ -1192,48 +1249,6 @@ TraceEventHandle TraceLog::AddTraceEventWithThreadIdAndTimestamp(
     thread_local_event_buffer = thread_local_event_buffer_.Get();
   }
 
-  // Check and update the current thread name only if the event is for the
-  // current thread to avoid locks in most cases.
-  if (thread_id == static_cast<int>(PlatformThread::CurrentId())) {
-    const char* new_name =
-        ThreadIdNameManager::GetInstance()->GetNameForCurrentThread();
-    // Check if the thread name has been set or changed since the previous
-    // call (if any), but don't bother if the new name is empty. Note this will
-    // not detect a thread name change within the same char* buffer address: we
-    // favor common case performance over corner case correctness.
-    static auto* current_thread_name = new ThreadLocalPointer<const char>();
-    if (new_name != current_thread_name->Get() && new_name && *new_name) {
-      current_thread_name->Set(new_name);
-
-      AutoLock thread_info_lock(thread_info_lock_);
-
-      auto existing_name = thread_names_.find(thread_id);
-      if (existing_name == thread_names_.end()) {
-        // This is a new thread id, and a new name.
-        thread_names_[thread_id] = new_name;
-      } else {
-        // This is a thread id that we've seen before, but potentially with a
-        // new name.
-        std::vector<StringPiece> existing_names = base::SplitStringPiece(
-            existing_name->second, ",", base::KEEP_WHITESPACE,
-            base::SPLIT_WANT_NONEMPTY);
-        if (!Contains(existing_names, new_name)) {
-          if (!existing_names.empty())
-            existing_name->second.push_back(',');
-          existing_name->second.append(new_name);
-        }
-      }
-    }
-  }
-
-#if defined(OS_WIN)
-  // This is done sooner rather than later, to avoid creating the event and
-  // acquiring the lock, which is not needed for ETW as it's already threadsafe.
-  if (*category_group_enabled & TraceCategory::ENABLED_FOR_ETW_EXPORT)
-    TraceEventETWExport::AddEvent(phase, category_group_enabled, name, id,
-                                  args);
-#endif  // OS_WIN
-
   if (*category_group_enabled & RECORDING_MODE) {
     auto trace_event_override =
         add_trace_event_override_.load(std::memory_order_relaxed);

base/trace_event/trace_log.h

@@ -231,6 +231,12 @@ class BASE_EXPORT TraceLog : public MemoryDumpProvider {
   // Called by TRACE_EVENT* macros, don't call this directly.
   // If |copy| is set, |name|, |arg_name1| and |arg_name2| will be deep copied
   // into the event; see "Memory scoping note" and TRACE_EVENT_COPY_XXX above.
+  bool ShouldAddAfterUpdatingState(char phase,
+                                   const unsigned char* category_group_enabled,
+                                   const char* name,
+                                   unsigned long long id,
+                                   int thread_id,
+                                   TraceArguments* args);
   TraceEventHandle AddTraceEvent(char phase,
                                  const unsigned char* category_group_enabled,
                                  const char* name,

services/tracing/public/cpp/BUILD.gn

@@ -33,6 +33,7 @@ if (!is_nacl && !is_ios) {
       "base_agent.h",
       "perfetto/dummy_producer.cc",
       "perfetto/dummy_producer.h",
+      "perfetto/event_context.h",
       "perfetto/interning_index.h",
       "perfetto/java_heap_profiler/hprof_buffer_android.cc",
       "perfetto/java_heap_profiler/hprof_buffer_android.h",
@@ -43,6 +44,9 @@ if (!is_nacl && !is_ios) {
       "perfetto/java_heap_profiler/hprof_parser_android.h",
       "perfetto/java_heap_profiler/java_heap_profiler_android.cc",
       "perfetto/java_heap_profiler/java_heap_profiler_android.h",
+      "perfetto/macros.h",
+      "perfetto/macros_internal.cc",
+      "perfetto/macros_internal.h",
       "perfetto/perfetto_config.cc",
       "perfetto/perfetto_config.h",
       "perfetto/perfetto_producer.cc",
@@ -58,8 +62,6 @@ if (!is_nacl && !is_ios) {
       "perfetto/system_trace_writer.h",
       "perfetto/task_runner.cc",
       "perfetto/task_runner.h",
-      "perfetto/thread_local_event_sink.cc",
-      "perfetto/thread_local_event_sink.h",
       "perfetto/trace_event_data_source.cc",
       "perfetto/trace_event_data_source.h",
       "perfetto/traced_value_proto_writer.cc",

services/tracing/public/cpp/perfetto/event_context.h

+// Copyright 2019 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef SERVICES_TRACING_PUBLIC_CPP_PERFETTO_EVENT_CONTEXT_H_
+#define SERVICES_TRACING_PUBLIC_CPP_PERFETTO_EVENT_CONTEXT_H_
+
+#include "third_party/perfetto/protos/perfetto/trace/trace_packet.pbzero.h"
+
+namespace perfetto {
+
+// An EventContext allows use of the TRACE_EVENT_BEGIN/TRACE_EVENT_END
+// functions to access the protozero TrackEvent to fill in typed arguments.
+//
+// TODO(nuskos): Switch chrome over to the perfetto client library's
+// EventContext.
+class EventContext {
+ public:
+  EventContext(protos::pbzero::TrackEvent* event) : event_(event) {}
+
+  protos::pbzero::TrackEvent* event() const { return event_; }
+
+ private:
+  protos::pbzero::TrackEvent* event_;
+};
+
+}  // namespace perfetto
+
+#endif  // SERVICES_TRACING_PUBLIC_CPP_PERFETTO_EVENT_CONTEXT_H_

services/tracing/public/cpp/perfetto/macros.h

+// Copyright 2019 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef SERVICES_TRACING_PUBLIC_CPP_PERFETTO_MACROS_H_
+#define SERVICES_TRACING_PUBLIC_CPP_PERFETTO_MACROS_H_
+
+#include "base/trace_event/trace_event.h"
+#include "services/tracing/public/cpp/perfetto/macros_internal.h"
+
+#if defined(TRACE_EVENT_BEGIN)
+#error "Another copy of perfetto tracing macros have been included"
+#endif
+
+// TODO(nuskos): This whole file should be removed and migrated to using the
+// Perfetto client library macros. However currently chromium isn't set up to
+// use the client library so in the meantime we've added support for features we
+// want sooner.
+
+namespace tracing {
+// The perfetto client library does not use event names for
+// TRACE_EVENT_PHASE_END. However currently Chrome expects all TraceEvents to
+// have event names. So until we move over to the client library we will use
+// this for all TRACE_EVENT_PHASE_END typed arguments.
+//
+// TODO(nuskos): remove this constant.
+constexpr char kTraceEventEndName[] = "";
+}  // namespace tracing
+
+// Begin a thread-scoped slice under |category| with the title |name|. Both
+// strings must be static constants. The track event is only recorded if
+// |category| is enabled for a tracing session.
+//
+// TODO(nuskos): Give a simple example once we have a typed event that doesn't
+// need interning.
+//   TRACE_EVENT_BEGIN("log", "LogMessage",
+//       [](perfetto::EventContext ctx) {
+//           auto* event = ctx.event();
+//           // Fill in some field in track_event.
+//       });
+#define TRACE_EVENT_BEGIN(category, name, ...)                              \
+  TRACING_INTERNAL_ADD_TRACE_EVENT(TRACE_EVENT_PHASE_BEGIN, category, name, \
+                                   ##__VA_ARGS__)
+
+// End a thread-scoped slice under |category|.
+#define TRACE_EVENT_END(category, ...)                              \
+  TRACING_INTERNAL_ADD_TRACE_EVENT(TRACE_EVENT_PHASE_END, category, \
+                                   kTraceEventEndName, ##__VA_ARGS__)
+
+#endif  // SERVICES_TRACING_PUBLIC_CPP_PERFETTO_MACROS_H_

services/tracing/public/cpp/perfetto/macros_internal.cc

+// Copyright 2019 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "services/tracing/public/cpp/perfetto/macros_internal.h"
+
+#include "base/trace_event/thread_instruction_count.h"
+
+namespace tracing {
+namespace internal {
+namespace {
+
+base::ThreadTicks ThreadNow() {
+  return base::ThreadTicks::IsSupported()
+             ? base::subtle::ThreadTicksNowIgnoringOverride()
+             : base::ThreadTicks();
+}
+
+base::trace_event::ThreadInstructionCount ThreadInstructionNow() {
+  return base::trace_event::ThreadInstructionCount::IsSupported()
+             ? base::trace_event::ThreadInstructionCount::Now()
+             : base::trace_event::ThreadInstructionCount();
+}
+
+}  // namespace
+
+base::Optional<base::trace_event::TraceEvent> CreateTraceEvent(
+    char phase,
+    const unsigned char* category_group_enabled,
+    const char* name) {
+  DCHECK(phase == TRACE_EVENT_PHASE_BEGIN || phase == TRACE_EVENT_PHASE_END);
+  DCHECK(category_group_enabled);
+  const int thread_id = static_cast<int>(base::PlatformThread::CurrentId());
+  auto* trace_log = base::trace_event::TraceLog::GetInstance();
+  DCHECK(trace_log);
+  if (!trace_log->ShouldAddAfterUpdatingState(phase, category_group_enabled,
+                                              name, trace_event_internal::kNoId,
+                                              thread_id, nullptr)) {
+    return base::nullopt;
+  }
+
+  base::TimeTicks now = TRACE_TIME_TICKS_NOW();
+  base::TimeTicks offset_event_timestamp = now - base::TimeDelta();
+  base::ThreadTicks thread_now = ThreadNow();
+  base::trace_event::ThreadInstructionCount thread_instruction_now =
+      ThreadInstructionNow();
+
+  return base::trace_event::TraceEvent(
+      thread_id, offset_event_timestamp, thread_now, thread_instruction_now,
+      phase, category_group_enabled, name, trace_event_internal::kGlobalScope,
+      trace_event_internal::kNoId, trace_event_internal::kNoId, nullptr,
+      TRACE_EVENT_FLAG_NONE);
+}
+
+}  // namespace internal
+}  // namespace tracing

services/tracing/public/cpp/perfetto/macros_internal.h

+// Copyright 2019 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef SERVICES_TRACING_PUBLIC_CPP_PERFETTO_MACROS_INTERNAL_H_
+#define SERVICES_TRACING_PUBLIC_CPP_PERFETTO_MACROS_INTERNAL_H_
+
+#include "base/component_export.h"
+#include "base/trace_event/trace_event.h"
+#include "services/tracing/public/cpp/perfetto/trace_event_data_source.h"
+
+namespace perfetto {
+class TrackEventContext;
+}
+
+namespace tracing {
+namespace internal {
+base::Optional<base::trace_event::TraceEvent> COMPONENT_EXPORT(TRACING_CPP)
+    CreateTraceEvent(char phase,
+                     const unsigned char* category_group_enabled,
+                     const char* name);
+
+// A simple function that will add the TraceEvent requested and will call the
+// |argument_func| after the track_event has been filled in.
+template <
+    typename TrackEventArgumentFunction = void (*)(perfetto::TrackEventContext)>
+static inline base::trace_event::TraceEventHandle AddTraceEvent(
+    char phase,
+    const unsigned char* category_group_enabled,
+    const char* name,
+    TrackEventArgumentFunction argument_func) {
+  base::trace_event::TraceEventHandle handle = {0, 0, 0};
+  auto maybe_event = CreateTraceEvent(phase, category_group_enabled, name);
+  if (!maybe_event) {
+    return handle;
+  }
+  TraceEventDataSource::OnAddTraceEvent(&maybe_event.value(),
+                                        /* thread_will_flush = */ false,
+                                        &handle, std::move(argument_func));
+  return handle;
+}
+}  // namespace internal
+}  // namespace tracing
+
+// Should not be called directly. Used by macros in
+// //services/tracing/perfetto/macros.h.
+#define TRACING_INTERNAL_ADD_TRACE_EVENT(phase, category, name, ...)     \
+  do {                                                                   \
+    INTERNAL_TRACE_EVENT_GET_CATEGORY_INFO(category);                    \
+    if (INTERNAL_TRACE_EVENT_CATEGORY_GROUP_ENABLED()) {                 \
+      tracing::internal::AddTraceEvent(                                  \
+          phase, INTERNAL_TRACE_EVENT_UID(category_group_enabled), name, \
+          ##__VA_ARGS__);                                                \
+    }                                                                    \
+  } while (false)
+
+#endif  // SERVICES_TRACING_PUBLIC_CPP_PERFETTO_MACROS_INTERNAL_H_

services/tracing/public/cpp/perfetto/thread_local_event_sink.cc

-// Copyright 2019 The Chromium Authors. All rights reserved.
-// Use of this source code is governed by a BSD-style license that can be
-// found in the LICENSE file.
-
-#include "services/tracing/public/cpp/perfetto/thread_local_event_sink.h"
-
-#include <atomic>
-#include <utility>
-
-#include "services/tracing/public/cpp/perfetto/trace_event_data_source.h"
-#include "third_party/perfetto/include/perfetto/ext/tracing/core/startup_trace_writer.h"
-
-namespace tracing {
-
-ThreadLocalEventSink::ThreadLocalEventSink(
-    std::unique_ptr<perfetto::StartupTraceWriter> trace_writer,
-    uint32_t session_id,
-    bool disable_interning)
-    : trace_writer_(std::move(trace_writer)),
-      session_id_(session_id),
-      disable_interning_(disable_interning) {
-  static std::atomic<uint32_t> g_sink_id_counter{0};
-  sink_id_ = ++g_sink_id_counter;
-}
-
-ThreadLocalEventSink::~ThreadLocalEventSink() {
-  // Subclass has already destroyed its message handles at this point.
-  TraceEventDataSource::GetInstance()->ReturnTraceWriter(
-      std::move(trace_writer_));
-}
-
-}  // namespace tracing

services/tracing/public/cpp/perfetto/thread_local_event_sink.h

-// Copyright 2019 The Chromium Authors. All rights reserved.
-// Use of this source code is governed by a BSD-style license that can be
-// found in the LICENSE file.
-
-#ifndef SERVICES_TRACING_PUBLIC_CPP_PERFETTO_THREAD_LOCAL_EVENT_SINK_H_
-#define SERVICES_TRACING_PUBLIC_CPP_PERFETTO_THREAD_LOCAL_EVENT_SINK_H_
-
-#include <memory>
-
-#include "base/component_export.h"
-#include "base/macros.h"
-#include "base/time/time.h"
-#include "base/trace_event/trace_event.h"
-
-namespace perfetto {
-class StartupTraceWriter;
-}
-
-namespace tracing {
-
-// Base class for a per-thread TraceEvent sink, which writes TraceEvents into
-// TracePacket protos in the SMB.
-class COMPONENT_EXPORT(TRACING_CPP) ThreadLocalEventSink {
- public:
-  ThreadLocalEventSink(
-      std::unique_ptr<perfetto::StartupTraceWriter> trace_writer,
-      uint32_t session_id,
-      bool disable_interning);
-
-  virtual ~ThreadLocalEventSink();
-
-  // TODO(oysteine): Adding trace events to Perfetto will stall in some
-  // situations, specifically when we overflow the buffer and need to make a
-  // sync call to flush it, and we're running on the same thread as the service.
-  // The short-term fix (while we're behind a flag) is to run the service on its
-  // own thread, the longer term fix is most likely to not go via Mojo in that
-  // specific case.
-  virtual void AddTraceEvent(base::trace_event::TraceEvent* trace_event,
-                             base::trace_event::TraceEventHandle* handle) = 0;
-
-  virtual void UpdateDuration(
-      base::trace_event::TraceEventHandle handle,
-      const base::TimeTicks& now,
-      const base::ThreadTicks& thread_now,
-      base::trace_event::ThreadInstructionCount thread_instruction_now) = 0;
-
-  virtual void Flush() = 0;
-
-  uint32_t session_id() const { return session_id_; }
-
- protected:
-  std::unique_ptr<perfetto::StartupTraceWriter> trace_writer_;
-  uint32_t session_id_;
-  bool disable_interning_;
-  uint32_t sink_id_;
-
- private:
-  DISALLOW_COPY_AND_ASSIGN(ThreadLocalEventSink);
-};
-
-}  // namespace tracing
-
-#endif  // SERVICES_TRACING_PUBLIC_CPP_PERFETTO_THREAD_LOCAL_EVENT_SINK_H_

services/tracing/public/cpp/perfetto/trace_event_data_source.cc

@@ -35,7 +35,6 @@
 #include "components/tracing/common/tracing_switches.h"
 #include "services/tracing/public/cpp/perfetto/perfetto_producer.h"
 #include "services/tracing/public/cpp/perfetto/perfetto_traced_process.h"
-#include "services/tracing/public/cpp/perfetto/thread_local_event_sink.h"
 #include "services/tracing/public/cpp/perfetto/traced_value_proto_writer.h"
 #include "services/tracing/public/cpp/perfetto/track_event_thread_local_event_sink.h"
 #include "services/tracing/public/cpp/trace_event_args_whitelist.h"
@@ -403,7 +402,7 @@ base::ThreadLocalStorage::Slot* ThreadLocalEventSinkSlot() {
       thread_local_event_sink_tls([](void* event_sink) {
         AutoThreadLocalBoolean thread_is_in_trace_event(
             TraceEventDataSource::GetThreadIsInTraceEventTLS());
-        delete static_cast<ThreadLocalEventSink*>(event_sink);
+        delete static_cast<TrackEventThreadLocalEventSink*>(event_sink);
       });
 
   return thread_local_event_sink_tls.get();
@@ -470,6 +469,53 @@ void TraceEventDataSource::UnregisterFromTraceLog() {
   DCHECK(!flush_complete_task_);
 }
 
+// static
+TrackEventThreadLocalEventSink* TraceEventDataSource::GetOrPrepareEventSink(
+    bool thread_will_flush) {
+  // Avoid re-entrancy, which can happen during PostTasks (the taskqueue can
+  // emit trace events). We discard the events in this case, as any PostTasking
+  // to deal with these events later would break the event ordering that the
+  // JSON traces rely on to merge 'A'/'B' events, as well as having to deal with
+  // updating duration of 'X' events which haven't been added yet.
+  if (GetThreadIsInTraceEventTLS()->Get()) {
+    return nullptr;
+  }
+
+  AutoThreadLocalBoolean thread_is_in_trace_event(GetThreadIsInTraceEventTLS());
+
+  auto* thread_local_event_sink = static_cast<TrackEventThreadLocalEventSink*>(
+      ThreadLocalEventSinkSlot()->Get());
+
+  // Make sure the sink was reset since the last tracing session. Normally, it
+  // is reset on Flush after the session is disabled. However, it may not have
+  // been reset if the current thread doesn't support flushing. In that case, we
+  // need to check here that it writes to the right buffer.
+  //
+  // Because we want to avoid locking for each event, we access |session_flags_|
+  // racily. It's OK if we don't see it change to the session immediately. In
+  // that case, the first few trace events may get lost, but we will eventually
+  // notice that we are writing to the wrong buffer once the change to
+  // |session_flags_| has propagated, and reset the sink. Note we will still
+  // acquire the |lock_| to safely recreate the sink in
+  // CreateThreadLocalEventSink().
+  if (thread_local_event_sink) {
+    SessionFlags new_session_flags =
+        GetInstance()->session_flags_.load(std::memory_order_relaxed);
+    if (new_session_flags.session_id != thread_local_event_sink->session_id()) {
+      delete thread_local_event_sink;
+      thread_local_event_sink = nullptr;
+    }
+  }
+
+  if (!thread_local_event_sink) {
+    thread_local_event_sink =
+        GetInstance()->CreateThreadLocalEventSink(thread_will_flush);
+    ThreadLocalEventSinkSlot()->Set(thread_local_event_sink);
+  }
+
+  return thread_local_event_sink;
+}
+
 bool TraceEventDataSource::IsEnabled() {
   base::AutoLock l(lock_);
   return is_enabled_;
@@ -822,8 +868,8 @@ TraceEventDataSource::CreateTraceWriterLocked() {
   return trace_writer;
 }
 
-ThreadLocalEventSink* TraceEventDataSource::CreateThreadLocalEventSink(
-    bool thread_will_flush) {
+TrackEventThreadLocalEventSink*
+TraceEventDataSource::CreateThreadLocalEventSink(bool thread_will_flush) {
   base::AutoLock lock(lock_);
   uint32_t session_id =
       session_flags_.load(std::memory_order_relaxed).session_id;
@@ -843,50 +889,8 @@ void TraceEventDataSource::OnAddTraceEvent(
     TraceEvent* trace_event,
     bool thread_will_flush,
     base::trace_event::TraceEventHandle* handle) {
-  // Avoid re-entrancy, which can happen during PostTasks (the taskqueue can
-  // emit trace events). We discard the events in this case, as any PostTasking
-  // to deal with these events later would break the event ordering that the
-  // JSON traces rely on to merge 'A'/'B' events, as well as having to deal with
-  // updating duration of 'X' events which haven't been added yet.
-  if (GetThreadIsInTraceEventTLS()->Get()) {
-    return;
-  }
-
-  AutoThreadLocalBoolean thread_is_in_trace_event(GetThreadIsInTraceEventTLS());
-
-  auto* thread_local_event_sink =
-      static_cast<ThreadLocalEventSink*>(ThreadLocalEventSinkSlot()->Get());
-
-  // Make sure the sink was reset since the last tracing session. Normally, it
-  // is reset on Flush after the session is disabled. However, it may not have
-  // been reset if the current thread doesn't support flushing. In that case, we
-  // need to check here that it writes to the right buffer.
-  //
-  // Because we want to avoid locking for each event, we access |session_flags_|
-  // racily. It's OK if we don't see it change to the session immediately. In
-  // that case, the first few trace events may get lost, but we will eventually
-  // notice that we are writing to the wrong buffer once the change to
-  // |session_flags_| has propagated, and reset the sink. Note we will still
-  // acquire the |lock_| to safely recreate the sink in
-  // CreateThreadLocalEventSink().
-  if (thread_local_event_sink) {
-    SessionFlags new_session_flags =
-        GetInstance()->session_flags_.load(std::memory_order_relaxed);
-    if (new_session_flags.session_id != thread_local_event_sink->session_id()) {
-      delete thread_local_event_sink;
-      thread_local_event_sink = nullptr;
-    }
-  }
-
-  if (!thread_local_event_sink) {
-    thread_local_event_sink =
-        GetInstance()->CreateThreadLocalEventSink(thread_will_flush);
-    ThreadLocalEventSinkSlot()->Set(thread_local_event_sink);
-  }
-
-  if (thread_local_event_sink) {
-    thread_local_event_sink->AddTraceEvent(trace_event, handle);
-  }
+  OnAddTraceEvent(trace_event, thread_will_flush, handle,
+                  [](perfetto::EventContext) {});
 }
 
 // static
@@ -901,8 +905,8 @@ void TraceEventDataSource::OnUpdateDuration(
 
   AutoThreadLocalBoolean thread_is_in_trace_event(GetThreadIsInTraceEventTLS());
 
-  auto* thread_local_event_sink =
-      static_cast<ThreadLocalEventSink*>(ThreadLocalEventSinkSlot()->Get());
+  auto* thread_local_event_sink = static_cast<TrackEventThreadLocalEventSink*>(
+      ThreadLocalEventSinkSlot()->Get());
   if (thread_local_event_sink) {
     thread_local_event_sink->UpdateDuration(handle, now, thread_now,
                                             thread_instruction_now);
@@ -911,8 +915,8 @@ void TraceEventDataSource::OnUpdateDuration(
 
 // static
 void TraceEventDataSource::FlushCurrentThread() {
-  auto* thread_local_event_sink =
-      static_cast<ThreadLocalEventSink*>(ThreadLocalEventSinkSlot()->Get());
+  auto* thread_local_event_sink = static_cast<TrackEventThreadLocalEventSink*>(
+      ThreadLocalEventSinkSlot()->Get());
   if (thread_local_event_sink) {
     // Prevent any events from being emitted while we're deleting
     // the sink (like from the TraceWriter being PostTask'ed for deletion).

services/tracing/public/cpp/perfetto/trace_event_data_source.h

@@ -20,6 +20,7 @@
 #include "base/timer/timer.h"
 #include "base/trace_event/trace_config.h"
 #include "services/tracing/public/cpp/perfetto/perfetto_traced_process.h"
+#include "services/tracing/public/cpp/perfetto/track_event_thread_local_event_sink.h"
 #include "third_party/perfetto/protos/perfetto/trace/chrome/chrome_metadata.pbzero.h"
 #include "third_party/perfetto/protos/perfetto/trace/chrome/chrome_trace_event.pbzero.h"
 
@@ -35,6 +36,7 @@ namespace perfetto {
 class StartupTraceWriter;
 class StartupTraceWriterRegistry;
 class TraceWriter;
+class EventContext;
 }
 
 namespace tracing {
@@ -185,6 +187,23 @@ class COMPONENT_EXPORT(TRACING_CPP) TraceEventDataSource
 
   bool IsEnabled();
 
+  static TrackEventThreadLocalEventSink* GetOrPrepareEventSink(
+      bool thread_will_flush);
+
+  template <
+      typename TrackEventArgumentFunction = void (*)(perfetto::EventContext)>
+  static void OnAddTraceEvent(base::trace_event::TraceEvent* trace_event,
+                              bool thread_will_flush,
+                              base::trace_event::TraceEventHandle* handle,
+                              TrackEventArgumentFunction func) {
+    auto* thread_local_event_sink = GetOrPrepareEventSink(thread_will_flush);
+    if (thread_local_event_sink) {
+      AutoThreadLocalBoolean thread_is_in_trace_event(
+          GetThreadIsInTraceEventTLS());
+      thread_local_event_sink->AddTraceEvent(trace_event, handle, func);
+    }
+  }
+
  private:
   friend class base::NoDestructor<TraceEventDataSource>;
 
@@ -203,7 +222,8 @@ class COMPONENT_EXPORT(TRACING_CPP) TraceEventDataSource
   void UnregisterFromTraceLog();
 
   std::unique_ptr<perfetto::StartupTraceWriter> CreateTraceWriterLocked();
-  ThreadLocalEventSink* CreateThreadLocalEventSink(bool thread_will_flush);
+  TrackEventThreadLocalEventSink* CreateThreadLocalEventSink(
+      bool thread_will_flush);
 
   // Callback from TraceLog, can be called from any thread.
   static void OnAddTraceEvent(base::trace_event::TraceEvent* trace_event,

services/tracing/public/cpp/perfetto/trace_event_data_source_unittest.cc

@@ -26,6 +26,7 @@
 #include "base/trace_event/trace_event.h"
 #include "base/trace_event/trace_log.h"
 #include "components/tracing/common/tracing_switches.h"
+#include "services/tracing/public/cpp/perfetto/macros.h"
 #include "services/tracing/public/cpp/perfetto/producer_client.h"
 #include "services/tracing/public/mojom/perfetto_service.mojom.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -35,6 +36,7 @@
 #include "third_party/perfetto/include/perfetto/protozero/scattered_stream_writer.h"
 #include "third_party/perfetto/protos/perfetto/trace/trace_packet.pb.h"
 #include "third_party/perfetto/protos/perfetto/trace/trace_packet.pbzero.h"
+#include "third_party/perfetto/protos/perfetto/trace/track_event/log_message.pbzero.h"
 #include "third_party/perfetto/protos/perfetto/trace/track_event/process_descriptor.pb.h"
 #include "third_party/perfetto/protos/perfetto/trace/track_event/thread_descriptor.pb.h"
 
@@ -1305,6 +1307,103 @@ TEST_F(TraceEventDataSourceTest, StartupTracingTimeout) {
   wait_for_stop.Run();
 }
 
+TEST_F(TraceEventDataSourceTest, TypedArgumentsTracingOff) {
+  TRACE_EVENT_BEGIN("log", "LogMessage", [](perfetto::EventContext ctx) {
+    ADD_FAILURE() << "lambda was called when tracing was off";
+  });
+
+  TRACE_EVENT_END("log", [](perfetto::EventContext ctx) {
+    ADD_FAILURE() << "lambda was called when tracing was off";
+  });
+}
+
+TEST_F(TraceEventDataSourceTest, TypedArgumentsTracingOnBegin) {
+  CreateTraceEventDataSource();
+
+  bool begin_called = false;
+
+  TRACE_EVENT_BEGIN("browser", "bar", [&](perfetto::EventContext ctx) {
+    begin_called = true;
+    ctx.event()->set_log_message()->set_body_iid(42);
+  });
+
+  EXPECT_TRUE(begin_called);
+
+  EXPECT_EQ(producer_client()->GetFinalizedPacketCount(), 2u);
+
+  auto* td_packet = producer_client()->GetFinalizedPacket();
+  ExpectThreadDescriptor(td_packet);
+
+  auto* e_packet = producer_client()->GetFinalizedPacket(1);
+  ExpectTraceEvent(e_packet, /*category_iid=*/1u, /*name_iid=*/1u,
+                   TRACE_EVENT_PHASE_BEGIN);
+
+  ExpectEventCategories(e_packet, {{1u, "browser"}});
+  ExpectEventNames(e_packet, {{1u, "bar"}});
+  ASSERT_TRUE(e_packet->track_event().has_log_message());
+  EXPECT_EQ(e_packet->track_event().log_message().body_iid(), 42u);
+}
+
+TEST_F(TraceEventDataSourceTest, TypedArgumentsTracingOnEnd) {
+  CreateTraceEventDataSource();
+
+  bool end_called = false;
+
+  TRACE_EVENT_END("browser", [&](perfetto::EventContext ctx) {
+    end_called = true;
+    ctx.event()->set_log_message()->set_body_iid(42);
+  });
+
+  EXPECT_TRUE(end_called);
+
+  EXPECT_EQ(producer_client()->GetFinalizedPacketCount(), 2u);
+
+  auto* td_packet = producer_client()->GetFinalizedPacket();
+  ExpectThreadDescriptor(td_packet);
+
+  auto* e_packet = producer_client()->GetFinalizedPacket(1);
+  ExpectTraceEvent(e_packet, /*category_iid=*/1u, /*name_iid=*/1u,
+                   TRACE_EVENT_PHASE_END);
+
+  ExpectEventCategories(e_packet, {{1u, "browser"}});
+  ExpectEventNames(e_packet, {{1u, kTraceEventEndName}});
+  ASSERT_TRUE(e_packet->track_event().has_log_message());
+  EXPECT_EQ(e_packet->track_event().log_message().body_iid(), 42u);
+}
+
+TEST_F(TraceEventDataSourceTest, TypedArgumentsTracingOnBeginAndEnd) {
+  CreateTraceEventDataSource();
+
+  TRACE_EVENT_BEGIN("browser", "bar", [&](perfetto::EventContext ctx) {
+    ctx.event()->set_log_message()->set_body_iid(42);
+  });
+  TRACE_EVENT_END("browser", [&](perfetto::EventContext ctx) {
+    ctx.event()->set_log_message()->set_body_iid(84);
+  });
+
+  EXPECT_EQ(producer_client()->GetFinalizedPacketCount(), 3u);
+
+  auto* td_packet = producer_client()->GetFinalizedPacket();
+  ExpectThreadDescriptor(td_packet);
+
+  auto* e_packet = producer_client()->GetFinalizedPacket(1);
+  ExpectTraceEvent(e_packet, /*category_iid=*/1u, /*name_iid=*/1u,
+                   TRACE_EVENT_PHASE_BEGIN);
+
+  ExpectEventCategories(e_packet, {{1u, "browser"}});
+  ExpectEventNames(e_packet, {{1u, "bar"}});
+  ASSERT_TRUE(e_packet->track_event().has_log_message());
+  EXPECT_EQ(e_packet->track_event().log_message().body_iid(), 42u);
+
+  e_packet = producer_client()->GetFinalizedPacket(2);
+  ExpectTraceEvent(e_packet, /*category_iid=*/1u, /*name_iid=*/2u,
+                   TRACE_EVENT_PHASE_END);
+
+  ExpectEventNames(e_packet, {{2u, kTraceEventEndName}});
+  ASSERT_TRUE(e_packet->track_event().has_log_message());
+  EXPECT_EQ(e_packet->track_event().log_message().body_iid(), 84u);
+}
+
 // TODO(eseckler): Add startup tracing unittests.
 
 }  // namespace

services/tracing/public/cpp/perfetto/track_event_thread_local_event_sink.h

@@ -10,15 +10,20 @@
 #include <string>
 #include <utility>
 
+#include "base/trace_event/trace_event.h"
+
 #include "base/component_export.h"
 #include "base/threading/thread_id_name_manager.h"
 #include "base/time/time.h"
 #include "base/trace_event/thread_instruction_count.h"
+#include "services/tracing/public/cpp/perfetto/event_context.h"
 #include "services/tracing/public/cpp/perfetto/interning_index.h"
-#include "services/tracing/public/cpp/perfetto/thread_local_event_sink.h"
+#include "third_party/perfetto/include/perfetto/ext/tracing/core/startup_trace_writer.h"
 #include "third_party/perfetto/include/perfetto/ext/tracing/core/trace_writer.h"
 #include "third_party/perfetto/include/perfetto/protozero/message_handle.h"
+#include "third_party/perfetto/protos/perfetto/trace/interned_data/interned_data.pbzero.h"
 #include "third_party/perfetto/protos/perfetto/trace/track_event/thread_descriptor.pbzero.h"
+#include "third_party/perfetto/protos/perfetto/trace/track_event/track_event.pbzero.h"
 
 namespace perfetto {
 class StartupTraceWriter;
@@ -28,9 +33,33 @@ namespace tracing {
 
 // ThreadLocalEventSink that emits TrackEvent protos.
 class COMPONENT_EXPORT(TRACING_CPP) TrackEventThreadLocalEventSink
-    : public ThreadLocalEventSink,
-      public base::ThreadIdNameManager::Observer {
+    : public base::ThreadIdNameManager::Observer {
  public:
+  enum class IndexType {
+    kName = 0,
+    kCategory = 1,
+    kAnnotationName = 2,
+    kSourceLocation = 3,
+    kLogMessage = 4
+  };
+  // IndexData is a temporary storage location for passing long updates to the
+  // interning indexes. Everything stored in it must have a lifetime that is
+  // at least as long as AddTraceEvent.
+  //
+  // In most cases this is easy since the provided |trace_event| is the source
+  // of most const char*s.
+  //
+  // This is important because when TRACE_EVENT_FLAG_COPY is set, the
+  // InternedIndexesUpdates are cleared within the same call to AddTraceEvent().
+  union IndexData {
+    const char* str_piece;
+    std::tuple<const char*, const char*, int> src_loc;
+    explicit IndexData(const char* str);
+    explicit IndexData(std::tuple<const char*, const char*, int>&& src);
+  };
+  using InternedIndexesUpdates =
+      std::vector<std::tuple<IndexType, IndexData, InterningIndexEntry>>;
+
   TrackEventThreadLocalEventSink(
       std::unique_ptr<perfetto::StartupTraceWriter> trace_writer,
       uint32_t session_id,
@@ -42,15 +71,59 @@ class COMPONENT_EXPORT(TRACING_CPP) TrackEventThreadLocalEventSink
   // (e.g. interning index entries and a ThreadDescriptor) to be emitted again.
   static void ClearIncrementalState();
 
-  // ThreadLocalEventSink implementation:
+  // Returns true if the current |trace_event| is part of a COMPLETE event
+  // phase. As part of this we will update any stack information to keep track
+  // of the information when we can pop something off the stack.
+  bool MaybeHandleCompleteEvent(base::trace_event::TraceEvent* trace_event,
+                                base::trace_event::TraceEventHandle* handle);
+
+  // If we need to perform an incremental reset we will do so, and also emit all
+  // the relevant descriptors to start a new fresh sequence.
+  void ResetIncrementalStateIfNeeded(
+      base::trace_event::TraceEvent* trace_event);
+
+  // Fills in all the fields in |track_event| that can be directly deduced from
+  // |trace_event|. Returns all the updates needed to be emitted into the
+  // |InternedData| field.
+  void PrepareTrackEvent(base::trace_event::TraceEvent* trace_event,
+                         base::trace_event::TraceEventHandle* handle,
+                         perfetto::protos::pbzero::TrackEvent* track_event);
+
+  // Given a list of updates to the indexes will fill in |interned_data| to
+  // reflect them.
+  void EmitStoredInternedData(
+      perfetto::protos::pbzero::InternedData* interned_data);
+
+  template <
+      typename TrackEventArgumentFunction = void (*)(perfetto::EventContext)>
   void AddTraceEvent(base::trace_event::TraceEvent* trace_event,
-                     base::trace_event::TraceEventHandle* handle) override;
-  void UpdateDuration(base::trace_event::TraceEventHandle handle,
-                      const base::TimeTicks& now,
-                      const base::ThreadTicks& thread_now,
-                      base::trace_event::ThreadInstructionCount
-                          thread_instruction_now) override;
-  void Flush() override;
+                     base::trace_event::TraceEventHandle* handle,
+                     TrackEventArgumentFunction arg_func) {
+    if (MaybeHandleCompleteEvent(trace_event, handle)) {
+      return;
+    }
+
+    ResetIncrementalStateIfNeeded(trace_event);
+
+    auto trace_packet = trace_writer_->NewTracePacket();
+    auto* track_event = trace_packet->set_track_event();
+    PrepareTrackEvent(trace_event, handle, track_event);
+
+    arg_func(perfetto::EventContext(track_event));
+
+    if (!pending_interning_updates_.empty()) {
+      EmitStoredInternedData(trace_packet->set_interned_data());
+    }
+  }
+
+  void UpdateDuration(
+      base::trace_event::TraceEventHandle handle,
+      const base::TimeTicks& now,
+      const base::ThreadTicks& thread_now,
+      base::trace_event::ThreadInstructionCount thread_instruction_now);
+  void Flush();
+
+  uint32_t session_id() const { return session_id_; }
 
   // ThreadIdNameManager::Observer implementation:
   void OnThreadNameChanged(const char* name) override;
@@ -80,6 +153,7 @@ class COMPONENT_EXPORT(TRACING_CPP) TrackEventThreadLocalEventSink
       interned_source_locations_;
   InterningIndex<TypeList<std::string>, SizeList<128>>
       interned_log_message_bodies_;
+  InternedIndexesUpdates pending_interning_updates_;
 
   static std::atomic<uint32_t> incremental_state_reset_id_;
 
@@ -98,6 +172,12 @@ class COMPONENT_EXPORT(TRACING_CPP) TrackEventThreadLocalEventSink
   uint32_t current_stack_depth_ = 0;
 
   const bool privacy_filtering_enabled_;
+
+  std::unique_ptr<perfetto::StartupTraceWriter> trace_writer_;
+  uint32_t session_id_;
+  bool disable_interning_;
+  uint32_t sink_id_;
+  DISALLOW_COPY_AND_ASSIGN(TrackEventThreadLocalEventSink);
 };
 
 }  // namespace tracing