refactor callbacks and retry strategy

lib/datadog/ci/test_retries/component.rb

@@ -13,6 +13,8 @@ module TestRetries
       # - retrying failed tests - improve success rate of CI pipelines
       # - retrying new tests - detect flaky tests as early as possible to prevent them from being merged
       class Component
+        FIBER_LOCAL_CURRENT_RETRY_STRATEGY_KEY = :__dd_current_retry_strategy
+
         attr_reader :retry_failed_tests_enabled, :retry_failed_tests_max_attempts,
           :retry_failed_tests_total_limit, :retry_failed_tests_count,
           :retry_new_tests_enabled, :retry_new_tests_duration_thresholds, :retry_new_tests_percentage_limit,
@@ -68,29 +70,15 @@ def configure(library_settings, test_session)
         end
 
         def with_retries(&block)
-          # @type var retry_strategy: Strategy::Base
-          retry_strategy = nil
-
-          test_finished_callback = lambda do |test_span|
-            if retry_strategy.nil?
-              # we always run test at least once and after first pass create a correct retry strategy
-              retry_strategy = build_strategy(test_span)
-            else
-              # after each retry we record the result, strategy will decide if we should retry again
-              retry_strategy&.record_retry(test_span)
-            end
-          end
-
-          # TODO: is there a better way to let test_visibility_component know that we are running under retries component?
-          test_visibility_component.set_test_finished_callback(test_finished_callback)
+          self.current_retry_strategy = nil
 
           loop do
             yield
 
-            break unless retry_strategy&.should_retry?
+            break unless current_retry_strategy&.should_retry?
           end
         ensure
-          test_visibility_component.remove_test_finished_callback
+          self.current_retry_strategy = nil
         end
 
         def build_strategy(test_span)
@@ -106,8 +94,30 @@ def build_strategy(test_span)
           end
         end
 
+        def record_test_finished(test_span)
+          if current_retry_strategy.nil?
+            # we always run test at least once and after the first pass create a correct retry strategy
+            self.current_retry_strategy = build_strategy(test_span)
+          else
+            # after each retry we record the result, strategy will decide if we should retry again
+            current_retry_strategy&.record_retry(test_span)
+          end
+        end
+
+        def record_test_span_duration(tracer_span)
+          # noop
+        end
+
         private
 
+        def current_retry_strategy
+          Thread.current[FIBER_LOCAL_CURRENT_RETRY_STRATEGY_KEY]
+        end
+
+        def current_retry_strategy=(strategy)
+          Thread.current[FIBER_LOCAL_CURRENT_RETRY_STRATEGY_KEY] = strategy
+        end
+
         def should_retry_failed_test?(test_span)
           @retry_failed_tests_enabled && !!test_span&.failed? && @retry_failed_tests_count < @retry_failed_tests_total_limit
         end

lib/datadog/ci/test_visibility/component.rb

@@ -19,8 +19,6 @@ module TestVisibility
       class Component
         attr_reader :test_suite_level_visibility_enabled
 
