pg_pipeline 0.2.1

data/lib/pg_pipeline/pool.rb

@@ -0,0 +1,420 @@
+# frozen_string_literal: true
+
+require "pg"
+require "async"
+require "async/semaphore"
+require "async/notification"
+
+require_relative "errors"
+require_relative "connection_driver"
+require_relative "server_caps"
+
+module PgPipeline
+  class Pool
+    DEFAULT_PIPELINE_SIZE = 4
+    DEFAULT_PINNED_SIZE = 2
+    DEFAULT_RECONNECT_INTERVAL = 0.5
+    DEFAULT_RECONNECT_BACKOFF_MAX = 30.0
+    DEFAULT_HEALTH_INTERVAL = 10.0
+    DEFAULT_HEALTH_TIMEOUT = 5.0
+    DISCARD_SEQUENCES_SERVER_VERSION = 90_400
+
+    attr_reader :reconnects, :pipeline_size
+
+    def initialize(connection_args,
+                   pipeline_size: DEFAULT_PIPELINE_SIZE,
+                   pinned_size: DEFAULT_PINNED_SIZE,
+                   max_pending: ConnectionDriver::DEFAULT_MAX_PENDING,
+                   max_in_flight: ConnectionDriver::DEFAULT_MAX_IN_FLIGHT,
+                   reconnect: true,
+                   reconnect_interval: DEFAULT_RECONNECT_INTERVAL,
+                   reconnect_backoff_max: DEFAULT_RECONNECT_BACKOFF_MAX,
+                   health_check: true,
+                   health_interval: DEFAULT_HEALTH_INTERVAL,
+                   health_timeout: DEFAULT_HEALTH_TIMEOUT,
+                   cancel_pinned_on_abort: true)
+      @connection_args = connection_args
+      @pipeline_size = PoolOps.positive_integer!(pipeline_size, :pipeline_size)
+      @pinned_size = PoolOps.nonnegative_integer!(pinned_size, :pinned_size)
+      @max_pending = max_pending
+      @max_in_flight = max_in_flight
+      @reconnect = reconnect
+      @reconnect_interval = Float(reconnect_interval)
+      @reconnect_backoff_max = Float(reconnect_backoff_max)
+      @health_check = health_check
+      @health_interval = Float(health_interval)
+      @health_timeout = Float(health_timeout)
+      @cancel_pinned_on_abort = cancel_pinned_on_abort
+
+      @drivers = []
+      @driver_backoff = []
+      @driver_attempts = []
+      @driver_last_health = []
+      @rr = 0
+      @reconnects = 0
+      @health_failures = 0
+      @supervisor_error = nil
+      @supervisor = nil
+
+      @pinned_free = []
+      @pinned_in_use = {}
+      @pinned_gate = nil
+      @pinned_error = nil
+      @pinned_active = 0
+      @pinned_owners = Hash.new(0)
+      @pinned_idle = Async::Notification.new
+
+      @started = false
+      @closing = false
+    end
+
+    def start(parent: Async::Task.current)
+      raise Error, "pool already started" if @started
+
+      begin
+        @pipeline_size.times { @drivers << start_pipeline_driver(parent) }
+        @driver_backoff = Array.new(@drivers.size, 0.0)
+        @driver_attempts = Array.new(@drivers.size, 0)
+        @driver_last_health = Array.new(@drivers.size, monotonic)
+        @pinned_gate = Async::Semaphore.new(@pinned_size) if @pinned_size.positive?
+      rescue Exception
+        cleanup_partial_start
+        raise
+      end
+
+      @started = true
+      @supervisor = parent.async { supervise } if @reconnect || @health_check
+      self
+    end
+
+    def closing? = @closing
+
+    def pipeline_driver
+      ensure_available!
+
+      driver, @rr = PoolOps.select_driver(@drivers, @rr)
+      raise NotDispatchedError, "no live pipeline connections; request was not dispatched" unless driver
+
+      driver
+    end
+
+    def with_pinned
+      ensure_available!
+      raise @pinned_error if @pinned_error
+      raise Error, "pinned pool is disabled (pinned_size=0)" if @pinned_size.zero?
+
+      @pinned_gate.acquire do
+        raise @pinned_error if @pinned_error
+        raise ShutdownError, "pool is closing" if @closing
+
+        owner = Async::Task.current
+        conn = nil
+        @pinned_active += 1
+        @pinned_owners[owner] += 1
+
+        begin
+          conn = @pinned_free.pop || PoolOps.new_connection(@connection_args)
+          raise @pinned_error if @pinned_error
+          raise ShutdownError, "pool is closing" if @closing
+
+          @pinned_in_use[conn] = owner
+          yield conn
+        ensure
+          @pinned_in_use.delete(conn) if conn
+          release_pinned(owner, conn)
+        end
+      end
+    end
+
+    def stats
+      {
+        pipeline: {
+          size: @pipeline_size,
+          live: @drivers.count(&:available?),
+          drivers: @drivers.map(&:stats)
+        },
+        pinned: {
+          size: @pinned_size,
+          active: @pinned_active,
+          free: @pinned_free.size,
+          in_use: @pinned_in_use.size
+        },
+        reconnects: @reconnects,
+        health_failures: @health_failures,
+        supervisor_error: @supervisor_error&.message,
+        closing: @closing,
+        pinned_error: @pinned_error&.message
+      }
+    end
+
+    def graceful_close
+      return unless @started
+
+      if @pinned_owners[Async::Task.current].positive?
+        raise Error, "cannot close the pool from inside Client#session/transaction"
+      end
+
+      @closing = true
+      @started = false
+      stop_supervisor
+      @drivers.each(&:graceful_close)
+      wait_for_pinned_idle
+      close_free_pinned
+      nil
+    end
+
+    def abort!
+      return unless @started || @closing
+
+      @closing = true
+      @started = false
+      stop_supervisor
+      cancel_in_use_pinned if @cancel_pinned_on_abort
+      @drivers.each(&:abort!)
+      close_free_pinned
+      nil
+    end
+
+    private
+
+    def supervise
+      parent = Async::Task.current
+      until @closing
+        begin
+          reap_and_replace(parent) if @reconnect
+          health_probe if @health_check
+          @supervisor_error = nil
+        rescue StandardError => e
+          # Keep the loop alive: one probe/reconnect failure must not permanently
+          # disable health checks and replacement for the worker lifetime.
+          @supervisor_error = e
+        end
+        parent.sleep(@reconnect_interval)
+      end
+    end
+
+    def reap_and_replace(parent)
+      now = monotonic
+      ensure_driver_slots!
+
+      @drivers.each_index do |index|
+        driver = @drivers[index]
+        next if driver.available?
+        next unless driver.dead?
+        next if now < (@driver_backoff[index] || 0.0)
+
+        begin
+          @drivers[index] = start_pipeline_driver(parent)
+          @driver_backoff[index] = 0.0
+          @driver_attempts[index] = 0
+          @driver_last_health[index] = monotonic
+          @reconnects += 1
+        rescue StandardError
+          @driver_attempts[index] = (@driver_attempts[index] || 0) + 1
+          @driver_backoff[index] = now + next_backoff(@driver_attempts[index])
+        end
+      end
+    end
+
+    def health_probe
+      now = monotonic
+      ensure_driver_slots!
+
+      @drivers.each_index do |index|
+        driver = @drivers[index]
+        next unless driver.available?
+        next unless driver.load.zero?
+        next if now - (@driver_last_health[index] || 0.0) < @health_interval
+
+        @driver_last_health[index] = now
+        next if driver.health_check(@health_timeout)
+
+        @health_failures += 1
+        driver.abort!
+      end
+    end
+
+    def ensure_driver_slots!
+      size = @drivers.size
+      @driver_backoff = Array.new(size, 0.0) if @driver_backoff.nil? || @driver_backoff.size != size
+      @driver_attempts = Array.new(size, 0) if @driver_attempts.nil? || @driver_attempts.size != size
+      @driver_last_health = Array.new(size, 0.0) if @driver_last_health.nil? || @driver_last_health.size != size
+    end
+
+    def next_backoff(attempts)
+      delay = @reconnect_interval * (2**(attempts - 1))
+      delay < @reconnect_backoff_max ? delay : @reconnect_backoff_max
+    end
+
+    def cancel_in_use_pinned
+      @pinned_in_use.keys.each do |conn|
+        conn.cancel if conn.respond_to?(:cancel)
+      rescue StandardError
+        nil
+      end
+    end
+
+    def stop_supervisor
+      @supervisor&.stop
+      @supervisor = nil
+    rescue StandardError
+      nil
+    end
+
+    def ensure_available!
+      raise Error, "pool not started" unless @started
+      raise ShutdownError, "pool is closing" if @closing
+    end
+
+    def start_pipeline_driver(parent)
+      conn = PoolOps.new_connection(@connection_args)
+      driver = ConnectionDriver.new(conn, max_pending: @max_pending, max_in_flight: @max_in_flight)
+      driver.start(parent: parent)
+    rescue Exception
+      PoolOps.safe_close(conn) if conn
+      raise
+    end
+
+    def release_pinned(owner, conn)
+      begin
+        if conn
+          if @closing
+            PoolOps.safe_close(conn)
+          else
+            recycled = recycle_pinned_connection(conn)
+            @pinned_free << recycled if recycled
+          end
+        end
+      ensure
+        @pinned_active -= 1
+        @pinned_owners[owner] -= 1
+        @pinned_owners.delete(owner) if @pinned_owners[owner].zero?
+        @pinned_idle.signal if @pinned_active.zero?
+      end
+    end
+
+    def recycle_pinned_connection(conn)
+      if conn.server_version < DISCARD_SEQUENCES_SERVER_VERSION
+        PoolOps.safe_close(conn)
+        return PoolOps.new_connection(@connection_args)
+      end
+
+      PoolOps.sanitize_pinned_connection(conn)
+      conn
+    rescue StandardError => cleanup_error
+      PoolOps.safe_close(conn)
+
+      begin
+        PoolOps.new_connection(@connection_args)
+      rescue StandardError => replacement_error
+        @pinned_error = ConnectionLostError.new(
+          "pinned connection reset failed (#{cleanup_error.class}: #{cleanup_error.message}) " \
+          "and replacement failed (#{replacement_error.class}: #{replacement_error.message})"
+        )
+        nil
+      end
+    end
+
+    def wait_for_pinned_idle
+      @pinned_idle.wait while @pinned_active.positive?
+    end
+
+    def cleanup_partial_start
+      @drivers.each do |driver|
+        driver.abort!
+      rescue StandardError
+        nil
+      end
+      @drivers.clear
+      close_free_pinned
+      @started = false
+      @closing = false
+    end
+
+    def close_free_pinned
+      @pinned_free.each { |conn| PoolOps.safe_close(conn) }
+      @pinned_free.clear
+    end
+
+    def monotonic = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+  end
+
+  module PoolOps
+    module_function
+
+    def positive_integer!(value, name)
+      integer = Integer(value)
+      raise ArgumentError, "#{name} must be >= 1" if integer < 1
+
+      integer
+    rescue ArgumentError, TypeError
+      raise ArgumentError, "#{name} must be an integer >= 1"
+    end
+
+    def nonnegative_integer!(value, name)
+      integer = Integer(value)
+      raise ArgumentError, "#{name} must be >= 0" if integer.negative?
+
+      integer
+    rescue ArgumentError, TypeError
+      raise ArgumentError, "#{name} must be an integer >= 0"
+    end
+
+    def select_driver(drivers, rr)
+      min_load = nil
+      count = 0
+      drivers.each do |driver|
+        next unless driver.available?
+
+        load = driver.load
+        if min_load.nil? || load < min_load
+          min_load = load
+          count = 1
+        elsif load == min_load
+          count += 1
+        end
+      end
+      return [nil, rr] if count.zero?
+
+      target = rr % count
+      index = 0
+      drivers.each do |driver|
+        next unless driver.available?
+        next unless driver.load == min_load
+
+        return [driver, rr + 1] if index == target
+
+        index += 1
+      end
+      [nil, rr]
+    end
+
+    def new_connection(connection_args)
+      connection_args.nil? ? PG::Connection.new : PG::Connection.new(connection_args)
+    end
+
+    def sanitize_pinned_connection(conn)
+      raise ConnectionLostError, "pinned connection is closed" if conn.finished?
+      raise ConnectionLostError, "pinned connection is bad" unless conn.status == PG::CONNECTION_OK
+
+      case conn.transaction_status
+      when PG::PQTRANS_IDLE
+        nil
+      when PG::PQTRANS_INTRANS, PG::PQTRANS_INERROR
+        conn.exec("ROLLBACK")
+      else
+        raise ConnectionLostError,
+              "pinned connection returned in unsafe transaction state #{conn.transaction_status}"
+      end
+
+      conn.exec("DISCARD ALL")
+      true
+    end
+
+    def safe_close(conn)
+      conn.close unless conn.finished?
+    rescue StandardError
+      nil
+    end
+  end
+end

