[POC] per test code coverage using rb_thread_add_event_hook in C extension

.gitignore

@@ -15,3 +15,8 @@ Gemfile-*.lock
 .ruby-version
 .DS_Store
 /test.rb
+
+# Native extension binaries
+lib/**/*.bundle
+lib/**/*.so
+lib/**/*.o

Gemfile

@@ -12,6 +12,9 @@ gem "pry"
 gem "rake"
 gem "os"
 
+# To compile native extensions
+gem "rake-compiler"
+
 gem "climate_control"
 
 gem "rspec"

Rakefile

@@ -121,3 +121,9 @@ end
 
 desc "CI task; it runs all tests for current version of Ruby"
 task ci: "test:all"
+
+require "rake/extensiontask"
+
+Rake::ExtensionTask.new("ddcov") do |ext|
+  ext.lib_dir = "lib/ddcov"
+end

datadog-ci.gemspec

@@ -35,8 +35,11 @@ Gem::Specification.new do |spec|
       NOTICE
       README.md
       lib/**/*
+      ext/**/*
     ]].select { |fn| File.file?(fn) } # We don't want directories, only files
 
+  spec.extensions << "ext/ddcov/extconf.rb"
+
   spec.require_paths = ["lib"]
 
   spec.add_dependency "msgpack"

ext/ddcov/ddcov.c

@@ -0,0 +1,185 @@
+#include <ruby.h>
+#include <ruby/debug.h>
+#include <stdio.h>
+
+// Utils
+static bool prefix(const char *pre, const char *str)
+{
+  return strncmp(pre, str, strlen(pre)) == 0;
+}
+
+// const
+#define DD_COVERAGE_TARGET_FILES 1
+#define DD_COVERAGE_TARGET_LINES 2
+
+// Data structure
+struct dd_cov_data
+{
+  char *root;
+  int mode;
+  VALUE coverage;
+};
+
+static void dd_cov_mark(void *ptr)
+{
+  // printf("MARK\n");
+  struct dd_cov_data *dd_cov_data = ptr;
+  rb_gc_mark_movable(dd_cov_data->coverage);
+}
+
+static void dd_cov_free(void *ptr)
+{
+  // printf("FREE\n");
+  struct dd_cov_data *dd_cov_data = ptr;
+
+  xfree(dd_cov_data);
+}
+
+static void dd_cov_compact(void *ptr)
+{
+  // printf("COMPACT\n");
+  struct dd_cov_data *dd_cov_data = ptr;
+  dd_cov_data->coverage = rb_gc_location(dd_cov_data->coverage);
+}
+
+const rb_data_type_t dd_cov_data_type = {
+    .wrap_struct_name = "dd_cov",
+    .function = {
+        .dmark = dd_cov_mark,
+        .dfree = dd_cov_free,
+        .dsize = NULL,
+        .dcompact = dd_cov_compact},
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY};
+
+static VALUE dd_cov_allocate(VALUE klass)
+{
+  // printf("ALLOCATE\n");
+  struct dd_cov_data *dd_cov_data;
+  VALUE obj = TypedData_Make_Struct(klass, struct dd_cov_data, &dd_cov_data_type, dd_cov_data);
+  dd_cov_data->coverage = rb_hash_new();
+  return obj;
+}
+
+// DDCov methods
+static VALUE dd_cov_initialize(int argc, VALUE *argv, VALUE self)
+{
+  VALUE opt;
+  int mode;
+
+  rb_scan_args(argc, argv, "10", &opt);
+  VALUE rb_root = rb_hash_lookup(opt, ID2SYM(rb_intern("root")));
+  if (!RTEST(rb_root))
+  {
+    rb_raise(rb_eArgError, "root is required");
+  }
+
+  VALUE rb_mode = rb_hash_lookup(opt, ID2SYM(rb_intern("mode")));
+  if (!RTEST(rb_mode) || rb_mode == ID2SYM(rb_intern("files")))
+  {
+    mode = DD_COVERAGE_TARGET_FILES;
+  }
+  else if (rb_mode == ID2SYM(rb_intern("lines")))
+  {
+    mode = DD_COVERAGE_TARGET_LINES;
+  }
+  else
+  {
+    rb_raise(rb_eArgError, "mode is invalid");
+  }
+
+  printf("MODE IS %d\n", mode);
+  struct dd_cov_data *dd_cov_data;
+  TypedData_Get_Struct(self, struct dd_cov_data, &dd_cov_data_type, dd_cov_data);
+
+  // printf("struct got\n");
+  dd_cov_data->root = StringValueCStr(rb_root);
+  dd_cov_data->mode = mode;
+  // printf("root set\n");
+
+  return Qnil;
+}
+
+static void dd_cov_update_line_coverage(rb_event_flag_t event, VALUE data, VALUE self, ID id, VALUE klass)
+{
+  // printf("EVENT HOOK FIRED\n");
+  // printf("FILE: %s\n", rb_sourcefile());
+  // printf("LINE: %d\n", rb_sourceline());
+  struct dd_cov_data *dd_cov_data;
+  TypedData_Get_Struct(data, struct dd_cov_data, &dd_cov_data_type, dd_cov_data);
+
+  const char *filename = rb_sourcefile();
+  if (filename == 0)
+  {
+    return;
+  }
+
+  if (!prefix(dd_cov_data->root, filename))
+  {
+    return;
+  }
+
+  unsigned long len_filename = strlen(filename);
+
+  VALUE rb_str_source_file = rb_str_new(filename, len_filename);
+
+  if (dd_cov_data->mode == DD_COVERAGE_TARGET_FILES)
+  {
+    rb_hash_aset(dd_cov_data->coverage, rb_str_source_file, Qtrue);
+    return;
+  }
+
+  if (dd_cov_data->mode == DD_COVERAGE_TARGET_LINES)
+  {
+    VALUE rb_lines = rb_hash_aref(dd_cov_data->coverage, rb_str_source_file);
+    if (rb_lines == Qnil)
+    {
+      rb_lines = rb_ary_new();
+      rb_hash_aset(dd_cov_data->coverage, rb_str_source_file, rb_lines);
+    }
+
+    VALUE line_number = INT2FIX(rb_sourceline());
+    if (rb_ary_includes(rb_lines, line_number) == Qfalse)
+    {
+      rb_ary_push(rb_lines, line_number);
+    }
+  }
+}
+
+static VALUE dd_cov_start(VALUE self)
+{
+  // get current thread
+  VALUE thval = rb_thread_current();
+
+  // add event hook
+  rb_thread_add_event_hook(thval, dd_cov_update_line_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_line_coverage);
+
+  struct dd_cov_data *dd_cov_data;
+  TypedData_Get_Struct(self, struct dd_cov_data, &dd_cov_data_type, dd_cov_data);
+
+  VALUE cov = dd_cov_data->coverage;
+
+  dd_cov_data->coverage = rb_hash_new();
+
+  return cov;
+}
+
+void Init_ddcov(void)
+{
+  VALUE cDDCov = rb_define_class("DDCov", rb_cObject);
+
+  rb_define_alloc_func(cDDCov, dd_cov_allocate);
+
+  rb_define_method(cDDCov, "initialize", dd_cov_initialize, -1);
+  rb_define_method(cDDCov, "start", dd_cov_start, 0);
+  rb_define_method(cDDCov, "stop", dd_cov_stop, 0);
+}

