Merge pull request #3973 from tgxworld/postgres_failover

FEATURE: AR adapter to failover to a replica DB server.

app/models/global_setting.rb

@@ -18,11 +18,14 @@ class GlobalSetting
 
   def self.database_config
     hash = {"adapter" => "postgresql"}
-    %w{pool timeout socket host port username password}.each do |s|
+    %w{pool timeout socket host port username password replica_host replica_port}.each do |s|
       if val = self.send("db_#{s}")
         hash[s] = val
       end
     end
+
+    hash["adapter"] = "postgresql_fallback" if hash["replica_host"]
+
     hostnames = [ hostname ]
     hostnames << backup_hostname if backup_hostname.present?
 

config/discourse_defaults.conf

@@ -43,6 +43,12 @@ db_password =
 # see: https://github.com/rails/rails/issues/21992
 db_prepared_statements = false
 
+# host address for db replica server
+db_replica_host =
+
+# port running replica db server, defaults to 5432 if not set
+db_replica_port =
+
 # hostname running the forum
 hostname = "www.example.com"
 

lib/active_record/connection_adapters/postgresql_fallback_adapter.rb

@@ -0,0 +1,117 @@
+require 'active_record/connection_adapters/abstract_adapter'
+require 'active_record/connection_adapters/postgresql_adapter'
+require 'discourse'
+
+class PostgreSQLFallbackHandler
+  include Singleton
+
+  attr_reader :running
+  attr_accessor :master
+
+  def initialize
+    @master = true
+    @running = false
+    @mutex = Mutex.new
+  end
+
+  def verify_master
+    @mutex.synchronize do
+      return if @running || recently_checked?
+      @running = true
+    end
+
+    Thread.new do
+      begin
+        logger.info "#{self.class}: Checking master server..."
+        connection = ActiveRecord::Base.postgresql_connection(config)
+
+        if connection.active?
+          connection.disconnect!
+          logger.info "#{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
+        @last_check = Time.zone.now
+        @running = false
+      end
+    end
+  end
+
+  private
+
+  def config
+    ActiveRecord::Base.configurations[Rails.env]
+  end
+
+  def logger
+    Rails.logger
+  end
+
+  def recently_checked?
+    if @last_check
+      Time.zone.now <= (@last_check + 5.seconds)
+    else
+      false
+    end
+  end
+end
+
+module ActiveRecord
+  module ConnectionHandling
+    def postgresql_fallback_connection(config)
+      fallback_handler = ::PostgreSQLFallbackHandler.instance
+      config = config.symbolize_keys
+
+      if !fallback_handler.master && !fallback_handler.running
+        connection = postgresql_connection(config.dup.merge({
+          host: config[:replica_host], port: config[:replica_port]
+        }))
+
+        verify_replica(connection)
+        Discourse.enable_readonly_mode
+      else
+        begin
+          connection = postgresql_connection(config)
+        rescue PG::ConnectionBad => e
+          fallback_handler.master = false
+          raise e
+        end
+      end
+
+      connection
+    end
+
+    private
+
+    def verify_replica(connection)
+      value = connection.raw_connection.exec("SELECT pg_is_in_recovery()").values[0][0]
+      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?
+        if !fallback_handler.master && !fallback_handler.running
+          fallback_handler.verify_master
+        end
+      end
+    end
+  end
+end

spec/components/active_record/connection_adapters/postgresql_fallback_adapter_spec.rb

@@ -0,0 +1,84 @@
+require 'rails_helper'
+require_dependency 'active_record/connection_adapters/postgresql_fallback_adapter'
+
+describe ActiveRecord::ConnectionHandling do
+  let(:replica_host) { "1.1.1.1" }
+  let(:replica_port) { "6432" }
+
+  let(:config) do
+    ActiveRecord::Base.configurations["test"].merge({
+      "adapter" => "postgresql_fallback",
+      "replica_host" => replica_host,
+      "replica_port" => replica_port
+    }).symbolize_keys!
+  end
+
+  after do
+    Discourse.disable_readonly_mode
+    ::PostgreSQLFallbackHandler.instance.master = true
+  end
+
+  describe "#postgresql_fallback_connection" do
+    it 'should return a PostgreSQL adapter' do
+      expect(ActiveRecord::Base.postgresql_fallback_connection(config))
+        .to be_an_instance_of(ActiveRecord::ConnectionAdapters::PostgreSQLAdapter)
+    end
+
+    context 'when master server is down' do
+      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)
+
+        ActiveRecord::Base.expects(:postgresql_connection).with(config.merge({
+          host: replica_host, port: replica_port
+        })).returns(@replica_connection)
+
+        expect { ActiveRecord::Base.postgresql_fallback_connection(config) }
+          .to raise_error(PG::ConnectionBad)
+
+        expect{ ActiveRecord::Base.postgresql_fallback_connection(config) }
+          .to change{ Discourse.readonly_mode? }.from(false).to(true)
+
+        ActiveRecord::Base.unstub(:postgresql_connection)
+
+        current_threads = Thread.list
+
+        expect{ ActiveRecord::Base.connection_pool.checkout }
+          .to change{ Thread.list.size }.by(1)
+
+        # Ensure that we don't try to connect back to the replica when a thread
+        # is running
+        begin
+          ActiveRecord::Base.postgresql_fallback_connection(config)
+        rescue PG::ConnectionBad => e
+          # This is expected if the thread finishes before the above is called.
+        end
+
+        # Wait for the thread to finish execution
+        threads = (Thread.list - current_threads).each(&:join)
+
+        expect(Discourse.readonly_mode?).to eq(false)
+
+        expect(ActiveRecord::Base.connection_pool.connections.count).to eq(0)
+
+        expect(ActiveRecord::Base.connection)
+          .to be_an_instance_of(ActiveRecord::ConnectionAdapters::PostgreSQLAdapter)
+      end
+    end
+
+    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
+
+        2.times do
+          expect { ActiveRecord::Base.postgresql_fallback_connection(config) }
+            .to raise_error(PG::ConnectionBad)
+        end
+      end
+    end
+  end
+end