From 872d49d5604e10543321d866dd670e682971b109 Mon Sep 17 00:00:00 2001
From: Ivan Giuliani <ivan@gocardless.com>
Date: Thu, 22 Apr 2021 19:24:05 +0100
Subject: [PATCH] Do not kill the worker thread on connection timeouts

Just like #66, we've seen instances of a que worker using fewer threads
than desired after failing to checking out a connection from AR. In some
cases, this can lead to a worker doing nothing (specifically, this can
happen if the que worker is configured to run with a single thread -- if
that thread dies, the worker will sit idle doing nothing forever).

This commit lets the ActiveRecord adapter handle the case where the
connection could not be estabilished for reasons other than a timeout,
which was already handled. Additionally, the internal
`UnavailableConnection` exception will also now include the original
exception as a nested `.cause` attribute to aid debugging.
---
 lib/que/adapters/active_record.rb |  9 +++++++--
 spec/lib/que/worker_spec.rb       | 15 +++++++++++++--
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/lib/que/adapters/active_record.rb b/lib/que/adapters/active_record.rb
index 0fbcb601..c6d5f46e 100644
--- a/lib/que/adapters/active_record.rb
+++ b/lib/que/adapters/active_record.rb
@@ -3,10 +3,15 @@
 module Que
   module Adapters
     class ActiveRecord < Base
+      AR_UNAVAILABLE_CONNECTION_ERRORS = [
+        ::ActiveRecord::ConnectionTimeoutError,
+        ::ActiveRecord::ConnectionNotEstablished,
+      ].freeze
+
       def checkout
         checkout_activerecord_adapter { |conn| yield conn.raw_connection }
-      rescue ::ActiveRecord::ConnectionTimeoutError => e
-        raise UnavailableConnection
+      rescue *AR_UNAVAILABLE_CONNECTION_ERRORS => e
+        raise UnavailableConnection.new(e)
       end
 
       def wake_worker_after_commit
diff --git a/spec/lib/que/worker_spec.rb b/spec/lib/que/worker_spec.rb
index a93a5c30..a3e24ed8 100644
--- a/spec/lib/que/worker_spec.rb
+++ b/spec/lib/que/worker_spec.rb
@@ -125,13 +125,24 @@
         end
       end
 
-      context "when we can't checkout a new connection" do
+      context "when we time out checking out a new connection" do
         it "rescues it and returns an error" do
           FakeJob.enqueue(1)
 
           expect(Que).
             to receive(:execute).with(:lock_job, ["default", 0]).
-              and_raise(ActiveRecord::ConnectionTimeoutError)
+            and_raise(ActiveRecord::ConnectionTimeoutError)
+          expect(subject).to eq(:postgres_error)
+        end
+      end
+
+      context "when we can't connect to postgres" do
+        it "rescues it and returns an error" do
+          FakeJob.enqueue(1)
+
+          expect(Que).
+            to receive(:execute).with(:lock_job, ["default", 0]).
+            and_raise(ActiveRecord::ConnectionNotEstablished)
           expect(subject).to eq(:postgres_error)
         end
       end