Add latest changes from gitlab-org/gitlab@15-6-stable-eev15.6.0-rc42

3 files changed, 111 insertions, 208 deletions

diff --git a/spec/lib/gitlab/sidekiq_middleware/duplicate_jobs/client_spec.rb b/spec/lib/gitlab/sidekiq_middleware/duplicate_jobs/client_spec.rb
index 4d12e4b3f6f..44c8df73463 100644
--- a/spec/lib/gitlab/sidekiq_middleware/duplicate_jobs/client_spec.rb
+++ b/spec/lib/gitlab/sidekiq_middleware/duplicate_jobs/client_spec.rb
@@ -2,7 +2,8 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::Client, :clean_gitlab_redis_queues do
+RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::Client, :clean_gitlab_redis_queues,
+:clean_gitlab_redis_shared_state do
   shared_context 'deduplication worker class' do |strategy, including_scheduled|
     let(:worker_class) do
       Class.new do
diff --git a/spec/lib/gitlab/sidekiq_middleware/duplicate_jobs/duplicate_job_spec.rb b/spec/lib/gitlab/sidekiq_middleware/duplicate_jobs/duplicate_job_spec.rb
index d240bf51e67..b6748d49739 100644
--- a/spec/lib/gitlab/sidekiq_middleware/duplicate_jobs/duplicate_job_spec.rb
+++ b/spec/lib/gitlab/sidekiq_middleware/duplicate_jobs/duplicate_job_spec.rb
@@ -11,8 +11,8 @@ RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::DuplicateJob, :clean_gi
 
   let(:wal_locations) do
     {
-      main: '0/D525E3A8',
-      ci: 'AB/12345'
+      'main' => '0/D525E3A8',
+      'ci' => 'AB/12345'
     }
   end
 
@@ -24,10 +24,6 @@ RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::DuplicateJob, :clean_gi
     "#{Gitlab::Redis::Queues::SIDEKIQ_NAMESPACE}:duplicate:#{queue}:#{hash}"
   end
 
-  let(:deduplicated_flag_key) do
-    "#{idempotency_key}:deduplicate_flag"
-  end
-
   describe '#schedule' do
     shared_examples 'scheduling with deduplication class' do |strategy_class|
       it 'calls schedule on the strategy' do
@@ -81,29 +77,26 @@ RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::DuplicateJob, :clean_gi
     end
   end
 
-  shared_examples 'tracking duplicates in redis' do
+  shared_examples 'with Redis cookies' do
+    let(:cookie_key) { "#{idempotency_key}:cookie:v2" }
+    let(:cookie) { get_redis_msgpack(cookie_key) }
+
     describe '#check!' do
       context 'when there was no job in the queue yet' do
         it { expect(duplicate_job.check!).to eq('123') }
 
         shared_examples 'sets Redis keys with correct TTL' do
           it "adds an idempotency key with correct ttl" do
-            expect { duplicate_job.check! }
-              .to change { read_idempotency_key_with_ttl(idempotency_key) }
-                    .from([nil, -2])
-                    .to(['123', be_within(1).of(expected_ttl)])
-          end
+            expected_cookie = {
+              'jid' => '123',
+              'offsets' => {},
+              'wal_locations' => {},
+              'existing_wal_locations' => wal_locations
+            }
 
-          context 'when wal locations is not empty' do
-            it "adds an existing wal locations key with correct ttl" do
-              expect { duplicate_job.check! }
-                .to change { read_idempotency_key_with_ttl(existing_wal_location_key(idempotency_key, :main)) }
-                      .from([nil, -2])
-                      .to([wal_locations[:main], be_within(1).of(expected_ttl)])
-                .and change { read_idempotency_key_with_ttl(existing_wal_location_key(idempotency_key, :ci)) }
-                      .from([nil, -2])
-                      .to([wal_locations[:ci], be_within(1).of(expected_ttl)])
-            end
+            duplicate_job.check!
+            expect(cookie).to eq(expected_cookie)
+            expect(redis_ttl(cookie_key)).to be_within(1).of(expected_ttl)
           end
         end
 
@@ -130,32 +123,23 @@ RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::DuplicateJob, :clean_gi
 
       context 'when there was already a job with same arguments in the same queue' do
         before do
-          set_idempotency_key(idempotency_key, 'existing-key')
-          wal_locations.each do |config_name, location|
-            set_idempotency_key(existing_wal_location_key(idempotency_key, config_name), location)
-          end
+          set_idempotency_key(cookie_key, existing_cookie.to_msgpack)
         end
 
