shoryuken 0.0.1

data/lib/shoryuken/client.rb

@@ -0,0 +1,23 @@
+
+module Shoryuken
+  class Client
+    @@queues = {}
+
+    def self.queues(name)
+      @@queues[name.to_s] ||= sqs.queues.named(name)
+    end
+
+    def self.receive_message(queue, options = {})
+      queues(queue).receive_message(Hash(options))
+    end
+
+    def self.reset!
+      # for test purposes
+      @@queues = {}
+    end
+
+    def self.sqs
+      @sqs ||= AWS::SQS.new
+    end
+  end
+end

data/lib/shoryuken/core_ext.rb

@@ -0,0 +1,47 @@
+begin
+  require 'active_support/core_ext/hash/keys'
+  require 'active_support/core_ext/hash/deep_merge'
+rescue LoadError
+  class Hash
+    def stringify_keys
+      keys.each do |key|
+        self[key.to_s] = delete(key)
+      end
+      self
+    end if !{}.respond_to?(:stringify_keys)
+
+    def symbolize_keys
+      keys.each do |key|
+        self[(key.to_sym rescue key) || key] = delete(key)
+      end
+      self
+    end if !{}.respond_to?(:symbolize_keys)
+
+    def deep_symbolize_keys
+      keys.each do |key|
+        value = delete(key)
+        self[(key.to_sym rescue key) || key] = value
+
+        value.deep_symbolize_keys if value.is_a? Hash
+      end
+      self
+    end if !{}.respond_to?(:deep_symbolize_keys)
+  end
+end
+
+begin
+  require 'active_support/core_ext/string/inflections'
+rescue LoadError
+  class String
+    def constantize
+      names = self.split('::')
+      names.shift if names.empty? || names.first.empty?
+
+      constant = Object
+      names.each do |name|
+        constant = constant.const_defined?(name) ? constant.const_get(name) : constant.const_missing(name)
+      end
+      constant
+    end
+  end if !"".respond_to?(:constantize)
+end

data/lib/shoryuken/fetcher.rb

@@ -0,0 +1,49 @@
+module Shoryuken
+  class Fetcher
+    include Celluloid
+    include Util
+
+    def initialize(manager)
+      @manager = manager
+    end
+
+    def receive_message(queue, limit)
+      # AWS limits the batch size by 10
+      limit = limit > 10 ? 10 : limit
+
+      Shoryuken::Client.receive_message queue, Shoryuken.options[:aws][:receive_message].to_h.merge(limit: limit)
+    end
+
+    def fetch(queue, available_processors)
+      watchdog('Fetcher#fetch died') do
+        started_at = Time.now
+
+        logger.info "Looking for new messages queue '#{queue}'"
+
+        begin
+          if (sqs_msgs = Array(receive_message(queue, available_processors))).any?
+            logger.info "Message found for queue '#{queue}'"
+
+            sqs_msgs.each do |sqs_msg|
+              @manager.async.rebalance_queue_weight!(queue)
+              @manager.async.assign(queue, sqs_msg)
+            end
+          else
+            logger.info "No message found for queue '#{queue}'"
+
+            @manager.async.pause_queue!(queue)
+          end
+
+          @manager.async.dispatch
+
+          logger.debug "Fetcher#fetch('#{queue}') completed in #{elapsed(started_at)} ms"
+        rescue => ex
+          logger.error("Error fetching message: #{ex}")
+          logger.error(ex.backtrace.first)
+
+          @manager.async.dispatch
+        end
+      end
+    end
+  end
+end

data/lib/shoryuken/launcher.rb

@@ -0,0 +1,41 @@
+module Shoryuken
+  class Launcher
+    include Celluloid
+    include Util
+
+    trap_exit :actor_died
+
+    attr_accessor :manager
+
+    def initialize
+      @manager = Shoryuken::Manager.new_link
+      @fetcher = Shoryuken::Fetcher.new_link(manager)
+
+      @done = false
+
+      manager.fetcher = @fetcher
+    end
+
+    def stop(options = {})
+      watchdog('Launcher#stop') do
+        @done = true
+        @fetcher.terminate if @fetcher.alive?
+
+        manager.async.stop(shutdown: !!options[:shutdown], timeout: Shoryuken.options[:timeout])
+        manager.wait(:shutdown)
+      end
+    end
+
+    def run
+      watchdog('Launcher#run') do
+        manager.async.start
+      end
+    end
+
+    def actor_died(actor, reason)
+      return if @done
+      Shoryuken.logger.warn 'Shoryuken died due to the following error, cannot recover, process exiting'
+      exit 1
+    end
+  end
+end

data/lib/shoryuken/logging.rb

