Accept exit codes from runner and mark jobs as allowed to fail

Accept exit codes from runner and mark jobs as allowed to fail

app/models/ci/build.rb

@@ -1031,6 +1031,14 @@ module Ci
       variables.any? { |variable| variable[:key] == 'CI_DEBUG_TRACE' && variable[:value].casecmp('true') == 0 }
     end
 
+    def conditionally_allow_failure!(exit_code)
+      return unless ::Gitlab::Ci::Features.allow_failure_with_exit_codes_enabled?
+
+      if options.dig(:allow_failure_criteria, :exit_codes).to_a.include?(exit_code)
+        update_columns(allow_failure: true)
+      end
+    end
+
     protected
 
     def run_status_commit_hooks!

app/services/ci/update_build_state_service.rb

@@ -111,6 +111,7 @@ module Ci
 
         Result.new(status: 200)
       when 'failed'
+        build.conditionally_allow_failure!(params[:exit_code])
         build.drop!(params[:failure_reason] || :unknown_failure)
 
         Result.new(status: 200)

lib/api/ci/runner.rb

@@ -180,6 +180,7 @@ module API
             optional :checksum, type: String, desc: %q(Job's trace CRC32 checksum)
             optional :bytesize, type: Integer, desc: %q(Job's trace size in bytes)
           end
+          optional :exit_code, type: Integer, desc: %q(Job's exit code)
         end
         put '/:id' do
           job = authenticate_job!

spec/models/ci/build_spec.rb

@@ -4823,4 +4823,59 @@ RSpec.describe Ci::Build do
       it { is_expected.to eq false }
     end
   end
+
+  describe '#conditionally_allow_failure!' do
+    let_it_be_with_refind(:job) do
+      create(:ci_build, pipeline: pipeline)
+    end
+
+    subject(:conditionally_allow_failure) do
+      job.conditionally_allow_failure!(1)
+    end
+
+    context 'when exit_codes are not defined' do
+      it 'does not change allow_failure' do
+        expect { conditionally_allow_failure }
+          .not_to change { job.reload.allow_failure }
+      end
+    end
+
+    context 'when exit_codes are defined but do not match' do
+      before do
+        job.options[:allow_failure_criteria] = { exit_codes: [2, 3, 4] }
+        job.save!
+      end
+
+      it 'does not change allow_failure' do
+        expect { conditionally_allow_failure }
+          .not_to change { job.reload.allow_failure }
+      end
+    end
+
+    context 'when exit_codes are defined and match' do
+      before do
+        job.options[:allow_failure_criteria] = { exit_codes: [1, 2, 3] }
+        job.save!
+      end
+
+      it 'does change allow_failure' do
+        expect { conditionally_allow_failure }
+          .to change { job.reload.allow_failure }
+      end
+    end
+
+    context 'when ci_allow_failure_with_exit_codes is disabled' do
+      before do
+        stub_feature_flags(ci_allow_failure_with_exit_codes: false)
+
+        job.options[:allow_failure_criteria] = { exit_codes: [1, 2, 3] }
+        job.save!
+      end
+
+      it 'does not change allow_failure' do
+        expect { conditionally_allow_failure }
+          .not_to change { job.reload.allow_failure }
+      end
+    end
+  end
 end

spec/requests/api/ci/runner/jobs_put_spec.rb

@@ -78,6 +78,33 @@ RSpec.describe API::Ci::Runner, :clean_gitlab_redis_shared_state do
           end
         end
 
+        context 'when an exit_code is provided' do
+          context 'when the exit_codes are acceptable' do
+            before do
+              job.options[:allow_failure] = { exit_codes: [1] }
+              job.save!
+            end
+
+            it 'accepts an exit code' do
+              update_job(state: 'failed', exit_code: 1)
+
+              expect(job.reload).to be_failed
+              expect(job.allow_failure).to be_truthy
+              expect(job).to be_unknown_failure
+            end
+          end
+
+          context 'when the exit_codes are not defined' do
+            it 'ignore the exit code' do
+              update_job(state: 'failed', exit_code: 1)
+
+              expect(job.reload).to be_failed
+              expect(job.allow_failure).to be_falsy
+              expect(job).to be_unknown_failure
+            end
+          end
+        end
+
         context 'when failure_reason is script_failure' do
           before do
             update_job(state: 'failed', failure_reason: 'script_failure')

spec/services/ci/update_build_state_service_spec.rb

@@ -82,8 +82,9 @@ RSpec.describe Ci::UpdateBuildStateService do
     let(:params) do
       {
         output: { checksum: 'crc32:12345678', bytesize: 123 },
+        state: 'failed',
         failure_reason: 'script_failure',
-        state: 'failed'
+        exit_code: 42
       }
     end
 
@@ -95,6 +96,15 @@ RSpec.describe Ci::UpdateBuildStateService do
         expect(result.status).to eq 200
       end
 
+      it 'updates the allow_failure flag' do
+        expect(build)
+          .to receive(:conditionally_allow_failure!)
+          .with(42)
+          .and_call_original
+
+        subject.execute
+      end
+
       it 'does not increment invalid trace metric' do
         execute_with_stubbed_metrics!
 
@@ -115,6 +125,15 @@ RSpec.describe Ci::UpdateBuildStateService do
         expect(build).to be_failed
       end
 
+      it 'updates the allow_failure flag' do
+        expect(build)
+          .to receive(:conditionally_allow_failure!)
+          .with(42)
+          .and_call_original
+
+        subject.execute
+      end
+
       it 'responds with 200 OK status' do
         result = subject.execute