shoryuken 5.0.4 → 5.2.3

data/gemfiles/rails_6_1.gemfile

@@ -0,0 +1,21 @@
+# This file was generated by Appraisal
+
+source "https://rubygems.org"
+
+group :test do
+  gem "activejob", "~> 6.1"
+  gem "aws-sdk-core", "~> 3"
+  gem "aws-sdk-sqs"
+  gem "codeclimate-test-reporter", require: nil
+  gem "httparty"
+  gem "multi_xml"
+  gem "simplecov"
+end
+
+group :development do
+  gem "appraisal", git: "https://github.com/thoughtbot/appraisal.git"
+  gem "pry-byebug", "3.9.0"
+  gem "rubocop"
+end
+
+gemspec path: "../"

data/lib/shoryuken/environment_loader.rb

@@ -70,12 +70,18 @@ module Shoryuken
             ::Rails.application.config.eager_load = true
           end
         end
-        require 'shoryuken/extensions/active_job_adapter' if Shoryuken.active_job?
+        if Shoryuken.active_job?
+          require 'shoryuken/extensions/active_job_extensions'
+          require 'shoryuken/extensions/active_job_adapter'
+          require 'shoryuken/extensions/active_job_concurrent_send_adapter'
+        end
         require File.expand_path('config/environment.rb')
       end
     end
 
     def prefix_active_job_queue_name(queue_name, weight)
+      return [queue_name, weight] if queue_name.start_with?('https://', 'arn:')
+
       queue_name_prefix = ::ActiveJob::Base.queue_name_prefix
       queue_name_delimiter = ::ActiveJob::Base.queue_name_delimiter
 

data/lib/shoryuken/extensions/active_job_adapter.rb

@@ -33,8 +33,12 @@ module ActiveJob
       def enqueue(job, options = {}) #:nodoc:
         register_worker!(job)
 
+        job.sqs_send_message_parameters.merge! options
+
         queue = Shoryuken::Client.queues(job.queue_name)
-        queue.send_message(message(queue, job, options))
+        send_message_params = message queue, job
+        job.sqs_send_message_parameters = send_message_params
+        queue.send_message send_message_params
       end
 
       def enqueue_at(job, timestamp) #:nodoc:
@@ -50,44 +54,47 @@ module ActiveJob
         delay
       end
 
-      def message(queue, job, options = {})
+      def message(queue, job)
         body = job.serialize
+        job_params = job.sqs_send_message_parameters
+
+        attributes = job_params[:message_attributes] || {}
 
-        msg = {}
+        msg = {
+          message_body: body,
+          message_attributes: attributes.merge(MESSAGE_ATTRIBUTES)
+        }
 
         if queue.fifo?
           # See https://github.com/phstc/shoryuken/issues/457
           msg[:message_deduplication_id] = Digest::SHA256.hexdigest(JSON.dump(body.except('job_id')))
         end
 
-        msg[:message_body] = body
-        msg[:message_attributes] = message_attributes
-
-        msg.merge(options)
+        msg.merge(job_params.except(:message_attributes))
       end
 
       def register_worker!(job)
         Shoryuken.register_worker(job.queue_name, JobWrapper)
       end
 
-      def message_attributes
-        @message_attributes ||= {
-          'shoryuken_class' => {
-            string_value: JobWrapper.to_s,
-            data_type: 'String'
-          }
-        }
-      end
-
       class JobWrapper #:nodoc:
         include Shoryuken::Worker
 
         shoryuken_options body_parser: :json, auto_delete: true
 
-        def perform(_sqs_msg, hash)
-          Base.execute hash
+        def perform(sqs_msg, hash)
+          receive_count = sqs_msg.attributes['ApproximateReceiveCount'].to_i
+          past_receives = receive_count - 1
+          Base.execute hash.merge({ 'executions' => past_receives })
         end
       end
+
+      MESSAGE_ATTRIBUTES = {
+        'shoryuken_class' => {
+          string_value: JobWrapper.to_s,
+          data_type: 'String'
+        }
+      }.freeze
     end
   end
 end

data/lib/shoryuken/extensions/active_job_extensions.rb

@@ -0,0 +1,38 @@
+module Shoryuken
+  module ActiveJobExtensions
+    # Adds an accessor for SQS SendMessage parameters on ActiveJob jobs
+    # (instances of ActiveJob::Base). Shoryuken ActiveJob queue adapters use
+    # these parameters when enqueueing jobs; other adapters can ignore them.
+    module SQSSendMessageParametersAccessor
+      extend ActiveSupport::Concern
+
+      included do
+        attr_accessor :sqs_send_message_parameters
+      end
+    end
+
+    # Initializes SQS SendMessage parameters on instances of ActiveJobe::Base
+    # to the empty hash, and populates it whenever `#enqueue` is called, such
+    # as when using ActiveJob::Base.set.
+    module SQSSendMessageParametersSupport
+      def initialize(*arguments)
+        super(*arguments)
+        self.sqs_send_message_parameters = {}
+      end
+      ruby2_keywords(:initialize) if respond_to?(:ruby2_keywords, true)
+
+      def enqueue(options = {})
+        sqs_options = options.extract! :message_attributes,
+                                       :message_system_attributes,
+                                       :message_deduplication_id,
+                                       :message_group_id
+        sqs_send_message_parameters.merge! sqs_options
+
+        super
+      end
+    end
+  end
+end
+
+ActiveJob::Base.include Shoryuken::ActiveJobExtensions::SQSSendMessageParametersAccessor
+ActiveJob::Base.prepend Shoryuken::ActiveJobExtensions::SQSSendMessageParametersSupport

