pw_trace_tokenized: Allow event callback to modify event

Change the event callback API to accept a pointer to
pw_trace_tokenized_TraceEvent, and move the event callback to before
inserting in the queue. This allows the platform to implement a callback
to modify run-time fields in the trace event.

Add example where the callback modifies the trace_id field to store the
Linux info, used for auto grouping trace events by tid.

Change-Id: Ib8dd8f69bf70e37c3bd53d69a2be1aec127cf7c4
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/126310
Reviewed-by: Rob Oliver <rgoliver@google.com>
Commit-Queue: Yvonne Yip <ykyyip@google.com>
Reviewed-by: Ted Pudlik <tpudlik@google.com>
Pigweed-Auto-Submit: Yvonne Yip <ykyyip@google.com>
diff --git a/pw_trace_tokenized/BUILD.bazel b/pw_trace_tokenized/BUILD.bazel
index 79451e8..26f1302 100644
--- a/pw_trace_tokenized/BUILD.bazel
+++ b/pw_trace_tokenized/BUILD.bazel
@@ -44,6 +44,12 @@
 )
 
 pw_cc_library(
+    name = "linux_config_overrides",
+    hdrs = ["linux_config_overrides.h"],
+    tags = ["manual"],
+)
+
+pw_cc_library(
     name = "pw_trace_tokenized",
     srcs = [
         "trace.cc",
@@ -268,3 +274,14 @@
         "//pw_trace:pw_trace_sample_app",
     ],
 )
+
+pw_cc_library(
+    name = "trace_tokenized_example_linux_group_by_tid",
+    srcs = ["example/linux_group_by_tid.cc"],
+    tags = ["manual"],
+    deps = [
+        ":pw_trace_example_to_file",
+        "//pw_log",
+        "//pw_trace",
+    ],
+)
diff --git a/pw_trace_tokenized/BUILD.gn b/pw_trace_tokenized/BUILD.gn
index 2aa8c46..7af4ddf 100644
--- a/pw_trace_tokenized/BUILD.gn
+++ b/pw_trace_tokenized/BUILD.gn
@@ -36,6 +36,19 @@
   public = [ "public/pw_trace_tokenized/config.h" ]
 }
 
+config("linux_include_config_overrides") {
+  cflags = [
+    "-include",
+    rebase_path("linux_config_overrides.h", root_build_dir),
+  ]
+  visibility = [ ":*" ]
+}
+
+pw_source_set("linux_config_overrides") {
+  public_configs = [ ":linux_include_config_overrides" ]
+  sources = [ "linux_config_overrides.h" ]
+}
+
 pw_test_group("tests") {
   tests = [
     ":trace_tokenized_test",
@@ -54,8 +67,9 @@
     ":core",
     "$dir_pw_tokenizer",
   ]
+  deps = []
   if (pw_trace_tokenizer_time != "") {
-    deps = [ "$pw_trace_tokenizer_time" ]
+    deps += [ "$pw_trace_tokenizer_time" ]
   }
 
   public = [ "public_overrides/pw_trace_backend/trace_backend.h" ]
@@ -162,6 +176,7 @@
     ":public_include_path",
   ]
   public_deps = [
+    "$dir_pw_log",
     "$dir_pw_status",
     "$dir_pw_tokenizer",
     dir_pw_span,
@@ -244,3 +259,17 @@
     ]
   }
 }
+
+if (current_os != "linux") {
+  group("trace_tokenized_example_linux_group_by_tid") {
+  }
+} else {
+  pw_executable("trace_tokenized_example_linux_group_by_tid") {
+    sources = [ "example/linux_group_by_tid.cc" ]
+    deps = [
+      ":trace_example_to_file",
+      "$dir_pw_log",
+      "$dir_pw_trace",
+    ]
+  }
+}
diff --git a/pw_trace_tokenized/docs.rst b/pw_trace_tokenized/docs.rst
index ca537e5..f58dc62 100644
--- a/pw_trace_tokenized/docs.rst
+++ b/pw_trace_tokenized/docs.rst
@@ -70,27 +70,35 @@
 provide hooks into tracing.
 
 The *event callbacks* are called when trace events occur, with the trace event