@@ -0,0 +1,47 @@
+require 'time'
+require 'logger'
+
+module Shoryuken
+  module Logging
+
+    class Pretty < Logger::Formatter
+      # Provide a call() method that returns the formatted message.
+      def call(severity, time, program_name, message)
+        "#{time.utc.iso8601} #{Process.pid} TID-#{Thread.current.object_id.to_s(36)}#{context} #{severity}: #{message}\n"
+      end
+
+      def context
+        c = Thread.current[:shoryuken_context]
+        c ? " #{c}" : ''
+      end
+    end
+
+    def self.with_context(msg)
+      begin
+        Thread.current[:shoryuken_context] = msg
+        yield
+      ensure
+        Thread.current[:shoryuken_context] = nil
+      end
+    end
+
+    def self.initialize_logger(log_target = STDOUT)
+      @logger = Logger.new(log_target)
+      @logger.level = Logger::INFO
+      @logger.formatter = Pretty.new
+      @logger
+    end
+
+    def self.logger
+      @logger || initialize_logger
+    end
+
+    def self.logger=(log)
+      @logger = (log ? log : Logger.new('/dev/null'))
+    end
+
+    def logger
+      shoryuken::Logging.logger
+    end
+  end
+end

data/lib/shoryuken/manager.rb

@@ -0,0 +1,220 @@
+require 'yaml'
+require 'aws-sdk'
+require 'celluloid'
+
+require 'shoryuken/version'
+require 'shoryuken/manager'
+require 'shoryuken/processor'
+require 'shoryuken/fetcher'
+
+module Shoryuken
+  class Manager
+    include Celluloid
+    include Util
+
+    attr_accessor :fetcher
+
+    trap_exit :processor_died
+
+    def initialize
+      @count  = Shoryuken.options[:concurrency] || 25
+      @queues = Shoryuken.queues.dup.uniq
+
+      @done = false
+
+      @busy  = []
+      @ready = @count.times.map { Processor.new_link(current_actor) }
+    end
+
+    def start
+      logger.info 'Starting'
+
+      dispatch
+    end
+
+    def stop(options = {})
+      watchdog('Manager#stop died') do
+        @done = true
+
+        @fetcher.terminate if @fetcher.alive?
+
+        logger.info { "Shutting down #{@ready.size} quiet workers" }
+
+        @ready.each do |processor|
+          processor.terminate if processor.alive?
+        end
+        @ready.clear
+
+        return after(0) { signal(:shutdown) } if @busy.empty?
+
+        if options[:shutdown]
+          hard_shutdown_in(options[:timeout])
+        else
+          soft_shutdown(options[:timeout])
+        end
+      end
+    end
+
+    def processor_done(queue, processor)
+      watchdog('Manager#processor_done died') do
+        logger.info "Process done for queue '#{queue}'"
+
+        @busy.delete processor
+
+        if stopped?
+          processor.terminate if processor.alive?
+        else
+          @ready << processor
+        end
+      end
+    end
+
+    def processor_died(processor, reason)
+      watchdog("Manager#processor_died died") do
+        logger.info "Process died, reason: #{reason}"
+
+        @busy.delete processor
+
+        unless stopped?
+          @ready << Processor.new_link(current_actor)
+        end
+      end
+    end
+
+    def stopped?
+      @done
+    end
+
+    def assign(queue, sqs_msg)
+      watchdog("Manager#assign died") do
+        logger.info "Assigning #{sqs_msg.id}"
+
+        processor = @ready.pop
+        @busy << processor
+
+        processor.async.process(queue, sqs_msg)
+      end
+    end
+
+    def rebalance_queue_weight!(queue)
+      watchdog('Manager#rebalance_queue_weight! died') do
+        if (original = original_queue_weight(queue)) > (current = current_queue_weight(queue))
+          if current + 1 > original
+            logger.info "Increasing queue '#{queue}' weight to #{current + 1}, original: #{original}"
+          else
+            logger.info "Queue '#{queue}' is back to its normal weight #{original}"
+          end
+
+          @queues << queue
+        end
+      end
+    end
+
+    def pause_queue!(queue)
+      return unless @queues.include? queue
+
+      logger.info "Pausing queue '#{queue}' for #{Shoryuken.options[:delay].to_f} seconds, because the queue is empty"
+
+      @queues.delete(queue)
+
+      after(Shoryuken.options[:delay].to_f) { async.restart_queue!(queue) }
+    end
+
+
+    def dispatch
+      return if stopped?
+
+      logger.debug { "Ready size: #{@ready.size}" }
+      logger.debug { "Busy size: #{@busy.size}" }
+      logger.debug { "Queues: #{@queues.inspect}" }
+
+      if @ready.empty?
+        logger.debug { 'Pausing fetcher, because all processors are busy' }
+
+        after(1) { dispatch }
+
+        return
+      end
+
+      if queue = next_queue
+        @fetcher.async.fetch(queue, @ready.size)
+      else
+        logger.debug { 'Pausing fetcher, because all queues are paused' }
+
+        @fetcher_paused = true
+      end
+    end
+
+    private
+
+    def restart_queue!(queue)
+      return if stopped?
+
+      unless @queues.include? queue
+        logger.info "Restarting queue '#{queue}'"
+
+        @queues << queue
+
+        if @fetcher_paused
+          logger.debug { 'Restarting fetcher' }
+
+          @fetcher_paused = false
+
+
+          dispatch
+        end
+      end
+    end
+
+    def current_queue_weight(queue)
+      queue_weight(@queues, queue)
+    end
+
+    def original_queue_weight(queue)
+      queue_weight(Shoryuken.queues, queue)
+    end
+
+    def queue_weight(queues, queue)
+      queues.count { |q| q == queue }
+    end
+
+    def next_queue
+      return nil if @queues.empty?
+
+      queue = @queues.shift
+      @queues << queue
+
+      queue
+    end
+
+    def soft_shutdown(delay)
+      logger.info { "Waiting for #{@busy.size} busy workers" }
+
+      if @busy.size > 0
+        after(delay) { soft_shutdown(delay) }
+      else
+        after(0) { signal(:shutdown) }
+      end
+    end
+
+    def hard_shutdown_in(delay)
+      logger.info { "Waiting for #{@busy.size} busy workers" }
+      logger.info { "Pausing up to #{delay} seconds to allow workers to finish..." }
+
+      after(delay) do
+        watchdog("Manager#hard_shutdown_in died") do
+          if @busy.size > 0
+            logger.info { "Hard shutting down #{@busy.size} busy workers" }
+
+            @busy.each do |processor|
+              t = processor.bare_object.actual_work_thread
+              t.raise Shutdown if processor.alive?
+            end
+          end
+
+          after(0) { signal(:shutdown) }
+        end
+      end
+    end
+  end
+end

