datadog-ci 1.13.0 → 1.15.0

data/lib/datadog/ci/test_visibility/component.rb

@@ -1,7 +1,10 @@
 # frozen_string_literal: true
 
+require "drb"
 require "rbconfig"
 
+require "datadog/core/utils/forking"
+
 require_relative "context"
 require_relative "telemetry"
 require_relative "total_coverage"
@@ -10,24 +13,35 @@ require_relative "../codeowners/parser"
 require_relative "../contrib/instrumentation"
 require_relative "../ext/test"
 require_relative "../git/local_repository"
+require_relative "../utils/file_storage"
+require_relative "../utils/stateful"
 
 require_relative "../worker"
 
 module Datadog
   module CI
     module TestVisibility
-      # Common behavior for CI tests
+      # Core functionality of the library: tracing tests' execution
       class Component
-        attr_reader :test_suite_level_visibility_enabled, :logical_test_session_name, :known_tests, :known_tests_enabled
+        include Core::Utils::Forking
+        include Datadog::CI::Utils::Stateful
+
+        FILE_STORAGE_KEY = "test_visibility_component_state"
+
+        attr_reader :test_suite_level_visibility_enabled, :logical_test_session_name,
+          :known_tests, :known_tests_enabled, :context_service_uri, :local_test_suites_mode
 
         def initialize(
           known_tests_client:,
           test_suite_level_visibility_enabled: false,
           codeowners: Codeowners::Parser.new(Git::LocalRepository.root).parse,
-          logical_test_session_name: nil
+          logical_test_session_name: nil,
+          context_service_uri: nil
         )
           @test_suite_level_visibility_enabled = test_suite_level_visibility_enabled
-          @context = Context.new
+
+          @context = Context.new(test_visibility_component: self)
+
           @codeowners = codeowners
           @logical_test_session_name = logical_test_session_name
 
@@ -36,46 +50,76 @@ module Datadog
           @known_tests_enabled = false
           @known_tests_client = known_tests_client
           @known_tests = Set.new
+
+          # this is used for parallel test runners such as parallel_tests
+          if context_service_uri
+            @context_service_uri = context_service_uri
+            @is_client_process = true
+          end
+
+          # This is used for parallel test runners such as parallel_tests.
+          # If true, then test suites are created in the worker process, not the parent.
+          #
+          # The only test runner that requires creating test suites in the remote process is rails test runner because
+          # it splits workload by test, not by test suite.
+          #
+          # Another test runner that splits workload by test is knapsack_pro, but we lack distributed test sessions/test suties
+          # support for that one (as of 2025-03).
+          @local_test_suites_mode = true
         end
 
         def configure(library_configuration, test_session)
           return unless test_suite_level_visibility_enabled
+          return unless library_configuration.known_tests_enabled?
 
-          if library_configuration.known_tests_enabled?
-            @known_tests_enabled = true
-            fetch_known_tests(test_session)
-          end
+          @known_tests_enabled = true
+          return if load_component_state
+
+          fetch_known_tests(test_session)
+          store_component_state if test_session.distributed
         end
 
-        def start_test_session(service: nil, tags: {}, total_tests_count: 0)
+        def start_test_session(service: nil, tags: {}, estimated_total_tests_count: 0, distributed: false, local_test_suites_mode: true)
           return skip_tracing unless test_suite_level_visibility_enabled
 
-          test_session = @context.start_test_session(service: service, tags: tags)
-          test_session.total_tests_count = total_tests_count
+          @local_test_suites_mode = local_test_suites_mode
+
+          start_drb_service
+
+          test_session = maybe_remote_context.start_test_session(service: service, tags: tags)
+          test_session.estimated_total_tests_count = estimated_total_tests_count
+          test_session.distributed = distributed
 
           on_test_session_started(test_session)
+
           test_session
         end
 
         def start_test_module(test_module_name, service: nil, tags: {})
           return skip_tracing unless test_suite_level_visibility_enabled
 
-          test_module = @context.start_test_module(test_module_name, service: service, tags: tags)
+          test_module = maybe_remote_context.start_test_module(test_module_name, service: service, tags: tags)
           on_test_module_started(test_module)
+
           test_module
         end
 
         def start_test_suite(test_suite_name, service: nil, tags: {})
           return skip_tracing unless test_suite_level_visibility_enabled
 
-          test_suite = @context.start_test_suite(test_suite_name, service: service, tags: tags)
+          context = @local_test_suites_mode ? @context : maybe_remote_context
+
+          test_suite = context.start_test_suite(test_suite_name, service: service, tags: tags)
           on_test_suite_started(test_suite)
           test_suite
         end
 
         def trace_test(test_name, test_suite_name, service: nil, tags: {}, &block)
