Improve the performance of Security::StoreScansWorker

ee/app/models/security/finding.rb

@@ -10,6 +10,7 @@
 module Security
   class Finding < ApplicationRecord
     include IgnorableColumns
+    include EachBatch
 
     self.table_name = 'security_findings'
 

ee/app/services/security/store_findings_metadata_service.rb

@@ -3,15 +3,18 @@
 module Security
   # This service class stores the findings metadata for all pipelines.
   class StoreFindingsMetadataService < ::BaseService
-    attr_reader :security_scan, :report
+    BATCH_SIZE = 50
 
-    def self.execute(security_scan, report)
-      new(security_scan, report).execute
+    attr_reader :security_scan, :report, :deduplicated_finding_uuids
+
+    def self.execute(security_scan, report, deduplicated_finding_uuids)
+      new(security_scan, report, deduplicated_finding_uuids).execute
     end
 
-    def initialize(security_scan, report)
+    def initialize(security_scan, report, deduplicated_finding_uuids)
       @security_scan = security_scan
       @report = report
+      @deduplicated_finding_uuids = deduplicated_finding_uuids
     end
 
     def execute
@@ -30,28 +33,39 @@ module Security
     end
 
     def store_findings
-      report_findings.each { |report_finding| store_finding!(report_finding) }
+      report_findings.each_slice(BATCH_SIZE) { |batch| store_finding_batch(batch) }
     end
 
     def report_findings
       report.findings.select(&:valid?)
     end
 
-    def store_finding!(report_finding)
-      security_scan.findings.create!(finding_data(report_finding))
+    def store_finding_batch(batch)
+      batch.map(&method(:finding_data))
+           .then(&method(:import_batch))
+    end
+
+    def import_batch(report_finding_data)
+      Security::Finding.insert_all(report_finding_data)
     end
 
     def finding_data(report_finding)
       {
+        scan_id: security_scan.id,
         severity: report_finding.severity,
         confidence: report_finding.confidence,
         uuid: report_finding.uuid,
         overridden_uuid: report_finding.overridden_uuid,
         project_fingerprint: report_finding.project_fingerprint,
-        scanner: persisted_scanner_for(report_finding.scanner)
+        scanner_id: persisted_scanner_for(report_finding.scanner).id,
+        deduplicated: deduplicated?(report_finding)
       }
     end
 
+    def deduplicated?(report_finding)
+      deduplicated_finding_uuids.include?(report_finding.uuid)
+    end
+
     def persisted_scanner_for(report_scanner)
       existing_scanners[report_scanner.key] ||= create_scanner!(report_scanner)
     end

ee/app/services/security/store_scan_service.rb

@@ -8,6 +8,8 @@
 # @param deduplicate [Boolean] attribute to force running deduplication logic.
 module Security
   class StoreScanService
+    DEDUPLICATE_BATCH_SIZE = 50
+
     def self.execute(artifact, known_keys, deduplicate)
       new(artifact, known_keys, deduplicate).execute
     end
@@ -49,8 +51,11 @@ module Security
     end
 
     def store_findings
-      StoreFindingsMetadataService.execute(security_scan, security_report)
-      deduplicate_findings? ? update_deduplicated_findings : register_finding_keys
+      StoreFindingsMetadataService.execute(security_scan, security_report, register_finding_keys).then do |result|
+        # If `StoreFindingsMetadataService` returns error, it means the findings
+        # have already been stored before so we may re-run the deduplication logic.
+        update_deduplicated_findings if result[:status] == :error && deduplicate_findings?
+      end
 
       deduplicate_findings?
     end
@@ -65,10 +70,19 @@ module Security
 
     def update_deduplicated_findings
       Security::Scan.transaction do
-        security_scan.findings.update_all(deduplicated: false)
+        mark_all_findings_as_duplicate
+        mark_unique_findings
+      end
+    end
+
+    def mark_all_findings_as_duplicate
+      security_scan.findings.deduplicated.each_batch(of: DEDUPLICATE_BATCH_SIZE) { |batch| batch.update_all(deduplicated: false) }
+    end
 
+    def mark_unique_findings
+      register_finding_keys.each_slice(DEDUPLICATE_BATCH_SIZE) do |batch|
         security_scan.findings
-                     .by_uuid(register_finding_keys)
+                     .by_uuid(batch)
                      .update_all(deduplicated: true)
       end
     end

ee/spec/services/security/store_findings_metadata_service_spec.rb

