shoryuken 0.0.1

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

@@ -0,0 +1,14 @@
+module Shoryuken
+  module Middleware
+    module Server
+      class AutoDelete
+        def call(worker, queue, sqs_msg)
+          yield
+
+          sqs_msg.delete if worker.class.get_shoryuken_options['auto_delete']
+        end
+      end
+    end
+  end
+end
+

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

@@ -0,0 +1,23 @@
+module Shoryuken
+  module Middleware
+    module Server
+      class Logging
+        include Util
+
+        def call(worker, queue, sqs_msg)
+          Shoryuken::Logging.with_context("#{worker.class.to_s} '#{queue}' '#{sqs_msg.id}'") do
+            begin
+              started_at = Time.now
+              logger.info { "started at #{started_at}" }
+              yield
+              logger.info { "completed in: #{elapsed(started_at)} ms" }
+            rescue Exception
+              logger.info { "failed in: #{elapsed(started_at)} ms" }
+              raise
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/shoryuken/processor.rb

@@ -0,0 +1,36 @@
+require 'multi_json'
+
+module Shoryuken
+  class Processor
+    include Celluloid
+    include Util
+
+    def initialize(manager)
+      @manager = manager
+    end
+
+    def process(queue, sqs_msg)
+      if worker_class = Shoryuken.workers[queue]
+        defer do
+          worker = worker_class.new
+
+          Shoryuken.server_middleware.invoke(worker, queue, sqs_msg) do
+            worker.perform(sqs_msg)
+          end
+        end
+      else
+        logger.error "Worker not found for queue '#{queue}'"
+      end
+
+      @manager.async.processor_done(queue, current_actor)
+    end
+
+    def self.default_middleware
+      Middleware::Chain.new do |m|
+        m.add Middleware::Server::Logging
+        m.add Middleware::Server::AutoDelete
+        # m.add Middleware::Server::RetryJobs
+      end
+    end
+  end
+end

data/lib/shoryuken/util.rb

@@ -0,0 +1,19 @@
+module Shoryuken
+  module Util
+    def watchdog(last_words)
+      yield
+    rescue => ex
+      logger.error last_words
+      logger.error ex
+      logger.error ex.backtrace.join("\n")
+    end
+
+    def logger
+      Shoryuken.logger
+    end
+
+    def elapsed(started_at)
+      (Time.now - started_at) * 1000
+    end
+  end
+end

data/lib/shoryuken/version.rb

@@ -0,0 +1,3 @@
+module Shoryuken
+  VERSION = '0.0.1'
+end

data/lib/shoryuken/worker.rb

@@ -0,0 +1,29 @@
+module Shoryuken
+  module Worker
+    def self.included(base)
+      base.extend(ClassMethods)
+    end
+
+    module ClassMethods
+      def shoryuken_options(opts = {})
+        @shoryuken_options = get_shoryuken_options.merge(stringify_keys(Hash(opts)))
+        queue = @shoryuken_options['queue']
+        queue = queue.call if queue.respond_to? :call
+
+
+        Shoryuken.register_worker(queue, self)
+      end
+
+      def get_shoryuken_options # :nodoc:
+        @shoryuken_options || { 'queue' => 'default' }
+      end
+
+      def stringify_keys(hash) # :nodoc:
+        hash.keys.each do |key|
+          hash[key.to_s] = hash.delete(key)
+        end
+        hash
+      end
+    end
+  end
+end

data/shoryuken.gemspec

@@ -0,0 +1,28 @@
+# coding: utf-8
+lib = File.expand_path('../lib', __FILE__)
+$LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
+require 'shoryuken/version'
+
+Gem::Specification.new do |spec|
+  spec.name          = "shoryuken"
+  spec.version       = Shoryuken::VERSION
+  spec.authors       = ["Pablo Cantero"]
+  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.files         = `git ls-files -z`.split("\x0")
+  spec.executables   = %w[shoryuken]
+  spec.test_files    = spec.files.grep(%r{^(test|spec|features)/})
+  spec.require_paths = ["lib"]
+
+  spec.add_development_dependency "bundler", "~> 1.6"
+  spec.add_development_dependency "rake"
+  spec.add_development_dependency "rspec"
+  spec.add_development_dependency "pry-byebug"
+
+  spec.add_dependency "aws-sdk"
+  spec.add_dependency "celluloid"
+  spec.add_dependency "multi_json"
+end

