shoryuken 0.0.5 → 1.0.0

data/lib/shoryuken/environment_loader.rb

@@ -0,0 +1,167 @@
+module Shoryuken
+  class EnvironmentLoader
+    attr_reader :options
+
+    def self.load(options)
+      new(options).load
+    end
+
+    def self.load_for_rails_console
+      load(config_file: (Rails.root + 'config' + 'shoryuken.yml'))
+    end
+
+    def initialize options
+      @options = options
+    end
+
+    def load
+      initialize_logger
+      load_rails if options[:rails]
+      Shoryuken.options.merge!(config_file_options)
+      merge_cli_defined_queues
+      Shoryuken.options.merge!(options)
+      parse_queues
+      require_workers
+      initialize_aws
+      validate_queues
+      validate_workers
+      patch_deprecated_workers
+    end
+
+    private
+
+    def config_file_options
+      if (path = options[:config_file])
+        unless File.exist?(path)
+          Shoryuken.logger.warn "Config file #{path} does not exist"
+          path = nil
+        end
+      end
+
+      return {} unless path
+
+      YAML.load(ERB.new(IO.read(path)).result).deep_symbolize_keys
+    end
+
+    def initialize_aws
+      # aws-sdk tries to load the credentials from the ENV variables: AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY
+      # when not explicit supplied
+      return if Shoryuken.options[:aws].empty?
+
+      shoryuken_keys = %i(
+        account_id
+        sns_endpoint
+        sqs_endpoint
+        receive_message)
+
+      aws_options = Shoryuken.options[:aws].reject do |k, v|
+        shoryuken_keys.include?(k)
+      end
+
+      credentials = Aws::Credentials.new(
+        aws_options.delete(:access_key_id),
+        aws_options.delete(:secret_access_key))
+
+      aws_options = aws_options.merge(credentials: credentials)
+
+      if callback = Shoryuken.aws_initialization_callback
+        Shoryuken.logger.info "Calling Shoryuken.on_aws_initialization block"
+        callback.call(aws_options)
+      end
+
+      Aws.config = aws_options
+    end
+
+    def initialize_logger
+      Shoryuken::Logging.initialize_logger(options[:logfile]) if options[:logfile]
+      Shoryuken.logger.level = Logger::DEBUG if options[:verbose]
+    end
+
+    def load_rails
+      # Adapted from: https://github.com/mperham/sidekiq/blob/master/lib/sidekiq/cli.rb
+
+      require 'rails'
+      if ::Rails::VERSION::MAJOR < 4
+        require File.expand_path("config/environment.rb")
+        ::Rails.application.eager_load!
+      else
+        # Painful contortions, see 1791 for discussion
+        require File.expand_path("config/application.rb")
+        ::Rails::Application.initializer "shoryuken.eager_load" do
+          ::Rails.application.config.eager_load = true
+        end
+        require 'shoryuken/extensions/active_job_adapter' if defined?(::ActiveJob)
+        require File.expand_path("config/environment.rb")
+      end
+
+      Shoryuken.logger.info "Rails environment loaded"
+    end
+
+    def merge_cli_defined_queues
+      cli_defined_queues = options.delete(:queues) || []
+
+      cli_defined_queues.each do |cli_defined_queue|
+        Shoryuken.options[:queues].delete_if do |config_file_queue|
+          config_file_queue[0] == cli_defined_queue[0]
+        end
+
+        Shoryuken.options[:queues] << cli_defined_queue
+      end
+    end
+
+    def parse_queue(queue, weight = nil)
+      [weight.to_i, 1].max.times { Shoryuken.queues << queue }
+    end
+
+    def parse_queues
+      Shoryuken.options[:queues].to_a.each do |queue_and_weight|
+        parse_queue(*queue_and_weight)
+      end
+    end
+
+    def patch_deprecated_workers
+      Shoryuken.worker_registry.queues.each do |queue|
+        Shoryuken.worker_registry.workers(queue).each do |worker_class|
+          if worker_class.instance_method(:perform).arity == 1
+            Shoryuken.logger.warn "[DEPRECATION] #{worker_class.name}#perform(sqs_msg) is deprecated. Please use #{worker_class.name}#perform(sqs_msg, body)"
+
+            worker_class.class_eval do
+              alias_method :deprecated_perform, :perform
+
+              def perform(sqs_msg, body = nil)
+                deprecated_perform(sqs_msg)
+              end
+            end
+          end
+        end
+      end
+    end
+
+    def require_workers
+      require Shoryuken.options[:require] if Shoryuken.options[:require]
+    end
+
+    def validate_queues
+      Shoryuken.logger.warn 'No queues supplied' if Shoryuken.queues.empty?
+
+      Shoryuken.queues.uniq.each do |queue|
+        begin
+          Shoryuken::Client.queues queue
+        rescue Aws::SQS::Errors::NonExistentQueue
+          Shoryuken.logger.warn "AWS Queue '#{queue}' does not exist"
+        end
+      end
+    end
+
+    def validate_workers
+      all_queues = Shoryuken.queues
+      queues_with_workers = Shoryuken.worker_registry.queues
+
+      unless defined?(::ActiveJob)
+        (all_queues - queues_with_workers).each do |queue|
+          Shoryuken.logger.warn "No worker supplied for '#{queue}'"
+        end
+      end
+    end
+  end
+end