data/lib/pg_pipeline/request.rb

@@ -0,0 +1,146 @@
+# frozen_string_literal: true
+
+require "async/notification"
+
+require_relative "errors"
+
+module PgPipeline
+  class Request
+    attr_reader :sql, :params, :condition
+    attr_accessor :state, :cancelled, :settled, :result_seen, :query_boundary_seen,
+                  :result, :error
+
+    def initialize(sql:, params: nil)
+      @sql = sql.to_s.dup.freeze
+      @params = RequestOps.snapshot_params(params)
+      @state = :new
+      @condition = Async::Notification.new
+      @cancelled = false
+      @settled = false
+      @result_seen = false
+      @query_boundary_seen = false
+      @result = nil
+      @error = nil
+    end
+
+    def queued! = RequestOps.transition!(self, :new, :queued)
+    def dispatched! = RequestOps.transition!(self, :queued, :dispatched)
+    def accept_result(result) = RequestOps.accept_result(self, result)
+    def record_error!(error, result: nil) = RequestOps.record_error!(self, error, result)
+    def query_boundary! = RequestOps.query_boundary!(self)
+    def finish! = RequestOps.finish!(self)
+    def reject!(error) = RequestOps.reject!(self, error)
+    def cancel! = RequestOps.cancel!(self)
+    def wait = RequestOps.wait(self)
+
+    def query_boundary_seen? = @query_boundary_seen
+    def cancelled? = @cancelled
+    def settled? = @settled
+  end
+
+  module RequestOps
+    module_function
+
+    def snapshot_params(params)
+      values = params.nil? ? [] : params
+      raise ArgumentError, "params must be an Array" unless values.is_a?(Array)
+
+      values.map { |value| snapshot_value(value) }.freeze
+    end
+
+    def snapshot_value(value)
+      case value
+      when String
+        value.dup.freeze
+      when Hash
+        value.each_with_object({}) do |(key, item), copy|
+          copy[key] = item.is_a?(String) ? item.dup.freeze : item
+        end.freeze
+      else
+        value
+      end
+    end
+
+    def transition!(req, from, to)
+      unless req.state == from
+        raise ProtocolError, "invalid request transition #{req.state.inspect} -> #{to.inspect}"
+      end
+
+      req.state = to
+    end
+
+    def accept_result(req, result)
+      assert_result_slot!(req)
+      req.result_seen = true
+
+      if req.cancelled
+        clear_result(result)
+      else
+        req.result = result
+      end
+    end
+
+    def record_error!(req, error, result)
+      assert_result_slot!(req)
+      req.result_seen = true
+
+      if req.cancelled
+        clear_result(result)
+      else
+        req.error ||= error
+      end
+    end
+
+    def query_boundary!(req)
+      raise ProtocolError, "query boundary before query result" unless req.result_seen
+      raise ProtocolError, "duplicate query boundary" if req.query_boundary_seen
+
+      req.query_boundary_seen = true
+    end
+
+    def finish!(req)
+      raise ProtocolError, "sync arrived before query boundary" unless req.query_boundary_seen
+      return if req.settled
+
+      req.settled = true
+      req.state = :done
+      req.condition.signal unless req.cancelled
+    end
+
+    def reject!(req, error)
+      return if req.settled
+
+      clear_result(req.result)
+      req.result = nil
+      req.error ||= error
+      req.settled = true
+      req.state = :done
+      req.condition.signal unless req.cancelled
+    end
+
+    def cancel!(req)
+      return if req.cancelled
+
+      req.cancelled = true
+      clear_result(req.result)
+      req.result = nil
+      req.error.clear_result! if req.error.respond_to?(:clear_result!)
+    end
+
+    def wait(req)
+      req.condition.wait unless req.settled
+      raise req.error if req.error
+
+      req.result
+    end
+
+    def assert_result_slot!(req)
+      raise ProtocolError, "result arrived after query boundary" if req.query_boundary_seen
+      raise ProtocolError, "multiple results for one pipeline unit" if req.result_seen
+    end
+
+    def clear_result(result)
+      result.clear if result.respond_to?(:clear)
+    end
+  end
+end