data/lib/shoryuken/middleware/chain.rb

@@ -0,0 +1,111 @@
+module Shoryuken
+  # Middleware is code configured to run before/after
+  # a message is processed.  It is patterned after Rack
+  # middleware. Middleware exists for the server
+  # side (when jobs are actually processed).
+  #
+  # To modify middleware for the server, just call
+  # with another block:
+  #
+  # Shoryuken.configure_server do |config|
+  #   config.server_middleware do |chain|
+  #     chain.add MyServerHook
+  #     chain.remove ActiveRecord
+  #   end
+  # end
+  #
+  # To insert immediately preceding another entry:
+  #
+  # Shoryuken.configure_server do |config|
+  #   config.server_middleware do |chain|
+  #     chain.insert_before ActiveRecord, MyServerHook
+  #   end
+  # end
+  #
+  # To insert immediately after another entry:
+  #
+  # Shoryuken.configure_server do |config|
+  #   config.server_middleware do |chain|
+  #     chain.insert_after ActiveRecord, MyServerHook
+  #   end
+  # end
+  #
+  # This is an example of a minimal server middleware:
+  #
+  # class MyServerHook
+  #   def call(worker_instance, queue, sqs_msg)
+  #     puts 'Before work'
+  #     yield
+  #     puts 'After work'
+  #   end
+  # end
+  #
+  module Middleware
+    class Chain
+      attr_reader :entries
+
+      def initialize
+        @entries = []
+        yield self if block_given?
+      end
+
+      def remove(klass)
+        entries.delete_if { |entry| entry.klass == klass }
+      end
+
+      def add(klass, *args)
+        entries << Entry.new(klass, *args) unless exists?(klass)
+      end
+
+      def insert_before(oldklass, newklass, *args)
+        i = entries.index { |entry| entry.klass == newklass }
+        new_entry = i.nil? ? Entry.new(newklass, *args) : entries.delete_at(i)
+        i = entries.find_index { |entry| entry.klass == oldklass } || 0
+        entries.insert(i, new_entry)
+      end
+
+      def insert_after(oldklass, newklass, *args)
+        i = entries.index { |entry| entry.klass == newklass }
+        new_entry = i.nil? ? Entry.new(newklass, *args) : entries.delete_at(i)
+        i = entries.find_index { |entry| entry.klass == oldklass } || entries.count - 1
+        entries.insert(i+1, new_entry)
+      end
+
+      def exists?(klass)
+        entries.any? { |entry| entry.klass == klass }
+      end
+
+      def retrieve
+        entries.map(&:make_new)
+      end
+
+      def clear
+        entries.clear
+      end
+
+      def invoke(*args, &final_action)
+        chain = retrieve.dup
+        traverse_chain = lambda do
+          if chain.empty?
+            final_action.call
+          else
+            chain.shift.call(*args, &traverse_chain)
+          end
+        end
+        traverse_chain.call
+      end
+    end
+
+    class Entry
+      attr_reader :klass
+      def initialize(klass, *args)
+        @klass = klass
+        @args  = args
+      end
+
+      def make_new
+        @klass.new(*@args)
+      end
+    end
+  end
+end