-data. Using the return flags, these callbacks can be used to adjust the trace
-behaviour at runtime in response to specific events. If requested (using
-``called_on_every_event``) the callback will be called on every trace event
-regardless if tracing is currently enabled or not. Using this, the application
-can trigger tracing on or off when specific traces or patterns of traces are
-observed, or can selectively filter traces to preserve the trace buffer.
+data, before the event is encoded or sent to the sinks. The callbacks may
+modify the run-time fields of the trace event, i.e. ``trace_id``,
+``data_buffer`` and ``data_size``. Using the return flags, these callbacks can
+be used to adjust the trace behaviour at runtime in response to specific events.
 
-The event callback is a single function which is provided the details of the
-trace as arguments, and returns ``pw_trace_TraceEventReturnFlags``, which can be
-used to change how the trace is handled.
+If requested (using ``called_on_every_event``) the callback will be called on
+every trace event regardless if tracing is currently enabled or not. Using this,
+the application can trigger tracing on or off when specific traces or patterns
+of traces are observed, or can selectively filter traces to preserve the trace
+buffer.
+
+The event callback is called in the context of the traced task. It must be
+ISR-safe to support tracing within ISRs. It must be lightweight to prevent
+performance issues in the trace tasks.
+
+The return flags ``pw_trace_TraceEventReturnFlags`` support the following
+behaviors:
+
+* ``PW_TRACE_EVENT_RETURN_FLAGS_SKIP_EVENT`` can be set true to skip this
+  sample.
+* ``PW_TRACE_EVENT_RETURN_FLAGS_DISABLE_AFTER_PROCESSING`` can be set true to
+  disable tracing after this sample.
 
 .. cpp:function:: pw_trace_TraceEventReturnFlags pw_trace_EventCallback( \
-    void* user_data, \
-    uint32_t trace_ref, \
-    pw_trace_EventType event_type, \
-    const char* module, \
-    uint32_t trace_id, \
-    uint8_t flags)
+    void* user_data,  \
+    pw_trace_tokenized_TraceEvent* event)
 .. cpp:function:: pw_Status pw_trace_RegisterEventCallback( \
     pw_trace_EventCallback callback, \
-    bool called_on_every_event, \
+    pw_trace_EventCallbackFlags flags, \
     void* user_data, \
     pw_trace_EventCallbackHandle* handle)
 .. cpp:function:: pw_Status pw_trace_UnregisterEventCallback( \
diff --git a/pw_trace_tokenized/example/filter.cc b/pw_trace_tokenized/example/filter.cc
index f76f64c..8461f56 100644
--- a/pw_trace_tokenized/example/filter.cc
+++ b/pw_trace_tokenized/example/filter.cc
@@ -35,15 +35,11 @@
 #include "pw_trace_tokenized/trace_tokenized.h"
 
 pw_trace_TraceEventReturnFlags TraceEventCallback(
-    void* /* user_data */,
-    uint32_t /* trace_ref */,
-    pw_trace_EventType /* event_type */,
-    const char* module,
-    uint32_t trace_id,
-    uint8_t /* flags */) {
+    void* /* user_data */, pw_trace_tokenized_TraceEvent* event) {
   // Filter out all traces from processing task, which aren't traceId 3
   static constexpr uint32_t kFilterId = 3;
-  return (strcmp("Processing", module) == 0 && trace_id != kFilterId)
+  return (strcmp("Processing", event->module) == 0 &&
+          event->trace_id != kFilterId)
              ? PW_TRACE_EVENT_RETURN_FLAGS_SKIP_EVENT
              : 0;
 }
diff --git a/pw_trace_tokenized/example/linux_group_by_tid.cc b/pw_trace_tokenized/example/linux_group_by_tid.cc
new file mode 100644
index 0000000..6c8a204
--- /dev/null
+++ b/pw_trace_tokenized/example/linux_group_by_tid.cc
@@ -0,0 +1,100 @@
+// Copyright 2023 The Pigweed Authors
+//
+// 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
+//
+//     https://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 <pthread.h>
+#include <sys/syscall.h>
+#include <unistd.h>
+
+#include <chrono>
+#include <cstdio>
+#include <thread>
+
+#include "pw_log/log.h"
+#include "pw_trace/trace.h"
+#include "pw_trace_tokenized/example/trace_to_file.h"
+
+// Example for annotating trace events with thread id.
+// The platform annotates instants and duration events with the thread id if the
+// caller does not explicitly provide a group. The thread id is written in
+// the trace_id field.
+//
+// This example requires linux_config_overrides.h to define
+// PW_TRACE_HAS_TRACE_ID. Set pw_trace_CONFIG to
+// "$dir_pw_trace_tokenized:linux_config_overrides" in target_toolchains.gni to
+// enable the override before building this example.
+//
+// TODO(ykyyip): update trace_tokenized.py to handle the trace_id.
+
+pw_trace_TraceEventReturnFlags TraceEventCallback(
+    void* /*user_data*/, pw_trace_tokenized_TraceEvent* event) {
+  // Instant and duration events with no group means group by pid/tid.
+  if ((event->event_type == PW_TRACE_EVENT_TYPE_INSTANT) ||
+      (event->event_type == PW_TRACE_EVENT_TYPE_DURATION_START) ||
+      (event->event_type == PW_TRACE_EVENT_TYPE_DURATION_END)) {
+    event->trace_id = syscall(__NR_gettid);
+  }
+  return PW_TRACE_EVENT_RETURN_FLAGS_NONE;
+}
+
+void ExampleTask(void* /*arg*/) {
+  int times_to_run = 10;
+  while (times_to_run--) {
+    PW_TRACE_START("Processing");
+    //  Fake processing time.
+    std::this_thread::sleep_for(std::chrono::milliseconds(42));
+    PW_TRACE_END("Processing");
+    //  Sleep for a random amount before running again.
+    int sleep_time = 1 + std::rand() % 20;
+    std::this_thread::sleep_for(std::chrono::milliseconds(sleep_time));
+  }
+}
+
+void RunThreadedTraceSampleApp() {
+  std::srand(std::time(nullptr));
+
+  // Start threads to show parallel processing.
+  int num_threads = 5;
+  while (num_threads--) {
+    PW_TRACE_INSTANT("CreateThread");
+    std::thread thread(ExampleTask, nullptr);
+    thread.detach();
+  }
+}
+
+int main(int argc, char** argv) {
+  if (argc != 2) {
+    PW_LOG_ERROR("Expected output file name as argument.\n");
+    return -1;
+  }
+
+  // Enable tracing.
+  PW_TRACE_SET_ENABLED(true);
+
+  // Dump trace data to the file passed in.
+  pw::trace::TraceToFile trace_to_file{argv[1]};
+
+  // Register platform callback
+  pw::trace::RegisterCallbackWhenCreated{TraceEventCallback};
+
+  PW_LOG_INFO("Running threaded trace example...\n");
+  RunThreadedTraceSampleApp();
+
+  // Sleep forever
+  while (true) {
+    std::this_thread::sleep_for(std::chrono::seconds(60));
+  }
+  return 0;
+}
diff --git a/pw_trace_tokenized/example/public/pw_trace_tokenized/example/trace_to_file.h b/pw_trace_tokenized/example/public/pw_trace_tokenized/example/trace_to_file.h
index 6b3a6f6..c392a38 100644
--- a/pw_trace_tokenized/example/public/pw_trace_tokenized/example/trace_to_file.h
+++ b/pw_trace_tokenized/example/public/pw_trace_tokenized/example/trace_to_file.h
@@ -20,7 +20,6 @@
 
 #include <fstream>
 
-#include "pw_trace/example/sample_app.h"
 #include "pw_trace_tokenized/trace_callback.h"
 
 namespace pw {
diff --git a/pw_trace_tokenized/example/trigger.cc b/pw_trace_tokenized/example/trigger.cc
index 7f84db4..cc7c393 100644
--- a/pw_trace_tokenized/example/trigger.cc
+++ b/pw_trace_tokenized/example/trigger.cc
@@ -53,17 +53,14 @@
 }  // namespace
 
 pw_trace_TraceEventReturnFlags TraceEventCallback(
-    void* /* user_data */,
-    uint32_t trace_ref,
-    pw_trace_EventType /* event_type */,
-    const char* /* module */,
-    uint32_t trace_id,
-    uint8_t /* flags */) {
-  if (trace_ref == kTriggerStartTraceRef && trace_id == kTriggerId) {
+    void* /* user_data */, pw_trace_tokenized_TraceEvent* event) {
+  if (event->trace_token == kTriggerStartTraceRef &&
+      event->trace_id == kTriggerId) {
     PW_LOG_INFO("Trace capture started!");
     PW_TRACE_SET_ENABLED(true);
   }
-  if (trace_ref == kTriggerEndTraceRef && trace_id == kTriggerId) {
+  if (event->trace_token == kTriggerEndTraceRef &&
+      event->trace_id == kTriggerId) {
     PW_LOG_INFO("Trace capture ended!");
     return PW_TRACE_EVENT_RETURN_FLAGS_DISABLE_AFTER_PROCESSING;
   }
diff --git a/pw_trace_tokenized/linux_config_overrides.h b/pw_trace_tokenized/linux_config_overrides.h
new file mode 100644
index 0000000..d4215d055a
--- /dev/null
+++ b/pw_trace_tokenized/linux_config_overrides.h
@@ -0,0 +1,30 @@
+// Copyright 2023 The Pigweed Authors
+//
+// 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
+//
+//     https://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.
+//==============================================================================
+//
+
+#pragma once
+
+// pw_trace config overrides when the host os is linux
+
+// When pw_trace_tokenized_ANNOTATION == linux_tid_trace_annotation,
+// the platform sets trace_id to the Linux thread id when the caller does
+// not provide a group.
+#define PW_TRACE_HAS_TRACE_ID(TRACE_TYPE)          \
+  ((TRACE_TYPE) == PW_TRACE_TYPE_INSTANT ||        \
+   (TRACE_TYPE) == PW_TRACE_TYPE_DURATION_START || \
+   (TRACE_TYPE) == PW_TRACE_TYPE_DURATION_END ||   \
+   (TRACE_TYPE) == PW_TRACE_TYPE_ASYNC_INSTANT ||  \
+   (TRACE_TYPE) == PW_TRACE_TYPE_ASYNC_START ||    \
+   (TRACE_TYPE) == PW_TRACE_TYPE_ASYNC_END)
diff --git a/pw_trace_tokenized/public/pw_trace_tokenized/internal/trace_tokenized_internal.h b/pw_trace_tokenized/public/pw_trace_tokenized/internal/trace_tokenized_internal.h
index 8427dc6..e2954dd 100644
--- a/pw_trace_tokenized/public/pw_trace_tokenized/internal/trace_tokenized_internal.h
+++ b/pw_trace_tokenized/public/pw_trace_tokenized/internal/trace_tokenized_internal.h
@@ -17,9 +17,11 @@
 #pragma once
 
 #include <stdbool.h>
+#include <stddef.h>
 #include <stdint.h>
 
 #include "pw_preprocessor/arguments.h"
+#include "pw_preprocessor/util.h"
 
 // Because __FUNCTION__ is not a string literal to the preprocessor it can't be
 // tokenized. So this backend redefines the implementation to instead use the
@@ -47,6 +49,7 @@
 #define PW_TRACE_TYPE_ASYNC_END PW_TRACE_EVENT_TYPE_ASYNC_END
 
 PW_EXTERN_C_START
+
 typedef enum {
   PW_TRACE_EVENT_TYPE_INVALID = 0,
   PW_TRACE_EVENT_TYPE_INSTANT = 1,
@@ -60,6 +63,16 @@
   PW_TRACE_EVENT_TYPE_DURATION_GROUP_END = 9,
 } pw_trace_EventType;
 
+typedef struct {
+  uint32_t trace_token;
+  pw_trace_EventType event_type;
+  const char* module;
+  uint8_t flags;
+  uint32_t trace_id;
+  size_t data_size;
+  const void* data_buffer;
+} pw_trace_tokenized_TraceEvent;
+
 // This should not be called directly, instead use the PW_TRACE_* macros.
 void pw_trace_TraceEvent(uint32_t trace_token,
                          pw_trace_EventType event_type,
diff --git a/pw_trace_tokenized/public/pw_trace_tokenized/trace_callback.h b/pw_trace_tokenized/public/pw_trace_tokenized/trace_callback.h
index f3d460a..48ac848 100644
--- a/pw_trace_tokenized/public/pw_trace_tokenized/trace_callback.h
+++ b/pw_trace_tokenized/public/pw_trace_tokenized/trace_callback.h
@@ -64,12 +64,7 @@
 
 typedef size_t pw_trace_EventCallbackHandle;
 typedef pw_trace_TraceEventReturnFlags (*pw_trace_EventCallback)(
-    void* user_data,
-    uint32_t trace_ref,
-    pw_trace_EventType event_type,
-    const char* module,
-    uint32_t trace_id,
-    uint8_t flags);
+    void* user_data, pw_trace_tokenized_TraceEvent* event);
 
 pw_Status pw_trace_RegisterEventCallback(
     pw_trace_EventCallback callback,
@@ -142,6 +137,7 @@
   };
   using EventCallback = pw_trace_EventCallback;
   using EventCallbackHandle = pw_trace_EventCallbackHandle;
+  using TraceEvent = pw_trace_tokenized_TraceEvent;
   struct EventCallbacks {
     void* user_data;
     EventCallback callback;
@@ -167,12 +163,7 @@
   pw::Status UnregisterAllEventCallbacks();
   EventCallbacks* GetEventCallback(EventCallbackHandle handle);
   pw_trace_TraceEventReturnFlags CallEventCallbacks(
-      CallOnEveryEvent called_on_every_event,
-      uint32_t trace_ref,
-      EventType event_type,
-      const char* module,
-      uint32_t trace_id,
-      uint8_t flags);
+      CallOnEveryEvent called_on_every_event, TraceEvent* event);
   size_t GetCalledOnEveryEventCount() const {
     return called_on_every_event_count_;
   }
diff --git a/pw_trace_tokenized/trace.cc b/pw_trace_tokenized/trace.cc
index 693c67f..5547c62 100644
--- a/pw_trace_tokenized/trace.cc
+++ b/pw_trace_tokenized/trace.cc
@@ -27,6 +27,8 @@
 TokenizedTraceImpl TokenizedTrace::instance_;
 CallbacksImpl Callbacks::instance_;
 
+using TraceEvent = pw_trace_tokenized_TraceEvent;
+
 void TokenizedTraceImpl::HandleTraceEvent(uint32_t trace_token,
                                           EventType event_type,
                                           const char* module,
@@ -40,16 +42,44 @@
     return;
   }
 
+  TraceEvent event = {
+      .trace_token = trace_token,
+      .event_type = event_type,
+      .module = module,
+      .flags = flags,
+      .trace_id = trace_id,
+      .data_size = data_size,
+      .data_buffer = data_buffer,
+  };
+
+  // Call any event callback which is registered to receive every event.
+  pw_trace_TraceEventReturnFlags ret_flags = 0;
+  ret_flags |= Callbacks::Instance().CallEventCallbacks(
+      CallbacksImpl::kCallOnEveryEvent, &event);
+  // Return if disabled.
+  if ((PW_TRACE_EVENT_RETURN_FLAGS_SKIP_EVENT & ret_flags) || !enabled_) {
+    return;
+  }
+
+  // Call any event callback not already called.
+  ret_flags |= Callbacks::Instance().CallEventCallbacks(
+      CallbacksImpl::kCallOnlyWhenEnabled, &event);
+  // Return if disabled (from a callback) or if a callback has indicated the
+  // sample should be skipped.
+  if ((PW_TRACE_EVENT_RETURN_FLAGS_SKIP_EVENT & ret_flags) || !enabled_) {
+    return;
+  }
+
   // Create trace event
   PW_TRACE_QUEUE_LOCK();
   if (!event_queue_
-           .TryPushBack(trace_token,
-                        event_type,
-                        module,
-                        trace_id,
-                        flags,
-                        data_buffer,
-                        data_size)
+           .TryPushBack(event.trace_token,
+                        event.event_type,
+                        event.module,
+                        event.trace_id,
+                        event.flags,
+                        event.data_buffer,
+                        event.data_size)
            .ok()) {
     // Queue full dropping sample
     // TODO(rgoliver): Allow other strategies, for example: drop oldest, try
@@ -66,6 +96,11 @@
     }
     PW_TRACE_UNLOCK();
   }
+
+  // Disable after processing if an event callback had set the flag.
+  if (PW_TRACE_EVENT_RETURN_FLAGS_DISABLE_AFTER_PROCESSING & ret_flags) {
+    enabled_ = false;
+  }
 }
 
 void TokenizedTraceImpl::HandleNextItemInQueue(
@@ -73,41 +108,11 @@
   // Get next item in queue
   uint32_t trace_token = event_block->trace_token;
   EventType event_type = event_block->event_type;
-  const char* module = event_block->module;
   uint32_t trace_id = event_block->trace_id;
-  uint8_t flags = event_block->flags;
   const std::byte* data_buffer =
       const_cast<const std::byte*>(event_block->data_buffer);
   size_t data_size = event_block->data_size;
 
-  // Call any event callback which is registered to receive every event.
-  pw_trace_TraceEventReturnFlags ret_flags = 0;
-  ret_flags |=
-      Callbacks::Instance().CallEventCallbacks(CallbacksImpl::kCallOnEveryEvent,
-                                               trace_token,
-                                               event_type,
-                                               module,
-                                               trace_id,
-                                               flags);
-  // Return if disabled.
-  if ((PW_TRACE_EVENT_RETURN_FLAGS_SKIP_EVENT & ret_flags) || !enabled_) {
-    return;
-  }
-
-  // Call any event callback not already called.
-  ret_flags |= Callbacks::Instance().CallEventCallbacks(
-      CallbacksImpl::kCallOnlyWhenEnabled,
-      trace_token,
-      event_type,
-      module,
-      trace_id,
-      flags);
-  // Return if disabled (from a callback) or if a callback has indicated the
-  // sample should be skipped.
-  if ((PW_TRACE_EVENT_RETURN_FLAGS_SKIP_EVENT & ret_flags) || !enabled_) {
-    return;
-  }
-
   // Create header to store trace info
   static constexpr size_t kMaxHeaderSize =
       sizeof(trace_token) + pw::varint::kMaxVarint64SizeBytes +  // time
@@ -139,30 +144,16 @@
       span<const std::byte>(header, header_size),
       span<const std::byte>(reinterpret_cast<const std::byte*>(data_buffer),
                             data_size));
-  // Disable after processing if an event callback had set the flag.
-  if (PW_TRACE_EVENT_RETURN_FLAGS_DISABLE_AFTER_PROCESSING & ret_flags) {
-    enabled_ = false;
-  }
 }
 
 pw_trace_TraceEventReturnFlags CallbacksImpl::CallEventCallbacks(
-    CallOnEveryEvent called_on_every_event,
-    uint32_t trace_ref,
-    EventType event_type,
-    const char* module,
-    uint32_t trace_id,
-    uint8_t flags) {
+    CallOnEveryEvent called_on_every_event, TraceEvent* event) {
   pw_trace_TraceEventReturnFlags ret_flags = 0;
   for (size_t i = 0; i < PW_TRACE_CONFIG_MAX_EVENT_CALLBACKS; i++) {
     if (event_callbacks_[i].callback &&
         event_callbacks_[i].called_on_every_event == called_on_every_event) {
       ret_flags |= Callbacks::Instance().GetEventCallback(i)->callback(
-          event_callbacks_[i].user_data,
-          trace_ref,
-          event_type,
-          module,
-          trace_id,
-          flags);
+          event_callbacks_[i].user_data, event);
     }
   }
   return ret_flags;
diff --git a/pw_trace_tokenized/trace_test.cc b/pw_trace_tokenized/trace_test.cc
index de03b3d..72457ab 100644
--- a/pw_trace_tokenized/trace_test.cc
+++ b/pw_trace_tokenized/trace_test.cc
@@ -88,21 +88,16 @@
   // buffer_ in the TraceSink callback, that way it only gets added to the
   // buffer if tracing is enabled and the sample was not surpressed.
   static pw_trace_TraceEventReturnFlags TraceEventCallback(
-      void* user_data,
-      uint32_t trace_ref,
-      pw_trace_EventType event_type,
-      const char* module,
-      uint32_t trace_id,
-      uint8_t /* flags */) {
+      void* user_data, pw_trace_tokenized_TraceEvent* event) {
     TraceTestInterface* test_interface =
         reinterpret_cast<TraceTestInterface*>(user_data);
     pw_trace_TraceEventReturnFlags ret = 0;
     if (test_interface->action_ != ActionOnEvent::None &&
-        (test_interface->event_match_.trace_ref == trace_ref ||
-         test_interface->event_match_.event_type == event_type ||
-         test_interface->event_match_.module == module ||
-         (trace_id != PW_TRACE_TRACE_ID_DEFAULT &&
-          test_interface->event_match_.trace_id == trace_id))) {
+        (test_interface->event_match_.trace_ref == event->trace_token ||
+         test_interface->event_match_.event_type == event->event_type ||
+         test_interface->event_match_.module == event->module ||
+         (event->trace_id != PW_TRACE_TRACE_ID_DEFAULT &&
+          test_interface->event_match_.trace_id == event->trace_id))) {
       if (test_interface->action_ == ActionOnEvent::Skip) {
         ret |= PW_TRACE_EVENT_RETURN_FLAGS_SKIP_EVENT;
       } else if (test_interface->action_ == ActionOnEvent::Enable) {
@@ -114,8 +109,8 @@
       }
     }
 
-    test_interface->current_trace_event_ =
-        TraceInfo{trace_ref, event_type, module, trace_id};
+    test_interface->current_trace_event_ = TraceInfo{
+        event->trace_token, event->event_type, event->module, event->trace_id};
     return ret;
   }