From 5a06c89c72bb6ab4570e86f0a0c3f092ac256000 Mon Sep 17 00:00:00 2001
From: Roman Samoilov <2270393+rsamoilov@users.noreply.github.com>
Date: Tue, 27 Aug 2024 19:44:34 +0100
Subject: [PATCH] Monitor the health of AR connections

---
 lib/rage/configuration.rb                     |  6 +++
 lib/rage/ext/active_record/connection_pool.rb | 40 +++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/lib/rage/configuration.rb b/lib/rage/configuration.rb
index 7faad39..84ff882 100644
--- a/lib/rage/configuration.rb
+++ b/lib/rage/configuration.rb
@@ -274,6 +274,12 @@ def manually_release_ar_connections?
       defined?(ActiveRecord) && ActiveRecord.version < Gem::Version.create("7.2.0")
     end
 
+    # whether we should manually reconnect closed AR connections;
+    # AR 7.1+ does this automatically while executing the query;
+    def should_manually_restore_ar_connections?
+      defined?(ActiveRecord) && ActiveRecord.version < Gem::Version.create("7.1.0")
+    end
+
     def inspect
       "#<#{self.class.name}>"
     end
diff --git a/lib/rage/ext/active_record/connection_pool.rb b/lib/rage/ext/active_record/connection_pool.rb
index cb6ed07..c2cac28 100644
--- a/lib/rage/ext/active_record/connection_pool.rb
+++ b/lib/rage/ext/active_record/connection_pool.rb
@@ -24,11 +24,30 @@ def to_a
     end
   end
 
+  # reconnect closed connections on checkout;
+  # only included with `Rage.config.should_manually_restore_ar_connections?`
+  module ConnectionWithVerify
+    def connection
+      conn = super
+
+      if conn.__needs_reconnect
+        conn.reconnect!
+        conn.__needs_reconnect = false
+      end
+
+      conn
+    end
+  end
+  if Rage.config.internal.should_manually_restore_ar_connections?
+    prepend ConnectionWithVerify
+  end
+
   def self.extended(instance)
     instance.class.alias_method :__checkout__, :checkout
     instance.class.alias_method :__remove__, :remove
 
     ActiveRecord::ConnectionAdapters::AbstractAdapter.attr_accessor(:__idle_since)
+    ActiveRecord::ConnectionAdapters::AbstractAdapter.attr_accessor(:__needs_reconnect)
   end
 
   def __init_rage_extension
@@ -65,6 +84,26 @@ def __init_rage_extension
       end
     end
 
+    # monitor connections health
+    if Rage.config.internal.should_manually_restore_ar_connections?
+      Iodine.run_every(1_000) do
+        i = 0
+        while i < @__connections.length
+          conn = @__connections[i]
+
+          unless conn.__needs_reconnect
+            needs_reconnect = !conn.active? rescue true
+            if needs_reconnect
+              conn.__needs_reconnect = true
+              conn.disconnect!
+            end
+          end
+
+          i += 1
+        end
+      end
+    end
+
     # resume blocked fibers once connections become available
     Iodine.subscribe("ext:ar-connection-released") do
       if @__blocked.length > 0 && @__connections.length > 0
@@ -135,6 +174,7 @@ def flush(minimum_idle = @__idle_timeout)
       conn = @__connections[i]
       if conn.__idle_since && current_time - conn.__idle_since >= minimum_idle
         conn.__idle_since = nil
+        conn.__needs_reconnect = true
         conn.disconnect!
       end
       i += 1