pg_pipeline 0.2.4 → 0.3.0

data/lib/pg_pipeline/pool.rb

@@ -1,11 +1,9 @@
 # frozen_string_literal: true
 
 require "pg"
-require "async"
-require "async/semaphore"
-require "async/notification"
 
 require_relative "errors"
+require_relative "runtime"
 require_relative "connection_driver"
 require_relative "prepared_statement"
 require_relative "server_caps"
@@ -19,7 +17,9 @@ module PgPipeline
     DEFAULT_HEALTH_INTERVAL = 10.0
     DEFAULT_HEALTH_TIMEOUT = 5.0
     DISCARD_SEQUENCES_SERVER_VERSION = 90_400
-    CANCEL_SIGNAL = Async::Cancel
+    CANCEL_SIGNAL = Runtime::Cancel
+    SUPERVISOR_JOIN_TIMEOUT = 5.0
+    MAX_PAUSE_FAILURES = 3
 
     attr_reader :reconnects, :pipeline_size
 
@@ -53,11 +53,13 @@ module PgPipeline
       @driver_attempts = []
       @driver_last_health = []
       @rr = 0
+      @rr_slot = [0]
       @reconnects = 0
       @health_failures = 0
       @supervisor_error = nil
       @supervisor = nil
-      @task_parent = nil
+      @supervisor_wake = Runtime::Notification.new
+      @pause_failures = 0
 
       @prepared_statements = {}
       @prepared_generation = 0
@@ -69,30 +71,28 @@ module PgPipeline
       @pinned_error = nil
       @pinned_active = 0
       @pinned_owners = Hash.new(0)
-      @pinned_idle = Async::Notification.new
+      @pinned_idle = Runtime::Notification.new
 
       @started = false
       @closing = false
       @closed = false
     end
 
-    def start(parent: nil)
+    def start
       raise Error, "pool already started" if @started
       if @closing || @closed
         raise ShutdownError, "pool is closing or was closed and cannot be restarted; create a new Pool"
       end
-
-      parent ||= Async::Task.current
-      @task_parent = parent
+      raise Error, "pool start requires an active Fiber scheduler" unless Fiber.scheduler
 
       begin
-        @pipeline_size.times { @drivers << start_pipeline_driver(parent) }
+        @pipeline_size.times { @drivers << start_pipeline_driver }
         @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?
+        @pinned_gate = Runtime::Semaphore.new(@pinned_size) if @pinned_size.positive?
         @started = true
-        @supervisor = parent.async { supervise } if @reconnect || @health_check
+        @supervisor = Runtime.spawn(name: :supervisor) { supervise } if @reconnect || @health_check
       rescue Exception
         cleanup_partial_start
         raise
@@ -106,7 +106,8 @@ module PgPipeline
     def pipeline_driver
       ensure_available!
 
-      driver, @rr = PoolOps.select_driver(@drivers, @rr)
+      driver = PoolOps.select_driver_into(@drivers, @rr, @rr_slot)
+      @rr = @rr_slot[0]
       raise NotDispatchedError, "no live pipeline connections; request was not dispatched" unless driver
 
       driver
@@ -172,9 +173,10 @@ module PgPipeline
       raise @pinned_error if @pinned_error
       raise Error, "pinned pool is disabled (pinned_size=0)" if @pinned_size.zero?
 
-      if @pinned_owners[Async::Task.current].positive?
+      owner = Fiber.current
+      if @pinned_owners[owner].positive?
         raise RecursiveCheckoutError,
-              "nested Client#session/transaction on the same task is not allowed; " \
+              "nested Client#session/transaction on the same fiber is not allowed; " \
               "use Transaction#savepoint for nested atomicity"
       end
 
@@ -182,7 +184,6 @@ module PgPipeline
         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
@@ -218,6 +219,7 @@ module PgPipeline
         reconnects: @reconnects,
         health_failures: @health_failures,
         supervisor_error: @supervisor_error&.message,
+        supervisor_alive: supervisor_alive?,
         closing: @closing,
         closed: @closed,
         pinned_error: @pinned_error&.message
@@ -227,7 +229,7 @@ module PgPipeline
     def graceful_close
       return unless @started
 
-      if @pinned_owners[Async::Task.current].positive?
+      if @pinned_owners[Fiber.current].positive?
         raise Error, "cannot close the pool from inside Client#session/transaction"
       end
 
@@ -282,7 +284,7 @@ module PgPipeline
     def supervise
       until @closing
         begin
-          reap_and_replace(@task_parent) if @reconnect
+          reap_and_replace if @reconnect
           health_probe if @health_check
           @supervisor_error = nil
         rescue StandardError => e
