Use AR Adapter callbacks to force reconnects

To prevent connection imbalances on PgBouncer, periodically reconnect
to the database when checking out a connection from the Rails connection
pool.

config/initializers/active_record_force_reconnects.rb

+# frozen_string_literal: true
+
+Gitlab::Database::ConnectionTimer.configure do |config|
+  config.interval = Rails.application.config_for(:database)[:force_reconnect_interval]
+end
+
+ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(Gitlab::Database::PostgresqlAdapter::ForceDisconnectableMixin)

lib/gitlab/database/connection_timer.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    class ConnectionTimer
+      DEFAULT_INTERVAL = 3600
+      RANDOMIZATION_INTERVAL = 600
+
+      class << self
+        def configure
+          yield self
+        end
+
+        def starting_now
+          # add a small amount of randomization to the interval, so reconnects don't all occur at once
+          new(interval_with_randomization, current_clock_value)
+        end
+
+        attr_writer :interval
+
+        def interval
+          @interval ||= DEFAULT_INTERVAL
+        end
+
+        def interval_with_randomization
+          interval + rand(RANDOMIZATION_INTERVAL) if interval.positive?
+        end
+
+        def current_clock_value
+          Concurrent.monotonic_time
+        end
+      end
+
+      attr_reader :interval, :starting_clock_value
+
+      def initialize(interval, starting_clock_value)
+        @interval = interval
+        @starting_clock_value = starting_clock_value
+      end
+
+      def expired?
+        interval&.positive? && self.class.current_clock_value > (starting_clock_value + interval)
+      end
+
+      def reset!
+        @starting_clock_value = self.class.current_clock_value
+      end
+    end
+  end
+end

lib/gitlab/database/postgresql_adapter/force_disconnectable_mixin.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module PostgresqlAdapter
+      module ForceDisconnectableMixin
+        extend ActiveSupport::Concern
+
+        prepended do
+          set_callback :checkin, :after, :force_disconnect_if_old!
+        end
+
+        def force_disconnect_if_old!
+          if force_disconnect_timer.expired?
+            disconnect!
+            reset_force_disconnect_timer!
+          end
+        end
+
+        def reset_force_disconnect_timer!
+          force_disconnect_timer.reset!
+        end
+
+        def force_disconnect_timer
+          @force_disconnect_timer ||= ConnectionTimer.starting_now
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/database/connection_timer_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Database::ConnectionTimer do
+  let(:current_clock_value) { 1234.56 }
+
+  before do
+    allow(described_class).to receive(:current_clock_value).and_return(current_clock_value)
+  end
+
+  describe '.starting_now' do
+    let(:default_interval) { described_class::DEFAULT_INTERVAL }
+    let(:random_value) { 120 }
+
+    before do
+      allow(described_class).to receive(:rand).and_return(random_value)
+    end
+
+    context 'when the configured interval is positive' do
+      before do
+        allow(described_class).to receive(:interval).and_return(default_interval)
+      end
+
+      it 'randomizes the interval of the created timer' do
+        timer = described_class.starting_now
+
+        expect(timer.interval).to eq(default_interval + random_value)
+      end
+    end
+
+    context 'when the configured interval is not positive' do
+      before do
+        allow(described_class).to receive(:interval).and_return(0)
+      end
+
+      it 'sets the interval of the created timer to nil' do
+        timer = described_class.starting_now
+
+        expect(timer.interval).to be_nil
+      end
+    end
+  end
+
+  describe '.expired?' do
+    context 'when the interval is positive' do
+      context 'when the interval has elapsed' do
+        it 'returns true' do
+          timer = described_class.new(20, current_clock_value - 30)
+
+          expect(timer).to be_expired
+        end
+      end
+
+      context 'when the interval has not elapsed' do
+        it 'returns false' do
+          timer = described_class.new(20, current_clock_value - 10)
+
+          expect(timer).not_to be_expired
+        end
+      end
+    end
+
+    context 'when the interval is not positive' do
+      context 'when the interval has elapsed' do
+        it 'returns false' do
+          timer = described_class.new(0, current_clock_value - 30)
+
+          expect(timer).not_to be_expired
+        end
+      end
+
+      context 'when the interval has not elapsed' do
+        it 'returns false' do
+          timer = described_class.new(0, current_clock_value + 10)
+
+          expect(timer).not_to be_expired
+        end
+      end
+    end
+
+    context 'when the interval is nil' do
+      it 'returns false' do
+        timer = described_class.new(nil, current_clock_value - 30)
+
+        expect(timer).not_to be_expired
+      end
+    end
+  end
+
+  describe '.reset!' do
+    it 'updates the timer clock value' do
+      timer = described_class.new(20, current_clock_value - 20)
+      expect(timer.starting_clock_value).not_to eql(current_clock_value)
+
+      timer.reset!
+      expect(timer.starting_clock_value).to eql(current_clock_value)
+    end
+  end
+end

spec/lib/gitlab/database/postgresql_adapter/force_disconnectable_mixin_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Database::PostgresqlAdapter::ForceDisconnectableMixin do
+  describe 'checking in a connection to the pool' do
+    let(:model) do
+      Class.new(ActiveRecord::Base) do
+        self.abstract_class = true
+
+        def self.name
+          'ForceDisconnectTestModel'
+        end
+      end
+    end
+    let(:config) { Rails.application.config_for(:database).merge(pool: 1) }
+    let(:pool) { model.establish_connection(config) }
+
+    it 'calls the force disconnect callback on checkin' do
+      connection = pool.connection
+
+      expect(pool.active_connection?).to be_truthy
+      expect(connection).to receive(:force_disconnect_if_old!).and_call_original
+
+      model.clear_active_connections!
+    end
+  end
+
+  describe 'disconnecting from the database' do
+    let(:connection) { ActiveRecord::Base.connection_pool.connection }
+    let(:timer) { connection.force_disconnect_timer }
+
+    context 'when the timer is expired' do
+      it 'disconnects from the database' do
+        allow(timer).to receive(:expired?).and_return(true)
+
+        expect(connection).to receive(:disconnect!).and_call_original
+        expect(timer).to receive(:reset!).and_call_original
+
+        connection.force_disconnect_if_old!
+      end
+    end
+
+    context 'when the timer is not expired' do
+      it 'does not disconnect from the database' do
+        allow(timer).to receive(:expired?).and_return(false)
+
+        expect(connection).not_to receive(:disconnect!)
+        expect(timer).not_to receive(:reset!)
+
+        connection.force_disconnect_if_old!
+      end
+    end
+  end
+end