shoryuken 0.0.4 → 0.0.5

data/lib/shoryuken/middleware/chain.rb

@@ -49,6 +49,10 @@ module Shoryuken
         yield self if block_given?
       end
 
+      def dup
+        self.class.new.tap { |new_chain| new_chain.entries.replace(entries) }
+      end
+
       def remove(klass)
         entries.delete_if { |entry| entry.klass == klass }
       end

data/lib/shoryuken/middleware/server/timing.rb

@@ -5,7 +5,7 @@ module Shoryuken
         include Util
 
         def call(worker, queue, sqs_msg, body)
-          Shoryuken::Logging.with_context("#{worker.class.to_s}/#{queue}/#{sqs_msg.id}") do
+          Shoryuken::Logging.with_context("#{worker_name(worker.class, sqs_msg, body)}/#{queue}/#{sqs_msg.id}") do
             begin
               started_at = Time.now
 

data/lib/shoryuken/processor.rb

@@ -9,22 +9,48 @@ module Shoryuken
       @manager = manager
     end
 
+    attr_accessor :proxy_id
+
     def process(queue, sqs_msg)
-      worker = Shoryuken.worker_loader.call(queue, sqs_msg)
+      @manager.async.real_thread(proxy_id, Thread.current)
+
+      worker = Shoryuken.worker_registry.fetch_worker(queue, sqs_msg)
+
+      timer = auto_visibility_timeout(queue, sqs_msg, worker.class)
 
-      defer do
-        body = get_body(worker.class, sqs_msg)
+      begin
+        defer do
+          body = get_body(worker.class, sqs_msg)
 
-        Shoryuken.server_middleware.invoke(worker, queue, sqs_msg, body) do
-          worker.perform(sqs_msg, body)
+          worker.class.server_middleware.invoke(worker, queue, sqs_msg, body) do
+            worker.perform(sqs_msg, body)
+          end
         end
-      end
 
-      @manager.async.processor_done(queue, current_actor)
+        @manager.async.processor_done(queue, current_actor)
+      ensure
+        timer.cancel if timer
+      end
     end
 
     private
 
+    def auto_visibility_timeout(queue, sqs_msg, worker_class)
+      if worker_class.auto_visibility_timeout?
+        timer = every(worker_class.visibility_timeout_heartbeat) do
+          begin
+            logger.debug "Extending message #{worker_name(worker_class, sqs_msg)}/#{queue}/#{sqs_msg.id} visibility timeout to #{worker_class.extended_visibility_timeout}"
+
+            sqs_msg.visibility_timeout = worker_class.extended_visibility_timeout
+          rescue => e
+            logger.error "Could not auto extend the message #{worker_class}/#{queue}/#{sqs_msg.id} visibility timeout. Error: #{e.message}"
+          end
+        end
+      end
+
+      timer
+    end
+
     def get_body(worker_class, sqs_msg)
       if sqs_msg.is_a? Array
         sqs_msg.map { |m| parse_body(worker_class, m) }

data/lib/shoryuken/util.rb

@@ -23,5 +23,18 @@ module Shoryuken
         queue_and_weights
       end.to_a
     end
+
+    def worker_name(worker_class, sqs_msg, body = nil)
+      if defined?(::ActiveJob) \
+          && !sqs_msg.is_a?(Array) \
+          && sqs_msg.message_attributes \
+          && sqs_msg.message_attributes['shoryuken_class'] \
+          && sqs_msg.message_attributes['shoryuken_class'][:string_value] == ActiveJob::QueueAdapters::ShoryukenAdapter::JobWrapper.to_s
+
+        "ActiveJob/#{body['job_class'].constantize}"
+      else
+        worker_class.to_s
+      end
+    end
   end
 end

data/lib/shoryuken/version.rb

@@ -1,3 +1,3 @@
 module Shoryuken
-  VERSION = '0.0.4'
+  VERSION = '0.0.5'
 end

data/lib/shoryuken/worker.rb

@@ -30,6 +30,12 @@ module Shoryuken
 
       alias_method :perform_at, :perform_in
 
