From 916e5e54be39f15e8292b9cf93fcf630aaec995f Mon Sep 17 00:00:00 2001
From: Jose Baiocchi <jbaiocchi@google.com>
Date: Thu, 6 Aug 2020 10:37:20 -0700
Subject: [PATCH] Cleanup TraceMe idioms

PiperOrigin-RevId: 325260871
Change-Id: Ifa970a3529c1c53a5b25899a61c8de940963d423
---
 tensorflow/core/profiler/lib/annotated_traceme.h   |  6 +++---
 .../python/profiler/internal/python_hooks.cc       | 14 +++++++++-----
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/tensorflow/core/profiler/lib/annotated_traceme.h b/tensorflow/core/profiler/lib/annotated_traceme.h
index c3257e2adbe..636b901e226 100644
--- a/tensorflow/core/profiler/lib/annotated_traceme.h
+++ b/tensorflow/core/profiler/lib/annotated_traceme.h
@@ -38,12 +38,12 @@ class AnnotatedTraceMe {
     bool annotation_enabled = ScopedAnnotation::IsEnabled();
     bool traceme_enabled = TraceMe::Active(level);
     if (TF_PREDICT_FALSE(annotation_enabled || traceme_enabled)) {
-      string label = name_generator();
+      string name = name_generator();
       if (annotation_enabled) {
-        scoped_annotation_.emplace(absl::string_view(label));
+        scoped_annotation_.emplace(absl::string_view(name));
       }
       if (TF_PREDICT_TRUE(traceme_enabled)) {
-        trace_me_.emplace(std::move(label), level);
+        trace_me_.emplace([name = std::move(name)] { return name; }, level);
       }
     }
   }
diff --git a/tensorflow/python/profiler/internal/python_hooks.cc b/tensorflow/python/profiler/internal/python_hooks.cc
index 33e182f8de0..ee2ad1e254b 100644
--- a/tensorflow/python/profiler/internal/python_hooks.cc
+++ b/tensorflow/python/profiler/internal/python_hooks.cc
@@ -120,8 +120,11 @@ void PythonHooks::ProfileFast(PyFrameObject* frame, int what, PyObject* arg) {
       function = py::reinterpret_borrow<py::str>(f_code->co_name);
     }
 
-    tracemes_[thread_id].push_back(absl::make_unique<TraceMe>(absl::StrCat(
-        "$", io::Basename(filename), ":", line_no, " ", function)));
+    tracemes_[thread_id].push_back(
+        absl::make_unique<TraceMe>([&filename, line_no, &function] {
+          return absl::StrCat("$", io::Basename(filename), ":", line_no, " ",
+                              function);
+        }));
   } else if (what == PyTrace_C_CALL && PyCFunction_Check(arg)) {
     // Python stack does not have a filename/line_no for native calls.
     auto* func = reinterpret_cast<PyCFunctionObject*>(arg);
@@ -139,9 +142,10 @@ void PythonHooks::ProfileFast(PyFrameObject* frame, int what, PyObject* arg) {
       filename = "<unknown>";
     }
 
-    string function(func->m_ml->ml_name);
-    tracemes_[thread_id].push_back(absl::make_unique<TraceMe>(
-        absl::StrCat(filename, " ", func->m_ml->ml_name)));
+    tracemes_[thread_id].push_back(
+        absl::make_unique<TraceMe>([&filename, func] {
+          return absl::StrCat(filename, " ", func->m_ml->ml_name);
+        }));
   } else if (what == PyTrace_RETURN || what == PyTrace_C_RETURN ||
              what == PyTrace_EXCEPTION || what == PyTrace_C_EXCEPTION) {
     auto& thread_tracemes = tracemes_[thread_id];