Fix, build and test the energy profiler daemon

CMakeLists.txt

@@ -178,6 +178,7 @@ if(NOT WIN32)
   add_subdirectory(profiling/chrome-tracing)
   add_subdirectory(profiling/space-time-stack)
   add_subdirectory(profiling/perfetto-connector)
+  add_subdirectory(profiling/energy-profiler)
 endif()
 
 # External lib connectors

profiling/energy-profiler/CMakeLists.txt

@@ -0,0 +1,10 @@
+# Periodic sampling thread for the energy profiling tools. It has no Kokkos
+# Tools entry points, so it is built as a static building block, not a tool.
+find_package(Threads REQUIRED)
+
+add_library(kp_energy_profiler_daemon STATIC daemon.cpp)
+set_target_properties(kp_energy_profiler_daemon
+  PROPERTIES POSITION_INDEPENDENT_CODE ON)
+target_include_directories(kp_energy_profiler_daemon
+  PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
+target_link_libraries(kp_energy_profiler_daemon PUBLIC Threads::Threads)

profiling/energy-profiler/daemon.cpp

@@ -2,29 +2,35 @@
 // SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project
 
 #include "daemon.hpp"
-#include <stdexcept>
+#include <condition_variable>
+#include <mutex>
 #include <thread>
 
 namespace KokkosTools::EnergyProfiler {
 void Daemon::start() {
   if (!running_) {
     running_ = true;
-    thread_  = std::thread(&Daemon::run, this);
+    thread_  = std::jthread([this](std::stop_token stop) { run(stop); });
   }
 }
 
 void Daemon::stop() {
   if (running_) {
-    running_ = false;
+    thread_.request_stop();
     thread_.join();
+    running_ = false;
   }
 }
 
-void Daemon::run() {
-  while (running_) {
-    auto next_run = std::chrono::high_resolution_clock::now() + interval_;
+void Daemon::run(std::stop_token stop) {
+  std::mutex mutex;
+  std::condition_variable_any wake;
+  std::unique_lock<std::mutex> lock(mutex);
+  while (!stop.stop_requested()) {
+    auto next_run = std::chrono::steady_clock::now() + interval_;
     func_();
-    std::this_thread::sleep_until(next_run);
+    // Returns early when a stop is requested.
+    wake.wait_until(lock, stop, next_run, [] { return false; });
   }
 }
 }  // namespace KokkosTools::EnergyProfiler

profiling/energy-profiler/daemon.hpp

@@ -4,14 +4,16 @@
 #ifndef KOKKOSP_ENERGY_PROFILER_DAEMON_HPP
 #define KOKKOSP_ENERGY_PROFILER_DAEMON_HPP
 
+#include <atomic>
 #include <chrono>
 #include <functional>
+#include <stop_token>
 #include <thread>
 
 namespace KokkosTools::EnergyProfiler {
 class Daemon {
  public:
-  Daemon(std::function<void()> func, const std::chrono::duration& interval)
+  Daemon(std::function<void()> func, std::chrono::nanoseconds interval)
       : interval_(interval), func_(std::move(func)){};
 
   void start();
@@ -20,11 +22,13 @@ class Daemon {
   auto& get_thread() { return thread_; }
 
  private:
-  void run();
-  std::chrono::duration interval_;
-  bool running_{false};
+  void run(std::stop_token stop);
+  std::chrono::nanoseconds interval_;
+  std::atomic<bool> running_{false};
   std::function<void()> func_;
-  std::thread thread_;
+  // Declared last so that it is destroyed first: the thread is stopped and
+  // joined before the members it uses go away.
+  std::jthread thread_;
 };
 }  // namespace KokkosTools::EnergyProfiler
 #endif

tests/CMakeLists.txt

@@ -129,3 +129,4 @@ target_link_libraries(test_common PUBLIC GTest::gtest GTest::gmock Kokkos::kokko
 add_subdirectory(space-time-stack)
 add_subdirectory(sampler)
 add_subdirectory(simple-kernel-timer)
+add_subdirectory(energy-profiler)

tests/energy-profiler/CMakeLists.txt

@@ -0,0 +1,5 @@
+kp_add_executable_and_test(
+    TARGET_NAME test_energy_profiler_daemon
+    SOURCE_FILE test_daemon.cpp
+)
+target_link_libraries(test_energy_profiler_daemon PRIVATE kp_energy_profiler_daemon)

tests/energy-profiler/test_daemon.cpp

@@ -0,0 +1,105 @@
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project
+
+#include <atomic>
+#include <chrono>
+#include <thread>
+
+#include "gtest/gtest.h"
+
+#include "daemon.hpp"
+
+using KokkosTools::EnergyProfiler::Daemon;
+using namespace std::chrono_literals;
+
+//! Sampling interval of the energy profiler (50 Hz).
+constexpr auto sampling_interval = 20ms;
+
+// Timing bounds are loose on purpose: shared CI runners can delay a thread by
+// tens of milliseconds.
+
+TEST(EnergyProfilerDaemonTest, calls_periodically) {
+  std::atomic<int> calls{0};
+  Daemon daemon([&] { ++calls; }, sampling_interval);
+  EXPECT_FALSE(daemon.is_running());
+
+  daemon.start();
+  EXPECT_TRUE(daemon.is_running());
+  std::this_thread::sleep_for(300ms);
+  daemon.stop();
+  EXPECT_FALSE(daemon.is_running());
+
+  // About 15 calls in 300 ms.
+  EXPECT_GE(calls.load(), 5);
+  EXPECT_LE(calls.load(), 20);
+}
+
+TEST(EnergyProfilerDaemonTest, repeated_start_and_stop) {
+  Daemon daemon([] {}, sampling_interval);
+  daemon.start();
+  daemon.start();
+  EXPECT_TRUE(daemon.is_running());
+
+  daemon.stop();
+  daemon.stop();
+  EXPECT_FALSE(daemon.is_running());
+}
+
+TEST(EnergyProfilerDaemonTest, restart_after_stop) {
+  std::atomic<int> calls{0};
+  Daemon daemon([&] { ++calls; }, sampling_interval);
+
+  daemon.start();
+  std::this_thread::sleep_for(100ms);
+  daemon.stop();
+  const int first_run = calls.load();
+
+  daemon.start();
+  std::this_thread::sleep_for(100ms);
+  daemon.stop();
+
+  EXPECT_GT(first_run, 0);
+  EXPECT_GT(calls.load(), first_run);
+}
+
+TEST(EnergyProfilerDaemonTest, stop_does_not_wait) {
+  Daemon daemon([] {}, 10s);
+  daemon.start();
+  std::this_thread::sleep_for(20ms);
+
+  const auto begin = std::chrono::steady_clock::now();
+  daemon.stop();
+  EXPECT_LT(std::chrono::steady_clock::now() - begin, 1s);
+}
+
+TEST(EnergyProfilerDaemonTest, destructor_stops_daemon) {
+  std::atomic<int> calls{0};
+  {
+    Daemon daemon([&] { ++calls; }, sampling_interval);
+    daemon.start();
+    std::this_thread::sleep_for(50ms);
+  }
+  const int after_destruction = calls.load();
+  std::this_thread::sleep_for(50ms);
+
+  EXPECT_GT(after_destruction, 0);
+  EXPECT_EQ(calls.load(), after_destruction);
+}
+
+TEST(EnergyProfilerDaemonTest, slow_function) {
+  std::atomic<int> calls{0};
+  Daemon daemon(
+      [&] {
+        ++calls;
+        std::this_thread::sleep_for(50ms);
+      },
+      sampling_interval);
+
+  daemon.start();
+  std::this_thread::sleep_for(300ms);
+  daemon.stop();
+
+  // One call every 50 ms instead of every 20 ms: about 6 in 300 ms.
+  EXPECT_GE(calls.load(), 2);
+  EXPECT_LE(calls.load(), 8);
+}