shoryuken 5.0.3 → 5.1.1

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/manager.rb

@@ -57,8 +57,13 @@ module Shoryuken
       @max_processors - busy
     end
 
-    def processor_done
+    def processor_done(queue)
       @busy_processors.decrement
+      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)
@@ -68,9 +73,10 @@ module Shoryuken
 
       @busy_processors.increment
 
-      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)

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.3'.freeze
+  VERSION = '5.1.1'.freeze
 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/shared_examples_for_active_job.rb

@@ -1,30 +1,44 @@
+require 'active_job'
+require 'shoryuken/extensions/active_job_extensions'
+
+# Stand-in for a job class specified by the user
+class TestJob < ActiveJob::Base; end
+
 # rubocop:disable Metrics/BlockLength
 RSpec.shared_examples 'active_job_adapters' do
-  let(:job) { double 'Job', id: '123', queue_name: 'queue' }
+  let(:job_sqs_send_message_parameters) { {} }
+  let(:job) do
+    job = TestJob.new
+    job.sqs_send_message_parameters = job_sqs_send_message_parameters
+    job
+  end
   let(:fifo) { false }
   let(:queue) { double 'Queue', fifo?: fifo }
 
   before do
     allow(Shoryuken::Client).to receive(:queues).with(job.queue_name).and_return(queue)
-    allow(job).to receive(:serialize).and_return(
-      'job_class' => 'Worker',
-      'job_id' => job.id,
-      'queue_name' => job.queue_name,
-      'arguments' => nil,
-      'locale' => nil
-    )
   end
 
   describe '#enqueue' do
     specify do
       expect(queue).to receive(:send_message) do |hash|
         expect(hash[:message_deduplication_id]).to_not be
+        expect(hash[:message_attributes]['shoryuken_class'][:string_value]).to eq(described_class::JobWrapper.to_s)
+        expect(hash[:message_attributes]['shoryuken_class'][:data_type]).to eq("String")
+        expect(hash[:message_attributes].keys).to eq(['shoryuken_class'])
       end
       expect(Shoryuken).to receive(:register_worker).with(job.queue_name, described_class::JobWrapper)
 
       subject.enqueue(job)
     end
 
+    it "should mutate the job's sqs_send_message_parameters reference to match those sent to the queue" do
+      expect(queue).to receive(:send_message) do |options|
+        expect(options).to be(job.sqs_send_message_parameters)
+      end
+      subject.enqueue(job)
+    end
+
     context 'when fifo' do
       let(:fifo) { true }
 
@@ -38,6 +52,209 @@ RSpec.shared_examples 'active_job_adapters' do
 
         subject.enqueue(job)
       end
