refactor

Author: KowalskiThomas

ddtrace/internal/datadog/profiling/dd_wrapper/CMakeLists.txt

@@ -47,6 +47,7 @@ add_library(
     src/ddup_interface.cpp
     src/profiler_stats.cpp
     src/profile.cpp
+    src/profile_borrow.cpp
     src/sample.cpp
     src/sample_manager.cpp
     src/static_sample_pool.cpp

ddtrace/internal/datadog/profiling/dd_wrapper/include/ddup_interface.hpp

@@ -1,14 +1,13 @@
 #pragma once
 
-#include <stddef.h>
-#include <stdint.h>
+#include <cstdint>
 #include <string_view>
 #include <unordered_map>
 
 // Forward decl of the return pointer
 namespace Datadog {
 class Sample;
-}
+} // namespace Datadog
 
 #ifdef __cplusplus
 extern "C"
@@ -68,6 +67,10 @@ extern "C"
                          int64_t line);
     void ddup_push_absolute_ns(Datadog::Sample* sample, int64_t timestamp_ns);
     void ddup_push_monotonic_ns(Datadog::Sample* sample, int64_t monotonic_ns);
+
+    void ddup_increment_sampling_event_count();
+    void ddup_increment_sample_count();
+
     void ddup_flush_sample(Datadog::Sample* sample);
     // Stack v2 specific flush, which reverses the locations
     void ddup_flush_sample_v2(Datadog::Sample* sample);

ddtrace/internal/datadog/profiling/dd_wrapper/include/profile.hpp

