From fe1e78ddf4d5315ecbd454d85df8fe531ac7c790 Mon Sep 17 00:00:00 2001
From: Guo Xiang Tan <tgx_world@hotmail.com>
Date: Fri, 20 Oct 2017 16:39:31 +0800
Subject: [PATCH] Make PostgreSQL failover work with distributed cache.

---
 .../postgresql_fallback_adapter.rb                   |  9 +++++----
 lib/distributed_cache.rb                             | 12 ++++++++++--
 .../postgresql_fallback_adapter_spec.rb              |  5 ++++-
 spec/components/distributed_cache_spec.rb            |  2 +-
 4 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/lib/active_record/connection_adapters/postgresql_fallback_adapter.rb b/lib/active_record/connection_adapters/postgresql_fallback_adapter.rb
index 94d8b987fcb..a508aafa94c 100644
--- a/lib/active_record/connection_adapters/postgresql_fallback_adapter.rb
+++ b/lib/active_record/connection_adapters/postgresql_fallback_adapter.rb
@@ -7,10 +7,12 @@ class PostgreSQLFallbackHandler
   include Singleton
 
   attr_reader :masters_down
+  attr_accessor :initialized
 
   def initialize
-    @masters_down = DistributedCache.new('masters_down')
+    @masters_down = DistributedCache.new('masters_down', namespace: false)
     @mutex = Mutex.new
+    @initialized = false
   end
 
   def verify_master
@@ -126,13 +128,12 @@ module ActiveRecord
       else
         begin
           connection = postgresql_connection(config)
-          fallback_handler.master_down = false
+          fallback_handler.initialized ||= true
         rescue PG::ConnectionBad => e
-          on_boot = fallback_handler.master_down?.nil?
           fallback_handler.master_down = true
           fallback_handler.verify_master
 
-          if on_boot
+          if !fallback_handler.initialized
             return postgresql_fallback_connection(config)
           else
             raise e
diff --git a/lib/distributed_cache.rb b/lib/distributed_cache.rb
index b88e4679c78..4a0a0348cae 100644
--- a/lib/distributed_cache.rb
+++ b/lib/distributed_cache.rb
@@ -105,10 +105,11 @@ class DistributedCache
 
   attr_reader :key
 
-  def initialize(key, manager = nil)
+  def initialize(key, manager: nil, namespace: true)
     @key = key
     @data = {}
     @manager = manager || DistributedCache.default_manager
+    @namespace = namespace
 
     @manager.ensure_subscribe!
     @manager.register(self)
@@ -142,7 +143,14 @@ class DistributedCache
   end
 
   def hash(db = nil)
-    db ||= RailsMultisite::ConnectionManagement.current_db
+    db ||= begin
+      if @namespace
+        RailsMultisite::ConnectionManagement.current_db
+      else
+        RailsMultisite::ConnectionManagement::DEFAULT
+      end
+    end
+
     @data[db] ||= ThreadSafe::Hash.new
   end
 
diff --git a/spec/components/active_record/connection_adapters/postgresql_fallback_adapter_spec.rb b/spec/components/active_record/connection_adapters/postgresql_fallback_adapter_spec.rb
index efec0c5d54e..c353d305289 100644
--- a/spec/components/active_record/connection_adapters/postgresql_fallback_adapter_spec.rb
+++ b/spec/components/active_record/connection_adapters/postgresql_fallback_adapter_spec.rb
@@ -62,6 +62,7 @@ describe ActiveRecord::ConnectionHandling do
 
       it 'should failover to a replica server' do
         RailsMultisite::ConnectionManagement.stubs(:all_dbs).returns(['default', multisite_db])
+        postgresql_fallback_handler.expects(:verify_master).at_least(3)
 
         [config, multisite_config].each do |configuration|
           ActiveRecord::Base.expects(:postgresql_connection).with(configuration).raises(PG::ConnectionBad)
@@ -113,13 +114,15 @@ describe ActiveRecord::ConnectionHandling do
 
     context 'when both master and replica server is down' do
       it 'should raise the right error' do
-        ActiveRecord::Base.expects(:postgresql_connection).with(config).raises(PG::ConnectionBad).once
+        ActiveRecord::Base.expects(:postgresql_connection).with(config).raises(PG::ConnectionBad)
 
         ActiveRecord::Base.expects(:postgresql_connection).with(config.dup.merge(
           host: replica_host,
           port: replica_port
         )).raises(PG::ConnectionBad).once
 
+        postgresql_fallback_handler.expects(:verify_master).twice
+
         2.times do
           expect { ActiveRecord::Base.postgresql_fallback_connection(config) }
             .to raise_error(PG::ConnectionBad)
diff --git a/spec/components/distributed_cache_spec.rb b/spec/components/distributed_cache_spec.rb
index 162e553cd3d..e1256346acb 100644
--- a/spec/components/distributed_cache_spec.rb
+++ b/spec/components/distributed_cache_spec.rb
@@ -14,7 +14,7 @@ describe DistributedCache do
   end
 
   def cache(name)
-    DistributedCache.new(name, @manager)
+    DistributedCache.new(name, manager: @manager)
   end
 
   let! :cache1 do