data/lib/shoryuken/extensions/active_job_adapter.rb

@@ -19,7 +19,8 @@ module ActiveJob
         def enqueue(job) #:nodoc:
           register_worker!(job)
 
-          Shoryuken::Client.send_message(job.queue_name, job.serialize, message_attributes: message_attributes)
+          queue = Shoryuken::Client.queues(job.queue_name)
+          queue.send_message(message(job))
         end
 
         def enqueue_at(job, timestamp) #:nodoc:
@@ -28,13 +29,20 @@ module ActiveJob
           delay = (timestamp - Time.current.to_f).round
           raise 'The maximum allowed delay is 15 minutes' if delay > 15.minutes
 
-          Shoryuken::Client.send_message(job.queue_name, job.serialize, delay_seconds: delay,
-                                                                        message_attributes: message_attributes)
+          queue = Shoryuken::Client.queues(job.queue_name)
+          queue.send_message(message(job, delay_seconds: delay))
         end
 
-
         private
 
+        def message(job, options = {})
+          body = job.serialize
+          body = JSON.dump(body) if body.is_a?(Hash)
+
+          { message_body: body,
+            message_attributes: message_attributes }.merge(options)
+        end
+
         def register_worker!(job)
           Shoryuken.register_worker(job.queue_name, JobWrapper)
         end

data/lib/shoryuken/fetcher.rb

@@ -9,16 +9,15 @@ module Shoryuken
       @manager = manager
     end
 
-    def receive_message(queue, limit)
+    def receive_messages(queue, limit)
       # AWS limits the batch size by 10
       limit = limit > FETCH_LIMIT ? FETCH_LIMIT : limit
 
       options = Shoryuken.options[:aws][:receive_message].to_h.dup
-      options[:limit] = limit
-      options[:message_attribute_names] ||= []
-      options[:message_attribute_names] << 'shoryuken_class'
+      options[:max_number_of_messages] = limit
+      options[:message_attribute_names] = %w(All)
 
-      Shoryuken::Client.receive_message queue, options
+      Shoryuken::Client.queues(queue).receive_messages options
     end
 
     def fetch(queue, available_processors)
@@ -31,7 +30,7 @@ module Shoryuken
           batch = Shoryuken.worker_registry.batch_receive_messages?(queue)
           limit = batch ? FETCH_LIMIT : available_processors
 
