From 0e3acf412e59335d779b5514f9c38b473377a1cf Mon Sep 17 00:00:00 2001
From: Natasha <natashamontero@gmail.com>
Date: Sat, 1 Apr 2023 23:46:39 -0400
Subject: [PATCH] Avoid setting status to jobs that are not
 Sidekiq::Status::Worker upon exceptions (#32)

---
 lib/sidekiq-status/server_middleware.rb       | 38 ++++++++++---------
 .../sidekiq-status/server_middleware_spec.rb  | 16 ++++++++
 spec/support/test_jobs.rb                     | 12 ++++++
 3 files changed, 48 insertions(+), 18 deletions(-)

diff --git a/lib/sidekiq-status/server_middleware.rb b/lib/sidekiq-status/server_middleware.rb
index c39a4cd..507fdd1 100644
--- a/lib/sidekiq-status/server_middleware.rb
+++ b/lib/sidekiq-status/server_middleware.rb
@@ -43,26 +43,28 @@ def call(worker, msg, queue)
         return
       end
 
-      # Determine job expiration
-      expiry = job_class.new.expiration || @expiration rescue @expiration
-
-      store_status worker.jid, :working,  expiry
-      yield
-      store_status worker.jid, :complete, expiry
-    rescue Worker::Stopped
-      store_status worker.jid, :stopped, expiry
-    rescue SystemExit, Interrupt
-      store_status worker.jid, :interrupted, expiry
-      raise
-    rescue Exception
-      status = :failed
-      if msg['retry']
-        if retry_attempt_number(msg) < retry_attempts_from(msg['retry'], DEFAULT_MAX_RETRY_ATTEMPTS)
-          status = :retrying
+      begin
+        # Determine job expiration
+        expiry = job_class.new.expiration || @expiration rescue @expiration
+
+        store_status worker.jid, :working,  expiry
+        yield
+        store_status worker.jid, :complete, expiry
+      rescue Worker::Stopped
+        store_status worker.jid, :stopped, expiry
+      rescue SystemExit, Interrupt
+        store_status worker.jid, :interrupted, expiry
+        raise
+      rescue Exception
+        status = :failed
+        if msg['retry']
+          if retry_attempt_number(msg) < retry_attempts_from(msg['retry'], DEFAULT_MAX_RETRY_ATTEMPTS)
+            status = :retrying
+          end
         end
+        store_status(worker.jid, status, expiry) if job_class && job_class.ancestors.include?(Sidekiq::Status::Worker)
+        raise
       end
-      store_status(worker.jid, status, expiry) if job_class && job_class.ancestors.include?(Sidekiq::Status::Worker)
-      raise
     end
 
     private
diff --git a/spec/lib/sidekiq-status/server_middleware_spec.rb b/spec/lib/sidekiq-status/server_middleware_spec.rb
index c53def9..4c2c807 100644
--- a/spec/lib/sidekiq-status/server_middleware_spec.rb
+++ b/spec/lib/sidekiq-status/server_middleware_spec.rb
@@ -61,6 +61,22 @@
         end
         expect(redis.hget("sidekiq:status:#{job_id}", :status)).to be_nil
       end
+
+      it "should not set any status on system exit signal" do
+        allow(SecureRandom).to receive(:hex).once.and_return(job_id)
+        start_server do
+          expect(ExitedNoStatusJob.perform_async).to eq(job_id)
+        end
+        expect(redis.hget("sidekiq:status:#{job_id}", :status)).to be_nil
+      end
+
+      it "should not set any status on interrupt signal" do
+        allow(SecureRandom).to receive(:hex).once.and_return(job_id)
+        start_server do
+          expect(InterruptedNoStatusJob.perform_async).to eq(job_id)
+        end
+        expect(redis.hget("sidekiq:status:#{job_id}", :status)).to be_nil
+      end
     end
 
     context "sets interrupted status" do
diff --git a/spec/support/test_jobs.rb b/spec/support/test_jobs.rb
index 393d787..6e36391 100644
--- a/spec/support/test_jobs.rb
+++ b/spec/support/test_jobs.rb
@@ -111,12 +111,24 @@ def perform
   end
 end
 
+class ExitedNoStatusJob < StubNoStatusJob
+  def perform
+    raise SystemExit
+  end
+end
+
 class InterruptedJob < StubJob
   def perform
     raise Interrupt
   end
 end
 
+class InterruptedNoStatusJob < StubNoStatusJob
+  def perform
+    raise Interrupt
+  end
+end
+
 class RetriedJob < StubJob
 
   sidekiq_options retry: true