@@ -290,9 +292,29 @@ module PgPipeline
         end
 
         break if @closing
+        break unless supervisor_pause
+      end
+    end
+
+    def supervisor_pause
+      @supervisor_wake.wait(supervisor_sleep_interval)
+      @pause_failures = 0
+      true
+    rescue StandardError => e
+      @supervisor_error = e
+      @pause_failures = (@pause_failures || 0) + 1
 
-        sleep(supervisor_sleep_interval)
+      unless ENV["PG_PIPELINE_SILENCE_WARNINGS"]
+        warn("pg_pipeline: supervisor pause failed (#{e.class}: #{e.message}) " \
+             "[#{@pause_failures}/#{MAX_PAUSE_FAILURES}]")
       end
+
+      @pause_failures < MAX_PAUSE_FAILURES
+    end
+
+    def supervisor_alive?
+      supervisor = @supervisor
+      !supervisor.nil? && !supervisor.finished?
     end
 
     def supervisor_sleep_interval
@@ -303,7 +325,7 @@ module PgPipeline
       interval.positive? ? interval : @reconnect_interval
     end
 
-    def reap_and_replace(parent)
+    def reap_and_replace
       now = monotonic
       ensure_driver_slots!
 
@@ -314,7 +336,7 @@ module PgPipeline
         next if now < (@driver_backoff[index] || 0.0)
 
         begin
-          @drivers[index] = start_pipeline_driver(parent)
+          @drivers[index] = start_pipeline_driver
           @driver_backoff[index] = 0.0
           @driver_attempts[index] = 0
           @driver_last_health[index] = monotonic
@@ -337,10 +359,16 @@ module PgPipeline
         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!
+        begin
+          next if driver.health_check(@health_timeout)
+
+          @health_failures += 1
+          driver.abort!
+        rescue StandardError => e
+          @health_failures += 1
+          @supervisor_error = e
+        end
       end
     end
 
@@ -374,10 +402,21 @@ module PgPipeline
     def stop_supervisor
       supervisor = @supervisor
       @supervisor = nil
-      supervisor&.stop
+      return nil unless supervisor
+
+      @supervisor_wake.signal
+
+      begin
+        supervisor.wait(SUPERVISOR_JOIN_TIMEOUT)
+      rescue Runtime::TimeoutError
+        unless ENV["PG_PIPELINE_SILENCE_WARNINGS"]
+          warn("pg_pipeline: supervisor did not exit within #{SUPERVISOR_JOIN_TIMEOUT}s and has been leaked")
+        end
+      rescue CANCEL_SIGNAL, StandardError
+        nil
+      end
+
       nil
-    rescue CANCEL_SIGNAL => e
-      e
     rescue StandardError
       nil
     end
@@ -411,11 +450,11 @@ module PgPipeline
       raise Error, "pool not started" unless @started
     end
 
-    def start_pipeline_driver(parent)
+    def start_pipeline_driver
       conn = PoolOps.new_connection(@connection_args)
       prepare_registered_statements(conn)
       driver = ConnectionDriver.new(conn, max_pending: @max_pending, max_in_flight: @max_in_flight)
-      driver.start(parent: parent)
+      driver.start
     rescue Exception
       PoolOps.safe_close(conn) if conn
       raise
@@ -445,28 +484,26 @@ module PgPipeline
     end
 
     def release_pinned(owner, conn)
-      begin
-        if conn
-          if @closing
-            PoolOps.safe_close(conn)
-          else
-            recycled = recycle_pinned_connection(conn)
-
-            if recycled
-              if @closing
-                PoolOps.safe_close(recycled)
-              else
-                @pinned_free << recycled
-              end
+      if conn
+        if @closing
+          PoolOps.safe_close(conn)
+        else
+          recycled = recycle_pinned_connection(conn)
+
+          if recycled
+            if @closing
+              PoolOps.safe_close(recycled)
+            else
+              @pinned_free << recycled
             end
           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
+    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
 
     def recycle_pinned_connection(conn)
@@ -511,7 +548,6 @@ module PgPipeline
       close_free_pinned
       @pinned_gate = nil
       @supervisor = nil
-      @task_parent = nil
       @started = false
       @closing = false
       @closed = false
@@ -557,28 +593,42 @@ module PgPipeline
       raise ArgumentError, "#{name} must be an integer >= 0"
     end
 
-    def select_driver(drivers, rr)
+    def select_driver_into(drivers, rr, slot)
       size = drivers.length
-      return [nil, rr] if size.zero?
+      if size.zero?
+        slot[0] = rr
+        return nil
+      end
 
+      start = rr % size
       best = nil
