PostgreSQL "active?" Rails adapter check async patch

gistfile1.rb

# This monkey patches the PostgreSQL adapter to use asynchronous
# communication when doing its "SELECT 1" check for is-connection-active.
# One symptom of this problem is stalled processing queues in Sidekiq.
# See also:
#
# https://github.com/rails/rails/issues/12867
#
# At the time of writing we have seen indefinite blocking down in the C
# library's "poll()" method due to the out-of-box Rails code using blocking
# I/O, in multithreaded environments - one thread gets stuck in the I/O and
# the others end up waiting on a semaphore to get access to the adapter.
# The hope is that the non-blocking I/O calls will not stall in this manner
# and we can thus implement timeouts to detect an apparently stalled
# keep-alive check.
#
# No attempt is made to do anything to tidy up a connection which appears
# to be in this stalled state - this patch just says "no longer active". It
# is likely that if the basic async solution's premise is sound, further
# work will be necessary to deal with the apparently stuck-in-I/O connection,
# else it's likely that the database-end connection limit will be reached as
# "stuck" connections build up.
#
# Adapted from:
#
# http://stackoverflow.com/questions/5893524/using-the-postgresql-gem-async
#
puts "Patching PostgreSQLAdapter"

module ActiveRecord
  module ConnectionAdapters
    class PostgreSQLAdapter < AbstractAdapter

      # Is this connection alive and ready for queries?
      def active?
        begin
          conn = @connection
          timeout = (ENV["PG_ASYNC_ACTIVE_TIMEOUT"] || "15").to_i # seconds

          # Grab a reference to the underlying socket so we know when the
          # connection is established
          socket = conn.socket_io # AKA IO.for_fd(conn.socket) - same thing

          # Track the progress of the connection, waiting for the socket to
          # become readable/writable before polling it

          poll_status = PG::PGRES_POLLING_WRITING
          until poll_status == PG::PGRES_POLLING_OK || poll_status == PG::PGRES_POLLING_FAILED

            # If the socket needs to read, wait until it becomes readable
            # to poll again
            case poll_status
            when PG::PGRES_POLLING_READING
              IO.select( [socket], nil, nil, timeout ) or raise "Asynchronous connection timed out!"

              # ...and the same for when the socket needs to write
            when PG::PGRES_POLLING_WRITING
              IO.select( nil, [socket], nil, timeout ) or raise "Asynchronous connection timed out!"
            end

            # Check to see if it's finished or failed yet
            poll_status = conn.connect_poll
          end

          raise "Connect failed: %s" % [ conn.error_message ] unless conn.status == PG::CONNECTION_OK

          # Send check-alive query
          conn.send_query( "SELECT 1" )

          # Fetch results until there aren't any more
          loop do
            # Buffer incoming data on the socket until a full result is ready.
            conn.consume_input
            while conn.is_busy
              IO.select( [socket], nil, nil, timeout ) or raise "Timeout waiting for response (indefinitely blocking 'SELECT 1' query potentially averted)."
              conn.consume_input
            end

            # Fetch the next result. If there isn't one, the query is finished
            result = conn.get_result or break
          end

          true

        rescue => e
          logger.error(">"*80)
          logger.error("#{Time.now} PostgreSQL gem 'active?' alert: #{e}")
          logger.error("="*80)
          logger.error(Kernel.caller.join("\n")) # Ruby >= 1.9, Kernel#caller -> backtrace
          logger.error("<"*80)
          false

        end

        # The original implementation in Rails/ActiveRecord v4.0.4 is:
        #
        # begin
        #   @connection.query 'SELECT 1'
        #   true
        # rescue PGError
        #   false
        # end
      end
    end
  end
end

This would go in config/initializers with an "early filename" (e.g. "00_<something>") so it gets executed soon in the application boot process.

Highly experimental!