Use a dedicated queue for each worker

3 files changed, 61 insertions, 13 deletions

diff --git a/lib/gitlab/sidekiq_config.rb b/lib/gitlab/sidekiq_config.rb
index dc9886732b5..c3d7814551c 100644
--- a/lib/gitlab/sidekiq_config.rb
+++ b/lib/gitlab/sidekiq_config.rb
@@ -1,16 +1,35 @@
 require 'yaml'
+require 'set'
 
 module Gitlab
   module SidekiqConfig
-    def self.redis_queues
-      @redis_queues ||= Sidekiq::Queue.all.map(&:name)
+    # This method is called by `bin/sidekiq-cluster` in EE, which runs outside
+    # of bundler/Rails context, so we cannot use any gem or Rails methods.
+    def self.worker_queues(rails_path = Rails.root.to_s)
+      @worker_queues ||= {}
+      @worker_queues[rails_path] ||= YAML.load_file(File.join(rails_path, 'app/workers/all_queues.yml'))
     end
 
     # This method is called by `bin/sidekiq-cluster` in EE, which runs outside
     # of bundler/Rails context, so we cannot use any gem or Rails methods.
-    def self.config_queues(rails_path = Rails.root.to_s)
+    def self.expand_queues(queues, all_queues = self.worker_queues)
+      return [] if queues.empty?
+
+      queues_set = all_queues.to_set
+
+      queues.flat_map do |queue|
+        [queue, *queues_set.grep(/\A#{queue}:/)]
+      end
+    end
+
+    def self.redis_queues
+      # Not memoized, because this can change during the life of the application
+      Sidekiq::Queue.all.map(&:name)
+    end
+
+    def self.config_queues
       @config_queues ||= begin
-        config = YAML.load_file(File.join(rails_path, 'config', 'sidekiq_queues.yml'))
+        config = YAML.load_file(Rails.root.join('config/sidekiq_queues.yml'))
         config[:queues].map(&:first)
       end
     end
@@ -23,14 +42,6 @@ module Gitlab
       @workers ||= find_workers(Rails.root.join('app', 'workers'))
     end
 
-    def self.default_queues
-      [ActionMailer::DeliveryJob.queue_name, 'default']
-    end
-
-    def self.worker_queues
-      @worker_queues ||= (workers.map(&:queue) + default_queues).uniq
-    end
-
     def self.find_workers(root)
       concerns = root.join('concerns').to_s
 
@@ -43,7 +54,7 @@ module Gitlab
         ns.camelize.constantize
       end
 
-      # Skip concerns
+      # Skip things that aren't workers
       workers.select { |w| w < Sidekiq::Worker }
     end
   end
diff --git a/lib/gitlab/sidekiq_versioning.rb b/lib/gitlab/sidekiq_versioning.rb
new file mode 100644
index 00000000000..9683214ec18
--- /dev/null
+++ b/lib/gitlab/sidekiq_versioning.rb
@@ -0,0 +1,25 @@
+module Gitlab
+  module SidekiqVersioning
+    def self.install!
+      Sidekiq::Manager.prepend SidekiqVersioning::Manager
+
+      # The Sidekiq client API always adds the queue to the Sidekiq queue
+      # list, but mail_room and gitlab-shell do not. This is only necessary
+      # for monitoring.
+      begin
+        queues = SidekiqConfig.worker_queues
+
+        if queues.any?
+          Sidekiq.redis do |conn|
+            conn.pipelined do
+              queues.each do |queue|
+                conn.sadd('queues', queue)
+              end
+            end
+          end
+        end
+      rescue ::Redis::BaseError, SocketError, Errno::ENOENT, Errno::EADDRNOTAVAIL, Errno::EAFNOSUPPORT, Errno::ECONNRESET, Errno::ECONNREFUSED
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/sidekiq_versioning/manager.rb b/lib/gitlab/sidekiq_versioning/manager.rb
new file mode 100644
index 00000000000..308be0fdf76
--- /dev/null
+++ b/lib/gitlab/sidekiq_versioning/manager.rb
@@ -0,0 +1,12 @@
+module Gitlab
+  module SidekiqVersioning
+    module Manager
+      def initialize(options = {})
+        options[:strict] = false
+        options[:queues] = SidekiqConfig.expand_queues(options[:queues])
+        Sidekiq.logger.info "Listening on queues #{options[:queues].uniq.sort}"
+        super
+      end
+    end
+  end
+end