data/lib/pg_pipeline/server_caps.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+
+require "pg"
+
+require_relative "errors"
+
+module PgPipeline
+  class ServerCaps
+    REQUIRED_LIBPQ_VERSION = 140_000
+    FAST_SYNC_LIBPQ_VERSION = 170_000
+    PROTOCOL_VERSION = 3
+
+    attr_reader :libpq_version, :protocol_version, :pipeline_api,
+                :fast_sync_api, :raw_pipeline_sync_api
+
+    def initialize(libpq_version:, protocol_version:, pipeline_api: true,
+                   fast_sync_api: false, raw_pipeline_sync_api: false)
+      @libpq_version = Integer(libpq_version)
+      @protocol_version = Integer(protocol_version)
+      @pipeline_api = pipeline_api
+      @fast_sync_api = fast_sync_api
+      @raw_pipeline_sync_api = raw_pipeline_sync_api
+    end
+
+    def self.from_connection(conn) = Caps.from_connection(conn)
+    def supported? = Caps.supported?(self)
+    def assert_supported! = Caps.assert_supported!(self)
+    def fast_sync? = Caps.fast_sync?(self)
+    def place_sync(conn) = Caps.place_sync(self, conn)
+  end
+
+  module Caps
+    module_function
+
+    def from_connection(conn)
+      ServerCaps.new(
+        libpq_version: PG.library_version,
+        protocol_version: conn.protocol_version,
+        pipeline_api: pipeline_api?(conn),
+        fast_sync_api: conn.respond_to?(:send_pipeline_sync),
+        raw_pipeline_sync_api: conn.respond_to?(:sync_pipeline_sync)
+      )
+    end
+
+    def pipeline_api?(conn)
+      conn.respond_to?(:enter_pipeline_mode) &&
+        conn.respond_to?(:exit_pipeline_mode) &&
+        conn.respond_to?(:pipeline_status) &&
+        (conn.respond_to?(:sync_pipeline_sync) || conn.respond_to?(:pipeline_sync))
+    end
+
+    def supported?(caps)
+      caps.libpq_version >= ServerCaps::REQUIRED_LIBPQ_VERSION &&
+        caps.protocol_version == ServerCaps::PROTOCOL_VERSION &&
+        caps.pipeline_api
+    end
+
+    def assert_supported!(caps)
+      return true if supported?(caps)
+
+      raise UnsupportedServerError,
+            "pg_pipeline requires libpq >= 14, PostgreSQL protocol v3, and " \
+            "ruby-pg pipeline bindings; libpq=#{caps.libpq_version}, " \
+            "protocol=#{caps.protocol_version}, pipeline_api=#{caps.pipeline_api}"
+    end
+
+    def fast_sync?(caps)
+      caps.libpq_version >= ServerCaps::FAST_SYNC_LIBPQ_VERSION && caps.fast_sync_api
+    end
+
+    def place_sync(caps, conn)
+      if fast_sync?(caps)
+        conn.send_pipeline_sync
+      elsif caps.raw_pipeline_sync_api
+        conn.sync_pipeline_sync
+      else
+        conn.pipeline_sync
+      end
+
+      nil
+    end
+  end
+end

