ActiveRecord Connection Lifetime controls for using postgres with an autoscaling pgbouncer service safely

activerecord-lifetime.rb

# frozen_string_literal: true

# Connection lifetime for ActiveRecord
#
# Make sure that connections to the database can only live for a certain number
# of seconds. Once lifetime is reached, the underlying connection will be
# reconnected. This is enforced when checking out a connection for use from the
# pool. Use in combination with idle_timeout to enforce connection lifetime on
# idle connections as well.
#
# This allows graceful failover of database connections when the connection
# targets underlying a database address might change over time. For example,
# when running a pgbouncer in front of postgres this allows pgbouncer itself to
# be scaled or updated without downtime by using a load balancer with a target
# deregistration connection draining timeout slightly longer than the
# connection lifetime. The difference should account for the maximum time a
# connection might be checked out, i.e. the maximum time of a web request.
#
# A potential future improvement would be to give opportunities for
# reconnection between transactions for transaction-pooled connections.
#
module ActiveRecord::ConnectionAdapters::Lifetime
  def initialize(...)
    super

    # Keep track of when this connection was opened.
    @connected_at = Concurrent.monotonic_time

    # And see if there is a configured time to live.
    @lifetime = self.class.type_cast_config_to_integer(
      @config.fetch(:lifetime, nil)
    )
  end

  attr_reader :connected_at, :lifetime

  def reconnect!(...)
    super

    # Make sure reconnect! resets the connection time.
    @connected_at = Concurrent.monotonic_time
  end

  def age
    Concurrent.monotonic_time - @connected_at
  end

  def old?
    lifetime && age > lifetime
  end

  def verify!(...)
    super

    # During checkout, reconnect if the connection is too old.
    reconnect! if old?
  end
end

activerecord-lifetime_spec.rb

require "rails_helper"

RSpec.describe ActiveRecord::ConnectionAdapters::Lifetime do
  # Informed by https://github.com/rails/rails/blob/main/activerecord/test/cases/connection_pool_test.rb
  let(:db_config_hash) { ActiveRecord::Base.configurations.configs_for(name: "primary").configuration_hash.merge(lifetime: lifetime) }
  let(:db_config) { ActiveRecord::Base.configurations.resolve(db_config_hash) }
  let(:pool_config) { ActiveRecord::ConnectionAdapters::PoolConfig.new(ActiveRecord::Base, db_config) }
  let(:connection_pool) { ActiveRecord::ConnectionAdapters::ConnectionPool.new(pool_config) }
  after { connection_pool.disconnect! }

  context "with a lifetime" do
    let(:lifetime) { 60 }

    it "does not reconnect young connections" do
      expect(connection_pool.connections.size).to eq(0)

      connection = connection_pool.checkout
      connected_at = connection.connected_at

      expect(connection).to be_active
      expect(connection_pool.connections.size).to eq(1)

      connection_pool.checkin(connection)

      expect(connection_pool.connections.size).to eq(1)
      expect(connection.connected_at).to eql(connected_at)

      connection = connection_pool.checkout

      expect(connection).to be_active
      expect(connection_pool.connections.size).to eq(1)
      expect(connection.connected_at).to eql(connected_at)
    end

    it "reconnects old connections during checkout" do
      expect(connection_pool.connections.size).to eq(0)

      connection = connection_pool.checkout

      expect(connection).to be_active
      expect(connection_pool.connections.size).to eq(1)

      connection.instance_variable_set(:@connected_at, Concurrent.monotonic_time - lifetime * 1.5)
      expect(connection).to be_active
      expect(connection.age).to be > lifetime
      expect(connection_pool.connections.size).to eq(1)

      connection_pool.checkin(connection)

      expect(connection_pool.connections.size).to eq(1)

      connection = connection_pool.checkout

      expect(connection).to be_active
      expect(connection_pool.connections.size).to eq(1)
      expect(connection.age).to be < lifetime
    end
  end

  context "without a lifetime" do
    let(:lifetime) { nil }

    it "does not affect connections" do
      expect(connection_pool.connections.size).to eq(0)

      connection = connection_pool.checkout
      connected_at = connection.connected_at

      expect(connection).to be_active
      expect(connection_pool.connections.size).to eq(1)

      connection_pool.checkin(connection)

      expect(connection_pool.connections.size).to eq(1)

      connection = connection_pool.checkout

      expect(connection).to be_active
      expect(connection_pool.connections.size).to eq(1)
      expect(connection.connected_at).to eql(connected_at)
    end
  end
end

config-initializers-activerecord-lifetime.rb

# frozen_string_literal: true

require "active_record/connection_adapters/lifetime"

ActiveSupport.on_load(:active_record) do
  ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend ActiveRecord::ConnectionAdapters::Lifetime

  # The Postgres adapter does a weaker "reset" when asked to reconnect by
  # default. The method docs suggest it should do what we want, but in practise
  # we're still seeing connections to database targets well after they should
  # no longer be receiving new connections through the load balancer. So let's
  # try throwing away the whole connection and recreating it.
  ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.class_eval do
    undef_method :reconnect!
    def reconnect!
      @lock.synchronize do
        super

        # @connection.close rescue nil
        disconnect!

        # @connection = PG::Connection.new ...
        connect
      end
    end
  end
end