+          test_suite = active_test_suite(test_suite_name)
+          tags[Ext::Test::TAG_SUITE] ||= test_suite_name
+
           if block
-            @context.trace_test(test_name, test_suite_name, service: service, tags: tags) do |test|
+            @context.trace_test(test_name, test_suite, service: service, tags: tags) do |test|
               subscribe_to_after_stop_event(test.tracer_span)
 
               on_test_started(test)
@@ -84,7 +128,7 @@ module Datadog
               res
             end
           else
-            test = @context.trace_test(test_name, test_suite_name, service: service, tags: tags)
+            test = @context.trace_test(test_name, test_suite, service: service, tags: tags)
             subscribe_to_after_stop_event(test.tracer_span)
             on_test_started(test)
             test
@@ -110,15 +154,19 @@ module Datadog
         end
 
         def active_test_session
-          @context.active_test_session
+          maybe_remote_context.active_test_session
         end
 
         def active_test_module
-          @context.active_test_module
+          maybe_remote_context.active_test_module
         end
 
         def active_test_suite(test_suite_name)
-          @context.active_test_suite(test_suite_name)
+          # when fetching test_suite to use as test's context, try local context instance first
+          local_test_suite = @context.active_test_suite(test_suite_name)
+          return local_test_suite if local_test_suite
+
+          maybe_remote_context.active_test_suite(test_suite_name)
         end
 
         def deactivate_test
@@ -146,9 +194,18 @@ module Datadog
           test_suite = active_test_suite(test_suite_name)
           on_test_suite_finished(test_suite) if test_suite
 
+          # deactivation always happens on the same process where test suite is located
           @context.deactivate_test_suite(test_suite_name)
         end
 
+        def total_tests_count
+          maybe_remote_context.total_tests_count
+        end
+
+        def tests_skipped_by_tia_count
+          maybe_remote_context.tests_skipped_by_tia_count
+        end
+
         def itr_enabled?
           test_optimisation.enabled?
         end
@@ -157,6 +214,10 @@ module Datadog
           # noop, there is no thread owned by test visibility component
         end
 
+        def client_process?
+          forked? || @is_client_process
+        end
+
         private
 
         # DOMAIN EVENTS
@@ -167,10 +228,6 @@ module Datadog
           # finds and instruments additional test libraries that we support (ex: selenium-webdriver)
           Contrib::Instrumentation.instrument_on_session_start
 
-          # sends internal telemetry events
-          Telemetry.test_session_started(test_session)
-          Telemetry.event_created(test_session)
-
           # sets logical test session name if none provided by the user
           override_logical_test_session_name!(test_session) if logical_test_session_name.nil?
 
@@ -179,19 +236,22 @@ module Datadog
           remote.configure(test_session)
         end
 
+        # intentionally empty
         def on_test_module_started(test_module)
-          Telemetry.event_created(test_module)
         end
 
         def on_test_suite_started(test_suite)
           set_codeowners(test_suite)
-
-          Telemetry.event_created(test_suite)
         end
 
         def on_test_started(test)
-          # sometimes test suite is not being assigned correctly
-          # fix it by fetching the one single running test suite from the global context
+          maybe_remote_context.incr_total_tests_count
+
+          # Sometimes test suite is not being assigned correctly.
+          # Fix it by fetching the one single running test suite from the process context.
+          #
+          # This is a hack to fix some edge cases that come from some minitest plugins,
+          # especially thoughtbot/shoulda-context.
           fix_test_suite!(test) if test.test_suite_id.nil?
           validate_test_suite_level_visibility_correctness(test)
 
@@ -208,14 +268,18 @@ module Datadog
         end
 
         def on_test_session_finished(test_session)
-          test_optimisation.write_test_session_tags(test_session)
+          test_optimisation.write_test_session_tags(test_session, maybe_remote_context.tests_skipped_by_tia_count)
 
           TotalCoverage.extract_lines_pct(test_session)
 
           Telemetry.event_finished(test_session)
+
+          Utils::FileStorage.cleanup
         end
 
         def on_test_module_finished(test_module)
+          @context.stop_all_test_suites
+
           Telemetry.event_finished(test_module)
         end
 
@@ -225,7 +289,7 @@ module Datadog
 
         def on_test_finished(test)
           test_optimisation.stop_coverage(test)
-          test_optimisation.count_skipped_test(test)
+          test_optimisation.on_test_finished(test, maybe_remote_context)
 
           Telemetry.event_finished(test)
 