data/spec/shoryuken/chain_spec.rb

@@ -0,0 +1,48 @@
+require 'spec_helper'
+
+describe Shoryuken::Middleware::Chain do
+  class CustomMiddleware
+    def initialize(name, recorder)
+      @name     = name
+      @recorder = recorder
+    end
+
+    def call(*args)
+      @recorder << [@name, 'before']
+      yield
+      @recorder << [@name, 'after']
+    end
+  end
+
+  it 'supports custom middleware' do
+    subject.add CustomMiddleware, 1, []
+
+    expect(CustomMiddleware).to eq subject.entries.last.klass
+  end
+
+  it 'executes middleware' do
+    recorder = []
+    subject.add CustomMiddleware, 'Pablo', recorder
+
+    final_action = nil
+    subject.invoke { final_action = true }
+    expect(final_action).to eq true
+    expect(recorder).to eq [%w[Pablo before], %w[Pablo after]]
+  end
+
+  class NonYieldingMiddleware
+    def call(*args)
+    end
+  end
+
+  it 'allows middleware to abruptly stop processing rest of chain' do
+    recorder = []
+    subject.add NonYieldingMiddleware
+    subject.add CustomMiddleware, 'Pablo', recorder
+
+    final_action = nil
+    subject.invoke { final_action = true }
+    expect(final_action).to eq nil
+    expect(recorder).to eq []
+  end
+end

data/spec/shoryuken/client_spec.rb

@@ -0,0 +1,22 @@
+require 'spec_helper'
+
+describe Shoryuken::Client do
+  let(:sqs)              { double 'SQS' }
+  let(:queue_collection) { double 'Queues Collection' }
+  let(:queue)            { double 'Queue' }
+  let(:queue_name)       { 'shoryuken' }
+
+  before do
+    allow(described_class).to receive(:sqs).and_return(sqs)
+    allow(sqs).to receive(:queues).and_return(queue_collection)
+  end
+
+  describe '.queues' do
+    it 'memoizes queues' do
+      expect(queue_collection).to receive(:named).once.with(queue_name).and_return(queue)
+
+      expect(Shoryuken::Client.queues(queue_name)).to eq queue
+      expect(Shoryuken::Client.queues(queue_name)).to eq queue
+    end
+  end
+end

data/spec/shoryuken/core_ext_spec.rb

@@ -0,0 +1,12 @@
+require 'spec_helper'
+
+describe 'core_ext' do
+  describe String do
+    describe '#constantize' do
+      class HelloWorld; end
+      it 'returns a class from a string' do
+        expect('HelloWorld'.constantize).to eq HelloWorld
+      end
+    end
+  end
+end

data/spec/shoryuken/fetcher_spec.rb

@@ -0,0 +1,36 @@
+require 'spec_helper'
+
+describe Shoryuken::Fetcher do
+  let(:manager)   { double Shoryuken::Manager }
+  let(:sqs_queue) { double 'sqs_queue' }
+  let(:queue)     { 'shoryuken' }
+  let(:sqs_msg)   { double 'SQS msg'}
+
+  subject { described_class.new(manager) }
+
+  before do
+    allow(manager).to receive(:async).and_return(manager)
+    allow(Shoryuken::Client).to receive(:queues).with(queue).and_return(sqs_queue)
+  end
+
+  describe '#fetch' do
+    it 'calls pause_queue! when not found' do
+      allow(sqs_queue).to receive(:receive_message).with(limit: 1).and_return([])
+
+      expect(manager).to receive(:pause_queue!).with(queue)
+      expect(manager).to receive(:dispatch)
+
+      subject.fetch(queue, 1)
+    end
+
+    it 'assigns messages' do
+      allow(sqs_queue).to receive(:receive_message).with(limit: 5).and_return(sqs_msg)
+
+      expect(manager).to receive(:rebalance_queue_weight!).with(queue)
+      expect(manager).to receive(:assign).with(queue, sqs_msg)
+      expect(manager).to receive(:dispatch)
+
+      subject.fetch(queue, 5)
+    end
+  end
+end