-        it { expect(duplicate_job.check!).to eq('existing-key') }
+        let(:existing_cookie) { { 'jid' => 'existing-jid' } }
 
-        it "does not change the existing key's TTL" do
-          expect { duplicate_job.check! }
-            .not_to change { read_idempotency_key_with_ttl(idempotency_key) }
-                  .from(['existing-key', -1])
-        end
+        it { expect(duplicate_job.check!).to eq('existing-jid') }
 
-        it "does not change the existing wal locations key's TTL" do
+        it "does not change the existing key's TTL" do
           expect { duplicate_job.check! }
-            .to not_change { read_idempotency_key_with_ttl(existing_wal_location_key(idempotency_key, :main)) }
-                  .from([wal_locations[:main], -1])
-            .and not_change { read_idempotency_key_with_ttl(existing_wal_location_key(idempotency_key, :ci)) }
-                  .from([wal_locations[:ci], -1])
+            .not_to change { redis_ttl(cookie_key) }
+                  .from(-1)
         end
 
         it 'sets the existing jid' do
           duplicate_job.check!
 
-          expect(duplicate_job.existing_jid).to eq('existing-key')
+          expect(duplicate_job.existing_jid).to eq('existing-jid')
         end
       end
     end
@@ -166,115 +150,90 @@ RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::DuplicateJob, :clean_gi
           { main: ::ActiveRecord::Base,
             ci: ::ActiveRecord::Base })
 
-        set_idempotency_key(existing_wal_location_key(idempotency_key, :main), existing_wal[:main])
-        set_idempotency_key(existing_wal_location_key(idempotency_key, :ci), existing_wal[:ci])
+        with_redis { |r| r.set(cookie_key, initial_cookie.to_msgpack, ex: expected_ttl) }
 
         # read existing_wal_locations
         duplicate_job.check!
       end
 
-      context "when the key doesn't exists in redis" do
-        let(:existing_wal) do
-          {
-            main: '0/D525E3A0',
-            ci: 'AB/12340'
-          }
-        end
+      let(:initial_cookie) do
+        {
+          'jid' => 'foobar',
+          'existing_wal_locations' => { 'main' => '0/D525E3A0', 'ci' => 'AB/12340' },
+          'offsets' => {},
+          'wal_locations' => {}
+        }
+      end
 
-        let(:new_wal_location_with_offset) do
-          {
-            # offset is relative to `existing_wal`
-            main: ['0/D525E3A8', '8'],
-            ci: ['AB/12345', '5']
-          }
-        end
+      let(:expected_ttl) { 123 }
+      let(:new_wal) do
+        {
+          # offset is relative to `existing_wal`
+          'main' => { location: '0/D525E3A8', offset: '8' },
+          'ci' => { location: 'AB/12345', offset: '5' }
+        }
+      end
 
-        let(:wal_locations) { new_wal_location_with_offset.transform_values(&:first) }
+      let(:wal_locations) { new_wal.transform_values { |v| v[:location] } }
 
-        it 'stores a wal location to redis with an offset relative to existing wal location' do
-          expect { duplicate_job.update_latest_wal_location! }
-            .to change { read_range_from_redis(wal_location_key(idempotency_key, :main)) }
-                  .from([])
-                  .to(new_wal_location_with_offset[:main])
-            .and change { read_range_from_redis(wal_location_key(idempotency_key, :ci)) }
-                  .from([])
-                  .to(new_wal_location_with_offset[:ci])
-        end
+      it 'stores a wal location to redis with an offset relative to existing wal location' do
+        duplicate_job.update_latest_wal_location!
+
+        expect(cookie['wal_locations']).to eq(wal_locations)
+        expect(cookie['offsets']).to eq(new_wal.transform_values { |v| v[:offset].to_i })
+        expect(redis_ttl(cookie_key)).to be_within(1).of(expected_ttl)
       end
+    end
 
-      context "when the key exists in redis" do
-        before do
-          rpush_to_redis_key(wal_location_key(idempotency_key, :main), *stored_wal_location_with_offset[:main])
-          rpush_to_redis_key(wal_location_key(idempotency_key, :ci), *stored_wal_location_with_offset[:ci])
+    describe 'UPDATE_WAL_COOKIE_SCRIPT' do
+      subject do
+        with_redis do |redis|
+          redis.eval(described_class::UPDATE_WAL_COOKIE_SCRIPT, keys: [cookie_key], argv: argv)
         end
