multi threaded code coverage support for datadog_cov

ext/datadog_cov/datadog_cov.c

@@ -1,7 +1,11 @@
 #include <ruby.h>
 #include <ruby/debug.h>
 
-const int PROFILE_FRAMES_BUFFER_SIZE = 1;
+#define PROFILE_FRAMES_BUFFER_SIZE 1
+
+// threading modes
+#define SINGLE_THREADED_COVERAGE_MODE 0
+#define MULTI_THREADED_COVERAGE_MODE 1
 
 char *ruby_strndup(const char *str, size_t size)
 {
@@ -27,6 +31,8 @@ struct dd_cov_data
 
   VALUE *frame_buffer;
   uintptr_t last_filename_ptr;
+
+  int threading_mode;
 };
 
 static void dd_cov_mark(void *ptr)
@@ -71,6 +77,7 @@ static VALUE dd_cov_allocate(VALUE klass)
   dd_cov_data->ignored_path_len = 0;
   dd_cov_data->last_filename_ptr = 0;
   dd_cov_data->frame_buffer = xcalloc(PROFILE_FRAMES_BUFFER_SIZE, sizeof(VALUE));
+  dd_cov_data->threading_mode = MULTI_THREADED_COVERAGE_MODE;
 
   return obj;
 }
@@ -88,9 +95,25 @@ static VALUE dd_cov_initialize(int argc, VALUE *argv, VALUE self)
   }
   VALUE rb_ignored_path = rb_hash_lookup(opt, ID2SYM(rb_intern("ignored_path")));
 
+  VALUE rb_threading_mode = rb_hash_lookup(opt, ID2SYM(rb_intern("threading_mode")));
+  int threading_mode;
+  if (!RTEST(rb_threading_mode) || rb_threading_mode == ID2SYM(rb_intern("multi")))
+  {
+    threading_mode = MULTI_THREADED_COVERAGE_MODE;
+  }
+  else if (rb_threading_mode == ID2SYM(rb_intern("single")))
+  {
+    threading_mode = SINGLE_THREADED_COVERAGE_MODE;
+  }
+  else
+  {
+    rb_raise(rb_eArgError, "mode is invalid");
+  }
+
   struct dd_cov_data *dd_cov_data;
   TypedData_Get_Struct(self, struct dd_cov_data, &dd_cov_data_type, dd_cov_data);
 
+  dd_cov_data->threading_mode = threading_mode;
   dd_cov_data->root_len = RSTRING_LEN(rb_root);
   dd_cov_data->root = ruby_strndup(RSTRING_PTR(rb_root), dd_cov_data->root_len);
 