+      def server_middleware
+        @server_chain ||= Shoryuken.server_middleware.dup
+        yield @server_chain if block_given?
+        @server_chain
+      end
+
       def shoryuken_options(opts = {})
         @shoryuken_options = get_shoryuken_options.merge(stringify_keys(Hash(opts)))
         queue = @shoryuken_options['queue']
@@ -41,8 +47,20 @@ module Shoryuken
         Shoryuken.register_worker(queue, self)
       end
 
+      def auto_visibility_timeout?
+        !!get_shoryuken_options['auto_visibility_timeout']
+      end
+
+      def visibility_timeout_heartbeat
+        extended_visibility_timeout - 5
+      end
+
+      def extended_visibility_timeout
+        Shoryuken::Client.visibility_timeout(get_shoryuken_options['queue'])
+      end
+
       def get_shoryuken_options # :nodoc:
-        @shoryuken_options || { 'queue' => 'default', 'delete' => false, 'auto_delete' => false, 'batch' => false }
+        @shoryuken_options || Shoryuken.default_worker_options
       end
 
       def stringify_keys(hash) # :nodoc:

data/lib/shoryuken/worker_registry.rb

@@ -0,0 +1,34 @@
+module Shoryuken
+  class WorkerRegistry
+    def batch_receive_messages?(queue)
+      # true if the workers for queue support batch processing of messages
+      fail NotImplementedError
+    end
+
+    def clear
+      # must remove all worker registrations
+      fail NotImplementedError
+    end
+
+    def fetch_worker(queue, message)
+      # must return an instance of the worker that handles
+      # message received on queue
+      fail NotImplementedError
+    end
+
+    def queues
+      # must return a list of all queues with registered workers
+      fail NotImplementedError
+    end
+
+    def register_worker(queue, clazz)
+      # must register the worker as a consumer of messages from queue
+      fail NotImplementedError
+    end
+
+    def workers(queue)
+      # must return the list of workers registered for queue, or []
+      fail NotImplementedError
+    end
+  end
+end

data/shoryuken.gemspec

@@ -10,7 +10,7 @@ Gem::Specification.new do |spec|
   spec.email         = ["pablo@pablocantero.com"]
   spec.description = spec.summary = %q{Shoryuken is a super efficient AWS SQS thread based message processor}
   spec.homepage      = "https://github.com/phstc/shoryuken"
-  spec.license       = "MIT"
+  spec.license       = "LGPL-3.0"
 
   spec.files         = `git ls-files -z`.split("\x0")
   spec.executables   = %w[shoryuken]
@@ -23,5 +23,5 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "pry-byebug"
 
   spec.add_dependency "aws-sdk-v1"
-  spec.add_dependency "celluloid", "~> 0.15.2"
+  spec.add_dependency "celluloid", "~> 0.16.0"
 end

data/spec/shoryuken/default_worker_registry_spec.rb

