shoryuken 2.1.1 → 2.1.2

data/lib/shoryuken/polling.rb

@@ -0,0 +1,206 @@
+module Shoryuken
+  module Polling
+    QueueConfiguration = Struct.new(:name, :options) do
+      def hash
+        name.hash
+      end
+
+      def ==(other)
+        case other
+        when String
+          if options.empty?
+            name == other
+          else
+            false
+          end
+        else
+          super
+        end
+      end
+
+      alias_method :eql?, :==
+
+      def to_s
+        if options.empty?
+          name
+        else
+          "#<QueueConfiguration #{name} options=#{options.inspect}>"
+        end
+      end
+    end
+
+    class BaseStrategy
+      include Util
+
+      def next_queue
+        fail NotImplementedError
+      end
+
+      def messages_found(queue, messages_found)
+        fail NotImplementedError
+      end
+
+      def active_queues
+        fail NotImplementedError
+      end
+
+      def ==(other)
+        case other
+        when Array
+          @queues == other
+        else
+          if other.respond_to?(:active_queues)
+            active_queues == other.active_queues
+          else
+            false
+          end
+        end
+      end
+
+      private
+
+      def delay
+        Shoryuken.options[:delay].to_f
+      end
+    end
+
+    class WeightedRoundRobin < BaseStrategy
+      
+      def initialize(queues)
+        @initial_queues = queues
+        @queues = queues.dup.uniq
+        @paused_queues = []
+      end
+
+      def next_queue
+        unpause_queues
+        queue = @queues.shift
+        return nil if queue.nil?
+
+        @queues << queue
+        QueueConfiguration.new(queue, {})
+      end
+
+      def messages_found(queue, messages_found)
+        if messages_found == 0
+          pause(queue)
+          return
+        end
+
+        maximum_weight = maximum_queue_weight(queue)
+        current_weight = current_queue_weight(queue)
+        if maximum_weight > current_weight
+          logger.info { "Increasing '#{queue}' weight to #{current_weight + 1}, max: #{maximum_weight}" }
+          @queues << queue
+        end
+      end
+
+      def active_queues
+        unparse_queues(@queues)
+      end
+
+      private
+
+      def pause(queue)
+        return unless @queues.delete(queue)
+        @paused_queues << [Time.now + delay, queue]
+        logger.debug "Paused '#{queue}'"
+      end
+ 
+      def unpause_queues
+        return if @paused_queues.empty?
+        return if Time.now < @paused_queues.first[0]
+        pause = @paused_queues.shift
+        @queues << pause[1]
+        logger.debug "Unpaused '#{pause[1]}'"
+      end
+
+      def current_queue_weight(queue)
+        queue_weight(@queues, queue)
+      end
+
+      def maximum_queue_weight(queue)
+        queue_weight(@initial_queues, queue)
+      end
+
+      def queue_weight(queues, queue)
+        queues.count { |q| q == queue }
+      end
+    end
+
+    class StrictPriority < BaseStrategy
+
+      def initialize(queues)
+        # Priority ordering of the queues, highest priority first
+        @queues = queues
+          .group_by { |q| q }
+          .sort_by { |_, qs| -qs.count }
+          .map(&:first)
+
+        # Pause status of the queues, default to past time (unpaused)
+        @paused_until = queues
+          .each_with_object(Hash.new) { |queue, h| h[queue] = Time.at(0) }
+
+        # Start queues at 0
+        reset_next_queue
+      end
+
+      def next_queue
+        next_queue = next_active_queue
+        next_queue.nil? ? nil : QueueConfiguration.new(next_queue, {})
+      end
+
+      def messages_found(queue, messages_found)
+        if messages_found == 0
+          pause(queue)
+        else
+          reset_next_queue
+        end
+      end
+
+      def active_queues
+        @queues
+          .reverse
+          .map.with_index(1)
+          .reject { |q, _| queue_paused?(q) }
+          .reverse
+      end
+
+      private
+
+      def next_active_queue
+        reset_next_queue if queues_unpaused_since?
+
+        size = @queues.length
+        size.times do
+          queue = @queues[@next_queue_index]
+          @next_queue_index = (@next_queue_index + 1) % size
+          return queue unless queue_paused?(queue)
+        end
+
+        return nil
+      end
+
+      def queues_unpaused_since?
+        last = @last_unpause_check
+        now = @last_unpause_check = Time.now
+
+        last && @paused_until.values.any? { |t| t > last && t <= now }
+      end
+
+      def reset_next_queue
+        @next_queue_index = 0
+      end
+
+      def queue_paused?(queue)
+        @paused_until[queue] > Time.now
+      end
+
+      def pause(queue)
+        return unless delay > 0
+        @paused_until[queue] = Time.now + delay
+        logger.debug "Paused '#{queue}'"
+      end
+    end
+  end
+end

