Merge pull request #4049 from tgxworld/support_multisite_in_postgresql_fallback

FEATURE: Support multisite in PostgreSQL fallback adapter.

lib/active_record/connection_adapters/postgresql_fallback_adapter.rb

@@ -5,65 +5,103 @@ require 'discourse'
 class PostgreSQLFallbackHandler
   include Singleton
 
-  attr_reader :running
-  attr_accessor :master
-
   def initialize
-    @master = true
-    @running = false
-    @mutex = Mutex.new
+    @master = {}
+    @running = {}
+    @mutex = {}
+    @last_check = {}
+
+    setup!
   end
 
   def verify_master
-    @mutex.synchronize do
-      return if @running || recently_checked?
-      @running = true
+    @mutex[namespace].synchronize do
+      return if running || recently_checked?
+      @running[namespace] = true
     end
 
+    current_namespace = namespace
     Thread.new do
-      begin
-        logger.warn "#{self.class}: Checking master server..."
-        connection = ActiveRecord::Base.postgresql_connection(config)
+      RailsMultisite::ConnectionManagement.with_connection(current_namespace) do
+        begin
+          logger.warn "#{log_prefix}: Checking master server..."
+          connection = ActiveRecord::Base.postgresql_connection(config)
 
-        if connection.active?
-          connection.disconnect!
-          logger.warn "#{self.class}: Master server is active. Reconnecting..."
-          ActiveRecord::Base.establish_connection(config)
-          Discourse.disable_readonly_mode
-          @master = true
-        end
-      rescue => e
-        if e.message.include?("could not connect to server")
-          logger.warn "#{self.class}: Connection to master PostgreSQL server failed with '#{e.message}'"
-        else
-          raise e
-        end
-      ensure
-        @mutex.synchronize do
-          @last_check = Time.zone.now
-          @running = false
+          if connection.active?
+            connection.disconnect!
+            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
+
+            Discourse.disable_readonly_mode
+            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
         end
       end
     end
   end
 
+  def master
+    @master[namespace]
+  end
+
+  def master=(args)
+    @master[namespace] = args
+  end
+
+  def running
+    @running[namespace]
+  end
+
+  def setup!
+    RailsMultisite::ConnectionManagement.all_dbs.each do |db|
+      @master[db] = true
+      @running[db] = false
+      @mutex[db] = Mutex.new
+    end
+  end
+
   private
 
   def config
-    ActiveRecord::Base.configurations[Rails.env]
+    ActiveRecord::Base.connection_config
   end
 
   def logger
     Rails.logger
   end
 
+  def log_prefix
+    "#{self.class} [#{namespace}]"
+  end
+
   def recently_checked?
-    if @last_check
-      Time.zone.now <= (@last_check + 5.seconds)
+    if @last_check[namespace]
+      Time.zone.now <= (@last_check[namespace] + 5.seconds)
     else
       false
     end
   end
+
+  def namespace
+    RailsMultisite::ConnectionManagement.current_db
+  end
 end
 
 module ActiveRecord

lib/discourse.rb

@@ -112,17 +112,22 @@ module Discourse
     end
   end
 
+  def self.last_read_only
+    @last_read_only ||= {}
+  end
+
   def self.recently_readonly?
-    return false unless @last_read_only
-    @last_read_only > 15.seconds.ago
+    read_only = last_read_only[$redis.namespace]
+    return false unless read_only
+    read_only > 15.seconds.ago
   end
 
   def self.received_readonly!
-    @last_read_only = Time.now
+    last_read_only[$redis.namespace] = Time.zone.now
   end
 
   def self.clear_readonly!
-    @last_read_only = nil
+    last_read_only[$redis.namespace] = nil
   end
 
   def self.disabled_plugin_names

spec/components/active_record/connection_adapters/postgresql_fallback_adapter_spec.rb

@@ -3,10 +3,10 @@ require_dependency 'active_record/connection_adapters/postgresql_fallback_adapte
 
 describe ActiveRecord::ConnectionHandling do
   let(:replica_host) { "1.1.1.1" }
-  let(:replica_port) { "6432" }
+  let(:replica_port) { 6432 }
 
   let(:config) do