-          if (sqs_msgs = Array(receive_message(queue, limit))).any?
+          if (sqs_msgs = Array(receive_messages(queue, limit))).any?
             logger.info "Found #{sqs_msgs.size} messages for '#{queue}'"
 
             if batch
@@ -63,7 +62,7 @@ module Shoryuken
 
     def patch_sqs_msgs!(sqs_msgs)
       sqs_msgs.instance_eval do
-        def id
+        def message_id
           "batch-with-#{size}-messages"
         end
       end

data/lib/shoryuken/manager.rb

@@ -32,6 +32,11 @@ module Shoryuken
       watchdog('Manager#stop died') do
         @done = true
 
+        if callback = Shoryuken.stop_callback
+          logger.info "Calling Shoryuken.on_stop block"
+          callback.call
+        end
+
         @fetcher.terminate if @fetcher.alive?
 
         logger.info { "Shutting down #{@ready.size} quiet workers" }
@@ -85,7 +90,7 @@ module Shoryuken
 
     def assign(queue, sqs_msg)
       watchdog("Manager#assign died") do
-        logger.info "Assigning #{sqs_msg.id}"
+        logger.info "Assigning #{sqs_msg.message_id}"
 
         processor = @ready.pop
         @busy << processor
@@ -187,7 +192,6 @@ module Shoryuken
 
       unless defined?(::ActiveJob) ||  !Shoryuken.worker_registry.workers(queue).empty?
         # when no worker registered pause the queue to avoid endless recursion
-
         logger.debug "Pausing '#{queue}' for #{Shoryuken.options[:delay].to_f} seconds, because no workers registered"
 
         after(Shoryuken.options[:delay].to_f) { async.restart_queue!(queue) }

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

@@ -7,7 +7,13 @@ module Shoryuken
 
           auto_delete = worker.class.get_shoryuken_options['delete'] || worker.class.get_shoryuken_options['auto_delete']
 
-          Shoryuken::Client.queues(queue).batch_delete(*Array(sqs_msg)) if auto_delete
+          if auto_delete
+            entries = [sqs_msg].flatten.map.with_index do |message, i|
+              { id: i.to_s, receipt_handle: message.receipt_handle }
+            end
+
+            Shoryuken::Client.queues(queue).delete_messages(entries: entries)
+          end
         end
       end
     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_name(worker.class, sqs_msg, body)}/#{queue}/#{sqs_msg.id}") do
+          Shoryuken::Logging.with_context("#{worker_name(worker.class, sqs_msg, body)}/#{queue}/#{sqs_msg.message_id}") do
             begin
               started_at = Time.now
 
@@ -15,7 +15,7 @@ module Shoryuken
 
               total_time = elapsed(started_at)
 
-              if (total_time / 1000.0) > (timeout = Shoryuken::Client.visibility_timeout(queue))
+              if (total_time / 1000.0) > (timeout = Shoryuken::Client.queues(queue).visibility_timeout)
                 logger.warn "exceeded the queue visibility timeout by #{total_time - (timeout * 1000)} ms"
               end
 

data/lib/shoryuken/processor.rb

@@ -37,13 +37,15 @@ module Shoryuken
 
     def auto_visibility_timeout(queue, sqs_msg, worker_class)
       if worker_class.auto_visibility_timeout?
-        timer = every(worker_class.visibility_timeout_heartbeat) do
+        queue_visibility_timeout = Shoryuken::Client.queues(queue).visibility_timeout
+
+        timer = every(queue_visibility_timeout - 5) do
           begin
-            logger.debug "Extending message #{worker_name(worker_class, sqs_msg)}/#{queue}/#{sqs_msg.id} visibility timeout to #{worker_class.extended_visibility_timeout}"
+            logger.debug "Extending message #{worker_name(worker_class, sqs_msg)}/#{queue}/#{sqs_msg.message_id} visibility timeout by #{queue_visibility_timeout}s."
 