+      end
 
-        let(:wal_locations) { new_wal_location_with_offset.transform_values(&:first) }
+      let(:argv) { ['c1', 1, 'loc1', 'c2', 2, 'loc2', 'c3', 3, 'loc3'] }
 
-        context "when the new offset is bigger then the existing one" do
-          let(:existing_wal) do
-            {
-              main: '0/D525E3A0',
-              ci: 'AB/12340'
-            }
-          end
+      it 'does not create the key' do
+        subject
 
-          let(:stored_wal_location_with_offset) do
-            {
-              # offset is relative to `existing_wal`
-              main: ['0/D525E3A3', '3'],
-              ci: ['AB/12342', '2']
-            }
-          end
+        expect(with_redis { |r| r.get(cookie_key) }).to eq(nil)
+      end
 
-          let(:new_wal_location_with_offset) do
-            {
-              # offset is relative to `existing_wal`
-              main: ['0/D525E3A8', '8'],
-              ci: ['AB/12345', '5']
-            }
-          end
+      context 'when the key exists' do
+        let(:existing_cookie) { { 'offsets' => {}, 'wal_locations' => {} } }
+        let(:expected_ttl) { 123 }
 
-          it 'updates a wal location to redis with an offset' do
-            expect { duplicate_job.update_latest_wal_location! }
-              .to change { read_range_from_redis(wal_location_key(idempotency_key, :main)) }
-                    .from(stored_wal_location_with_offset[:main])
-                    .to(new_wal_location_with_offset[:main])
-              .and change { read_range_from_redis(wal_location_key(idempotency_key, :ci)) }
-                    .from(stored_wal_location_with_offset[:ci])
-                    .to(new_wal_location_with_offset[:ci])
-          end
+        before do
+          with_redis { |r| r.set(cookie_key, existing_cookie.to_msgpack, ex: expected_ttl) }
         end
 
-        context "when the old offset is not bigger then the existing one" do
-          let(:existing_wal) do
-            {
-              main: '0/D525E3A0',
-              ci: 'AB/12340'
-            }
-          end
+        it 'updates all connections' do
+          subject
 
-          let(:stored_wal_location_with_offset) do
-            {
-              # offset is relative to `existing_wal`
-              main: ['0/D525E3A8', '8'],
-              ci: ['AB/12345', '5']
-            }
-          end
+          expect(cookie['wal_locations']).to eq({ 'c1' => 'loc1', 'c2' => 'loc2', 'c3' => 'loc3' })
+          expect(cookie['offsets']).to eq({ 'c1' => 1, 'c2' => 2, 'c3' => 3 })
+        end
+
+        it 'preserves the ttl' do
+          subject
 
-          let(:new_wal_location_with_offset) do
+          expect(redis_ttl(cookie_key)).to be_within(1).of(expected_ttl)
+        end
+
+        context 'and low offsets' do
+          let(:existing_cookie) do
             {
-              # offset is relative to `existing_wal`
-              main: ['0/D525E3A2', '2'],
-              ci: ['AB/12342', '2']
+              'offsets' => { 'c1' => 0, 'c2' => 2 },
+              'wal_locations' => { 'c1' => 'loc1old', 'c2' => 'loc2old' }
             }
           end
 
-          it "does not update a wal location to redis with an offset" do
-            expect { duplicate_job.update_latest_wal_location! }
-              .to not_change { read_range_from_redis(wal_location_key(idempotency_key, :main)) }
-                    .from(stored_wal_location_with_offset[:main])
-              .and not_change { read_range_from_redis(wal_location_key(idempotency_key, :ci)) }
-                    .from(stored_wal_location_with_offset[:ci])
+          it 'updates only some connections' do
+            subject
+
+            expect(cookie['wal_locations']).to eq({ 'c1' => 'loc1', 'c2' => 'loc2old', 'c3' => 'loc3' })
+            expect(cookie['offsets']).to eq({ 'c1' => 1, 'c2' => 2, 'c3' => 3 })
           end
         end
       end
@@ -283,11 +242,11 @@ RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::DuplicateJob, :clean_gi
     describe '#latest_wal_locations' do
       context 'when job was deduplicated and wal locations were already persisted' do
         before do