data/lib/shoryuken/processor.rb

@@ -12,17 +12,18 @@ module Shoryuken
     attr_accessor :proxy_id
 
     def process(queue, sqs_msg)
-      @manager.async.real_thread(proxy_id, Thread.current)
-
       worker = Shoryuken.worker_registry.fetch_worker(queue, sqs_msg)
+      body = get_body(worker.class, sqs_msg)
 
-        body = get_body(worker.class, sqs_msg)
+      worker.class.server_middleware.invoke(worker, queue, sqs_msg, body) do
+        worker.perform(sqs_msg, body)
+      end
 
-        worker.class.server_middleware.invoke(worker, queue, sqs_msg, body) do
-          worker.perform(sqs_msg, body)
-        end
+      @manager.async.processor_done(queue, current_actor)
+    end
 
-        @manager.async.processor_done(queue, current_actor)
+    def running_thread
+      Thread.current
     end
 
     private

data/lib/shoryuken/version.rb

@@ -1,3 +1,3 @@
 module Shoryuken
-  VERSION = '2.1.1'
+  VERSION = '2.1.2'
 end

data/spec/shoryuken/fetcher_spec.rb

@@ -3,69 +3,35 @@ require 'shoryuken/manager'
 require 'shoryuken/fetcher'
 
 describe Shoryuken::Fetcher do
-  let(:manager)    { double Shoryuken::Manager }
-  let(:queue)      { double Shoryuken::Queue }
+  let(:queue)      { instance_double('Shoryuken::Queue') }
   let(:queue_name) { 'default' }
+  let(:queue_config) { Shoryuken::Polling::QueueConfiguration.new(queue_name, {}) }
 
   let(:sqs_msg) do
-    double Shoryuken::Message,
+    double(Shoryuken::Message,
       queue_url: queue_name,
       body: 'test',
-      message_id: 'fc754df7-9cc2-4c41-96ca-5996a44b771e'
-  end
-
-  subject { described_class.new(manager) }
-
-  before do
-    allow(manager).to receive(:async).and_return(manager)
-    allow(Shoryuken::Client).to receive(:queues).with(queue_name).and_return(queue)
+      message_id: 'fc754df79cc24c4196ca5996a44b771e',
+          )
   end
 
+  subject { described_class.new }
 
   describe '#fetch' do
-    it 'calls pause when no message' do
-      allow(queue).to receive(:receive_messages).with(max_number_of_messages: 1, attribute_names: ['All'], message_attribute_names: ['All']).and_return([])
-
-      expect(manager).to receive(:pause_queue!).with(queue_name)
-      expect(manager).to receive(:dispatch)
-
-      subject.fetch(queue_name, 1)
-    end
-
-    it 'assigns messages' do
-      allow(queue).to receive(:receive_messages).with(max_number_of_messages: 5, attribute_names: ['All'], message_attribute_names: ['All']).and_return(sqs_msg)
-
-      expect(manager).to receive(:rebalance_queue_weight!).with(queue_name)
-      expect(manager).to receive(:assign).with(queue_name, sqs_msg)
-      expect(manager).to receive(:dispatch)
-
-      subject.fetch(queue_name, 5)
+    it 'calls Shoryuken::Client to receive messages' do
+      expect(Shoryuken::Client).to receive(:queues).with(queue_name).and_return(queue)
+      expect(queue).to receive(:receive_messages).
+        with(max_number_of_messages: 1, attribute_names: ['All'], message_attribute_names: ['All']).
+        and_return([])
+      subject.fetch(queue_config, 1)
     end
 
