From 88f1a04474bf3986413c516b367d445fb83c35e3 Mon Sep 17 00:00:00 2001
From: Junwhan Ahn <junwhan@google.com>
Date: Tue, 3 Feb 2026 01:01:14 -0800
Subject: [PATCH] Always run the name generators passed to TraceMe on debug
 builds to help catch bugs

Today, the name generators are run only when there is an active tracing session, which means that none of the code inside the name generators are tested at all. Since users sometimes perform non-trivial work (such as accessing pointers that may be dangling) inside the name generators, it is useful to test such code by always invoking the name generators in debug builds where the extra overhead is not a concern.

PiperOrigin-RevId: 864742934
---
 tsl/profiler/lib/traceme.h | 34 ++++++++++++++++++++++++++++------
 1 file changed, 28 insertions(+), 6 deletions(-)

diff --git a/tsl/profiler/lib/traceme.h b/tsl/profiler/lib/traceme.h
index 7c9a1f4ac..b58c74b31 100644
--- a/tsl/profiler/lib/traceme.h
+++ b/tsl/profiler/lib/traceme.h
@@ -155,6 +155,12 @@ class TraceMe {
                          TraceMeRecorder::CheckFilter(filter_mask))) {
       name_.Emplace(std::forward<NameGeneratorT>(name_generator)());
       start_time_ = GetCurrentTimeNanos();
+    } else {
+#ifndef NDEBUG
+      // Always invoke the name generator on debug builds to catch bugs in name
+      // generators that would otherwise only be caught when tracing is enabled.
+      std::forward<NameGeneratorT>(name_generator)();
+#endif
     }
 #endif
   }
@@ -212,12 +218,16 @@ class TraceMe {
       std::enable_if_t<std::is_invocable_v<MetadataGeneratorT>, bool> = true>
   void AppendMetadata(MetadataGeneratorT&& metadata_generator) {
 #if !defined(IS_MOBILE_PLATFORM)
-    if (TF_PREDICT_FALSE(start_time_ != kUntracedActivity)) {
-      if (TF_PREDICT_TRUE(TraceMeRecorder::Active())) {
-        traceme_internal::AppendMetadata(
-            &name_.value,
-            std::forward<MetadataGeneratorT>(metadata_generator)());
-      }
+    if (TF_PREDICT_FALSE(start_time_ != kUntracedActivity) &&
+        TF_PREDICT_TRUE(TraceMeRecorder::Active())) {
+      traceme_internal::AppendMetadata(
+          &name_.value, std::forward<MetadataGeneratorT>(metadata_generator)());
+    } else {
+#ifndef NDEBUG
+      // Always invoke the name generator on debug builds to catch bugs in name
+      // generators that would otherwise only be caught when tracing is enabled.
+      std::forward<MetadataGeneratorT>(metadata_generator)();
+#endif
     }
 #endif
   }
@@ -239,6 +249,12 @@ class TraceMe {
       TraceMeRecorder::Record({std::forward<NameGeneratorT>(name_generator)(),
                                GetCurrentTimeNanos(), -activity_id});
       return activity_id;
+    } else {
+#ifndef NDEBUG
+      // Always invoke the name generator on debug builds to catch bugs in name
+      // generators that would otherwise only be caught when tracing is enabled.
+      std::forward<NameGeneratorT>(name_generator)();
+#endif
     }
 #endif
     return kUntracedActivity;
@@ -300,6 +316,12 @@ class TraceMe {
       int64_t now = GetCurrentTimeNanos();
       TraceMeRecorder::Record({std::forward<NameGeneratorT>(name_generator)(),
                                /*start_time=*/now, /*end_time=*/now});
+    } else {
+#ifndef NDEBUG
+      // Always invoke the name generator on debug builds to catch bugs in name
+      // generators that would otherwise only be caught when tracing is enabled.
+      std::forward<NameGeneratorT>(name_generator)();
+#endif
     }
 #endif
   }