Merge branch '341022-follow-up-idempotent-lb-workers' into 'master'

Address followup items

See merge request gitlab-org/gitlab!70780

lib/gitlab/database.rb

@@ -61,6 +61,7 @@ module Gitlab
     def self.databases
       @databases ||= database_base_models
         .transform_values { |connection_class| Connection.new(connection_class) }
+        .with_indifferent_access
         .freeze
     end
 

lib/gitlab/sidekiq_middleware/duplicate_jobs/duplicate_job.rb

@@ -64,9 +64,9 @@ module Gitlab
 
           Sidekiq.redis do |redis|
             redis.multi do |multi|
-              redis.set(idempotency_key, jid, ex: expiry, nx: true)
-              read_wal_locations = check_existing_wal_locations!(redis, expiry)
-              read_jid = redis.get(idempotency_key)
+              multi.set(idempotency_key, jid, ex: expiry, nx: true)
+              read_wal_locations = check_existing_wal_locations!(multi, expiry)
+              read_jid = multi.get(idempotency_key)
             end
           end
 
@@ -81,9 +81,9 @@ module Gitlab
           return unless job_wal_locations.present?
 
           Sidekiq.redis do |redis|
-            redis.multi do
+            redis.multi do |multi|
               job_wal_locations.each do |connection_name, location|
-                redis.eval(LUA_SET_WAL_SCRIPT, keys: [wal_location_key(connection_name)], argv: [location, pg_wal_lsn_diff(connection_name).to_i, WAL_LOCATION_TTL])
+                multi.eval(LUA_SET_WAL_SCRIPT, keys: [wal_location_key(connection_name)], argv: [location, pg_wal_lsn_diff(connection_name).to_i, WAL_LOCATION_TTL])
               end
             end
           end
@@ -96,9 +96,9 @@ module Gitlab
             read_wal_locations = {}
 
             Sidekiq.redis do |redis|
-              redis.multi do
+              redis.multi do |multi|
                 job_wal_locations.keys.each do |connection_name|
-                  read_wal_locations[connection_name] = redis.lindex(wal_location_key(connection_name), 0)
+                  read_wal_locations[connection_name] = multi.lindex(wal_location_key(connection_name), 0)
                 end
               end
             end
@@ -110,8 +110,8 @@ module Gitlab
         def delete!
           Sidekiq.redis do |redis|
             redis.multi do |multi|
-              redis.del(idempotency_key)
-              delete_wal_locations!(redis)
+              multi.del(idempotency_key)
+              delete_wal_locations!(multi)
             end
           end
         end
@@ -147,7 +147,7 @@ module Gitlab
 
         private
 
-        attr_accessor :existing_wal_locations
+        attr_writer :existing_wal_locations
         attr_reader :queue_name, :job
         attr_writer :existing_jid
 
@@ -155,6 +155,31 @@ module Gitlab
           @worker_klass ||= worker_class_name.to_s.safe_constantize
         end
 
+        def delete_wal_locations!(redis)
+          job_wal_locations.keys.each do |connection_name|
+            redis.del(wal_location_key(connection_name))
+            redis.del(existing_wal_location_key(connection_name))
+          end
+        end
+
+        def check_existing_wal_locations!(redis, expiry)
+          read_wal_locations = {}
+
+          job_wal_locations.each do |connection_name, location|
+            key = existing_wal_location_key(connection_name)
+            redis.set(key, location, ex: expiry, nx: true)
+            read_wal_locations[connection_name] = redis.get(key)
+          end
+
+          read_wal_locations
+        end
+
+        def job_wal_locations
+          return {} unless preserve_wal_location?
+
+          job['wal_locations'] || {}
+        end
+
         def pg_wal_lsn_diff(connection_name)
           Gitlab::Database.databases[connection_name].pg_wal_lsn_diff(job_wal_locations[connection_name], existing_wal_locations[connection_name])
         end
@@ -179,12 +204,6 @@ module Gitlab
           job['jid']
         end
 
-        def job_wal_locations
-          return {} unless preserve_wal_location?
-
-          job['wal_locations'] || {}
-        end
-
         def existing_wal_location_key(connection_name)
           "#{idempotency_key}:#{connection_name}:existing_wal_location"
         end
@@ -209,23 +228,8 @@ module Gitlab
           "#{worker_class_name}:#{Sidekiq.dump_json(arguments)}"
         end
 
-        def delete_wal_locations!(redis)
-          job_wal_locations.keys.each do |connection_name|
-            redis.del(wal_location_key(connection_name))
-            redis.del(existing_wal_location_key(connection_name))
-          end
-        end
-
-        def check_existing_wal_locations!(redis, expiry)
-          read_wal_locations = {}
-
-          job_wal_locations.each do |connection_name, location|
-            key = existing_wal_location_key(connection_name)
-            redis.set(key, location, ex: expiry, nx: true)
-            read_wal_locations[connection_name] = redis.get(key)
-          end
-
-          read_wal_locations
+        def existing_wal_locations
+          @existing_wal_locations ||= {}
         end
 
         def preserve_wal_location?

lib/gitlab/sidekiq_middleware/duplicate_jobs/strategies/deduplicates_when_scheduling.rb

@@ -4,11 +4,15 @@ module Gitlab
   module SidekiqMiddleware
     module DuplicateJobs
       module Strategies
-        module DeduplicatesWhenScheduling
+        class DeduplicatesWhenScheduling < Base
+          extend ::Gitlab::Utils::Override
+
+          override :initialize
           def initialize(duplicate_job)
             @duplicate_job = duplicate_job
           end
 
+          override :schedule
           def schedule(job)
             if deduplicatable_job? && check! && duplicate_job.duplicate?
               job['duplicate-of'] = duplicate_job.existing_jid
@@ -25,6 +29,7 @@ module Gitlab
             yield
           end
 
+          override :perform
           def perform(job)
             update_job_wal_location!(job)
           end

lib/gitlab/sidekiq_middleware/duplicate_jobs/strategies/until_executed.rb

@@ -7,11 +7,7 @@ module Gitlab
         # This strategy takes a lock before scheduling the job in a queue and
         # removes the lock after the job has executed preventing a new job to be queued
         # while a job is still executing.
-        class UntilExecuted < Base
-          extend ::Gitlab::Utils::Override
-
-          include DeduplicatesWhenScheduling
-
+        class UntilExecuted < DeduplicatesWhenScheduling
           override :perform
           def perform(job)
             super

lib/gitlab/sidekiq_middleware/duplicate_jobs/strategies/until_executing.rb

@@ -7,11 +7,7 @@ module Gitlab
         # This strategy takes a lock before scheduling the job in a queue and
         # removes the lock before the job starts allowing a new job to be queued
         # while a job is still executing.
-        class UntilExecuting < Base
-          extend ::Gitlab::Utils::Override
-
-          include DeduplicatesWhenScheduling
-
+        class UntilExecuting < DeduplicatesWhenScheduling
           override :perform
           def perform(job)
             super

spec/lib/gitlab/database_spec.rb

@@ -15,6 +15,13 @@ RSpec.describe Gitlab::Database do
     end
   end
 
+  describe '.databases' do
+    it 'stores connections as a HashWithIndifferentAccess' do
+      expect(described_class.databases.has_key?('main')).to be true
+      expect(described_class.databases.has_key?(:main)).to be true
+    end
+  end
+
   describe '.default_pool_size' do
     before do
       allow(Gitlab::Runtime).to receive(:max_threads).and_return(7)