From ffc727a859d9ae7631fbc7647392efa05032211a Mon Sep 17 00:00:00 2001
From: xdje42 <xdje42@gmail.com>
Date: Thu, 24 Oct 2024 02:22:58 -0700
Subject: [PATCH] Verify RegisterProfilerManager doesn't overwrite an existing
 registration (#1837)

* Verify RegisterProfilerManager doesn't overwrite an existing registration

Tested:
Add a second registration to test/profiler_manager_test.cc and
verify the test crashes as expected.

* Verify RegisterProfilerManager doesn't overwrite an existing registration

Tested:
Configure with:
cmake -GNinja -DCMAKE_BUILD_TYPE=Debug -DBENCHMARK_DOWNLOAD_DEPENDENCIES=on
Then run:
ctest -R profiler_manager_gtest
Before change test fails (expected), after change test passes (expected)

---------

Co-authored-by: dominic <510002+dmah42@users.noreply.github.com>
---
 src/benchmark.cc               |  4 ++++
 test/CMakeLists.txt            |  1 +
 test/profiler_manager_gtest.cc | 42 ++++++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+)
 create mode 100644 test/profiler_manager_gtest.cc

diff --git a/src/benchmark.cc b/src/benchmark.cc
index 26050774..0ea90aeb 100644
--- a/src/benchmark.cc
+++ b/src/benchmark.cc
@@ -668,6 +668,10 @@ void RegisterMemoryManager(MemoryManager* manager) {
 }
 
 void RegisterProfilerManager(ProfilerManager* manager) {
+  // Don't allow overwriting an existing manager.
+  if (manager != nullptr) {
+    BM_CHECK_EQ(internal::profiler_manager, nullptr);
+  }
   internal::profiler_manager = manager;
 }
 
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index 815b5818..321e24d9 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -254,6 +254,7 @@ if (BENCHMARK_ENABLE_GTEST_TESTS)
   add_gtest(perf_counters_gtest)
   add_gtest(time_unit_gtest)
   add_gtest(min_time_parse_gtest)
+  add_gtest(profiler_manager_gtest)
 endif(BENCHMARK_ENABLE_GTEST_TESTS)
 
 ###############################################################################
diff --git a/test/profiler_manager_gtest.cc b/test/profiler_manager_gtest.cc
new file mode 100644
index 00000000..434e4eca
--- /dev/null
+++ b/test/profiler_manager_gtest.cc
@@ -0,0 +1,42 @@
+#include <memory>
+
+#include "benchmark/benchmark.h"
+#include "gtest/gtest.h"
+
+namespace {
+
+class TestProfilerManager : public benchmark::ProfilerManager {
+ public:
+  void AfterSetupStart() override { ++start_called; }
+  void BeforeTeardownStop() override { ++stop_called; }
+
+  int start_called = 0;
+  int stop_called = 0;
+};
+
+void BM_empty(benchmark::State& state) {
+  for (auto _ : state) {
+    auto iterations = state.iterations();
+    benchmark::DoNotOptimize(iterations);
+  }
+}
+BENCHMARK(BM_empty);
+
+TEST(ProfilerManager, ReregisterManager) {
+#if GTEST_HAS_DEATH_TEST
+  // Tests only runnable in debug mode (when BM_CHECK is enabled).
+#ifndef NDEBUG
+#ifndef TEST_BENCHMARK_LIBRARY_HAS_NO_ASSERTIONS
+  ASSERT_DEATH_IF_SUPPORTED(
+      {
+        std::unique_ptr<TestProfilerManager> pm(new TestProfilerManager());
+        benchmark::RegisterProfilerManager(pm.get());
+        benchmark::RegisterProfilerManager(pm.get());
+      },
+      "RegisterProfilerManager");
+#endif
+#endif
+#endif
+}
+
+}  // namespace