-        FIBER_LOCAL_TEST_FINISHED_CALLBACK_KEY = :__dd_test_finished_callback
-
         def initialize(
           test_suite_level_visibility_enabled: false,
           codeowners: Codeowners::Parser.new(Git::LocalRepository.root).parse
@@ -130,20 +128,6 @@ def deactivate_test_suite(test_suite_name)
           @context.deactivate_test_suite(test_suite_name)
         end
 
-        # sets fiber-local callback to be called after test is finished
-
-        def test_finished_callback
-          Thread.current[FIBER_LOCAL_TEST_FINISHED_CALLBACK_KEY]
-        end
-
-        def set_test_finished_callback(callback)
-          Thread.current[FIBER_LOCAL_TEST_FINISHED_CALLBACK_KEY] = callback
-        end
-
-        def remove_test_finished_callback
-          Thread.current[FIBER_LOCAL_TEST_FINISHED_CALLBACK_KEY] = nil
-        end
-
         def itr_enabled?
           test_optimisation.enabled?
         end
@@ -212,11 +196,11 @@ def on_test_finished(test)
 
           Telemetry.event_finished(test)
 
-          test_finished_callback&.call(test)
+          test_retries.record_test_finished(test)
         end
 
         def on_after_test_span_finished(tracer_span)
-          # noop
+          test_retries.record_test_span_duration(tracer_span)
         end
 
         # HELPERS
@@ -287,6 +271,10 @@ def test_optimisation
           Datadog.send(:components).test_optimisation
         end
 
+        def test_retries
+          Datadog.send(:components).test_retries
+        end
+
         def git_tree_upload_worker
           Datadog.send(:components).git_tree_upload_worker
         end

sig/datadog/ci/test_retries/component.rbs

@@ -2,6 +2,8 @@ module Datadog
   module CI
     module TestRetries
       class Component
+        FIBER_LOCAL_CURRENT_RETRY_STRATEGY_KEY: Symbol
+
         attr_reader retry_failed_tests_enabled: bool
 
         attr_reader retry_failed_tests_max_attempts: Integer
@@ -32,8 +34,16 @@ module Datadog
 
         def build_strategy: (Datadog::CI::Test test) -> Datadog::CI::TestRetries::Strategy::Base
 
+        def record_test_finished: (Datadog::CI::Test test) -> void
+
+        def record_test_span_duration: (Datadog::Tracing::SpanOperation span) -> void
+
         private
 
+        def current_retry_strategy: () -> Datadog::CI::TestRetries::Strategy::Base?
+
+        def current_retry_strategy=: (Datadog::CI::TestRetries::Strategy::Base? strategy) -> void
+
         def should_retry_failed_test?: (Datadog::CI::Test test) -> bool
 
         def test_visibility_component: () -> Datadog::CI::TestVisibility::Component

sig/datadog/ci/test_visibility/component.rbs

@@ -7,8 +7,6 @@ module Datadog
         @codeowners: Datadog::CI::Codeowners::Matcher
         @context: Datadog::CI::TestVisibility::Context
 
-        FIBER_LOCAL_TEST_FINISHED_CALLBACK_KEY: Symbol
-
         attr_reader test_suite_level_visibility_enabled: bool
 
         def initialize: (?test_suite_level_visibility_enabled: bool, ?codeowners: Datadog::CI::Codeowners::Matcher) -> void
@@ -41,12 +39,6 @@ module Datadog
 
         def deactivate_test_suite: (String test_suite_name) -> void
 
-        def test_finished_callback: () -> Proc?
-
-        def set_test_finished_callback: (Proc callback) -> void
-
-        def remove_test_finished_callback: () -> void
-
         def itr_enabled?: () -> bool
 
         def shutdown!: () -> void
@@ -87,6 +79,8 @@ module Datadog
 
         def test_optimisation: () -> Datadog::CI::TestOptimisation::Component
 
+        def test_retries: () -> Datadog::CI::TestRetries::Component
+
         def git_tree_upload_worker: () -> Datadog::CI::Worker
 
         def remote: () -> Datadog::CI::Remote::Component

spec/datadog/ci/test_visibility/component_spec.rb

@@ -797,51 +797,4 @@
       end
     end
   end
-
-  describe "#set_test_finished_callback" do
-    include_context "CI mode activated"
-
-    let(:callback) { spy(:callback) }
-
-    it "sets the test finished callback which will be executed when any test is finished" do
-      test_visibility.set_test_finished_callback(callback)
-
-      ci_test = test_visibility.trace_test("my test", "my suite")
-      test_visibility.deactivate_test
-
-      expect(callback).to have_received(:call).with(ci_test)
-    end
-
-    it "only fires callback on the same thread where it was set" do
-      test_visibility.set_test_finished_callback(callback)
-
-      t = Thread.new do
-        test_visibility.trace_test("my test", "my suite")
-        test_visibility.deactivate_test
-      end
-      t.join
-
-      expect(callback).to_not have_received(:call)
-    end
-  end
-
-  describe "#remove_test_finished_callback" do
-    include_context "CI mode activated"
-
-    let(:callback) { spy(:callback) }
-
-    it "removes the callback" do
-      test_visibility.set_test_finished_callback(callback)
-
-      test_visibility.trace_test("my test", "my suite")
-      test_visibility.deactivate_test
-
-      test_visibility.remove_test_finished_callback
-
-      test_visibility.trace_test("my test", "my suite")
-      test_visibility.deactivate_test
-
-      expect(callback).to have_received(:call).once
-    end
-  end
 end