@@ -15,9 +15,13 @@ extern "C"
 
 namespace Datadog {
 
+class ProfileBorrow;
+
 // Serves to collect individual samples, as well as lengthen the scope of string data
 class Profile
 {
+    friend class ProfileBorrow;
+
   private:
     // Serialization for static state
     // - string table
@@ -45,6 +49,10 @@ class Profile
 
     Datadog::ProfilerStats profiler_stats{};
 
+    // Internal access methods - not for direct use
+    ddog_prof_Profile& profile_borrow_internal();
+    void profile_release();
+
   public:
     // State management
     void one_time_init(SampleType type, unsigned int _max_nframes);
@@ -53,11 +61,8 @@ class Profile
 
     // Getters
     size_t get_sample_type_length();
-    ddog_prof_Profile& profile_borrow();
-    ProfilerStats& profiler_stats_borrow();
 
-    void profiler_stats_release();
-    void profile_release();
+    ProfileBorrow borrow();
 
     // constref getters
     const ValueIndex& val();

ddtrace/internal/datadog/profiling/dd_wrapper/include/profile_borrow.hpp

@@ -0,0 +1,33 @@
+#pragma once
+
+#include "profile.hpp"
+
+namespace Datadog {
+
+// Forward declaration
+class Profile;
+
+// RAII wrapper for borrowing both profile and stats under a single lock
+class ProfileBorrow
+{
+  private:
+    Profile* profile_ptr;
+
+  public:
+    explicit ProfileBorrow(Profile& profile);
+    ~ProfileBorrow();
+
+    // Disable copy
+    ProfileBorrow(const ProfileBorrow&) = delete;
+    ProfileBorrow& operator=(const ProfileBorrow&) = delete;
+
+    // Enable move
+    ProfileBorrow(ProfileBorrow&& other) noexcept;
+    ProfileBorrow& operator=(ProfileBorrow&& other) noexcept;
+
+    // Accessors
+    ddog_prof_Profile& profile();
+    ProfilerStats& stats();
+};
+
+} // namespace Datadog

ddtrace/internal/datadog/profiling/dd_wrapper/include/sample.hpp

@@ -2,7 +2,7 @@
 
 #include "libdatadog_helpers.hpp"
 #include "profile.hpp"
-#include "profiler_stats.hpp"
+#include "profile_borrow.hpp"
 #include "types.hpp"
 
 #include <string>
@@ -135,9 +135,7 @@ class Sample
     // Flushes the current buffer, clearing it
     bool flush_sample(bool reverse_locations = false);
 
-    static ddog_prof_Profile& profile_borrow();
-    static ProfilerStats& profiler_stats_borrow();
-    static void profile_release();
+    static ProfileBorrow profile_borrow();
     static void postfork_child();
     Sample(SampleType _type_mask, unsigned int _max_nframes);
 

ddtrace/internal/datadog/profiling/dd_wrapper/src/ddup_interface.cpp

@@ -303,6 +303,20 @@ ddup_push_monotonic_ns(Datadog::Sample* sample, int64_t monotonic_ns) // cppchec
     sample->push_monotonic_ns(monotonic_ns);
 }
 
+void
+ddup_increment_sampling_event_count() // cppcheck-suppress unusedFunction
+{
+    auto borrowed = Datadog::Sample::profile_borrow();
+    borrowed.stats().increment_sampling_event_count();
+}
+
+void
+ddup_increment_sample_count() // cppcheck-suppress unusedFunction
+{
+    auto borrowed = Datadog::Sample::profile_borrow();
+    borrowed.stats().increment_sample_count();
+}
+
 void
 ddup_flush_sample(Datadog::Sample* sample) // cppcheck-suppress unusedFunction
 {
@@ -350,18 +364,19 @@ ddup_upload() // cppcheck-suppress unusedFunction
     //  be modified.  It gets released and cleared after uploading.
     // * Uploading cancels inflight uploads. There are better ways to do this, but this is what
     //   we have for now.
-    uploader.upload(Datadog::Sample::profile_borrow(), Datadog::Sample::profiler_stats_borrow());
-    Datadog::Sample::profiler_stats_borrow().reset_state();
-    Datadog::Sample::profile_release();
-    return true;
+    auto borrowed = Datadog::Sample::profile_borrow();
+    bool success = uploader.upload(borrowed.profile(), borrowed.stats());
+    borrowed.stats().reset_state();
+    return success;
 }
 
 void
 ddup_profile_set_endpoints(
   std::unordered_map<int64_t, std::string_view> span_ids_to_endpoints) // cppcheck-suppress unusedFunction
 {
     static bool already_warned = false; // cppcheck-suppress threadsafety-threadsafety
-    ddog_prof_Profile& profile = Datadog::Sample::profile_borrow();
+    auto borrowed = Datadog::Sample::profile_borrow();
+    ddog_prof_Profile& profile = borrowed.profile();
     for (const auto& [span_id, trace_endpoint] : span_ids_to_endpoints) {
         ddog_CharSlice trace_endpoint_slice = Datadog::to_slice(trace_endpoint);
         auto res = ddog_prof_Profile_set_endpoint(&profile, span_id, trace_endpoint_slice);
@@ -375,14 +390,14 @@ ddup_profile_set_endpoints(
             ddog_Error_drop(&err);
         }
     }
-    Datadog::Sample::profile_release();
 }
 
 void
 ddup_profile_add_endpoint_counts(std::unordered_map<std::string_view, int64_t> trace_endpoints_to_counts)
 {
     static bool already_warned = false; // cppcheck-suppress threadsafety-threadsafety
-    ddog_prof_Profile& profile = Datadog::Sample::profile_borrow();
+    auto borrowed = Datadog::Sample::profile_borrow();
+    ddog_prof_Profile& profile = borrowed.profile();
     for (const auto& [trace_endpoint, count] : trace_endpoints_to_counts) {
         ddog_CharSlice trace_endpoint_slice = Datadog::to_slice(trace_endpoint);
         auto res = ddog_prof_Profile_add_endpoint_count(&profile, trace_endpoint_slice, count);
@@ -396,5 +411,4 @@ ddup_profile_add_endpoint_counts(std::unordered_map<std::string_view, int64_t> t
             ddog_Error_drop(&err);
         }
     }
-    Datadog::Sample::profile_release();
 }

ddtrace/internal/datadog/profiling/dd_wrapper/src/profile.cpp

@@ -1,6 +1,8 @@
 #include "profile.hpp"
 
 #include "libdatadog_helpers.hpp"
+#include "profile_borrow.hpp"
+#include "profiler_stats.hpp"
 
 #include <functional>
 #include <iostream>
@@ -129,11 +131,16 @@ Datadog::Profile::get_sample_type_length()
     return samplers.size();
 }
 
