Spawn a single thread that checks for PostgreSQL fallback.

lib/active_record/connection_adapters/postgresql_fallback_adapter.rb

@@ -6,23 +6,46 @@ class PostgreSQLFallbackHandler
   include Singleton
 
   def initialize
-    @master = {}
-    @running = {}
-    @mutex = {}
-    @last_check = {}
-
-    setup!
+    @masters_down = {}
+    @mutex = Mutex.new
   end
 
   def verify_master
-    @mutex[namespace].synchronize do
-      return if running || recently_checked?
-      @running[namespace] = true
-    end
+    synchronize { return if @thread && @thread.alive? }
 
-    current_namespace = namespace
-    Thread.new do
-      RailsMultisite::ConnectionManagement.with_connection(current_namespace) do
+    @thread = Thread.new do
+      while true do
+        begin
+          thread = Thread.new { initiate_fallback_to_master }
+          thread.join
+          break if synchronize { @masters_down.empty? }
+          sleep 10
+        ensure
+          thread.kill
+        end
+      end
+    end
+  end
+
+  def master_down?
+    synchronize { @masters_down[namespace] }
+  end
+
+  def master_down=(args)
+    synchronize { @masters_down[namespace] = args }
+  end
+
+  def master_up(namespace)
+    synchronize { @masters_down.delete(namespace) }
+  end
+
+  def running?
+    synchronize { @thread.alive? }
+  end
+
+  def initiate_fallback_to_master
+    @masters_down.keys.each do |key|
+      RailsMultisite::ConnectionManagement.with_connection(key) do
         begin
           logger.warn "#{log_prefix}: Checking master server..."
           connection = ActiveRecord::Base.postgresql_connection(config)
@@ -32,54 +55,20 @@ class PostgreSQLFallbackHandler
             ActiveRecord::Base.clear_all_connections!
             logger.warn "#{log_prefix}: Master server is active. Reconnecting..."
 
-            if namespace == RailsMultisite::ConnectionManagement::DEFAULT
-              ActiveRecord::Base.establish_connection(config)
-            else
-              RailsMultisite::ConnectionManagement.establish_connection(db: namespace)
-            end
-
+            self.master_up(key)
             Discourse.disable_readonly_mode
-            self.master = true
           end
         rescue => e
-          if e.message.include?("could not connect to server")
-            logger.warn "#{log_prefix}: Connection to master PostgreSQL server failed with '#{e.message}'"
-          else
-            raise e
-          end
-        ensure
-          @mutex[namespace].synchronize do
-            @last_check[namespace] = Time.zone.now
-            @running[namespace] = false
-          end
+          byebug
+          logger.warn "#{log_prefix}: Connection to master PostgreSQL server failed with '#{e.message}'"
         end
       end
     end
   end
 
-  def master
-    @master[namespace]
-  end
-
-  def master=(args)
-    @master[namespace] = args
-  end
-
-  def running
-    @running[namespace]
-  end
-
+  # Use for testing
   def setup!
-    RailsMultisite::ConnectionManagement.all_dbs.each do |db|
-      @master[db] = true
-      @running[db] = false
-      @mutex[db] = Mutex.new
-      @last_check[db] = nil
-    end
-  end
-
-  def verify?
-    !master && !running
+    @masters_down = {}
   end
 
   private
@@ -96,17 +85,13 @@ class PostgreSQLFallbackHandler
     "#{self.class} [#{namespace}]"
   end
 
-  def recently_checked?
-    if @last_check[namespace]
-      Time.zone.now <= (@last_check[namespace] + 5.seconds)
-    else
-      false
-    end
-  end
-
   def namespace
     RailsMultisite::ConnectionManagement.current_db
   end
+
+  def synchronize
+    @mutex.synchronize { yield }
+  end
 end
 
 module ActiveRecord
@@ -115,7 +100,7 @@ module ActiveRecord
       fallback_handler = ::PostgreSQLFallbackHandler.instance
       config = config.symbolize_keys
 
-      if fallback_handler.verify?
+      if fallback_handler.master_down?
         connection = postgresql_connection(config.dup.merge({
           host: config[:replica_host], port: config[:replica_port]
         }))
@@ -126,7 +111,8 @@ module ActiveRecord
         begin
           connection = postgresql_connection(config)
         rescue PG::ConnectionBad => e
-          fallback_handler.master = false
+          fallback_handler.master_down = true
+          fallback_handler.verify_master
           raise e
         end
       end
@@ -141,20 +127,4 @@ module ActiveRecord
       raise "Replica database server is not in recovery mode." if value == 'f'
     end
   end
-
-  module ConnectionAdapters
-    class PostgreSQLAdapter
-      set_callback :checkout, :before, :switch_back?
-
-      private
-
-      def fallback_handler
-        @fallback_handler ||= ::PostgreSQLFallbackHandler.instance
-      end
-
-      def switch_back?
-        fallback_handler.verify_master if fallback_handler.verify?
-      end
-    end
-  end
 end