+
+      context 'with message_deduplication_id' do
+        context 'when message_deduplication_id is specified in options' do
+          it 'should enqueue a message with the deduplication_id specified in options' do
+            expect(queue).to receive(:send_message) do |hash|
+              expect(hash[:message_deduplication_id]).to eq('options-dedupe-id')
+            end
+            subject.enqueue(job, message_deduplication_id: 'options-dedupe-id')
+          end
+        end
+
+        context 'when message_deduplication_id is specified on the job' do
+          let(:job_sqs_send_message_parameters) { { message_deduplication_id: 'job-dedupe-id' } }
+
+          it 'should enqueue a message with the deduplication_id specified on the job' do
+            expect(queue).to receive(:send_message) do |hash|
+              expect(hash[:message_deduplication_id]).to eq('job-dedupe-id')
+            end
+            subject.enqueue job
+          end
+        end
+
+        context 'when message_deduplication_id is specified on the job and also in options' do
+          let(:job_sqs_send_message_parameters) { { message_deduplication_id: 'job-dedupe-id' } }
+
+          it 'should enqueue a message with the deduplication_id specified in options' do
+            expect(queue).to receive(:send_message) do |hash|
+              expect(hash[:message_deduplication_id]).to eq('options-dedupe-id')
+            end
+            subject.enqueue(job, message_deduplication_id: 'options-dedupe-id')
+          end
+        end
+      end
+    end
+
+    context 'with message_group_id' do
+      context 'when message_group_id is specified in options' do
+        it 'should enqueue a message with the group_id specified in options' do
+          expect(queue).to receive(:send_message) do |hash|
+            expect(hash[:message_group_id]).to eq('options-group-id')
+          end
+          subject.enqueue(job, message_group_id: 'options-group-id')
+        end
+      end
+
+      context 'when message_group_id is specified on the job' do
+        let(:job_sqs_send_message_parameters) { { message_group_id: 'job-group-id' } }
+
+        it 'should enqueue a message with the group_id specified on the job' do
+          expect(queue).to receive(:send_message) do |hash|
+            expect(hash[:message_group_id]).to eq('job-group-id')
+          end
+          subject.enqueue job
+        end
+      end
+
+      context 'when message_group_id is specified on the job and also in options' do
+        let(:job_sqs_send_message_parameters) { { message_group_id: 'job-group-id' } }
+
+        it 'should enqueue a message with the group_id specified in options' do
+          expect(queue).to receive(:send_message) do |hash|
+            expect(hash[:message_group_id]).to eq('options-group-id')
+          end
+          subject.enqueue(job, message_group_id: 'options-group-id')
+        end
+      end
+    end
+
+    context 'with additional message attributes' do
+      it 'should combine with activejob attributes' do
+        custom_message_attributes = {
+          'tracer_id' => {
+            string_value: SecureRandom.hex,
+            data_type: 'String'
+          }
+        }
+
+        expect(queue).to receive(:send_message) do |hash|
+          expect(hash[:message_attributes]['shoryuken_class'][:string_value]).to eq(described_class::JobWrapper.to_s)
+          expect(hash[:message_attributes]['shoryuken_class'][:data_type]).to eq("String")
+          expect(hash[:message_attributes]['tracer_id'][:string_value]).to eq(custom_message_attributes['tracer_id'][:string_value])
+          expect(hash[:message_attributes]['tracer_id'][:data_type]).to eq("String")
+        end
+        expect(Shoryuken).to receive(:register_worker).with(job.queue_name, described_class::JobWrapper)
+
+        subject.enqueue(job, message_attributes: custom_message_attributes)
+      end
+
+      context 'when message_attributes are specified on the job' do
+        let(:job_sqs_send_message_parameters) do
+          {
+            message_attributes: {
+              'tracer_id' => {
+                data_type: 'String',
+                string_value: 'job-value'
+              }
+            }
+          }
+        end
+
+        it 'should enqueue a message with the message_attributes specified on the job' do
+          expect(queue).to receive(:send_message) do |hash|
+            expect(hash[:message_attributes]['tracer_id']).to eq({ data_type: 'String', string_value: 'job-value' })
+            expect(hash[:message_attributes]['shoryuken_class']).to eq({ data_type: 'String', string_value: described_class::JobWrapper.to_s })
+          end
+          subject.enqueue job
+        end
+      end
+
+      context 'when message_attributes are specified on the job and also in options' do
+        let(:job_sqs_send_message_parameters) do
+          {
+            message_attributes: {
+              'tracer_id' => {
+                data_type: 'String',
+                string_value: 'job-value'
+              }
+            }
+          }
+        end
+
+        it 'should enqueue a message with the message_attributes speficied in options' do
+          custom_message_attributes = {
+            'options_tracer_id' => {
+              string_value: 'options-value',
+              data_type: 'String'
+            }
+          }
+
+          expect(queue).to receive(:send_message) do |hash|
+            expect(hash[:message_attributes]['tracer_id']).to be_nil
+            expect(hash[:message_attributes]['options_tracer_id']).to eq({ data_type: 'String', string_value: 'options-value' })
+            expect(hash[:message_attributes]['shoryuken_class']).to eq({ data_type: 'String', string_value: described_class::JobWrapper.to_s })
+          end
+          subject.enqueue job, message_attributes: custom_message_attributes
+        end
+      end
+    end
+  end
+
+  context 'with message_system_attributes' do
+    context 'when message_system_attributes are specified in options' do
+      it 'should enqueue a message with message_system_attributes specified in options' do
+        system_attributes = {
+          'AWSTraceHeader' => {
+            string_value: 'trace_id',
+            data_type: 'String'
+          }
+        }
+        expect(queue).to receive(:send_message) do |hash|
+          expect(hash[:message_system_attributes]['AWSTraceHeader'][:string_value]).to eq('trace_id')
+          expect(hash[:message_system_attributes]['AWSTraceHeader'][:data_type]).to eq('String')
+        end
+        subject.enqueue(job, message_system_attributes: system_attributes)
+      end
+    end
+
+    context 'when message_system_attributes are specified on the job' do
+      let(:job_sqs_send_message_parameters) do
+        {
+          message_system_attributes: {
+            'AWSTraceHeader' => {
+              string_value: 'job-value',
+              data_type: 'String'
+            }
+          }
+        }
+      end
+
+      it 'should enqueue a message with the message_system_attributes specified on the job' do
+        expect(queue).to receive(:send_message) do |hash|
+          expect(hash[:message_system_attributes]['AWSTraceHeader']).to eq({ data_type: 'String', string_value: 'job-value' })
+        end
+        subject.enqueue job
+      end
+    end
+
+    context 'when message_system_attributes are specified on the job and also in options' do
+      let(:job_sqs_send_message_parameters) do
+        {
+          message_system_attributes: {
+            'job_trace_header' => {
+              string_value: 'job-value',
+              data_type: 'String'
+            }
+          }
+        }
+      end
+
+      it 'should enqueue a message with the message_system_attributes speficied in options' do
+        custom_message_attributes = {
+          'options_trace_header' => {
+            string_value: 'options-value',
+            data_type: 'String'
+          }
+        }
+
+        expect(queue).to receive(:send_message) do |hash|
+          expect(hash[:message_system_attributes]['job_trace_header']).to be_nil
+          expect(hash[:message_system_attributes]['options_trace_header']).to eq({ data_type: 'String', string_value: 'options-value' })
+        end
+        subject.enqueue job, message_system_attributes: custom_message_attributes
+      end
     end
   end
 
@@ -59,4 +276,4 @@ RSpec.shared_examples 'active_job_adapters' do
     end
   end
 end
-# rubocop:enable Metrics/BlockLength
+# rubocop:enable Metrics/BlockLength