From d3fea44790e07d1a88d0563ae6d1c6a63490fc48 Mon Sep 17 00:00:00 2001
From: "Ben Sheldon [he/him]" <bensheldon@github.com>
Date: Mon, 15 Jul 2024 09:32:03 -0500
Subject: [PATCH] Add keepalive SQL query to Notifier

---
 lib/good_job/notifier.rb           |  7 +++++++
 spec/lib/good_job/notifier_spec.rb | 14 ++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/lib/good_job/notifier.rb b/lib/good_job/notifier.rb
index 2a8621b4..6584425e 100644
--- a/lib/good_job/notifier.rb
+++ b/lib/good_job/notifier.rb
@@ -27,6 +27,8 @@ class Notifier
     RECONNECT_INTERVAL = 5
     # Number of consecutive connection errors before reporting an error
     CONNECTION_ERRORS_REPORTING_THRESHOLD = 6
+    # Interval for emitting a noop SQL query to keep the connection alive
+    KEEPALIVE_INTERVAL = 10
 
     # Connection errors that will wait {RECONNECT_INTERVAL} before reconnecting
     CONNECTION_ERRORS = %w[
@@ -78,6 +80,7 @@ def initialize(*recipients, enable_listening: true, capsule: GoodJob.capsule, ex
       @enable_listening = enable_listening
       @task = nil
       @capsule = capsule
+      @last_keepalive_time = Time.current
 
       start
       self.class.instances << self
@@ -269,6 +272,10 @@ def wait_for_notify
         raw_connection.wait_for_notify(WAIT_INTERVAL) do |channel, _pid, payload|
           yield(channel, payload)
         end
+        if Time.current - @last_keepalive_time >= KEEPALIVE_INTERVAL
+          raw_connection.async_exec("SELECT 1")
+          @last_keepalive_time = Time.current
+        end
       elsif @enable_listening && raw_connection.respond_to?(:jdbc_connection)
         raw_connection.execute_query("SELECT 1")
         notifications = raw_connection.jdbc_connection.getNotifications
diff --git a/spec/lib/good_job/notifier_spec.rb b/spec/lib/good_job/notifier_spec.rb
index c6f88881..264c443c 100644
--- a/spec/lib/good_job/notifier_spec.rb
+++ b/spec/lib/good_job/notifier_spec.rb
@@ -134,6 +134,20 @@
 
       notifier.shutdown
     end
+
+    it 'executes a noop SQL query every 10 seconds to keep the connection alive' do
+      stub_const("GoodJob::Notifier::KEEPALIVE_INTERVAL", 0.1)
+      stub_const("GoodJob::Notifier::WAIT_INTERVAL", 0.1)
+
+      notifier = described_class.new(enable_listening: true)
+      original_keepalive = notifier.instance_variable_get(:@last_keepalive_time)
+
+      expect(notifier).to be_listening(timeout: 2)
+      sleep 0.2
+      expect(notifier.instance_variable_get(:@last_keepalive_time)).to be > original_keepalive
+
+      notifier.shutdown
+    end
   end
 
   describe '#shutdown' do