-      best_index = nil
-      best_load = nil
+      best_index = 0
+      best_load = 0
+      offset = 0
 
-      size.times do |offset|
-        index = (rr + offset) % size
+      while offset < size
+        index = start + offset
+        index -= size if index >= size
         driver = drivers[index]
+        offset += 1
         next unless driver.available?
 
         load = driver.load
-        if best.nil? || load < best_load
-          best = driver
-          best_index = index
-          best_load = load
-        end
+        next unless best.nil? || load < best_load
+
+        best = driver
+        best_index = index
+        best_load = load
       end
 
-      best ? [best, (best_index + 1) % size] : [nil, rr]
+      slot[0] = best ? (best_index + 1) % size : rr
+      best
+    end
+
+    def select_driver(drivers, rr)
+      slot = [rr]
+      driver = select_driver_into(drivers, rr, slot)
+      [driver, slot[0]]
     end
 
     def new_connection(connection_args)

data/lib/pg_pipeline/prepared_statement.rb

@@ -16,7 +16,7 @@ module PgPipeline
       freeze
     end
 
-    def query(params = []) = PreparedStatementOps.query(self, params)
+    def query(params = RequestOps::EMPTY_PARAMS) = PreparedStatementOps.query(self, params)
     alias call query
 
     def inspect

data/lib/pg_pipeline/request.rb

@@ -22,10 +22,16 @@ module PgPipeline
       @waiter_scheduler = nil
     end
 
+    def self.build(sql, params)
+      request = allocate
+      request.__send__(:init_query, sql, params)
+      request
+    end
+
     def self.prepare(statement) = PrepareRequest.new(statement)
 
     def self.prepared_query(statement, params: nil)
-      PreparedQueryRequest.new(statement, params: params)
+      PreparedQueryRequest.build(statement, params)
     end
 
     def operation = :query
@@ -43,6 +49,23 @@ module PgPipeline
     def query_boundary_seen? = @query_boundary_seen
     def cancelled? = @cancelled
     def settled? = @settled
+
+    private
+
+    def init_query(sql, params)
+      @sql = RequestOps.snapshot_sql(sql)
+      @params = RequestOps.snapshot_params(params)
+      @state = :new
+      @cancelled = false
+      @settled = false
+      @result_seen = false
+      @query_boundary_seen = false
+      @result = nil
+      @error = nil
+      @waiter = nil
+      @waiter_scheduler = nil
+      self
+    end
   end
 
   class PrepareRequest < Request
@@ -65,7 +88,21 @@ module PgPipeline
       @statement_name = RequestOps.snapshot_name(statement.physical_name)
     end
 
+    def self.build(statement, params)
+      request = allocate
+      request.__send__(:init_prepared, statement, params)
+      request
+    end
+
     def operation = :prepared_query
+
+    private
+
+    def init_prepared(statement, params)
+      init_query(statement.sql, params)
+      @statement_name = RequestOps.snapshot_name(statement.physical_name)
+      self
+    end
   end
 
   module RequestOps
@@ -76,11 +113,34 @@ module PgPipeline
       value.frozen? ? value : value.dup.freeze
     end
 
+    EMPTY_PARAMS = [].freeze
+
     def snapshot_params(params)
-      values = params.nil? ? [] : params
-      raise ArgumentError, "params must be an Array" unless values.is_a?(Array)
+      return EMPTY_PARAMS if params.nil?
+      raise ArgumentError, "params must be an Array" unless params.is_a?(Array)
+      return EMPTY_PARAMS if params.empty?
+
+      return params if params.frozen? && immutable_values?(params)
 
-      values.map { |value| snapshot_value(value) }.freeze
+      params.map { |value| snapshot_value(value) }.freeze
+    end
+
+    def immutable_values?(values)
+      index = 0
+      size = values.size
+      while index < size
+        value = values[index]
+        case value
+        when Integer, Float, Symbol, NilClass, TrueClass, FalseClass
+          nil
+        when String
+          return false unless value.frozen?
+        else
+          return false
+        end
+        index += 1
+      end
+      true
     end
 
     def snapshot_value(value)
@@ -195,13 +255,11 @@ module PgPipeline
       req.waiter = waiter
       req.waiter_scheduler = scheduler
 
-      begin
-        scheduler.block(req, nil) until req.settled
-      ensure
-        if req.waiter.equal?(waiter)
-          req.waiter = nil
-          req.waiter_scheduler = nil
-        end
+      scheduler.block(req, nil) until req.settled
+    ensure
+      if waiter && req.waiter.equal?(waiter)
+        req.waiter = nil
+        req.waiter_scheduler = nil
       end
     end
 

