Merge branch 'move_lb_call_ci_trace' into 'master'

Move LB call in Ci::Trace to Core [RUN AS-IF-FOSS]

See merge request gitlab-org/gitlab!62749

ee/lib/ee/gitlab/ci/trace.rb

-# frozen_string_literal: true
-
-module EE
-  module Gitlab
-    module Ci
-      module Trace
-        extend ::Gitlab::Utils::Override
-
-        override :destroy_stream
-        def destroy_stream(build)
-          if consistent_archived_trace?(build)
-            ::Gitlab::Database::LoadBalancing::Sticking
-              .stick('ci/build/trace', build.id)
-          end
-
-          yield
-        end
-
-        override :read_trace_artifact
-        def read_trace_artifact(build)
-          if consistent_archived_trace?(build)
-            ::Gitlab::Database::LoadBalancing::Sticking
-              .unstick_or_continue_sticking('ci/build/trace', build.id)
-          end
-
-          yield
-        end
-
-        def consistent_archived_trace?(build)
-          ::Feature.enabled?(:gitlab_ci_archived_trace_consistent_reads, build.project, default_enabled: false)
-        end
-      end
-    end
-  end
-end

lib/gitlab/ci/trace.rb

@@ -14,6 +14,8 @@ module Gitlab
       UPDATE_FREQUENCY_DEFAULT = 60.seconds
       UPDATE_FREQUENCY_WHEN_BEING_WATCHED = 3.seconds
 
+      LOAD_BALANCING_STICKING_NAMESPACE = 'ci/build/trace'
+
       ArchiveError = Class.new(StandardError)
       AlreadyArchivedError = Class.new(StandardError)
       LockedError = Class.new(StandardError)
@@ -296,25 +298,31 @@ module Gitlab
         read_trace_artifact(job) { job.job_artifacts_trace }
       end
 
-      ##
-      # Overridden in EE
-      #
-      def destroy_stream(job)
+      def destroy_stream(build)
+        if consistent_archived_trace?(build)
+          ::Gitlab::Database::LoadBalancing::Sticking
+            .stick(LOAD_BALANCING_STICKING_NAMESPACE, build.id)
+        end
+
         yield
       end
 
-      ##
-      # Overriden in EE
-      #
-      def read_trace_artifact(job)
+      def read_trace_artifact(build)
+        if consistent_archived_trace?(build)
+          ::Gitlab::Database::LoadBalancing::Sticking
+            .unstick_or_continue_sticking(LOAD_BALANCING_STICKING_NAMESPACE, build.id)
+        end
+
         yield
       end
 
+      def consistent_archived_trace?(build)
+        ::Feature.enabled?(:gitlab_ci_archived_trace_consistent_reads, build.project, default_enabled: false)
+      end
+
       def being_watched_cache_key
         "gitlab:ci:trace:#{job.id}:watched"
       end
     end
   end
 end
-
-::Gitlab::Ci::Trace.prepend_mod_with('Gitlab::Ci::Trace')

spec/support/shared_examples/lib/gitlab/ci/ci_trace_shared_examples.rb

@@ -29,6 +29,34 @@ RSpec.shared_examples 'common trace features' do
     end
   end
 
+  describe '#read' do
+    context 'gitlab_ci_archived_trace_consistent_reads feature flag enabled' do
+      before do
+        stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: trace.job.project)
+      end
+
+      it 'calls ::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking' do
+        expect(::Gitlab::Database::LoadBalancing::Sticking).to receive(:unstick_or_continue_sticking)
+          .with(described_class::LOAD_BALANCING_STICKING_NAMESPACE, trace.job.id)
+          .and_call_original
+
+        trace.read { |stream| stream }
+      end
+    end
+
+    context 'gitlab_ci_archived_trace_consistent_reads feature flag disabled' do
+      before do
+        stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: false)
+      end
+
+      it 'does not call ::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking' do
+        expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:unstick_or_continue_sticking)
+
+        trace.read { |stream| stream }
+      end
+    end
+  end
+
   describe '#extract_coverage' do
     let(:regex) { '\(\d+.\d+\%\) covered' }
 
@@ -253,6 +281,52 @@ RSpec.shared_examples 'common trace features' do
   describe '#archive!' do
     subject { trace.archive! }
 
+    context 'when live trace chunks exists' do
+      before do
+        # Build a trace_chunk manually
+        # It is possible to do so with trace.set but only if ci_enable_live_trace FF is enabled
+        #
+        # We need the job to have a trace_chunk because we only use #stick in
+        # the case where trace_chunks exist.
+        stream = Gitlab::Ci::Trace::Stream.new do
+          Gitlab::Ci::Trace::ChunkedIO.new(trace.job)
+        end
+
+        stream.set(+"12\n34")
+      end
+
+      # We check the before setup actually sets up job trace_chunks
+      it 'has job trace_chunks' do
+        expect(trace.job.trace_chunks).to be_present
+      end
+
+      context 'gitlab_ci_archived_trace_consistent_reads feature flag enabled' do
+        before do
+          stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: trace.job.project)
+        end
+
+        it 'calls ::Gitlab::Database::LoadBalancing::Sticking.stick' do
+          expect(::Gitlab::Database::LoadBalancing::Sticking).to receive(:stick)
+            .with(described_class::LOAD_BALANCING_STICKING_NAMESPACE, trace.job.id)
+            .and_call_original
+
+          subject
+        end
+      end
+
+      context 'gitlab_ci_archived_trace_consistent_reads feature flag disabled' do
+        before do
+          stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: false)
+        end
+
+        it 'does not call ::Gitlab::Database::LoadBalancing::Sticking.stick' do
+          expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:stick)
+
+          subject
+        end
+      end
+    end
+
     context 'when build status is success' do
       let!(:build) { create(:ci_build, :success, :trace_live) }