-            sqs_msg.visibility_timeout = worker_class.extended_visibility_timeout
+            sqs_msg.visibility_timeout = queue_visibility_timeout
           rescue => e
-            logger.error "Could not auto extend the message #{worker_class}/#{queue}/#{sqs_msg.id} visibility timeout. Error: #{e.message}"
+            logger.error "Could not auto extend the message #{worker_class}/#{queue}/#{sqs_msg.message_id} visibility timeout. Error: #{e.message}"
           end
         end
       end

data/lib/shoryuken/sns_arn.rb

@@ -0,0 +1,27 @@
+module Shoryuken
+  class SnsArn
+    def initialize topic
+      @topic = topic
+    end
+
+    def to_s
+      @arn ||= "arn:aws:sns:#{region}:#{account_id}:#{@topic}"
+    end
+
+    private
+
+    def account_id
+      Shoryuken::Client.account_id.tap do |account_id|
+        if account_id.to_s.empty?
+          fail "To generate SNS ARNs, you must assign an :account_id in your Shoryuken::Client."
+        end
+      end
+    end
+
+    def region
+      Aws.config.fetch(:region) do
+        fail "To generate SNS ARNs, you must include a :region in your AWS config."
+      end
+    end
+  end
+end

data/lib/shoryuken/topic.rb

@@ -0,0 +1,17 @@
+module Shoryuken
+  class Topic
+    def initialize(name, sns)
+      @name, @sns = name, sns
+    end
+
+    def arn
+      @arn ||= Client.sns_arn.new(@name).to_s
+    end
+
+    def send_message(body, options = {})
+      body = JSON.dump(body) if body.is_a?(Hash)
+
+      @sns.publish(topic_arn: arn, message: body)
+    end
+  end
+end

data/lib/shoryuken/version.rb

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

data/lib/shoryuken/worker.rb

@@ -13,7 +13,10 @@ module Shoryuken
           data_type: 'String'
         }
 
-        Shoryuken::Client.send_message(get_shoryuken_options['queue'], body, options)
+        body = JSON.dump(body) if body.is_a?(Hash)
+        options[:message_body] = body
+
+        Shoryuken::Client.queues(get_shoryuken_options['queue']).send_message(options)
       end
 
       def perform_in(interval, body, options = {})
@@ -51,14 +54,6 @@ module Shoryuken
         !!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 || Shoryuken.default_worker_options
       end

data/shoryuken.gemspec

@@ -21,7 +21,10 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "rake"
   spec.add_development_dependency "rspec"
   spec.add_development_dependency "pry-byebug"
+  spec.add_development_dependency "nokogiri"
+  spec.add_development_dependency "dotenv"
 
-  spec.add_dependency "aws-sdk-v1"
+  spec.add_dependency "aws-sdk-core", "2.0.21"
+  spec.add_dependency "aws-sdk-resources", "2.0.21.pre"
   spec.add_dependency "celluloid", "~> 0.16.0"
 end

data/spec/integration/launcher_spec.rb

@@ -5,83 +5,41 @@ require 'shoryuken/launcher'
 describe Shoryuken::Launcher do
   describe 'Consuming messages', slow: :true do
     before do
-      Shoryuken.options[:aws] ||= {}
-      Shoryuken.options[:aws][:receive_message] ||= {}
-      Shoryuken.options[:aws][:receive_message][:wait_time_seconds] = 5
-
-      Shoryuken.queues << 'shoryuken'
-      Shoryuken.queues << 'shoryuken_command'
-
-      Shoryuken.register_worker 'shoryuken',          StandardWorker
-      Shoryuken.register_worker 'shoryuken_command',  CommandWorker
-
-      subject.run
+      Shoryuken.options[:aws][:receive_message] = { wait_time_seconds: 5 }
 
       StandardWorker.received_messages = 0
-    end
-
-    after { subject.stop }
-
-    class CommandWorker
-      include Shoryuken::Worker
 