data/lib/pg_pipeline/session.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+require_relative "errors"
+
+module PgPipeline
+  class Session
+    attr_reader :owner_fiber
+
+    def initialize(conn)
+      @conn = conn
+      @active = true
+      @owner_fiber = Fiber.current
+    end
+
+    def query(sql, params = []) = SessionOps.query(self, sql, params)
+    def exec(sql) = SessionOps.exec(self, sql)
+    def prepare(name, sql, param_types = nil) = SessionOps.prepare(self, name, sql, param_types)
+    def exec_prepared(name, params = []) = SessionOps.exec_prepared(self, name, params)
+    def active? = @active
+
+    private
+
+    attr_accessor :conn, :active
+  end
+
+  module SessionOps
+    module_function
+
+    def ensure_active!(session)
+      conn = connection(session)
+      raise Error, "session handle is no longer active" unless session.active? && conn
+      return if Fiber.current.equal?(session.owner_fiber)
+
+      raise Error, "session handle is fiber-local and cannot be used from another fiber"
+    end
+
+    def query(session, sql, params)
+      ensure_active!(session)
+      connection(session).exec_params(sql, params)
+    end
+
+    def exec(session, sql)
+      ensure_active!(session)
+      connection(session).exec(sql)
+    end
+
+    def prepare(session, name, sql, param_types)
+      ensure_active!(session)
+      if param_types
+        connection(session).prepare(name, sql, param_types)
+      else
+        connection(session).prepare(name, sql)
+      end
+    end
+
+    def exec_prepared(session, name, params)
+      ensure_active!(session)
+      connection(session).exec_prepared(name, params)
+    end
+
+    def close!(session)
+      session.__send__(:active=, false)
+      session.__send__(:conn=, nil)
+      nil
+    end
+
+    def connection(session)
+      session.__send__(:conn)
+    end
+  end
+end