+Datadog::ProfileBorrow
+Datadog::Profile::borrow()
+{
+    return ProfileBorrow(*this);
+}
+
 ddog_prof_Profile&
-Datadog::Profile::profile_borrow()
+Datadog::Profile::profile_borrow_internal()
 {
-    // We could wrap this in an object for better RAII, but since this
-    // sequence is only used in a single place, we'll hold off on that sidequest.
+    // Note: Caller is responsible for ensuring profile_release() is called
     profile_mtx.lock();
     return cur_profile;
 }
@@ -221,5 +228,6 @@ Datadog::Profile::postfork_child()
 {
     new (&profile_mtx) std::mutex();
     // Reset the profile to clear any samples collected in the parent process
+    profiler_stats.reset_state();
     reset_profile();
 }

ddtrace/internal/datadog/profiling/dd_wrapper/src/profile_borrow.cpp

@@ -0,0 +1,50 @@
+#include "profile_borrow.hpp"
+#include "profile.hpp"
+
+Datadog::ProfileBorrow::ProfileBorrow(Profile& profile)
+  : profile_ptr(&profile)
+{
+    // Lock the mutex on construction
+    profile_ptr->profile_borrow_internal();
+}
+
+Datadog::ProfileBorrow::~ProfileBorrow()
+{
+    if (profile_ptr) {
+        profile_ptr->profile_release();
+    }
+}
+
+Datadog::ProfileBorrow::ProfileBorrow(ProfileBorrow&& other) noexcept
+  : profile_ptr(other.profile_ptr)
+{
+    other.profile_ptr = nullptr;
+}
+
+Datadog::ProfileBorrow&
+Datadog::ProfileBorrow::operator=(ProfileBorrow&& other) noexcept
+{
+    if (this != &other) {
+        // Release current lock if any
+        if (profile_ptr) {
+            profile_ptr->profile_release();
+        }
+
+        // Take ownership from other
+        profile_ptr = other.profile_ptr;
+        other.profile_ptr = nullptr;
+    }
+    return *this;
+}
+
+ddog_prof_Profile&
+Datadog::ProfileBorrow::profile()
+{
+    return profile_ptr->cur_profile;
+}
+
+Datadog::ProfilerStats&
+Datadog::ProfileBorrow::stats()
+{
+    return profile_ptr->profiler_stats;
+}

ddtrace/internal/datadog/profiling/dd_wrapper/src/sample.cpp

@@ -523,22 +523,10 @@ Datadog::Sample::is_timeline_enabled() const
     return timeline_enabled;
 }
 
-ddog_prof_Profile&
+Datadog::ProfileBorrow
 Datadog::Sample::profile_borrow()
 {
-    return profile_state.profile_borrow();
-}
-
-Datadog::ProfilerStats&
-Datadog::Sample::profiler_stats_borrow()
-{
-    return profile_state.profiler_stats_borrow();
-}
-
-void
-Datadog::Sample::profile_release()
-{
-    profile_state.profile_release();
+    return profile_state.borrow();
 }
 
 void

ddtrace/internal/datadog/profiling/stack_v2/include/stack_renderer.hpp

@@ -1,13 +1,7 @@
 #pragma once
 
-#include <chrono>
-#include <fstream>
-#include <memory>
-#include <mutex>
 #include <string>
 #include <string_view>
-#include <thread>
-#include <vector>
 
 #include "python_headers.hpp"