-      @@received_messages = 0
+      queue = "test_shoryuken#{StandardWorker}_#{SecureRandom.uuid}"
 
-      shoryuken_options queue: 'shoryuken_command', auto_delete: true
+      Shoryuken::Client.sqs.create_queue queue_name: queue
 
-      def perform(sqs_msg, body)
-        @@received_messages = Array(sqs_msg).size
-      end
+      Shoryuken.queues << queue
 
-      def self.received_messages
-        @@received_messages
-      end
+      StandardWorker.get_shoryuken_options['queue'] = queue
 
-      def self.received_messages=(received_messages)
-        @@received_messages = received_messages
-      end
+      Shoryuken.register_worker queue, StandardWorker
     end
 
-    class StandardWorker
-      include Shoryuken::Worker
-
-      @@received_messages = 0
-
-      shoryuken_options queue: 'shoryuken', auto_delete: true
+    after do
+      queue_url = Shoryuken::Client.sqs.get_queue_url(queue_name: StandardWorker.get_shoryuken_options['queue']).queue_url
 
-      def perform(sqs_msg, body)
-        @@received_messages = Array(sqs_msg).size
-      end
-
-      def self.received_messages
-        @@received_messages
-      end
-
-      def self.received_messages=(received_messages)
-        @@received_messages = received_messages
-      end
+      Shoryuken::Client.sqs.delete_queue queue_url: queue_url
     end
 
     it 'consumes as a command worker' do
-      CommandWorker.perform_async('Yo')
+      StandardWorker.perform_async('Yo')
 
-      10.times do
-        break if CommandWorker.received_messages > 0
-        sleep 1
-      end
+      poll_queues_until { StandardWorker.received_messages > 0 }
 
-      expect(CommandWorker.received_messages).to eq 1
+      expect(StandardWorker.received_messages).to eq 1
     end
 
     it 'consumes a message' do
       StandardWorker.get_shoryuken_options['batch'] = false
 
-      Shoryuken::Client.queues('shoryuken').send_message('Yo')
+      Shoryuken::Client.queues(StandardWorker.get_shoryuken_options['queue']).send_message(message_body: 'Yo')
 
-      10.times do
-        break if StandardWorker.received_messages > 0
-        sleep 1
-      end
+      poll_queues_until { StandardWorker.received_messages > 0 }
 
       expect(StandardWorker.received_messages).to eq 1
     end
@@ -89,15 +47,48 @@ describe Shoryuken::Launcher do
     it 'consumes a batch' do
       StandardWorker.get_shoryuken_options['batch'] = true
 
-      Shoryuken::Client.queues('shoryuken').batch_send *(['Yo'] * 10)
+      entries = 10.times.map { |i| { id: SecureRandom.uuid, message_body: i.to_s } }
 
-      10.times do
-        break if StandardWorker.received_messages > 0
-        sleep 1
-      end
+      Shoryuken::Client.queues(StandardWorker.get_shoryuken_options['queue']).send_messages(entries: entries)
+
+      # Give the messages a chance to hit the queue so they are all available at the same time
+      sleep 2
+
+      poll_queues_until { StandardWorker.received_messages > 0 }
 
-      # the fetch result is uncertain, should be greater than 1, but hard to tell the exact size
       expect(StandardWorker.received_messages).to be > 1
     end
+
+    def poll_queues_until
+      subject.run
+
+      Timeout::timeout(10) do
+        begin
+          sleep 0.5
+        end until yield
+      end
+    ensure
+      subject.stop
+    end
+
+    class StandardWorker
+      include Shoryuken::Worker
+
+      @@received_messages = 0
+
+      shoryuken_options auto_delete: true
+
+      def perform(sqs_msg, body)
+        @@received_messages += Array(sqs_msg).size
+      end
+
+      def self.received_messages
+        @@received_messages
+      end
+
+      def self.received_messages=(received_messages)
+        @@received_messages = received_messages
+      end
+    end
   end
 end