-          rpush_to_redis_key(wal_location_key(idempotency_key, :main), wal_locations[:main], 1024)
-          rpush_to_redis_key(wal_location_key(idempotency_key, :ci), wal_locations[:ci], 1024)
+          cookie = { 'wal_locations' => { 'main' => 'abc', 'ci' => 'def' } }.to_msgpack
+          set_idempotency_key(cookie_key, cookie)
         end
 
-        it { expect(duplicate_job.latest_wal_locations).to eq(wal_locations) }
+        it { expect(duplicate_job.latest_wal_locations).to eq({ 'main' => 'abc', 'ci' => 'def' }) }
       end
 
       context 'when job is not deduplication and wal locations were not persisted' do
@@ -302,60 +261,22 @@ RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::DuplicateJob, :clean_gi
 
       context 'when the key exists in redis' do
         before do
-          set_idempotency_key(idempotency_key, 'existing-jid')
-          set_idempotency_key(deduplicated_flag_key, 1)
-          wal_locations.each do |config_name, location|
-            set_idempotency_key(existing_wal_location_key(idempotency_key, config_name), location)
-            set_idempotency_key(wal_location_key(idempotency_key, config_name), location)
-          end
+          set_idempotency_key(cookie_key, "garbage")
         end
 
         shared_examples 'deleting the duplicate job' do
           shared_examples 'deleting keys from redis' do |key_name|
             it "removes the #{key_name} from redis" do
               expect { duplicate_job.delete! }
-                .to change { read_idempotency_key_with_ttl(key) }
-                      .from([from_value, -1])
-                      .to([nil, -2])
+                .to change { with_redis { |r| r.get(key) } }
+                      .from(from_value)
+                      .to(nil)
             end
           end
 
-          shared_examples 'does not delete key from redis' do |key_name|
-            it "does not remove the #{key_name} from redis" do
-              expect { duplicate_job.delete! }
-                .to not_change { read_idempotency_key_with_ttl(key) }
-                      .from([from_value, -1])
-            end
-          end
-
-          it_behaves_like 'deleting keys from redis', 'idempotent key' do
-            let(:key) { idempotency_key }
-            let(:from_value) { 'existing-jid' }
-          end
-
-          it_behaves_like 'deleting keys from redis', 'deduplication counter key' do
-            let(:key) { deduplicated_flag_key }
-            let(:from_value) { '1' }
-          end
-
-          it_behaves_like 'deleting keys from redis', 'existing wal location keys for main database' do
-            let(:key) { existing_wal_location_key(idempotency_key, :main) }
-            let(:from_value) { wal_locations[:main] }
-          end
-
-          it_behaves_like 'deleting keys from redis', 'existing wal location keys for ci database' do
-            let(:key) { existing_wal_location_key(idempotency_key, :ci) }
-            let(:from_value) { wal_locations[:ci] }
-          end
-
-          it_behaves_like 'deleting keys from redis', 'latest wal location keys for main database' do
-            let(:key) { wal_location_key(idempotency_key, :main) }
-            let(:from_value) { wal_locations[:main] }
-          end
-
-          it_behaves_like 'deleting keys from redis', 'latest wal location keys for ci database' do
-            let(:key) { wal_location_key(idempotency_key, :ci) }
-            let(:from_value) { wal_locations[:ci] }
+          it_behaves_like 'deleting keys from redis', 'cookie key' do
+            let(:key) { cookie_key }
+            let(:from_value) { "garbage" }
           end
         end
 
@@ -387,15 +308,14 @@ RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::DuplicateJob, :clean_gi
     describe '#set_deduplicated_flag!' do
       context 'when the job is reschedulable' do
         before do
+          duplicate_job.check! # ensure cookie exists
           allow(duplicate_job).to receive(:reschedulable?) { true }
         end
 
         it 'sets the key in Redis' do
           duplicate_job.set_deduplicated_flag!
 
-          flag = with_redis { |redis| redis.get(deduplicated_flag_key) }
-
-          expect(flag).to eq(described_class::DEDUPLICATED_FLAG_VALUE.to_s)
+          expect(cookie['deduplicated']).to eq('1')
         end
 
         it 'sets, gets and cleans up the deduplicated flag' do
@@ -415,11 +335,10 @@ RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::DuplicateJob, :clean_gi
         end
 
         it 'does not set the key in Redis' do
+          duplicate_job.check!
           duplicate_job.set_deduplicated_flag!
 
