Make build trace chunks comparable by chunk index

This commit implements Comparable concern in a build trace chunk class,
making chunks comparable by a chunk_index if these are associated with
the same build.

app/models/ci/build_trace_chunk.rb

@@ -3,6 +3,7 @@
 module Ci
   class BuildTraceChunk < ApplicationRecord
     extend ::Gitlab::Ci::Model
+    include ::Comparable
     include ::FastDestroyAll
     include ::Checksummable
     include ::Gitlab::ExclusiveLeaseHelpers
@@ -143,6 +144,12 @@ module Ci
         build.trace_chunks.maximum(:chunk_index).to_i == chunk_index
     end
 
+    def <=>(other)
+      return unless self.build_id == other.build_id
+
+      self.chunk_index <=> other.chunk_index
+    end
+
     private
 
     def get_data

lib/gitlab/ci/trace/checksum.rb

@@ -43,10 +43,8 @@ module Gitlab
           end
         end
 
-        def chunks_count
-          strong_memoize(:chunks_count) do
-            build.trace_chunks.maximum(:chunk_index)
-          end
+        def last_chunk
+          strong_memoize(:last_chunk) { trace_chunks.max }
         end
 
         ##
@@ -62,14 +60,16 @@ module Gitlab
         # metadata on the database level too.
         #
         def trace_chunks
-          build.trace_chunks.persisted
-            .select(::Ci::BuildTraceChunk.metadata_attributes)
+          strong_memoize(:trace_chunks) do
+            build.trace_chunks.persisted
+              .select(::Ci::BuildTraceChunk.metadata_attributes)
+          end
         end
 
         private
 
         def chunk_size(chunk)
-          if chunk.chunk_index == chunks_count
+          if chunk == last_chunk
             chunk.size
           else
             ::Ci::BuildTraceChunk::CHUNK_SIZE

spec/lib/gitlab/ci/trace/checksum_spec.rb

@@ -94,6 +94,25 @@ RSpec.describe Gitlab::Ci::Trace::Checksum do
     end
   end
 
+  describe '#last_chunk' do
+    context 'when there are no chunks' do
+      it 'returns nil' do
+        expect(subject.last_chunk).to be_nil
+      end
+    end
+
+    context 'when there are multiple chunks' do
+      before do
+        create_chunk(index: 1, data: '1234')
+        create_chunk(index: 0, data: 'abcd')
+      end
+
+      it 'returns chunk with the highest index' do
+        expect(subject.last_chunk.chunk_index).to eq 1
+      end
+    end
+  end
+
   def create_chunk(index:, data:)
     create(:ci_build_trace_chunk, :persisted, build: build,
                                               chunk_index: index,

spec/lib/gitlab/exclusive_lease_helpers_spec.rb

@@ -25,13 +25,17 @@ RSpec.describe Gitlab::ExclusiveLeaseHelpers, :clean_gitlab_redis_shared_state d
       let!(:lease) { stub_exclusive_lease(unique_key, 'uuid') }
 
       it 'calls the given block' do
-        expect { |b| class_instance.in_lock(unique_key, &b) }.to yield_with_args(false)
+        expect { |b| class_instance.in_lock(unique_key, &b) }
+          .to yield_with_args(false, an_instance_of(described_class::SleepingLock))
       end
 
       it 'calls the given block continuously' do
-        expect { |b| class_instance.in_lock(unique_key, &b) }.to yield_with_args(false)
-        expect { |b| class_instance.in_lock(unique_key, &b) }.to yield_with_args(false)
-        expect { |b| class_instance.in_lock(unique_key, &b) }.to yield_with_args(false)
+        expect { |b| class_instance.in_lock(unique_key, &b) }
+          .to yield_with_args(false, an_instance_of(described_class::SleepingLock))
+        expect { |b| class_instance.in_lock(unique_key, &b) }
+          .to yield_with_args(false, an_instance_of(described_class::SleepingLock))
+        expect { |b| class_instance.in_lock(unique_key, &b) }
+          .to yield_with_args(false, an_instance_of(described_class::SleepingLock))
       end
 
       it 'cancels the exclusive lease after the block' do
@@ -74,7 +78,8 @@ RSpec.describe Gitlab::ExclusiveLeaseHelpers, :clean_gitlab_redis_shared_state d
             expect(lease).to receive(:try_obtain).exactly(3).times { nil }
             expect(lease).to receive(:try_obtain).once { unique_key }
 
-            expect { |b| class_instance.in_lock(unique_key, &b) }.to yield_with_args(true)
+            expect { |b| class_instance.in_lock(unique_key, &b) }
+              .to yield_with_args(true, an_instance_of(described_class::SleepingLock))
           end
         end
       end

spec/models/ci/build_trace_chunk_spec.rb

@@ -779,4 +779,62 @@ RSpec.describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
       it_behaves_like 'deletes all build_trace_chunk and data in redis'
     end
   end
+
+  describe 'comparable build trace chunks' do
+    describe '#<=>' do
+      context 'when chunks are associated with different builds' do
+        let(:first) { create(:ci_build_trace_chunk, build: build, chunk_index: 1) }
+        let(:second) { create(:ci_build_trace_chunk, chunk_index: 1) }
+
+        it 'returns nil' do
+          expect(first <=> second).to be_nil
+        end
+      end
+
+      context 'when there are two chunks with different indexes' do
+        let(:first) { create(:ci_build_trace_chunk, build: build, chunk_index: 1) }
+        let(:second) { create(:ci_build_trace_chunk, build: build, chunk_index: 0) }
+
+        it 'indicates the the first one is greater than then second' do
+          expect(first <=> second).to eq 1
+        end
+      end
+
+      context 'when there are two chunks with the same index within the same build' do
+        let(:chunk) { create(:ci_build_trace_chunk) }
+
+        it 'indicates the these are equal' do
+          expect(chunk <=> chunk).to be_zero # rubocop:disable Lint/UselessComparison
+        end
+      end
+    end
+
+    describe '#==' do
+      context 'when chunks have the same index' do
+        let(:chunk) { create(:ci_build_trace_chunk) }
+
+        it 'indicates that the chunks are equal' do
+          expect(chunk).to eq chunk
+        end
+      end
+
+      context 'when chunks have different indexes' do
+        let(:first) { create(:ci_build_trace_chunk, build: build, chunk_index: 1) }
+        let(:second) { create(:ci_build_trace_chunk, build: build, chunk_index: 0) }
+
+        it 'indicates that the chunks are not equal' do
+          expect(first).not_to eq second
+        end
+      end
+
+      context 'when chunks are associated with different builds' do
+        let(:first) { create(:ci_build_trace_chunk, build: build, chunk_index: 1) }
+        let(:second) { create(:ci_build_trace_chunk, chunk_index: 1) }
+
+        it 'indicates that the chunks are not equal' do
+          expect(first).not_to eq second
+        end
+      end
+    end
+  end
 end