Ping less patch for MySQL.

ping_less.rb

# frozen_string_literal: true

# By default, ActiveRecord will issue a PING command to the database to check
# if it is active at various points. This overhead is unnecessary; we instead
# attempt to issue queries without checking the connection, then if the query
# returns an error and the connection was closed, try to reconnect.
# This also allows for reconnection during a single UoW, improving resiliency
# under transient connection failure (e.g. ProxySQL restarts).
#
# To avoid amplifying load when a database is intermittently down, the attempt
# to reconnect is allowed only when verification is requested by ActiveRecord,
# and automatically at most once per VERIFICATION_INTERVAL.
module MysqlAdapterPingLess
  VERIFICATION_INTERVAL = 1.minute.to_f

  def verify!
    if connected?
      @__verify_connection = true
    else
      reconnect!
    end
  end

  def execute(sql, name = nil)
    start_time = Time.now.to_f
    super # Try executing the query.

  rescue ::ActiveRecord::StatementInvalid
    now = Time.now.to_f
    query_time = now - start_time

    raise_unless_verifying_connection(now)
    raise unless @connection.closed?
    raise if query_time > pingless_query_timeout
    raise if current_transaction.open?
    @__last_verify_time = now

    begin
      reconnect!
    rescue => e
      raise translate_exception_class(e, sql, [])
    end

    super # Retry the query.
  ensure
    @__verify_connection = false
  end

  def quote_string(string)
    super
  rescue Mysql2::Error
    now = Time.now.to_f

    raise_unless_verifying_connection(now)
    raise unless @connection.closed?
    raise if current_transaction.open?
    @__last_verify_time = now

    reconnect!
    super
  ensure
    @__verify_connection = false
  end

  private

  def connected?
    @connection && @connection.socket.present?
  rescue Mysql2::Error
    false
  end

  def raise_unless_verifying_connection(now)
    allow_automatic_verify = !@__last_verify_time || now - @__last_verify_time > VERIFICATION_INTERVAL
    raise if !@__verify_connection && !allow_automatic_verify
  end

  def pingless_query_timeout
    5
  end
end

ActiveRecord::ConnectionAdapters::Mysql2Adapter.prepend(MysqlAdapterPingLess)

ping_less_test.rb

require 'test_helper'

class ActiveRecordPingLessTest < ActiveSupport::TestCase
  setup do
    # Emulate connection pool middleware by putting all active connections back
    # in the pool.
    ActiveRecord::Base.clear_active_connections!
  end

  teardown do
    Semian[:mysql_readonly].reset
  end

  test "active? calls ping" do
    with_pings(1) do
      connection.active?
    end
  end

  test "checking out connection does not ping it" do
    with_pings(0) do
      connection
    end
  end

  test "checking out already established connection does not require a connection to the database" do
    connection

    toxiproxy_mysql_down("master_readonly") do
      connection
    end
  end

  test "reconnects and retries on first query after checkout" do
    connection.expects(:reconnect!)

    toxiproxy_mysql_down("master_readonly") do
      assert_raise ActiveRecord::StatementInvalid do
        connection.execute("SELECT 1")
      end
    end
  end

  test "reconnects and retries automatically at most once within VERIFICATION_INTERVAL" do
    connection.execute("SELECT 1") # force connection to open
    connection.expects(:reconnect!).once

    toxiproxy_mysql_down("master_readonly") do
      assert_raise ActiveRecord::StatementInvalid do
        connection.execute("SELECT 1")
      end

      Timecop.travel(MysqlAdapterPingLess::VERIFICATION_INTERVAL - 1.second) do
        assert_raise ActiveRecord::StatementInvalid do
          connection.execute("SELECT 1")
        end
      end
    end
  end

  test "reconnects and retries automatically after VERIFICATION_INTERVAL" do
    connection.execute("SELECT 1") # force connection to open
    connection.expects(:reconnect!).twice

    toxiproxy_mysql_down("master_readonly") do
      assert_raise ActiveRecord::StatementInvalid do
        connection.execute("SELECT 1")
      end

      Timecop.travel(MysqlAdapterPingLess::VERIFICATION_INTERVAL + 1.second) do
        assert_raise ActiveRecord::StatementInvalid do
          connection.execute("SELECT 1")
        end
      end
    end
  end

  test "does not reconnect if connection is still open" do
    connection.execute("SELECT 1") # force connection to open
    connection.expects(:reconnect!).never

    Timecop.travel(MysqlAdapterPingLess::VERIFICATION_INTERVAL + 1.second) do
      assert_raise ActiveRecord::StatementInvalid do
        connection.execute("THIS IS NOT A VALID QUERY")
      end
    end
  end

  test "does not reconnect in a transaction" do
    connection.execute("SELECT 1") # force connection to open
    connection.expects(:reconnect!).never

    connection.begin_transaction

    toxiproxy_mysql_down("master_readonly") do
      assert_raise ActiveRecord::StatementInvalid do
        connection.execute("SELECT 1")
      end

      assert_raise Mysql2::Error do
        connection.quote("'; DROP DATABASE WALRUS_PARTY_DB")
      end
    end

    assert !connection.current_transaction.open?
  end

  test "reconnect for string quoting" do
    conn = connection
    conn.raw_connection.close

    conn.quote("'; DROP DATABASE WALRUS_PARTY_DB")

    ActiveRecord::Base.clear_active_connections!

    connection.quote("'; DROP DATABASE WALRUS_PARTY_DB")
  end

  test "reconnect when checking out disconnected connection" do
    semain_resource = connection.semian_resource
    toxiproxy_mysql_down("master") do
      assert_raise(ActiveRecord::StatementInvalid) do
        connection.execute("SELECT 1")
      end
      semain_resource.reset
      ActiveRecord::Base.clear_active_connections!
      error = assert_raise(Mysql2::Error) do
        connection
      end
      assert_match(/Can't connect to MySQL server/, error.message)
    end
  end

  test "does not retry slow queries over timeout" do
    with_pt_kill(0.2, /^SELECT SLEEP/) do
      stub_pingless_timeout(0.1)
      assert_raise(ActiveRecord::StatementInvalid) do
        connection.execute("SELECT SLEEP(5)")
      end
    end
  end

  test "does retry slow queries under timeout" do
    with_pt_kill(0.2, /^SELECT SLEEP/) do
      connection.execute("SELECT SLEEP(1)")
    end
  end

  private

  def connection
    ActiveRecord::Base.connection(model_check: false)
  end

  def with_pings(n)
    Mysql2::Client.any_instance.expects(:ping).times(n).returns(true)
    yield
  ensure
    Mysql2::Client.any_instance.unstub(:ping)
  end

  def with_pt_kill(timeout, match)
    pool = connection.pool
    c = pool.checkout

    t = Thread.new do
      sleep timeout
      processlist = c.execute("SHOW PROCESSLIST").to_a
      process = processlist.find do |(_id, _user, _host, _db, _command, _time, _state, info)|
        info =~ match
      end
      raise "target query not found" if process.empty?
      c.execute("KILL #{process[0]}")
    end

    yield

    t.join
  ensure
    ActiveRecord::Base.clear_all_connections!
  end

  def stub_pingless_timeout(n)
    ActiveRecord::ConnectionAdapters::Mysql2Adapter.any_instance.stubs(:pingless_query_timeout).returns(n)
  end
end