-    it 'assigns messages in batch' do
-      TestWorker.get_shoryuken_options['batch'] = true
-
-      allow(queue).to receive(:receive_messages).with(max_number_of_messages: described_class::FETCH_LIMIT, attribute_names: ['All'], message_attribute_names: ['All']).and_return(sqs_msg)
-
-      expect(manager).to receive(:rebalance_queue_weight!).with(queue_name)
-      expect(manager).to receive(:assign).with(queue_name, [sqs_msg])
-      expect(manager).to receive(:dispatch)
-
-      subject.fetch(queue_name, 5)
-    end
-
-    context 'when worker not found' do
-      let(:queue_name) { 'notfound' }
-
-      it 'ignores batch' do
-        allow(queue).to receive(:receive_messages).with(max_number_of_messages: 5, attribute_names: ['All'], message_attribute_names: ['All']).and_return(sqs_msg)
-
-        expect(manager).to receive(:rebalance_queue_weight!).with(queue_name)
-        expect(manager).to receive(:assign).with(queue_name, sqs_msg)
-        expect(manager).to receive(:dispatch)
-
-        subject.fetch(queue_name, 5)
-      end
+    it 'maxes messages to receive to 10 (SQS limit)' do
+      allow(Shoryuken::Client).to receive(:queues).with(queue_name).and_return(queue)
+      expect(queue).to receive(:receive_messages).
+        with(max_number_of_messages: 10, attribute_names: ['All'], message_attribute_names: ['All']).
+        and_return([])
+      subject.fetch(queue_config, 20)
     end
   end
 end

data/spec/shoryuken/manager_spec.rb

@@ -1,11 +1,37 @@
 require 'spec_helper'
 require 'shoryuken/manager'
 
+RSpec::Matchers.define :queue_config_of do |expected|
+  match do |actual|
+    actual.name == expected
+  end
+end
+
 RSpec.describe Shoryuken::Manager do
-  subject do
+  let(:queue) { 'default' }
+  let(:queues) { [queue] }
+  let(:polling_strategy) { Shoryuken::Polling::WeightedRoundRobin.new(queues) }
+  let(:fetcher) { Shoryuken::Fetcher.new }
+  let(:condvar) do
     condvar = double(:condvar)
     allow(condvar).to receive(:signal).and_return(nil)
-    Shoryuken::Manager.new(condvar)
+    condvar
+  end
+  let(:async_manager) { instance_double(described_class.name) }
+  let(:concurrency) { 1 }
+
+  subject { Shoryuken::Manager.new(condvar) }
+
+  before(:each) do
+    Shoryuken.options[:concurrency] = concurrency
+    subject.fetcher = fetcher
+    subject.polling_strategy = polling_strategy
+    allow_any_instance_of(described_class).to receive(:async).and_return(async_manager)
+  end
+
+  after(:each) do
+    Shoryuken.options[:concurrency] = 1
+    TestWorker.get_shoryuken_options['batch'] = false
   end
 
   describe 'Invalid concurrency setting' do
@@ -14,109 +40,52 @@ RSpec.describe Shoryuken::Manager do
       expect { Shoryuken::Manager.new(nil) }
         .to raise_error(ArgumentError, 'Concurrency value -1 is invalid, it needs to be a positive number')
     end
-
   end
 
-  describe 'Auto Scaling' do
-    it 'decreases weight' do
-      queue1 = 'shoryuken'
-      queue2 = 'uppercut'
-
-      Shoryuken.queues.clear
-      # [shoryuken, 2]
-      # [uppercut,  1]
-      Shoryuken.queues << queue1
-      Shoryuken.queues << queue1
-      Shoryuken.queues << queue2
-
-      expect(subject.instance_variable_get('@queues')).to eq [queue1, queue2]
-
-      subject.pause_queue!(queue1)
-
-      expect(subject.instance_variable_get('@queues')).to eq [queue2]
+  describe '#dispatch' do
+    it 'pauses when there are no active queues' do
+      expect(polling_strategy).to receive(:next_queue).and_return(nil)
+      expect_any_instance_of(described_class).to receive(:after)
+      subject.dispatch
     end
 