@@ -7,18 +7,20 @@ RSpec.describe Security::StoreFindingsMetadataService do
   let_it_be(:project) { security_scan.project }
   let_it_be(:security_finding_1) { build(:ci_reports_security_finding) }
   let_it_be(:security_finding_2) { build(:ci_reports_security_finding) }
-  let_it_be(:security_finding_3) { build(:ci_reports_security_finding, uuid: nil) }
+  let_it_be(:security_finding_3) { build(:ci_reports_security_finding) }
+  let_it_be(:security_finding_4) { build(:ci_reports_security_finding, uuid: nil) }
+  let_it_be(:deduplicated_finding_uuids) { [security_finding_1.uuid, security_finding_3.uuid] }
   let_it_be(:security_scanner) { build(:ci_reports_security_scanner) }
   let_it_be(:report) do
     build(
       :ci_reports_security_report,
-      findings: [security_finding_1, security_finding_2, security_finding_3],
+      findings: [security_finding_1, security_finding_2, security_finding_3, security_finding_4],
       scanners: [security_scanner]
     )
   end
 
   describe '#execute' do
-    let(:service_object) { described_class.new(security_scan, report) }
+    let(:service_object) { described_class.new(security_scan, report, deduplicated_finding_uuids) }
 
     subject(:store_findings) { service_object.execute }
 
@@ -27,6 +29,10 @@ RSpec.describe Security::StoreFindingsMetadataService do
         create(:security_finding, scan: security_scan)
       end
 
+      it 'returns error message' do
+        expect(store_findings).to eq({ status: :error, message: "Findings are already stored!" })
+      end
+
       it 'does not create new findings in database' do
         expect { store_findings }.not_to change { Security::Finding.count }
       end
@@ -38,11 +44,14 @@ RSpec.describe Security::StoreFindingsMetadataService do
       end
 
       it 'creates the security finding entries in database' do
-        expect { store_findings }.to change { security_scan.findings.count }.by(2)
-                                 .and change { security_scan.findings.first&.severity }.to(security_finding_1.severity.to_s)
-                                 .and change { security_scan.findings.first&.confidence }.to(security_finding_1.confidence.to_s)
-                                 .and change { security_scan.findings.first&.uuid }.to(security_finding_1.uuid)
-                                 .and change { security_scan.findings.last&.uuid }.to(security_finding_2.uuid)
+        store_findings
+
+        expect(security_scan.findings.reload.as_json(only: [:uuid, :deduplicated]))
+          .to match_array([
+            { "uuid" => security_finding_1.uuid, "deduplicated" => true },
+            { "uuid" => security_finding_2.uuid, "deduplicated" => false },
+            { "uuid" => security_finding_3.uuid, "deduplicated" => true }
+          ])
       end
 
       context 'when the scanners already exist in the database' do

ee/spec/services/security/store_scan_service_spec.rb

@@ -56,7 +56,7 @@ RSpec.describe Security::StoreScanService do
     subject(:store_scan) { service_object.execute }
 
     before do
-      allow(Security::StoreFindingsMetadataService).to receive(:execute)
+      allow(Security::StoreFindingsMetadataService).to receive(:execute).and_return(status: :success)
 
       known_keys.add(finding_key)
     end
@@ -170,12 +170,24 @@ RSpec.describe Security::StoreScanService do
         context 'when the `deduplicate` param is set as true' do
           let(:deduplicate) { true }
 
-          it 'does not change the deduplicated flag of duplicated finding false' do
-            expect { store_scan }.not_to change { duplicated_security_finding.reload.deduplicated }.from(false)
+          context 'when the `StoreFindingsMetadataService` returns success' do
+            it 'does not run the re-deduplicate logic' do
+              expect { store_scan }.not_to change { unique_security_finding.reload.deduplicated }.from(false)
+            end
           end
 
-          it 'sets the deduplicated flag of unique finding as true' do
-            expect { store_scan }.to change { unique_security_finding.reload.deduplicated }.to(true)
+          context 'when the `StoreFindingsMetadataService` returns error' do
+            before do
+              allow(Security::StoreFindingsMetadataService).to receive(:execute).and_return({ status: :error })
+            end
+
+            it 'does not change the deduplicated flag of duplicated finding from false' do
+              expect { store_scan }.not_to change { duplicated_security_finding.reload.deduplicated }.from(false)
+            end
+
+            it 'sets the deduplicated flag of unique finding as true' do
+              expect { store_scan }.to change { unique_security_finding.reload.deduplicated }.to(true)
+            end
           end
         end
       end