ext/ddcov/extconf.rb

@@ -0,0 +1,5 @@
+require "mkmf"
+
+extension_name = "ddcov"
+dir_config(extension_name)
+create_makefile(extension_name)

lib/datadog/ci/itr/coverage/collector.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require "coverage"
+
+require_relative "../../utils/git"
+require_relative "../../../../ddcov/ddcov"
+
+module Datadog
+  module CI
+    module Itr
+      module Coverage
+        class Collector
+          def initialize(mode: :files, enabled: true)
+            # TODO: make this thread local
+            # modes available: :files, :lines
+            @mode = mode
+            @enabled = enabled
+
+            if @enabled
+              @ddcov = DDCov.new(root: Utils::Git.root, mode: mode)
+            end
+          end
+
+          def setup
+            if @enabled
+              p "RUNNING WITH CODE COVERAGE ENABLED AND MODE #{@mode}"
+            else
+              p "RUNNING WITH CODE COVERAGE DISABLED"
+            end
+          end
+
+          def start
+            @ddcov.start if @enabled
+          end
+
+          def stop
+            @ddcov.stop if @enabled
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/ci/itr/coverage/filter.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+require_relative "../../utils/git"
+
+module Datadog
+  module CI
+    module Itr
+      module Coverage
+        # not filter, but rather filter and transformer
+        class Filter
+          def self.call(raw_result)
+            new.call(raw_result)
+          end
+
+          def initialize(root: Utils::Git.root)
+            @regex = /\A#{Regexp.escape(root + File::SEPARATOR)}/i.freeze
+          end
+
+          def call(raw_result)
+            return nil if raw_result.nil?
+
+            # p "RAW"
+            # p raw_result.count
+
+            raw_result.filter_map do |path, coverage|
+              next unless path =~ @regex
+
+              [path, coverage]
+            end
+          end
+
+          private
+
+          def convert_lines_to_bitmap(lines)
+            bitmap = []
+            current = 0
+            bit = 1 << 63
+            lines.each do |line|
+              if !line.nil? && line > 0
+                current |= bit
+              end
+              bit >>= 1
+              if bit == 0
+                bitmap << current
+                current = 0
+                bit = 1 << 63
+              end
+            end
+            bitmap << current
+            lines
+          end
+        end
+      end
+    end
+  end
+end