add an after_requeue callback to RequeueExceptions

lib/qless/middleware/requeue_exceptions.rb

@@ -35,13 +35,18 @@ def raise_if_exhausted_requeues(error, requeues)
         end
       end
 
-      def requeue_on(*exceptions, options)
+      def requeue_on(*exceptions, options, &block)
+        after_requeue_callbacks << block if block_given?
         RequeueableException.from_splat_and_options(
           *exceptions, options).each do |exc|
           requeueable_exceptions[exc.klass] = exc
         end
       end
 
+      def after_requeue_callbacks
+        @after_requeue_callbacks ||= []
+      end
+
       def around_perform(job)
         super
       rescue *requeueable_exceptions.keys => e
@@ -55,6 +60,8 @@ def around_perform(job)
 
         requeues_by_exception[config.klass.name] += 1
         job.requeue(job.queue_name, delay: config.delay, data: job.data)
+
+        after_requeue_callbacks.each { |callback| callback.call(e, job) }
       end
 
       def requeueable_exceptions

spec/unit/middleware/requeue_exceptions_spec.rb

@@ -28,6 +28,11 @@ def around_perform(job)
       matched_exception_2 = KeyError
       unmatched_exception = RegexpError
 
+      let(:requeue_on_args) do
+        [matched_exception_1, matched_exception_2, MessageSpecificException,
+          delay_range: delay_range, max_attempts: max_attempts]
+      end
+
       module MessageSpecificException
         def self.===(other)
           ArgumentError === other && other.message.include?("foo")
@@ -36,16 +41,32 @@ def self.===(other)
 
       before do
         container.extend(RequeueExceptions)
-        container.requeue_on(matched_exception_1, matched_exception_2,
-                             MessageSpecificException,
-                             delay_range: delay_range,
-                             max_attempts: max_attempts)
+        container.requeue_on(*requeue_on_args)
+      end
+
+      def add_requeue_callback
+        callback = ->(error, job) { callback_catcher << [error, job] }
+        container.after_requeue_callbacks << callback
+      end
+
+      def callback_catcher
+        @callback_catcher ||= []
       end
 
       def perform
         container.around_perform(job)
       end
 
+      describe '.requeue_on' do
+        it 'accepts a block to set an after requeue callback' do
+          container.extend(RequeueExceptions)
+
+          expect {
+            container.requeue_on(*requeue_on_args) { |*| true }
+          }.to change {container.after_requeue_callbacks.size }.from(0).to(1)
+        end
+      end
+
       context 'when no exception is raised' do
         before { container.perform = -> { } }
 
@@ -62,6 +83,16 @@ def perform
           job.should_not_receive(:requeue)
           expect { perform }.to raise_error(unmatched_exception)
         end
+
+        context 'when an after requeue callback is set' do
+          before { add_requeue_callback }
+
+          it 'does not call the callback' do
+            expect { perform }.to raise_error(unmatched_exception)
+
+            expect(callback_catcher.size).to eq(0)
+          end
+        end
       end
 
       shared_context "requeues on matching exception" do |exception, exception_name|
@@ -129,6 +160,16 @@ def perform
 
           expect { perform }.to raise_error(exception)
         end
+
+        context 'when an after requeue callback is set' do
+          before { add_requeue_callback }
+
+          it 'calls the callback' do
+            expect {
+              perform
+            }.to change { callback_catcher.size }.from(0).to(1)
+          end
+        end
       end
 
       context "when a matched exception is raised" do