@@ -258,7 +322,7 @@ module Datadog
         def fix_test_suite!(test)
           return unless test_suite_level_visibility_enabled
 
-          test_suite = @context.single_active_test_suite
+          test_suite = maybe_remote_context.single_active_test_suite
           unless test_suite
             Datadog.logger.debug do
               "Trying to fix test suite for test [#{test.name}] but no single test suite is running."
@@ -370,6 +434,42 @@ module Datadog
         def test_management
           Datadog.send(:components).test_management
         end
+
+        # DISTRIBUTED RUBY CONTEXT
+        def start_drb_service
+          return if @context_service_uri
+          return if client_process?
+
+          @context_service = DRb.start_service("drbunix:", @context)
+          @context_service_uri = @context_service.uri
+        end
+
+        # depending on whether we are in a forked process or not, returns either the global context or its DRbObject
+        def maybe_remote_context
+          return @context unless client_process?
+          return @context_client if defined?(@context_client)
+
+          # at least once per fork we must stop the running DRb server that was copied from the parent process
+          # otherwise, client will be confused thinking it's server which leads to terrible bugs
+          @context_service&.stop_service
+
+          @context_client = DRbObject.new_with_uri(@context_service_uri)
+        end
+
+        # Implementation of Stateful interface
+        def serialize_state
+          {
+            known_tests: @known_tests
+          }
+        end
+
+        def restore_state(state)
+          @known_tests = state[:known_tests]
+        end
+
+        def storage_key
+          FILE_STORAGE_KEY
+        end
       end
     end
   end

data/lib/datadog/ci/test_visibility/context.rb

@@ -4,8 +4,9 @@ require "datadog/tracing"
 require "datadog/tracing/contrib/component"
 require "datadog/tracing/trace_digest"
 
-require_relative "store/global"
-require_relative "store/local"
+require_relative "store/process"
+require_relative "store/fiber_local"
+require_relative "telemetry"
 
 require_relative "../ext/app_types"
 require_relative "../ext/environment"
@@ -27,24 +28,38 @@ module Datadog
       # Its responsibility includes building domain models for test visibility as well.
       # Internally it uses Datadog::Tracing module to create spans.
       class Context
-        def initialize
-          @local_context = Store::Local.new
-          @global_context = Store::Global.new
+        attr_reader :total_tests_count, :tests_skipped_by_tia_count
+
+        def initialize(test_visibility_component:)
+          @test_visibility_component = test_visibility_component
+
+          @fiber_local_context = Store::FiberLocal.new
+          @process_context = Store::Process.new
+
+          @mutex = Mutex.new
+
+          @total_tests_count = 0
+          @tests_skipped_by_tia_count = 0
         end
 
         def start_test_session(service: nil, tags: {})
-          @global_context.fetch_or_activate_test_session do
+          @process_context.fetch_or_activate_test_session do
             tracer_span = start_datadog_tracer_span(
               "test.session", build_tracing_span_options(service, Ext::AppTypes::TYPE_TEST_SESSION)
             )
             set_session_context(tags, tracer_span)
 
-            build_test_session(tracer_span, tags)
+            test_session = build_test_session(tracer_span, tags)
+
+            Telemetry.test_session_started(test_session)
+            Telemetry.event_created(test_session)
+
+            test_session
           end
         end
 
         def start_test_module(test_module_name, service: nil, tags: {})
-          @global_context.fetch_or_activate_test_module do
+          @process_context.fetch_or_activate_test_module do
             set_inherited_globals(tags)
             set_session_context(tags)
 
@@ -53,12 +68,16 @@ module Datadog
             )
             set_module_context(tags, tracer_span)
 
-            build_test_module(tracer_span, tags)
+            test_module = build_test_module(tracer_span, tags)
+
+            Telemetry.event_created(test_module)
+
+            test_module
           end
         end
 
         def start_test_suite(test_suite_name, service: nil, tags: {})
-          @global_context.fetch_or_activate_test_suite(test_suite_name) do
+          @process_context.fetch_or_activate_test_suite(test_suite_name) do
             set_inherited_globals(tags)
             set_session_context(tags)
             set_module_context(tags)
@@ -66,17 +85,21 @@ module Datadog
             tracer_span = start_datadog_tracer_span(
               test_suite_name, build_tracing_span_options(service, Ext::AppTypes::TYPE_TEST_SUITE)
             )
-            set_suite_context(tags, span: tracer_span)
+            set_suite_context(tags, test_suite: tracer_span)
 