-    ActiveRecord::Base.configurations["test"].merge({
+    ActiveRecord::Base.configurations[Rails.env].merge({
       "adapter" => "postgresql_fallback",
       "replica_host" => replica_host,
       "replica_port" => replica_port
@@ -14,8 +14,7 @@ describe ActiveRecord::ConnectionHandling do
   end
 
   after do
-    Discourse.disable_readonly_mode
-    ::PostgreSQLFallbackHandler.instance.master = true
+    ::PostgreSQLFallbackHandler.instance.setup!
   end
 
   describe "#postgresql_fallback_connection" do
@@ -25,17 +24,39 @@ describe ActiveRecord::ConnectionHandling do
     end
 
     context 'when master server is down' do
+      let(:multisite_db) { "database_2" }
+
+      let(:multisite_config) do
+        {
+          host: 'localhost1',
+          port: 5432,
+          replica_host: replica_host,
+          replica_port: replica_port
+        }
+      end
+
       before do
         @replica_connection = mock('replica_connection')
       end
 
-      it 'should failover to a replica server' do
-        ActiveRecord::Base.expects(:postgresql_connection).with(config).raises(PG::ConnectionBad)
-        ActiveRecord::Base.expects(:verify_replica).with(@replica_connection)
+      after do
+        with_multisite_db(multisite_db) { Discourse.disable_readonly_mode }
+        Discourse.disable_readonly_mode
+        ActiveRecord::Base.establish_connection(ActiveRecord::Base.configurations[Rails.env])
+      end
 
-        ActiveRecord::Base.expects(:postgresql_connection).with(config.merge({
-          host: replica_host, port: replica_port
-        })).returns(@replica_connection)
+      it 'should failover to a replica server' do
+        RailsMultisite::ConnectionManagement.stubs(:all_dbs).returns(['default', multisite_db])
+        ::PostgreSQLFallbackHandler.instance.setup!
+
+        [config, multisite_config].each do |configuration|
+          ActiveRecord::Base.expects(:postgresql_connection).with(configuration).raises(PG::ConnectionBad)
+          ActiveRecord::Base.expects(:verify_replica).with(@replica_connection)
+
+          ActiveRecord::Base.expects(:postgresql_connection).with(configuration.merge({
+            host: replica_host, port: replica_port
+          })).returns(@replica_connection)
+        end
 
         expect { ActiveRecord::Base.postgresql_fallback_connection(config) }
           .to raise_error(PG::ConnectionBad)
@@ -43,6 +64,14 @@ describe ActiveRecord::ConnectionHandling do
         expect{ ActiveRecord::Base.postgresql_fallback_connection(config) }
           .to change{ Discourse.readonly_mode? }.from(false).to(true)
 
+        with_multisite_db(multisite_db) do
+          expect { ActiveRecord::Base.postgresql_fallback_connection(multisite_config) }
+            .to raise_error(PG::ConnectionBad)
+
+          expect{ ActiveRecord::Base.postgresql_fallback_connection(multisite_config) }
+            .to change{ Discourse.readonly_mode? }.from(false).to(true)
+        end
+
         ActiveRecord::Base.unstub(:postgresql_connection)
 
         current_threads = Thread.list
@@ -59,7 +88,7 @@ describe ActiveRecord::ConnectionHandling do
         end
 
         # Wait for the thread to finish execution
-        threads = (Thread.list - current_threads).each(&:join)
+        (Thread.list - current_threads).each(&:join)
 
         expect(Discourse.readonly_mode?).to eq(false)
 
@@ -72,7 +101,11 @@ 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).raises(PG::ConnectionBad).twice
+        ActiveRecord::Base.expects(:postgresql_connection).with(config).raises(PG::ConnectionBad).once
+
+        ActiveRecord::Base.expects(:postgresql_connection).with(config.dup.merge({
+          host: replica_host, port: replica_port
+        })).raises(PG::ConnectionBad).once
 
         2.times do
           expect { ActiveRecord::Base.postgresql_fallback_connection(config) }
@@ -81,4 +114,10 @@ describe ActiveRecord::ConnectionHandling do
       end
     end
   end
+
+  def with_multisite_db(dbname)
+    RailsMultisite::ConnectionManagement.expects(:current_db).returns(dbname).at_least_once
+    yield
+    RailsMultisite::ConnectionManagement.unstub(:current_db)
+  end
 end

spec/components/discourse_spec.rb

@@ -111,8 +111,12 @@ describe Discourse do
     end
 
     it "returns true when the key is present in redis" do
-      $redis.expects(:get).with(Discourse.readonly_mode_key).returns("1")
-      expect(Discourse.readonly_mode?).to eq(true)
+      begin
+        $redis.set(Discourse.readonly_mode_key, 1)
+        expect(Discourse.readonly_mode?).to eq(true)
+      ensure
+        $redis.del(Discourse.readonly_mode_key)
+      end
     end
 
     it "returns true when Discourse is recently read only" do
@@ -121,6 +125,13 @@ describe Discourse do
     end
   end
 
+  context ".received_readonly!" do
+    it "sets the right time" do
+      time = Discourse.received_readonly!
+      expect(Discourse.last_read_only['default']).to eq(time)
+    end
+  end
+
   context "#handle_exception" do
 
     class TempSidekiqLogger < Sidekiq::ExceptionHandler::Logger