@@ -153,31 +176,45 @@ static void dd_cov_update_coverage(rb_event_flag_t event, VALUE data, VALUE self
 
 static VALUE dd_cov_start(VALUE self)
 {
-  // get current thread
-  VALUE thval = rb_thread_current();
 
   struct dd_cov_data *dd_cov_data;
   TypedData_Get_Struct(self, struct dd_cov_data, &dd_cov_data_type, dd_cov_data);
 
-  if (dd_cov_data->root_len != 0)
+  if (dd_cov_data->root_len == 0)
+  {
+    return self;
+  }
+
+  if (dd_cov_data->threading_mode == SINGLE_THREADED_COVERAGE_MODE)
   {
-    // add event hook
+    VALUE thval = rb_thread_current();
     rb_thread_add_event_hook(thval, dd_cov_update_coverage, RUBY_EVENT_LINE, self);
   }
+  else
+  {
+    rb_add_event_hook(dd_cov_update_coverage, RUBY_EVENT_LINE, self);
+  }
 
   return self;
 }
 
 static VALUE dd_cov_stop(VALUE self)
 {
-  // get current thread
-  VALUE thval = rb_thread_current();
-  // remove event hook for the current thread
-  rb_thread_remove_event_hook(thval, dd_cov_update_coverage);
-
   struct dd_cov_data *dd_cov_data;
   TypedData_Get_Struct(self, struct dd_cov_data, &dd_cov_data_type, dd_cov_data);
 
+  if (dd_cov_data->threading_mode == SINGLE_THREADED_COVERAGE_MODE)
+  {
+    // get current thread
+    VALUE thval = rb_thread_current();
+    // remove event hook for the current thread
+    rb_thread_remove_event_hook(thval, dd_cov_update_coverage);
+  }
+  else
+  {
+    rb_remove_event_hook(dd_cov_update_coverage);
+  }
+
   VALUE res = dd_cov_data->coverage;
 
   dd_cov_data->coverage = rb_hash_new();

sig/datadog/ci/itr/coverage/ddcov.rbs

@@ -3,8 +3,9 @@ module Datadog
     module ITR
       module Coverage
         class DDCov
+          type threading_mode = :multi | :single
 
-          def initialize: (root: String, ignored_path: String?) -> void
+          def initialize: (root: String, ignored_path: String?, ?threading_mode: threading_mode?) -> void
 
           def start: () -> void
 

spec/ddcov/ddcov_spec.rb

@@ -161,50 +161,75 @@ def absolute_path(path)
 
       context "multi threaded execution" do
         def thread_local_cov
-          Thread.current[:datadog_ci_cov] ||= described_class.new(root: root)
+          Thread.current[:datadog_ci_cov] ||= described_class.new(root: root, threading_mode: threading_mode)
         end
 
-        it "collects coverage for each thread separately" do
-          t1_queue = Queue.new
-          t2_queue = Queue.new
+        context "in single threaded coverage mode" do
+          let(:threading_mode) { :single }
 
-          t1 = Thread.new do
-            cov = thread_local_cov
-            cov.start
+          it "collects coverage for each thread separately" do
+            t1_queue = Queue.new
+            t2_queue = Queue.new
 
-            t1_queue << :ready
-            expect(t2_queue.pop).to be(:ready)
+            t1 = Thread.new do
+              cov = thread_local_cov
+              cov.start
 
-            expect(calculator.add(1, 2)).to eq(3)
-            expect(calculator.multiply(1, 2)).to eq(2)
+              t1_queue << :ready
+              expect(t2_queue.pop).to be(:ready)
 
-            t1_queue << :done
-            expect(t2_queue.pop).to be :done
+              expect(calculator.add(1, 2)).to eq(3)
+              expect(calculator.multiply(1, 2)).to eq(2)
 
-            coverage = cov.stop
-            expect(coverage.size).to eq(2)
-            expect(coverage.keys).to include(absolute_path("calculator/operations/add.rb"))
-            expect(coverage.keys).to include(absolute_path("calculator/operations/multiply.rb"))
+              t1_queue << :done
+              expect(t2_queue.pop).to be :done
+
+              coverage = cov.stop
+              expect(coverage.size).to eq(2)
+              expect(coverage.keys).to include(absolute_path("calculator/operations/add.rb"))
+              expect(coverage.keys).to include(absolute_path("calculator/operations/multiply.rb"))
+            end
+
+            t2 = Thread.new do
+              cov = thread_local_cov
+              cov.start
+
+              t2_queue << :ready
+              expect(t1_queue.pop).to be(:ready)
+
+              expect(calculator.subtract(1, 2)).to eq(-1)
+
+              t2_queue << :done
+              expect(t1_queue.pop).to be :done
+
+              coverage = cov.stop
+              expect(coverage.size).to eq(1)
+              expect(coverage.keys).to include(absolute_path("calculator/operations/subtract.rb"))
+            end
+
+            [t1, t2].each(&:join)
           end
+        end
+
+        context "in multi threaded code coverage mode" do
+          let(:threading_mode) { :multi }
 
-          t2 = Thread.new do
+          it "collects coverage for background threads" do
             cov = thread_local_cov
             cov.start
 
-            t2_queue << :ready
-            expect(t1_queue.pop).to be(:ready)
-
-            expect(calculator.subtract(1, 2)).to eq(-1)
+            t = Thread.new do
+              expect(calculator.add(1, 2)).to eq(3)
+            end
 
-            t2_queue << :done
-            expect(t1_queue.pop).to be :done
+            expect(calculator.multiply(1, 2)).to eq(2)
+            t.join
 
             coverage = cov.stop
-            expect(coverage.size).to eq(1)
-            expect(coverage.keys).to include(absolute_path("calculator/operations/subtract.rb"))
+            expect(coverage.size).to eq(2)
+            expect(coverage.keys).to include(absolute_path("calculator/operations/add.rb"))
+            expect(coverage.keys).to include(absolute_path("calculator/operations/multiply.rb"))
           end
-
-          [t1, t2].each(&:join)
         end
       end
     end