Add latest changes from gitlab-org/gitlab@master

1 files changed, 79 insertions, 5 deletions

diff --git a/lib/gitlab/sidekiq_config/cli_methods.rb b/lib/gitlab/sidekiq_config/cli_methods.rb
index 0676e9df9c5..8f19b557d24 100644
--- a/lib/gitlab/sidekiq_config/cli_methods.rb
+++ b/lib/gitlab/sidekiq_config/cli_methods.rb
@@ -18,17 +18,39 @@ module Gitlab
         result
       end.freeze
 
-      def worker_queues(rails_path = Rails.root.to_s)
+      QUERY_OR_OPERATOR = '|'
+      QUERY_AND_OPERATOR = '&'
+      QUERY_CONCATENATE_OPERATOR = ','
+      QUERY_TERM_REGEX = %r{^(\w+)(!?=)([\w#{QUERY_CONCATENATE_OPERATOR}]+)}.freeze
+
+      QUERY_PREDICATES = {
+        feature_category: :to_sym,
+        has_external_dependencies: lambda { |value| value == 'true' },
+        latency_sensitive: lambda { |value| value == 'true' },
+        name: :to_s,
+        resource_boundary: :to_sym
+      }.freeze
+
+      QueryError = Class.new(StandardError)
+      InvalidTerm = Class.new(QueryError)
+      UnknownOperator = Class.new(QueryError)
+      UnknownPredicate = Class.new(QueryError)
+
+      def all_queues(rails_path = Rails.root.to_s)
         @worker_queues ||= {}
 
         @worker_queues[rails_path] ||= QUEUE_CONFIG_PATHS.flat_map do |path|
           full_path = File.join(rails_path, path)
-          queues = File.exist?(full_path) ? YAML.load_file(full_path) : []
 
-          # https://gitlab.com/gitlab-org/gitlab/issues/199230
-          queues.map { |queue| queue.is_a?(Hash) ? queue[:name] : queue }
+          File.exist?(full_path) ? YAML.load_file(full_path) : []
         end
       end
+      # rubocop:enable Gitlab/ModuleWithInstanceVariables
+
+      def worker_queues(rails_path = Rails.root.to_s)
+        # https://gitlab.com/gitlab-org/gitlab/issues/199230
+        worker_names(all_queues(rails_path))
+      end
 
       def expand_queues(queues, all_queues = self.worker_queues)
         return [] if queues.empty?
@@ -40,12 +62,64 @@ module Gitlab
         end
       end
 
+      def query_workers(query_string, queues)
+        worker_names(queues.select(&query_string_to_lambda(query_string)))
+      end
+
       def clear_memoization!
         if instance_variable_defined?('@worker_queues')
           remove_instance_variable('@worker_queues')
         end
       end
-      # rubocop:enable Gitlab/ModuleWithInstanceVariables
+
+      private
+
+      def worker_names(workers)
+        workers.map { |queue| queue.is_a?(Hash) ? queue[:name] : queue }
+      end
+
+      def query_string_to_lambda(query_string)
+        or_clauses = query_string.split(QUERY_OR_OPERATOR).map do |and_clauses_string|
+          and_clauses_predicates = and_clauses_string.split(QUERY_AND_OPERATOR).map do |term|
+            predicate_for_term(term)
+          end
+
+          lambda { |worker| and_clauses_predicates.all? { |predicate| predicate.call(worker) } }
+        end
+
+        lambda { |worker| or_clauses.any? { |predicate| predicate.call(worker) } }
+      end
+
+      def predicate_for_term(term)
+        match = term.match(QUERY_TERM_REGEX)
+
+        raise InvalidTerm.new("Invalid term: #{term}") unless match
+
+        _, lhs, op, rhs = *match
+
+        predicate_for_op(op, predicate_factory(lhs, rhs.split(QUERY_CONCATENATE_OPERATOR)))
+      end
+
+      def predicate_for_op(op, predicate)
+        case op
+        when '='
+          predicate
+        when '!='
+          lambda { |worker| !predicate.call(worker) }
+        else
+          # This is unreachable because InvalidTerm will be raised instead, but
+          # keeping it allows to guard against that changing in future.
+          raise UnknownOperator.new("Unknown operator: #{op}")
+        end
+      end
+
+      def predicate_factory(lhs, values)
+        values_block = QUERY_PREDICATES[lhs.to_sym]
+
+        raise UnknownPredicate.new("Unknown predicate: #{lhs}") unless values_block
+
+        lambda { |queue| values.map(&values_block).include?(queue[lhs.to_sym]) }
+      end
     end
   end
 end