-            build_test_suite(tracer_span, tags)
+            test_suite = build_test_suite(tracer_span, tags)
+
+            Telemetry.event_created(test_suite)
+
+            test_suite
           end
         end
 
-        def trace_test(test_name, test_suite_name, service: nil, tags: {}, &block)
+        def trace_test(test_name, test_suite, service: nil, tags: {}, &block)
           set_inherited_globals(tags)
           set_session_context(tags)
           set_module_context(tags)
-          set_suite_context(tags, name: test_suite_name)
+          set_suite_context(tags, test_suite: test_suite)
 
           tags[Ext::Test::TAG_NAME] = test_name
           tags[Ext::Test::TAG_TYPE] ||= Ext::Test::Type::TEST
@@ -93,14 +116,14 @@ module Datadog
             start_datadog_tracer_span(test_name, span_options) do |tracer_span|
               test = build_test(tracer_span, tags)
 
-              @local_context.activate_test(test) do
+              @fiber_local_context.activate_test(test) do
                 block.call(test)
               end
             end
           else
             tracer_span = start_datadog_tracer_span(test_name, span_options)
             test = build_test(tracer_span, tags)
-            @local_context.activate_test(test)
+            @fiber_local_context.activate_test(test)
             test
           end
         end
@@ -129,39 +152,51 @@ module Datadog
         end
 
         def active_test
-          @local_context.active_test
+          @fiber_local_context.active_test
         end
 
         def active_test_session
-          @global_context.active_test_session
+          @process_context.active_test_session
         end
 
         def active_test_module
-          @global_context.active_test_module
+          @process_context.active_test_module
         end
 
         def active_test_suite(test_suite_name)
-          @global_context.active_test_suite(test_suite_name)
+          @process_context.active_test_suite(test_suite_name)
         end
 
         def single_active_test_suite
-          @global_context.fetch_single_test_suite
+          @process_context.fetch_single_test_suite
+        end
+
+        def stop_all_test_suites
+          @process_context.stop_all_test_suites
         end
 
         def deactivate_test
-          @local_context.deactivate_test
+          @fiber_local_context.deactivate_test
         end
 
         def deactivate_test_session
-          @global_context.deactivate_test_session!
+          @process_context.deactivate_test_session!
         end
 
         def deactivate_test_module
-          @global_context.deactivate_test_module!
+          @process_context.deactivate_test_module!
         end
 
         def deactivate_test_suite(test_suite_name)
-          @global_context.deactivate_test_suite!(test_suite_name)
+          @process_context.deactivate_test_suite!(test_suite_name)
+        end
+
+        def incr_total_tests_count
+          @mutex.synchronize { @total_tests_count += 1 }
+        end
+
+        def incr_tests_skipped_by_tia_count
+          @mutex.synchronize { @tests_skipped_by_tia_count += 1 }
         end
 
         private
@@ -212,37 +247,32 @@ module Datadog
 
         # PROPAGATING CONTEXT FROM TOP-LEVEL TO THE LOWER LEVELS
         def set_inherited_globals(tags)
-          # this code achieves the same as @global_context.inheritable_session_tags.merge(tags)
-          # but without allocating a new hash
-          @global_context.inheritable_session_tags.each do |key, value|
+          # Copy inheritable tags from the test session context to the provided tags
+          test_session_context&.inheritable_tags&.each do |key, value|
             tags[key] = value unless tags.key?(key)
           end
         end
 
         def set_session_context(tags, test_session = nil)
-          test_session ||= active_test_session
+          # we need to call TestVisibility::Component here because active test session might be remote
+          test_session ||= test_session_context
           tags[Ext::Test::TAG_TEST_SESSION_ID] = test_session.id.to_s if test_session
         end
 
         def set_module_context(tags, test_module = nil)
-          test_module ||= active_test_module
+          # we need to call TestVisibility::Component here because active test module might be remote
+          test_module ||= test_module_context
           if test_module
             tags[Ext::Test::TAG_TEST_MODULE_ID] = test_module.id.to_s
             tags[Ext::Test::TAG_MODULE] = test_module.name
           end
         end
 
-        def set_suite_context(tags, span: nil, name: nil)
-          return if span.nil? && name.nil?
-
-          test_suite = span || active_test_suite(name)
+        def set_suite_context(tags, test_suite: nil)
+          return if test_suite.nil?
 
-          if test_suite
-            tags[Ext::Test::TAG_TEST_SUITE_ID] = test_suite.id.to_s
-            tags[Ext::Test::TAG_SUITE] = test_suite.name
-          else
-            tags[Ext::Test::TAG_SUITE] = name
-          end
+          tags[Ext::Test::TAG_TEST_SUITE_ID] = test_suite.id.to_s
+          tags[Ext::Test::TAG_SUITE] = test_suite.name
         end
 
         # INTERACTIONS WITH TRACING
