Redesign crash tracker to behave as regular object

ext/datadog_profiling_native_extension/crash_tracker.c

@@ -2,28 +2,36 @@
 #include <datadog/common.h>
 #include <libdatadog_helpers.h>
 
-static VALUE _native_start_crashtracker(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _self);
+static VALUE _native_start_or_update_on_fork(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _self);
+static VALUE _native_stop(DDTRACE_UNUSED VALUE _self);
 
 // Used to report Ruby VM crashes.
 // Once initialized, segfaults will be reported automatically using libdatadog.
 
 void crash_tracker_init(VALUE profiling_module) {
   VALUE crash_tracker_class = rb_define_class_under(profiling_module, "CrashTracker", rb_cObject);
 
-  rb_define_singleton_method(crash_tracker_class, "_native_start_crashtracker", _native_start_crashtracker, -1);
+  rb_define_singleton_method(crash_tracker_class, "_native_start_or_update_on_fork", _native_start_or_update_on_fork, -1);
+  rb_define_singleton_method(crash_tracker_class, "_native_stop", _native_stop, 0);
 }
 
-static VALUE _native_start_crashtracker(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _self) {
+static VALUE _native_start_or_update_on_fork(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _self) {
   VALUE options;
   rb_scan_args(argc, argv, "0:", &options);
 
   VALUE exporter_configuration = rb_hash_fetch(options, ID2SYM(rb_intern("exporter_configuration")));
   VALUE path_to_crashtracking_receiver_binary = rb_hash_fetch(options, ID2SYM(rb_intern("path_to_crashtracking_receiver_binary")));
   VALUE tags_as_array = rb_hash_fetch(options, ID2SYM(rb_intern("tags_as_array")));
+  VALUE action = rb_hash_fetch(options, ID2SYM(rb_intern("action")));
+  VALUE start_action = ID2SYM(rb_intern("start"));
+  VALUE update_on_fork_action = ID2SYM(rb_intern("update_on_fork"));
 
   ENFORCE_TYPE(exporter_configuration, T_ARRAY);
   ENFORCE_TYPE(tags_as_array, T_ARRAY);
   ENFORCE_TYPE(path_to_crashtracking_receiver_binary, T_STRING);
+  ENFORCE_TYPE(action, T_SYMBOL);
+
+  if (action != start_action && action != update_on_fork_action) rb_raise(rb_eArgError, "Unexpected action: %+"PRIsVALUE, action);
 
   VALUE version = ddtrace_version();
   ddog_Endpoint endpoint = endpoint_from(exporter_configuration);
@@ -45,13 +53,26 @@ static VALUE _native_start_crashtracker(int argc, VALUE *argv, DDTRACE_UNUSED VA
     .tags = &tags,
   };
 
-  ddog_prof_Profile_Result result = ddog_prof_crashtracker_init(config, metadata);
+  ddog_prof_Profile_Result result =
+    action == start_action ?
+      ddog_prof_crashtracker_init(config, metadata) :
+      ddog_prof_crashtracker_update_on_fork(config, metadata);
 
   // Clean up before potentially raising any exceptions
   ddog_Vec_Tag_drop(tags);
 
   if (result.tag == DDOG_PROF_PROFILE_RESULT_ERR) {
-    rb_raise(rb_eRuntimeError, "Failed to initialize the crash tracker: %"PRIsVALUE, get_error_details_and_drop(&result.err));
+    rb_raise(rb_eRuntimeError, "Failed to start/update the crash tracker: %"PRIsVALUE, get_error_details_and_drop(&result.err));
+  }
+
+  return Qtrue;
+}
+
+static VALUE _native_stop(DDTRACE_UNUSED VALUE _self) {
+  ddog_prof_Profile_Result result = ddog_prof_crashtracker_shutdown();
+
+  if (result.tag == DDOG_PROF_PROFILE_RESULT_ERR) {
+    rb_raise(rb_eRuntimeError, "Failed to stop the crash tracker: %"PRIsVALUE, get_error_details_and_drop(&result.err));
   }
 
   return Qtrue;

lib/datadog/profiling/crash_tracker.rb

@@ -17,41 +17,61 @@ module Profiling
     #
     # Methods prefixed with _native_ are implemented in `crash_tracker.c`
     class CrashTracker
-      def self.build_crash_tracker(
+      private
+
+      attr_reader :exporter_configuration, :tags_as_array, :path_to_crashtracking_receiver_binary
+
+      public
+
+      def initialize(
         exporter_configuration:,
         tags:,
         path_to_crashtracking_receiver_binary: Libdatadog.path_to_crashtracking_receiver_binary
       )
+        @exporter_configuration = exporter_configuration
+        @tags_as_array = tags.to_a
+        @path_to_crashtracking_receiver_binary = path_to_crashtracking_receiver_binary
+      end
+
+      def start
+        start_or_update_on_fork(action: :start)
+      end
+
+      def reset_after_fork
+        start_or_update_on_fork(action: :update_on_fork)
+      end
+
+      def stop
+        begin
+          self.class._native_stop
+          Datadog.logger.debug('Crash tracking stopped successfully')
+        rescue => e
+          Datadog.logger.error("Failed to stop crash tracking: #{e.message}")
+        end
+      end
+
+      private
+
+      def start_or_update_on_fork(action:)
         unless path_to_crashtracking_receiver_binary
           Datadog.logger.warn(
-            'Cannot enable profiling crash tracking as no path_to_crashtracking_receiver_binary was found'
+            "Cannot #{action} profiling crash tracking as no path_to_crashtracking_receiver_binary was found"
           )
           return
         end
 
         begin
-          new(
+          self.class._native_start_or_update_on_fork(
+            action: action,
             exporter_configuration: exporter_configuration,
             path_to_crashtracking_receiver_binary: path_to_crashtracking_receiver_binary,
-            tags_as_array: tags.to_a,
-          ).tap {
-            Datadog.logger.debug('Crash tracker enabled')
-          }
+            tags_as_array: tags_as_array,
+          )
+          Datadog.logger.debug("Crash tracking #{action} successful")
         rescue => e
-          Datadog.logger.error("Failed to initialize crash tracking: #{e.message}")
-          nil
+          Datadog.logger.error("Failed to #{action} crash tracking: #{e.message}")
         end
       end
-
-      private
-
-      def initialize(exporter_configuration:, path_to_crashtracking_receiver_binary:, tags_as_array:)
-        self.class._native_start_crashtracker(
-          exporter_configuration: exporter_configuration,
-          path_to_crashtracking_receiver_binary: path_to_crashtracking_receiver_binary,
-          tags_as_array: tags_as_array,
-        )
-      end
     end
   end
 end

sig/datadog/profiling/crash_tracker.rbs

@@ -1,25 +1,38 @@
 module Datadog
   module Profiling
     class CrashTracker
-      def self.build_crash_tracker: (
-        exporter_configuration: [:agentless | :agent, untyped],
-        tags: ::Hash[::String, ::String],
-        ?path_to_crashtracking_receiver_binary: ::String,
-      ) -> CrashTracker?
+      type exporter_configuration_array = [:agentless | :agent, untyped]
 
       private
 
+      attr_reader exporter_configuration: exporter_configuration_array
+      attr_reader tags_as_array: ::Array[[::String, ::String]]
+      attr_reader path_to_crashtracking_receiver_binary: ::String
+
+      public
+
       def initialize: (
-        exporter_configuration: [:agentless | :agent, untyped],
+        exporter_configuration: exporter_configuration_array,
+        tags: ::Hash[::String, ::String],
         path_to_crashtracking_receiver_binary: ::String,
-        tags_as_array: ::Array[[::String, ::String]],
       ) -> void
 
-      def self._native_start_crashtracker: (
-        exporter_configuration: [:agentless | :agent, untyped],
-        path_to_crashtracking_receiver_binary: ::String,
+      def start: -> void
+      def stop: -> void
+      def reset_after_fork: -> void
+
+      private
+
+      def start_or_update_on_fork: (action: :start | :update_on_fork) -> void
+
+      def self._native_start_or_update_on_fork: (
+        action: :start | :update_on_fork,
+        exporter_configuration: exporter_configuration_array,
         tags_as_array: ::Array[[::String, ::String]],
+        path_to_crashtracking_receiver_binary: ::String,
       ) -> void
+
+      def self._native_stop: -> void
     end
   end
 end