data/spec/shoryuken/integration/launcher_spec.rb

@@ -0,0 +1,36 @@
+require 'spec_helper'
+
+describe Shoryuken::Launcher do
+  describe 'Consuming messages', slow: :true do
+    before do
+      subject.run
+
+      $received_messages = 0
+    end
+
+    after do
+      subject.stop
+    end
+
+    class ShoryukenWorker
+      include Shoryuken::Worker
+
+      shoryuken_options queue: 'shoryuken', auto_delete: true
+
+      def perform(sqs_msg)
+        $received_messages += 1
+      end
+    end
+
+    it 'consumes a message' do
+      Shoryuken::Client.queues('shoryuken').send_message('Yo')
+
+      10.times do
+        break if $received_messages > 0
+        sleep 0.5
+      end
+
+      expect($received_messages).to eq 1
+    end
+  end
+end

data/spec/shoryuken/manager_spec.rb

@@ -0,0 +1,73 @@
+require 'spec_helper'
+
+describe Shoryuken::Manager do
+  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]
+    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]
+    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]
+    end
+  end
+end

data/spec/shoryuken/processor_spec.rb

@@ -0,0 +1,93 @@
+require 'spec_helper'
+
+describe Shoryuken::Processor do
+  let(:manager)   { double Shoryuken::Manager }
+  let(:sqs_queue) { double 'Queue' }
+  let(:queue)     { 'yo' }
+  let(:sqs_msg)   { double AWS::SQS::ReceivedMessage, id: 'fc754df7-9cc2-4c41-96ca-5996a44b771e' }
+
+  subject { described_class.new(manager) }
+
+  before do
+    allow(manager).to receive(:async).and_return(manager)
+    allow(Shoryuken::Client).to receive(:queues).with(queue).and_return(sqs_queue)
+  end
+
+  describe '#process' do
+    class YoWorker
+      include Shoryuken::Worker
+
+      shoryuken_options queue: 'yo'
+
+      def perform(sqs_msg); end
+    end
+
+    it 'skips when worker not found' do
+      queue = 'notfound'
+
+      expect(manager).to receive(:processor_done).with(queue, subject)
+
+      expect(sqs_msg).to_not receive(:delete)
+
+      subject.process(queue, sqs_msg)
+    end
+
+    context 'when custom middleware' do
+      class WorkerCalledMiddleware
+        def call(worker, queue, sqs_msg)
+          worker.called(sqs_msg, queue)
+          yield
+        end
+      end
+
+      before do
+        Shoryuken.configure_server do |config|
+          config.server_middleware do |chain|
+            chain.add WorkerCalledMiddleware
+          end
+        end
+      end
+
+      after do
+        Shoryuken.configure_server do |config|
+          config.server_middleware do |chain|
+            chain.remove WorkerCalledMiddleware
+          end
+        end
+      end
+
+      it 'invokes middleware' do
+        expect(manager).to receive(:processor_done).with(queue, subject)
+
+        expect_any_instance_of(YoWorker).to receive(:perform).with(sqs_msg)
+        expect_any_instance_of(YoWorker).to receive(:called).with(sqs_msg, queue)
+
+        subject.process(queue, sqs_msg)
+      end
+    end
+
+    it 'performs with auto delete' do
+      YoWorker.get_shoryuken_options['auto_delete'] = true
+
+      expect(manager).to receive(:processor_done).with(queue, subject)
+
+      expect_any_instance_of(YoWorker).to receive(:perform).with(sqs_msg)
+
+      expect(sqs_msg).to receive(:delete)
+
+      subject.process(queue, sqs_msg)
+    end
+
+    it 'performs without auto delete' do
+      YoWorker.get_shoryuken_options['auto_delete'] = false
+
+      expect(manager).to receive(:processor_done).with(queue, subject)
+
+      expect_any_instance_of(YoWorker).to receive(:perform).with(sqs_msg)
+
+      expect(sqs_msg).to_not receive(:delete)
+
+      subject.process(queue, sqs_msg)
+    end
+  end
+end