From 60c6b155f626f40e2971cda10aa4c3565897874a Mon Sep 17 00:00:00 2001
From: MediaPipe Team <mediapipe-team@google.com>
Date: Thu, 29 Dec 2022 10:16:10 -0800
Subject: [PATCH] Save an integer id in graph profiler objects to distinguish
 between different profiler instances during benchmarking.

PiperOrigin-RevId: 498409363
---
 .../framework/profiler/graph_profiler.cc      |  1 +
 mediapipe/framework/profiler/graph_profiler.h |  9 +++++++
 .../framework/profiler/graph_profiler_test.cc | 26 +++++++++++++++++++
 3 files changed, 36 insertions(+)

diff --git a/mediapipe/framework/profiler/graph_profiler.cc b/mediapipe/framework/profiler/graph_profiler.cc
index f14acfc7..6aead525 100644
--- a/mediapipe/framework/profiler/graph_profiler.cc
+++ b/mediapipe/framework/profiler/graph_profiler.cc
@@ -194,6 +194,7 @@ void GraphProfiler::Initialize(
         "Calculator \"$0\" has already been added.", node_name);
   }
   profile_builder_ = std::make_unique<GraphProfileBuilder>(this);
+  graph_id_ = ++next_instance_id_;
 
   is_initialized_ = true;
 }
diff --git a/mediapipe/framework/profiler/graph_profiler.h b/mediapipe/framework/profiler/graph_profiler.h
index 23caed4e..6358cb05 100644
--- a/mediapipe/framework/profiler/graph_profiler.h
+++ b/mediapipe/framework/profiler/graph_profiler.h
@@ -237,6 +237,9 @@ class GraphProfiler : public std::enable_shared_from_this<ProfilingContext> {
     return validated_graph_;
   }
 
+  // Gets a numerical identifier for this GraphProfiler object.
+  uint64_t GetGraphId() { return graph_id_; }
+
  private:
   // This can be used to add packet info for the input streams to the graph.
   // It treats the stream defined by |stream_name| as a stream produced by a
@@ -357,6 +360,12 @@ class GraphProfiler : public std::enable_shared_from_this<ProfilingContext> {
   class GraphProfileBuilder;
   std::unique_ptr<GraphProfileBuilder> profile_builder_;
 
+  // The globally incrementing identifier for all graphs in a process.
+  static inline std::atomic_int next_instance_id_ = 0;
+
+  // A unique identifier for this object. Only unique within a process.
+  uint64_t graph_id_;
+
   // For testing.
   friend GraphProfilerTestPeer;
 };
diff --git a/mediapipe/framework/profiler/graph_profiler_test.cc b/mediapipe/framework/profiler/graph_profiler_test.cc
index 81ba90cd..75d1c7eb 100644
--- a/mediapipe/framework/profiler/graph_profiler_test.cc
+++ b/mediapipe/framework/profiler/graph_profiler_test.cc
@@ -442,6 +442,32 @@ TEST_F(GraphProfilerTestPeer, InitializeMultipleTimes) {
                "Cannot initialize .* multiple times.");
 }
 
+// Tests that graph identifiers are not reused, even after destruction.
+TEST_F(GraphProfilerTestPeer, InitializeMultipleProfilers) {
+  auto raw_graph_config = R"(
+    profiler_config {
+      enable_profiler: true
+    }
+    input_stream: "input_stream"
+    node {
+      calculator: "DummyTestCalculator"
+      input_stream: "input_stream"
+    })";
+  const int n_iterations = 100;
+  absl::flat_hash_set<int> seen_ids;
+  for (int i = 0; i < n_iterations; ++i) {
+    std::shared_ptr<ProfilingContext> profiler =
+        std::make_shared<ProfilingContext>();
+    auto graph_config = CreateGraphConfig(raw_graph_config);
+    mediapipe::ValidatedGraphConfig validated_graph;
+    QCHECK_OK(validated_graph.Initialize(graph_config));
+    profiler->Initialize(validated_graph);
+
+    int id = profiler->GetGraphId();
+    ASSERT_THAT(seen_ids, testing::Not(testing::Contains(id)));
+    seen_ids.insert(id);
+  }
+}
 // Tests that Pause(), Resume(), and Reset() works.
 TEST_F(GraphProfilerTestPeer, PauseResumeReset) {
   InitializeProfilerWithGraphConfig(R"(