Validate Geo queues exist before queue draining checks

ee/lib/gitlab/geo/geo_tasks.rb

@@ -66,6 +66,8 @@ def enable_maintenance_mode
       # Note that since some non-Geo cron jobs are enabled, empty queues will be a transient state.
       # It is a sufficient check when the site is in Maintenance Mode.
       def drain_non_geo_queues
+        validate_geo_queues_exist!
+
         puts 'Sidekiq Queues: Disabling all non-Geo cron jobs'
 
         disable_non_geo_cron_jobs
@@ -106,6 +108,8 @@ def wait_until_replicated_and_verified
       #
       # Note that we need this check because e.g. Geo update events may be enqueued in Redis.
       def drain_geo_secondary_queues
+        validate_geo_queues_exist!
+
         puts "Sidekiq Queues: Waiting for all Geo queues to be empty"
 
         poll_selected_queues_until_empty do |queue|
@@ -195,6 +199,28 @@ def poll_selected_queues_until_empty
         # rubocop:enable Cop/RedisQueueUsage
       end
 
+      # It's possible to configure GitLab Sidekiq queues and their names. If there are no queues
+      # with Geo in the name, then raise, because this code assumes that it can easily know if Geo
+      # or non-Geo jobs are drained, based on queue name.
+      def validate_geo_queues_exist!
+        any_geo_queues = false
+
+        # Watch Sidekiq Queues on all shards
+        # rubocop:disable Cop/RedisQueueUsage -- valid usage
+        Gitlab::Redis::Queues.instances.each_value do |inst|
+          Sidekiq::Client.via(inst.sidekiq_redis) do # scope all Sidekiq operations to use shard's redis pool
+            # rubocop:disable Cop/SidekiqApiUsage -- valid usage
+            any_geo_queues ||= Sidekiq::Queue.all.any? { |queue| geo_queue?(queue) }
+            # rubocop:enable Cop/SidekiqApiUsage
+          end
+
+          break if any_geo_queues
+        end
+        # rubocop:enable Cop/RedisQueueUsage
+
+        raise "No Geo queues detected. Unable to check if Geo or non-Geo jobs are drained" unless any_geo_queues
+      end
+
       def geo_queue?(queue)
         queue.name.include?('geo')
       end

ee/spec/lib/gitlab/geo/geo_tasks_spec.rb

@@ -124,6 +124,10 @@
   end
 
   describe '.drain_non_geo_queues' do
+    before do
+      allow(described_class).to receive(:validate_geo_queues_exist!)
+    end
+
     it 'disables all non-Geo Sidekiq cron jobs' do
       cronjob1 = instance_double(Sidekiq::Cron::Job)
       cronjob2 = instance_double(Sidekiq::Cron::Job)
@@ -189,6 +193,7 @@
 
   describe '.drain_geo_secondary_queues' do
     it 'waits until all Geo queues are empty' do
+      allow(described_class).to receive(:validate_geo_queues_exist!)
       queue1 = instance_double(Sidekiq::Queue, name: 'geo_foo')
       queue2 = instance_double(Sidekiq::Queue, name: 'bar')
       queue3 = instance_double(Sidekiq::Queue, name: 'geo_baz')
@@ -206,6 +211,7 @@
     end
 
     it 'outputs what it is doing' do
+      allow(described_class).to receive(:validate_geo_queues_exist!)
       expected_output = <<~MSG
         Sidekiq Queues: Waiting for all Geo queues to be empty
         Sidekiq Queues: Geo queues empty
@@ -315,4 +321,66 @@
       end
     end
   end
+
+  describe '.validate_geo_queues_exist!' do
+    let(:geo_queue) { instance_double(Sidekiq::Queue, name: 'geo:sync') }
+    let(:non_geo_queue) { instance_double(Sidekiq::Queue, name: 'default') }
+
+    before do
+      allow(Gitlab::Redis::Queues).to receive(:instances)
+        .and_return({ main: Gitlab::Redis::Queues })
+      allow(Gitlab::Redis::Queues).to receive(:sidekiq_redis)
+      allow(Sidekiq::Client).to receive(:via).and_yield
+    end
+
+    context 'when Geo queues exist' do
+      before do
+        allow(Sidekiq::Queue).to receive(:all).and_return([geo_queue, non_geo_queue])
+      end
+
+      it 'does not raise an error' do
+        expect { described_class.validate_geo_queues_exist! }.not_to raise_error
+      end
+    end
+
+    context 'when no Geo queues exist' do
+      before do
+        allow(Sidekiq::Queue).to receive(:all).and_return([non_geo_queue])
+      end
+
+      it 'raises an error' do
+        expect { described_class.validate_geo_queues_exist! }.to raise_error(
+          RuntimeError, "No Geo queues detected. Unable to check if Geo or non-Geo jobs are drained")
+      end
+    end
+
+    context 'when multiple Redis instances are present' do
+      before do
+        allow(Gitlab::Redis::Queues).to receive(:instances)
+          .and_return({ foo: Gitlab::Redis::Queues, bar: Gitlab::Redis::Queues })
+        allow(Gitlab::Redis::Queues).to receive(:sidekiq_redis)
+      end
+
+      it 'checks all instances and stops when a Geo queue is found' do
+        expect(Sidekiq::Client).to receive(:via).ordered.and_yield
+        expect(Sidekiq::Queue).to receive(:all).and_return([non_geo_queue])
+
+        expect(Sidekiq::Client).to receive(:via).ordered.and_yield
+        expect(Sidekiq::Queue).to receive(:all).and_return([geo_queue])
+
+        expect { described_class.validate_geo_queues_exist! }.not_to raise_error
+      end
+    end
+
+    context 'when no Redis instances are present' do
+      before do
+        allow(Gitlab::Redis::Queues).to receive(:instances).and_return({})
+      end
+
+      it 'raises an error' do
+        expect { described_class.validate_geo_queues_exist! }.to raise_error(
+          RuntimeError, "No Geo queues detected. Unable to check if Geo or non-Geo jobs are drained")
+      end
+    end
+  end
 end