data/lib/shoryuken/launcher.rb

@@ -71,6 +71,7 @@ module Shoryuken
     def create_managers
       Shoryuken.groups.map do |group, options|
         Shoryuken::Manager.new(
+          group,
           Shoryuken::Fetcher.new(group),
           Shoryuken.polling_strategy(group).new(options[:queues], Shoryuken.delay(group)),
           options[:concurrency],

data/lib/shoryuken/manager.rb

@@ -6,7 +6,8 @@ module Shoryuken
     # See https://github.com/phstc/shoryuken/issues/348#issuecomment-292847028
     MIN_DISPATCH_INTERVAL = 0.1
 
-    def initialize(fetcher, polling_strategy, concurrency, executor)
+    def initialize(group, fetcher, polling_strategy, concurrency, executor)
+      @group            = group
       @fetcher          = fetcher
       @polling_strategy = polling_strategy
       @max_processors   = concurrency
@@ -16,6 +17,7 @@ module Shoryuken
     end
 
     def start
+      fire_utilization_update_event
       dispatch_loop
     end
 
@@ -57,8 +59,15 @@ module Shoryuken
       @max_processors - busy
     end
 
-    def processor_done
+    def processor_done(queue)
       @busy_processors.decrement
+      fire_utilization_update_event
+
+      client_queue = Shoryuken::Client.queues(queue)
+      return unless client_queue.fifo?
+      return unless @polling_strategy.respond_to?(:message_processed)
+
+      @polling_strategy.message_processed(queue)
     end
 
     def assign(queue_name, sqs_msg)
@@ -67,10 +76,12 @@ module Shoryuken
       logger.debug { "Assigning #{sqs_msg.message_id}" }
 
       @busy_processors.increment
+      fire_utilization_update_event
 
-      Concurrent::Promise.execute(
-        executor: @executor
-      ) { Processor.process(queue_name, sqs_msg) }.then { processor_done }.rescue { processor_done }
+      Concurrent::Promise
+        .execute(executor: @executor) { Processor.process(queue_name, sqs_msg) }
+        .then { processor_done(queue_name) }
+        .rescue { processor_done(queue_name) }
     end
 
     def dispatch_batch(queue)
@@ -108,5 +119,13 @@ module Shoryuken
 
       @running.make_false
     end
+
+    def fire_utilization_update_event
+      fire_event :utilization_update, false, {
+        group: @group,
+        max_processors: @max_processors,
+        busy_processors: busy
+      }
+    end
   end
 end

data/lib/shoryuken/options.rb

@@ -9,6 +9,7 @@ module Shoryuken
       lifecycle_events: {
         startup: [],
         dispatch: [],
+        utilization_update: [],
         quiet: [],
         shutdown: []
       }

data/lib/shoryuken/polling/base.rb

@@ -40,6 +40,8 @@ module Shoryuken
         fail NotImplementedError
       end
 
+      def message_processed(_queue); end
+
       def active_queues
         fail NotImplementedError
       end

data/lib/shoryuken/polling/strict_priority.rb

@@ -38,6 +38,11 @@ module Shoryuken
           .reverse
       end
 
+      def message_processed(queue)
+        logger.debug "Unpausing #{queue}"
+        @paused_until[queue] = Time.now
+      end
+
       private
 
       def next_active_queue
@@ -70,6 +75,7 @@ module Shoryuken
 
       def pause(queue)
         return unless delay > 0
+
         @paused_until[queue] = Time.now + delay
         logger.debug "Paused #{queue}"
       end

data/lib/shoryuken/polling/weighted_round_robin.rb

@@ -35,10 +35,20 @@ module Shoryuken
         unparse_queues(@queues)
       end
 
+      def message_processed(queue)
+        return if @paused_queues.empty?
+
+        logger.debug "Unpausing #{queue}"
+        @paused_queues.reject! { |_time, name| name == queue }
+        @queues << queue
+        @queues.uniq!
+      end
+
       private
 
       def pause(queue)
         return unless @queues.delete(queue)
+
         @paused_queues << [Time.now + delay, queue]
         logger.debug "Paused #{queue}"
       end
@@ -46,6 +56,7 @@ module Shoryuken
       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]}"

data/lib/shoryuken/queue.rb

@@ -8,9 +8,9 @@ module Shoryuken
 
     attr_accessor :name, :client, :url
 