@@ -268,11 +298,42 @@ module Datadog
         end
 
         def build_tracing_span_options(service, type, other_options = {})
-          other_options[:service] = service || @global_context.service
+          other_options[:service] = service || test_session_context&.service
           other_options[:type] = type
 
           other_options
         end
+
+        # one of:
+        #   1. Currrent test session from the Store::Process
+        #   2. Readonly copy of the remote test session (if test session was started by a parent process and local copy was created)
+        #   3. Remote test session as DRb::DRbObject link (in this case also local copy will be created)
+        def test_session_context
+          local_test_session = @process_context.active_test_session
+          return local_test_session if local_test_session
+
+          local_readonly_test_session = @process_context.readonly_test_session
+          return local_readonly_test_session if local_readonly_test_session
+
+          remote_test_session = @test_visibility_component.active_test_session
+          @process_context.set_readonly_test_session(remote_test_session)
+
+          remote_test_session
+        end
+
+        # works similar to test_session_context
+        def test_module_context
+          local_test_module = @process_context.active_test_module
+          return local_test_module if local_test_module
+
+          local_readonly_test_module = @process_context.readonly_test_module
+          return local_readonly_test_module if local_readonly_test_module
+
+          remote_test_module = @test_visibility_component.active_test_module
+          @process_context.set_readonly_test_module(remote_test_module)
+
+          remote_test_module
+        end
       end
     end
   end

data/lib/datadog/ci/test_visibility/null_component.rb

@@ -8,7 +8,7 @@ module Datadog
         def configure(_, _)
         end
 
-        def start_test_session(service: nil, tags: {}, total_tests_count: 0)
+        def start_test_session(service: nil, tags: {}, estimated_total_tests_count: 0)
           skip_tracing
         end
 
@@ -63,6 +63,10 @@ module Datadog
         def logical_test_session_name
         end
 
+        def client_process?
+          false
+        end
+
         private
 
         def skip_tracing(block = nil)

data/lib/datadog/ci/test_visibility/store/{local.rb → fiber_local.rb}

@@ -4,7 +4,7 @@ module Datadog
   module CI
     module TestVisibility
       module Store
-        class Local
+        class FiberLocal
           def initialize
             @key = :datadog_ci_active_test
 

data/lib/datadog/ci/test_visibility/store/{global.rb → process.rb}

@@ -1,11 +1,16 @@
 # frozen_string_literal: true
 
+require_relative "../../readonly_test_session"
+require_relative "../../readonly_test_module"
+
 module Datadog
   module CI
     module TestVisibility
       module Store
         # This context is shared between threads and represents the current test session and test module.
-        class Global
+        class Process
+          attr_reader :readonly_test_session, :readonly_test_module
+
           def initialize
             # we are using Monitor instead of Mutex because it is reentrant
             @mutex = Monitor.new
@@ -13,6 +18,11 @@ module Datadog
             @test_session = nil
             @test_module = nil
             @test_suites = {}
+
+            # small copies of id, name and some tags: store them in the current process to set session/module context
+            # for any spans faster
+            @readonly_test_session = nil
+            @readonly_test_module = nil
           end
 
           def fetch_or_activate_test_suite(test_suite_name, &block)
@@ -53,20 +63,10 @@ module Datadog
             @mutex.synchronize { @test_suites[test_suite_name] }
           end
 
-          def service
-            @mutex.synchronize do
-              @test_session&.service
-            end
-          end
-
-          def inheritable_session_tags
+          def stop_all_test_suites
             @mutex.synchronize do
-              test_session = @test_session
-              if test_session
-                test_session.inheritable_tags
-              else
-                {}
-              end
+              @test_suites.each_value(&:finish)
+              @test_suites.clear
             end
           end
 
@@ -81,6 +81,18 @@ module Datadog
           def deactivate_test_suite!(test_suite_name)
             @mutex.synchronize { @test_suites.delete(test_suite_name) }
           end
+
+          def set_readonly_test_session(remote_test_session)
+            return if remote_test_session.nil?
+
+            @readonly_test_session = Datadog::CI::ReadonlyTestSession.new(remote_test_session)
+          end
+
+          def set_readonly_test_module(remote_test_module)
+            return if remote_test_module.nil?
+
+            @readonly_test_module = Datadog::CI::ReadonlyTestModule.new(remote_test_module)
+          end
         end
       end
     end