data/lib/pg_pipeline/runtime/notification.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module PgPipeline
+  module Runtime
+    class Notification
+      def initialize
+        @waiters = []
+      end
+
+      def wait(timeout = nil)
+        Runtime.with_waiter(self, @waiters) do |waiter|
+          Runtime.with_timeout(timeout) { Runtime.park(self, waiter) { false } }
+          true
+        rescue TimeoutError
+          false
+        end
+      end
+
+      def signal
+        waiter = @waiters.shift or return false
+        Runtime.wake_dequeued(waiter, self)
+        true
+      end
+
+      def signal_all
+        pending, @waiters = @waiters, []
+        pending.each { |waiter| Runtime.wake_dequeued(waiter, self) }
+        nil
+      end
+    end
+  end
+end

data/lib/pg_pipeline/runtime/queue.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module PgPipeline
+  module Runtime
+    class Queue
+      def initialize
+        @queue = ::Thread::Queue.new
+      end
+
+      def enqueue(item)
+        @queue.push(item)
+        item
+      rescue ClosedQueueError
+        nil
+      end
+
+      def dequeue
+        @queue.pop
+      end
+
+      def close
+        @queue.close
+        nil
+      end
+
+      def closed? = @queue.closed?
+      def empty? = @queue.empty?
+      def size = @queue.size
+    end
+  end
+end

data/lib/pg_pipeline/runtime/semaphore.rb

@@ -0,0 +1,138 @@
+# frozen_string_literal: true
+
+module PgPipeline
+  module Runtime
+    class Semaphore
+      def initialize(limit)
+        @limit = Integer(limit)
+        raise ArgumentError, "limit must be >= 0" if @limit.negative?
+
+        @available = @limit
+        @waiting = WaitList.new
+      end
+
+      def acquire
+        wait
+        return @available unless block_given?
+
+        begin
+          yield
+        ensure
+          release
+        end
+      end
+
+      def release
+        if (node = @waiting.shift)
+          node.grant!
+          node.resume
+        else
+          @available += 1
+        end
+        @available
+      end
+
+      private
+
+      def wait
+        return (@available -= 1) if @available.positive?
+
+        task = Runtime.current_task
+        node = FiberNode.new(Fiber.current, self)
+        @waiting.push(node)
+        task&.enter_block(node.waiter)
+
+        begin
+          until node.granted?
+            node.suspend
+            abandon_wait!(node, task) if task&.cancelled?
+          end
+          abandon_wait!(node, task) if task&.cancelled?
+        ensure
+          task&.exit_block
+          @waiting.remove(node) unless node.granted? || node.list.nil?
+        end
+      end
+
+      def abandon_wait!(node, task)
+        node.granted? ? release : @waiting.remove(node)
+        task.raise_if_cancelled!
+      end
+
+      class WaitList
+        attr_reader :size
+
+        def initialize
+          @head = @tail = nil
+          @size = 0
+        end
+
+        def first = @head
+        def empty? = @head.nil?
+
+        def push(node)
+          raise ArgumentError, "node already queued" if node.list
+
+          node.list = self
+          node.prev = @tail
+          node.next = nil
+          @tail ? (@tail.next = node) : (@head = node)
+          @tail = node
+          @size += 1
+          node
+        end
+
+        def shift
+          node = @head or return nil
+          remove(node)
+        end
+
+        def remove(node)
+          return nil unless node.list.equal?(self)
+
+          node.prev ? (node.prev.next = node.next) : (@head = node.next)
+          node.next ? (node.next.prev = node.prev) : (@tail = node.prev)
+          node.list = node.prev = node.next = nil
+          @size -= 1
+          node
+        end
+      end
+
+      class FiberNode
+        attr_accessor :list, :prev, :next
+        attr_reader :fiber, :waiter, :blocker
+
+        def initialize(fiber, blocker)
+          @fiber = fiber
+          @blocker = blocker
+          @list = @prev = @next = nil
+          @granted = false
+          @scheduler = Fiber.scheduler
+          @waiter = {fiber: fiber, scheduler: @scheduler, ready: false, blocker: blocker, queued: true}
+        end
+
+        def granted? = @granted
+
+        def grant!
+          @granted = true
+          @waiter[:ready] = true
+          @waiter[:queued] = false
+          self
+        end
+
+        def suspend
+          @scheduler.block(@blocker, nil)
+        end
+
+        def resume
+          fiber = @fiber
+          return unless fiber.alive?
+
+          @scheduler.unblock(@blocker, fiber)
+        end
+      end
+
+      private_constant :WaitList, :FiberNode
+    end
+  end
+end