-    def initialize(client, name_or_url)
+    def initialize(client, name_or_url_or_arn)
       self.client = client
-      set_name_and_url(name_or_url)
+      set_name_and_url(name_or_url_or_arn)
     end
 
     def visibility_timeout
@@ -50,32 +50,60 @@ module Shoryuken
       # Make sure the memoization work with boolean to avoid multiple calls to SQS
       # see https://github.com/phstc/shoryuken/pull/529
       return @_fifo if defined?(@_fifo)
+
       @_fifo = queue_attributes.attributes[FIFO_ATTR] == 'true'
+      @_fifo
     end
 
     private
 
-    def set_by_name(name)
+    def initialize_fifo_attribute
+      # calling fifo? will also initialize it
+      fifo?
+    end
+
+    def set_by_name(name) # rubocop:disable Naming/AccessorMethodName
       self.name = name
       self.url  = client.get_queue_url(queue_name: name).queue_url
     end
 
-    def set_by_url(url)
+    def set_by_url(url) # rubocop:disable Naming/AccessorMethodName
       self.name = url.split('/').last
       self.url  = url
     end
 
-    def set_name_and_url(name_or_url)
-      if name_or_url.include?('://')
-        set_by_url(name_or_url)
+    def arn_to_url(arn_str)
+      *, region, account_id, resource = arn_str.split(':')
+
+      required = [region, account_id, resource].map(&:to_s)
+      valid = required.none?(&:empty?)
+
+      abort "Invalid ARN: #{arn_str}. A valid ARN must include: region, account_id and resource." unless valid
+
+      "https://sqs.#{region}.amazonaws.com/#{account_id}/#{resource}"
+    end
+
+    def set_name_and_url(name_or_url_or_arn) # rubocop:disable Naming/AccessorMethodName
+      if name_or_url_or_arn.include?('://')
+        set_by_url(name_or_url_or_arn)
+
+        # anticipate the fifo? checker for validating the queue URL
+        initialize_fifo_attribute
+        return
+      end
+
+      if name_or_url_or_arn.start_with?('arn:')
+        url = arn_to_url(name_or_url_or_arn)
+        set_by_url(url)
 
         # anticipate the fifo? checker for validating the queue URL
-        return fifo?
+        initialize_fifo_attribute
+        return
       end
 
-      set_by_name(name_or_url)
-    rescue Aws::Errors::NoSuchEndpointError, Aws::SQS::Errors::NonExistentQueue => ex
-      raise ex, "The specified queue #{name_or_url} does not exist."
+      set_by_name(name_or_url_or_arn)
+    rescue Aws::Errors::NoSuchEndpointError, Aws::SQS::Errors::NonExistentQueue => e
+      raise e, "The specified queue #{name_or_url_or_arn} does not exist."
     end
 
     def queue_attributes

data/lib/shoryuken/version.rb

@@ -1,3 +1,3 @@
 module Shoryuken
-  VERSION = '5.0.4'.freeze
+  VERSION = '5.2.3'.freeze
 end

data/lib/shoryuken.rb

@@ -89,6 +89,7 @@ module Shoryuken
 end
 
 if Shoryuken.active_job?
+  require 'shoryuken/extensions/active_job_extensions'
   require 'shoryuken/extensions/active_job_adapter'
   require 'shoryuken/extensions/active_job_concurrent_send_adapter'
 end

data/shoryuken.gemspec

@@ -18,7 +18,6 @@ Gem::Specification.new do |spec|
   spec.require_paths = ['lib']
 
   spec.add_development_dependency 'dotenv'
-  spec.add_development_dependency 'pry-byebug'
   spec.add_development_dependency 'rake'
   spec.add_development_dependency 'rspec'
 

data/spec/integration/launcher_spec.rb

@@ -4,10 +4,37 @@ require 'shoryuken/launcher'
 require 'securerandom'
 
 RSpec.describe Shoryuken::Launcher do
-  describe 'Consuming messages', slow: true do
+  let(:sqs_client) do
+    Aws::SQS::Client.new(
+      region: 'us-east-1',
+      endpoint: 'http://localhost:5000',
+      access_key_id: 'fake',
+      secret_access_key: 'fake'
+    )
+  end
+
+  let(:executor) do
+    # We can't use Concurrent.global_io_executor in these tests since once you
+    # shut down a thread pool, you can't start it back up. Instead, we create
+    # one new thread pool executor for each spec. We use a new
+    # CachedThreadPool, since that most closely resembles
+    # Concurrent.global_io_executor
+    Concurrent::CachedThreadPool.new auto_terminate: true
+  end
+
+  describe 'Consuming messages' do
     before do
       Aws.config[:stub_responses] = false
-      Aws.config[:region] = 'us-east-1'
+
+      allow(Shoryuken).to receive(:launcher_executor).and_return(executor)
+
+      Shoryuken.configure_client do |config|
+        config.sqs_client = sqs_client
+      end
+
+      Shoryuken.configure_server do |config|
+        config.sqs_client = sqs_client
+      end
 
       StandardWorker.received_messages = 0