@@ -0,0 +1,83 @@
+require 'spec_helper'
+
+describe Shoryuken::DefaultWorkerRegistry do
+  class RegistryTestWorker
+    include Shoryuken::Worker
+
+    shoryuken_options queue: 'registry-test'
+
+    def perform(sqs_msg, body); end
+  end
+
+  subject do
+    Shoryuken.worker_registry
+  end
+
+  before do
+    subject.register_worker 'registry-test', RegistryTestWorker
+  end
+
+  describe 'a registry containing workers is cleared' do
+    it 'removes all registrations' do
+      queue = 'some-other-queue'
+
+      registry = described_class.new
+
+      worker_class = Class.new do
+        include Shoryuken::Worker
+
+        shoryuken_options queue: queue
+
+        def perform(sqs_msg, body); end
+      end
+
+      registry.register_worker(queue, worker_class)
+
+      expect(registry.workers(queue)).to eq([worker_class])
+
+      registry.clear
+
+      expect(registry.workers(queue)).to eq([])
+    end
+  end
+
+  describe 'a registry with workers is handling messages' do
+    def build_message queue, explicit_worker = nil
+      attributes = {}
+      attributes['shoryuken_class'] = {
+        string_value: explicit_worker.to_s,
+        data_type: 'String' } if explicit_worker
+
+      double AWS::SQS::ReceivedMessage,
+        body: 'test',
+        message_attributes: attributes,
+        message_id: SecureRandom.uuid
+    end
+
+    context 'a batch of messages is being processed' do
+      it 'returns an instance of the worker registered for that queue' do
+        batch = [build_message('default', RegistryTestWorker)]
+        expect(subject.fetch_worker('default', batch)).to be_instance_of(TestWorker)
+      end
+    end
+
+    context 'a single message is being processed' do
+      context 'a worker class name is included in the message attributes' do
+        it 'returns an instance of that worker' do
+          message = build_message('default', RegistryTestWorker)
+          expect(subject.fetch_worker('default', message)).to be_instance_of(RegistryTestWorker)
+        end
+      end
+
+      context 'a worker class name is not included in the message attributes' do
+        it 'returns an instance of the worker registered for that queue' do
+          message = build_message('default')
+          expect(subject.fetch_worker('default', message)).to be_instance_of(TestWorker)
+
+          message = build_message('registry-test')
+          expect(subject.fetch_worker('registry-test', message)).to be_instance_of(RegistryTestWorker)
+        end
+      end
+    end
+  end
+end

data/spec/shoryuken/manager_spec.rb

@@ -2,6 +2,12 @@ require 'spec_helper'
 require 'shoryuken/manager'
 
 describe Shoryuken::Manager do
+  subject do
+    condvar = double(:condvar)
+    allow(condvar).to receive(:signal).and_return(nil)
+    Shoryuken::Manager.new(condvar)
+  end
+
   describe 'Auto Scaling' do
     it 'decreases weight' do
       queue1 = 'shoryuken'

data/spec/shoryuken/processor_spec.rb

@@ -12,11 +12,12 @@ describe Shoryuken::Processor do
 
   before do
     allow(manager).to receive(:async).and_return(manager)
+    allow(manager).to receive(:real_thread)
     allow(Shoryuken::Client).to receive(:queues).with(queue).and_return(sqs_queue)
   end
 
   describe '#process' do
-    it 'parsers the body into JSON' do
+    it 'parses the body into JSON' do
       TestWorker.get_shoryuken_options['body_parser'] = :json
 
       body = { 'test' => 'hi' }
@@ -28,7 +29,7 @@ describe Shoryuken::Processor do
       subject.process(queue, sqs_msg)
     end
 
-    it 'parsers the body calling the proc' do
+    it 'parses the body calling the proc' do
       TestWorker.get_shoryuken_options['body_parser'] = Proc.new { |sqs_msg| "*#{sqs_msg.body}*" }
 
       expect_any_instance_of(TestWorker).to receive(:perform).with(sqs_msg, '*test*')
@@ -38,7 +39,7 @@ describe Shoryuken::Processor do
       subject.process(queue, sqs_msg)
     end
 
-    it 'parsers the body as text' do
+    it 'parses the body as text' do
       TestWorker.get_shoryuken_options['body_parser'] = :text
 
       body = 'test'
@@ -50,7 +51,7 @@ describe Shoryuken::Processor do
       subject.process(queue, sqs_msg)
     end
 
-    it 'parsers calling `.parse`' do
+    it 'parses calling `.parse`' do
       TestWorker.get_shoryuken_options['body_parser'] = JSON
 
       body = { 'test' => 'hi' }
@@ -77,7 +78,7 @@ describe Shoryuken::Processor do
     end
 
     context 'when `object_type: nil`' do
-      it 'parsers the body as text' do
+      it 'parses the body as text' do
         TestWorker.get_shoryuken_options['body_parser'] = nil
 
         body = 'test'
@@ -91,6 +92,8 @@ describe Shoryuken::Processor do
     end
 
     context 'when custom middleware' do
+      let(:queue) { 'worker_called_middleware' }
+
       class WorkerCalledMiddleware
         def call(worker, queue, sqs_msg, body)
           # called is defined with `allow(...).to receive(...)`
