lib/sidekiq/throttled/registry.rb in sidekiq-throttled-0.10.0.beta vs lib/sidekiq/throttled/registry.rb in sidekiq-throttled-0.10.0

@@ -1,9 +1,10 @@
 # frozen_string_literal: true
 
 # internal
 require "sidekiq/throttled/strategy"
+require "sidekiq/throttled/utils"
 
 module Sidekiq
   module Throttled
     # Registred strategies.
     #
@@ -11,10 +12,12 @@
     module Registry
       @strategies = {}
       @aliases    = {}
 
       class << self
+        include Utils
+
         # Adds strategy to the registry.
         #
         # @note prints a warning to STDERR upon duplicate strategy name
         # @param (see Strategy#initialize)
         # @return [Strategy]
@@ -47,16 +50,19 @@
         #   @param [#to_s] name
         #   @return [Strategy, nil] registred strategy
         #
         # @overload get(name, &block)
         #   Yields control to the block if requested strategy was found.
+        #   @param [#to_s] name
         #   @yieldparam [Strategy] strategy
         #   @yield [strategy] Gives found strategy to the block
         #   @return result of a block
         def get(name)
-          strategy = @strategies[name.to_s] || @aliases[name.to_s]
+          strategy = find(name.to_s) || find_by_class(name)
+
           return yield strategy if strategy && block_given?
+
           strategy
         end
 
         # @overload each()
         #   @return [Enumerator]
@@ -83,9 +89,37 @@
         def each_with_static_keys
           return to_enum(__method__) unless block_given?
           @strategies.each do |name, strategy|
             yield(name, strategy) unless strategy.dynamic?
           end
+        end
+
+        private
+
+        # Find strategy by it's name.
+        #
+        # @param name [String]
+        # @return [Strategy, nil]
+        def find(name)
+          @strategies[name] || @aliases[name]
+        end
+
+        # Find strategy by class or it's parents.
+        #
+        # @param name [Class, #to_s]
+        # @return [Strategy, nil]
+        def find_by_class(name)
+          return unless Throttled.configuration.inherit_strategies?
+
+          const = name.is_a?(Class) ? name : constantize(name)
+          return unless const.is_a?(Class)
+
+          const.ancestors.each do |m|
+            strategy = find(m.name)
+            return strategy if strategy
+          end
+
+          nil
         end
       end
     end
   end
 end