-          flag = with_redis { |redis| redis.get(deduplicated_flag_key) }
-
-          expect(flag).to be_nil
+          expect(cookie['deduplicated']).to eq(nil)
         end
 
         it 'does not set the deduplicated flag' do
@@ -445,43 +364,24 @@ RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::DuplicateJob, :clean_gi
         expect(duplicate_job.duplicate?).to be(false)
       end
 
-      it 'returns false if the existing jid is different from the job jid' do
-        set_idempotency_key(idempotency_key, 'a different jid')
+      it 'returns true if the existing jid is different from the job jid' do
+        set_idempotency_key(cookie_key, { 'jid' => 'a different jid' }.to_msgpack)
         duplicate_job.check!
 
         expect(duplicate_job.duplicate?).to be(true)
       end
     end
 
-    def existing_wal_location_key(idempotency_key, connection_name)
-      "#{idempotency_key}:#{connection_name}:existing_wal_location"
-    end
-
-    def wal_location_key(idempotency_key, connection_name)
-      "#{idempotency_key}:#{connection_name}:wal_location"
-    end
-
-    def set_idempotency_key(key, value = '1')
+    def set_idempotency_key(key, value)
       with_redis { |r| r.set(key, value) }
     end
 
-    def rpush_to_redis_key(key, wal, offset)
-      with_redis { |r| r.rpush(key, [wal, offset]) }
-    end
-
-    def read_idempotency_key_with_ttl(key)
-      with_redis do |redis|
-        redis.pipelined do |p|
-          p.get(key)
-          p.ttl(key)
-        end
-      end
+    def get_redis_msgpack(key)
+      MessagePack.unpack(with_redis { |redis| redis.get(key) })
     end
 
-    def read_range_from_redis(key)
-      with_redis do |redis|
-        redis.lrange(key, 0, -1)
-      end
+    def redis_ttl(key)
+      with_redis { |redis| redis.ttl(key) }
     end
   end
 
@@ -497,7 +397,7 @@ RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::DuplicateJob, :clean_gi
       duplicate_job.check!
     end
 
-    it_behaves_like 'tracking duplicates in redis'
+    it_behaves_like 'with Redis cookies'
   end
 
   context 'when both multi-store feature flags are off' do
@@ -517,7 +417,7 @@ RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::DuplicateJob, :clean_gi
       duplicate_job.check!
     end
 
-    it_behaves_like 'tracking duplicates in redis'
+    it_behaves_like 'with Redis cookies'
   end
 
   describe '#scheduled?' do
@@ -562,6 +462,7 @@ RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::DuplicateJob, :clean_gi
 
       context 'with deduplicated flag' do
         before do
+          duplicate_job.check! # ensure cookie exists
           duplicate_job.set_deduplicated_flag!
         end
 
@@ -578,6 +479,7 @@ RSpec.describe Gitlab::SidekiqMiddleware::DuplicateJobs::DuplicateJob, :clean_gi
 
       context 'with deduplicated flag' do
         before do
+          duplicate_job.check! # ensure cookie exists
           duplicate_job.set_deduplicated_flag!
         end
 
diff --git a/spec/lib/gitlab/sidekiq_middleware/server_metrics_spec.rb b/spec/lib/gitlab/sidekiq_middleware/server_metrics_spec.rb
index 54a1723afbc..1a53a9b8701 100644
--- a/spec/lib/gitlab/sidekiq_middleware/server_metrics_spec.rb
+++ b/spec/lib/gitlab/sidekiq_middleware/server_metrics_spec.rb
@@ -10,7 +10,7 @@ RSpec.describe Gitlab::SidekiqMiddleware::ServerMetrics do
 
       describe '.initialize_process_metrics' do
         it 'sets concurrency metrics' do
-          expect(concurrency_metric).to receive(:set).with({}, Sidekiq.options[:concurrency].to_i)
+          expect(concurrency_metric).to receive(:set).with({}, Sidekiq[:concurrency].to_i)
 
           described_class.initialize_process_metrics
         end
@@ -65,7 +65,7 @@ RSpec.describe Gitlab::SidekiqMiddleware::ServerMetrics do
           end
 
           it 'sets the concurrency metric' do
-            expect(concurrency_metric).to receive(:set).with({}, Sidekiq.options[:concurrency].to_i)
+            expect(concurrency_metric).to receive(:set).with({}, Sidekiq[:concurrency].to_i)
 
             described_class.initialize_process_metrics
           end