-    it 'increases weight' do
-      queue1 = 'shoryuken'
-      queue2 = 'uppercut'
-
-      Shoryuken.queues.clear
-      # [shoryuken, 3]
-      # [uppercut,  1]
-      Shoryuken.queues << queue1
-      Shoryuken.queues << queue1
-      Shoryuken.queues << queue1
-      Shoryuken.queues << queue2
-
-      expect(subject.instance_variable_get('@queues')).to eq [queue1, queue2]
-      subject.pause_queue!(queue1)
-      expect(subject.instance_variable_get('@queues')).to eq [queue2]
-
-      subject.rebalance_queue_weight!(queue1)
-      expect(subject.instance_variable_get('@queues')).to eq [queue2, queue1]
-
-      subject.rebalance_queue_weight!(queue1)
-      expect(subject.instance_variable_get('@queues')).to eq [queue2, queue1, queue1]
-
-      subject.rebalance_queue_weight!(queue1)
-      expect(subject.instance_variable_get('@queues')).to eq [queue2, queue1, queue1, queue1]
+    it 'calls dispatch_batch if worker wants batches' do
+      TestWorker.get_shoryuken_options['batch'] = true
+      expect_any_instance_of(described_class).to receive(:dispatch_batch).with(queue_config_of(queue))
+      expect_any_instance_of(described_class).to receive(:async).and_return(async_manager)
+      expect(async_manager).to receive(:dispatch)
+      subject.dispatch
     end
 
-    it 'adds queue back' do
-      queue1 = 'shoryuken'
-      queue2 = 'uppercut'
-
-      Shoryuken.queues.clear
-      # [shoryuken, 2]
-      # [uppercut,  1]
-      Shoryuken.queues << queue1
-      Shoryuken.queues << queue1
-      Shoryuken.queues << queue2
-
-      Shoryuken.options[:delay] = 0.1
-
-      fetcher = double('Fetcher').as_null_object
-      subject.fetcher = fetcher
-
-      subject.pause_queue!(queue1)
-      expect(subject.instance_variable_get('@queues')).to eq [queue2]
-
-      sleep 0.5
-
-      expect(subject.instance_variable_get('@queues')).to eq [queue2, queue1]
+    it 'calls dispatch_single_messages if worker wants single messages' do
+      expect_any_instance_of(described_class).to receive(:dispatch_single_messages).
+        with(queue_config_of(queue))
+      expect(async_manager).to receive(:dispatch)
+      subject.dispatch
     end
   end
 
-  describe '#next_queue' do
-    it 'returns queues' do
-      queue1 = 'shoryuken'
-      queue2 = 'uppercut'
-
-      Shoryuken.queues.clear
-
-      Shoryuken.register_worker queue1, TestWorker
-      Shoryuken.register_worker queue2, TestWorker
-
-      Shoryuken.queues << queue1
-      Shoryuken.queues << queue2
-
-      expect(subject.send :next_queue).to eq queue1
-      expect(subject.send :next_queue).to eq queue2
+  describe '#dispatch_batch' do
+    it 'assings batch as a single message' do
+      q = polling_strategy.next_queue
+      messages = [1, 2, 3]
+      expect(fetcher).to receive(:fetch).with(q, described_class::BATCH_LIMIT).and_return(messages)
+      expect_any_instance_of(described_class).to receive(:assign).with(q.name, messages)
+      subject.send(:dispatch_batch, q)
     end
+  end
 
-    it 'skips when no worker' do
-      queue1 = 'shoryuken'
-      queue2 = 'uppercut'
-
-      Shoryuken.queues.clear
-
-      Shoryuken.register_worker queue2, TestWorker
-
-      Shoryuken.queues << queue1
-      Shoryuken.queues << queue2
-
-      expect(subject.send :next_queue).to eq queue2
-      expect(subject.send :next_queue).to eq queue2
+  describe '#dispatch_single_messages' do
+    let(:concurrency) { 3 }
+
+    it 'assings messages from batch one by one' do
+      q = polling_strategy.next_queue
+      messages = [1, 2, 3]
+      expect(fetcher).to receive(:fetch).with(q, concurrency).and_return(messages)
+      expect_any_instance_of(described_class).to receive(:assign).with(q.name, 1)
+      expect_any_instance_of(described_class).to receive(:assign).with(q.name, 2)
+      expect_any_instance_of(described_class).to receive(:assign).with(q.name, 3)
+      subject.send(:dispatch_single_messages, q)
     end
   end
 end