@@ -100,6 +103,14 @@ describe Shoryuken::Processor do
       end
 
       before do
+        class WorkerCalledMiddlewareWorker
+          include Shoryuken::Worker
+
+          shoryuken_options queue: 'worker_called_middleware'
+
+          def perform(sqs_msg, body); end
+        end
+
         Shoryuken.configure_server do |config|
           config.server_middleware do |chain|
             chain.add WorkerCalledMiddleware
@@ -118,8 +129,8 @@ describe Shoryuken::Processor do
       it 'invokes middleware' do
         expect(manager).to receive(:processor_done).with(queue, subject)
 
-        expect_any_instance_of(TestWorker).to receive(:perform).with(sqs_msg, sqs_msg.body)
-        expect_any_instance_of(TestWorker).to receive(:called).with(sqs_msg, queue)
+        expect_any_instance_of(WorkerCalledMiddlewareWorker).to receive(:perform).with(sqs_msg, sqs_msg.body)
+        expect_any_instance_of(WorkerCalledMiddlewareWorker).to receive(:called).with(sqs_msg, queue)
 
         subject.process(queue, sqs_msg)
       end
@@ -158,7 +169,7 @@ describe Shoryuken::Processor do
       } }
 
       it 'performs without delete' do
-        Shoryuken.workers.clear # unregister TestWorker
+        Shoryuken.worker_registry.clear # unregister TestWorker
 
         expect(manager).to receive(:processor_done).with(queue, subject)
 
@@ -170,4 +181,50 @@ describe Shoryuken::Processor do
       end
     end
   end
+
+  describe '#auto_visibility_timeout' do
+    let(:heartbeat)          { sqs_queue.visibility_timeout - 5 }
+    let(:visibility_timeout) { sqs_queue.visibility_timeout }
+
+    before do
+      TestWorker.get_shoryuken_options['auto_visibility_timeout'] = true
+
+      allow(sqs_queue).to receive(:visibility_timeout).and_return(15)
+    end
+
+    context 'when the worker takes a long time', slow: true do
+      it 'extends the message invisibility to prevent it from being dequeued concurrently' do
+        TestWorker.get_shoryuken_options['body_parser'] = Proc.new do |sqs_msg|
+          sleep visibility_timeout
+          'test'
+        end
+
+        expect(sqs_msg).to receive(:visibility_timeout=).with(visibility_timeout).once
+        expect(manager).to receive(:processor_done).with(queue, subject)
+
+        allow(sqs_msg).to receive(:body).and_return('test')
+
+        subject.process(queue, sqs_msg)
+      end
+    end
+
+    context 'when the worker takes a short time' do
+      it 'does not extend the message invisibility' do
+        expect(sqs_msg).to receive(:visibility_timeout=).never
+        expect(manager).to receive(:processor_done).with(queue, subject)
+
+        allow(sqs_msg).to receive(:body).and_return('test')
+
+        subject.process(queue, sqs_msg)
+      end
+    end
+
+    context 'when the worker fails' do
+      it 'does not extend the message invisibility' do
+        expect(sqs_msg).to receive(:visibility_timeout=).never
+        expect_any_instance_of(TestWorker).to receive(:perform).and_raise 'worker failed'
+        expect { subject.process(queue, sqs_msg) }.to raise_error
+      end
+    end
+  end
 end

data/spec/shoryuken/util_spec.rb

@@ -14,4 +14,14 @@ describe 'Shoryuken::Util' do
       expect(subject.unparse_queues(queues)).to eq([['queue1', 2], ['queue2', 1], ['queue3', 1], ['queue4', 3]])
     end
   end
+
+  describe '#worker_name' do
+    let(:sqs_msg) { double AWS::SQS::ReceivedMessage, id: 'fc754df7-9cc2-4c41-96ca-5996a44b771e', message_attributes: {} }
+
+    it 'returns Shoryuken worker name' do
+      expect(subject.worker_name(TestWorker, sqs_msg)).to eq 'TestWorker'
+    end
+
+    it 'returns ActiveJob worker name'
+  end
 end