pg_pipeline 0.2.5 → 0.3.1

data/lib/pg_pipeline/connection_driver.rb

@@ -1,10 +1,9 @@
 # frozen_string_literal: true
 
 require "pg"
-require "async"
-require "async/queue"
 
 require_relative "errors"
+require_relative "runtime"
 require_relative "bounded_queue"
 require_relative "server_caps"
 require_relative "request"
@@ -21,7 +20,7 @@ module PgPipeline
                   :owner_task, :reader_task, :writer_task
 
     attr_writer :readable_events, :results_read, :units_completed, :flush_calls,
-                :flush_incomplete, :dispatches
+                :flush_incomplete, :dispatches, :leaked_watchers
 
     def readable_events = @readable_events || 0
     def results_read = @results_read || 0
@@ -29,6 +28,7 @@ module PgPipeline
     def flush_calls = @flush_calls || 0
     def flush_incomplete = @flush_incomplete || 0
     def dispatches = @dispatches || 0
+    def leaked_watchers = @leaked_watchers || 0
 
     def initialize(conn, max_pending: DEFAULT_MAX_PENDING, max_in_flight: DEFAULT_MAX_IN_FLIGHT)
       @max_pending = DriverOps.positive_integer!(max_pending, :max_pending)
@@ -38,37 +38,19 @@ module PgPipeline
       @caps = ServerCaps.from_connection(conn)
       @caps.assert_supported!
       DriverOps.warn_flush_coupling_once(@caps)
-
       @requests = BoundedQueue.new(@max_pending)
-      @events = Async::Queue.new
-      @reader_rearm = Async::Queue.new
-      @writer_commands = Async::Queue.new
-
+      @events = Runtime::Queue.new
+      @reader_rearm = Runtime::Queue.new
+      @writer_commands = Runtime::Queue.new
+
+      @accepting = @running = @draining = @needs_flush = @writer_armed = @request_event_pending = false
+      @readable_events = @results_read = @units_completed = @flush_calls =
+        @flush_incomplete = @dispatches = @leaked_watchers = @submitting = 0
+      @socket = @owner_task = @reader_task = @writer_task = @dispatching = nil
       @inflight = []
-      @dispatching = nil
-      @submitting = 0
-
-      @accepting = false
-      @running = false
-      @draining = false
-      @needs_flush = false
-      @writer_armed = false
-      @request_event_pending = false
-
-      @readable_events = 0
-      @results_read = 0
-      @units_completed = 0
-      @flush_calls = 0
-      @flush_incomplete = 0
-      @dispatches = 0
-
-      @socket = nil
-      @owner_task = nil
-      @reader_task = nil
-      @writer_task = nil
-    end
-
-    def start(parent: Async::Task.current) = DriverOps.start(self, parent)
+    end
+
+    def start = DriverOps.start(self)
     def submit(request) = DriverOps.submit(self, request)
     def load = @requests.size + @inflight.size + @submitting + (@dispatching ? 1 : 0)
     def available? = @accepting && @running
@@ -91,7 +73,8 @@ module PgPipeline
         dispatches: dispatches,
         units_per_readable: DriverOps.ratio(units_completed, readable_events),
         results_per_readable: DriverOps.ratio(results_read, readable_events),
-        flush_calls_per_unit: DriverOps.ratio(flush_calls, units_completed)
+        flush_calls_per_unit: DriverOps.ratio(flush_calls, units_completed),
+        leaked_watchers: leaked_watchers
       }
     end
 
@@ -99,21 +82,19 @@ module PgPipeline
       return true unless available?
 
       probe = Request.build("SELECT 1", nil)
-      begin
-        submit(probe)
-        Async::Task.current.with_timeout(timeout) { probe.wait }
-        true
-      rescue Async::TimeoutError
-        DriverOps.abort_timed_out_health_probe(self, probe)
-      rescue QueryError, PipelineAbortedError
-        true
-      rescue ShutdownError, NotDispatchedError
-        true
-      rescue ConnectionLostError, PG::Error
-        false
-      ensure
-        probe.cancel! unless probe.settled?
-      end
+      submit(probe)
+      Runtime.with_timeout(timeout) { probe.wait }
+      true
+    rescue Runtime::TimeoutError
+      DriverOps.abort_timed_out_health_probe(self, probe)
+    rescue QueryError, PipelineAbortedError
+      true
+    rescue ShutdownError, NotDispatchedError
+      true
+    rescue ConnectionLostError, PG::Error
+      false
+    ensure
+      probe.cancel! if probe && !probe.settled?
     end
 
     def graceful_close = DriverOps.graceful_close(self)
@@ -124,6 +105,10 @@ module PgPipeline
   end
 
   module DriverOps
+    WATCHER_JOIN_TIMEOUT = 2.0
+    OWNER_JOIN_TIMEOUT = 5.0
+    WATCHER_POLL_INTERVAL = 0.25
+
     module_function
 
     def ratio(numerator, denominator)
@@ -155,25 +140,25 @@ module PgPipeline
       raise ArgumentError, "#{name} must be an integer >= 1"
     end
 
-    def start(d, parent)
+    def start(d)
       raise Error, "driver already started" if d.running
+      raise Error, "driver start requires an active Fiber scheduler" unless Fiber.scheduler
 
       d.conn.setnonblocking(true)
       d.conn.enter_pipeline_mode
-
       d.socket = d.conn.socket_io
       d.accepting = true
       d.running = true
 
-      d.reader_task = parent.async { reader_watcher(d) }
-      d.writer_task = parent.async { writer_watcher(d) }
-      d.owner_task = parent.async { owner_loop(d) }
+      d.reader_task = Runtime.spawn(name: :reader) { reader_watcher(d) }
+      d.writer_task = Runtime.spawn(name: :writer) { writer_watcher(d) }
+      d.owner_task = Runtime.spawn(name: :owner) { owner_loop(d) }
       d
     rescue Exception
       d.accepting = false
       d.running = false
-      stop_watchers(d)
-      safe_close_conn(d)
+      teardown_watchers(d)
+      join_owner(d)
       raise
     end
 
@@ -200,7 +185,7 @@ module PgPipeline
       d.accepting = false
       d.requests.close(ShutdownError.new("driver is closing"))
       d.events.enqueue(:close)
-      d.owner_task.wait
+      join_owner(d)
       nil
     end
 
@@ -210,7 +195,7 @@ module PgPipeline
       d.accepting = false
       d.requests.close(not_dispatched_error(error))
       d.events.enqueue([:abort, error])
-      d.owner_task.wait unless Async::Task.current.equal?(d.owner_task)
+      join_owner(d)
       nil
     end
 
@@ -232,7 +217,12 @@ module PgPipeline
     end
 
     def owner_loop(d)
-      process_event(d, d.events.dequeue) while d.running
+      while d.running
+        event = d.events.dequeue
+        break if event.nil?
+
+        process_event(d, event)
+      end
     rescue StandardError => e
       fatal_close(d, ConnectionLostError.new("driver crashed: #{e.class}: #{e.message}"))
     ensure
@@ -241,7 +231,6 @@ module PgPipeline
 
     def process_event(d, event)
       input_changed = false
-
       case event
       when :requests
         d.request_event_pending = false
@@ -275,6 +264,7 @@ module PgPipeline
 
     def notify_requests(d)
       return if d.request_event_pending
+      return unless d.running
 
       d.request_event_pending = true
       d.events.enqueue(:requests)
@@ -282,7 +272,6 @@ module PgPipeline
 
     def pump_requests(d)
       dispatched = false
-
       while d.inflight.size < d.max_in_flight && !d.requests.empty?
         request = d.requests.dequeue
         break unless request
@@ -363,7 +352,6 @@ module PgPipeline
       return unless d.running
 
       d.flush_calls += 1
-
       if d.conn.sync_flush
         d.needs_flush = false
       else
@@ -377,6 +365,7 @@ module PgPipeline
 
     def arm_writer(d)
       return if d.writer_armed
+      return unless d.running
 
       d.writer_armed = true
       d.writer_commands.enqueue(:wait_writable)
@@ -404,7 +393,6 @@ module PgPipeline
           end
 
           status = result.result_status
-
           case status
           when PG::PGRES_TUPLES_OK
             ensure_before_query_boundary!(request, status)
@@ -465,15 +453,11 @@ module PgPipeline
     def finish_graceful_close(d)
       d.accepting = false
       d.running = false
-
-      begin
-        d.conn.exit_pipeline_mode
-      rescue PG::Error => e
-        fail_all(d, ConnectionLostError.new("failed to exit pipeline mode: #{e.message}"))
-      ensure
-        stop_watchers(d)
-        safe_close_conn(d)
-      end
+      d.conn.exit_pipeline_mode
+    rescue PG::Error => e
+      fail_all(d, ConnectionLostError.new("failed to exit pipeline mode: #{e.message}"))
+    ensure
+      teardown_watchers(d)
     end
 
     def fatal_close(d, error)
@@ -483,8 +467,8 @@ module PgPipeline
       d.running = false
       d.requests.close(not_dispatched_error(error))
       fail_all(d, error)
-      stop_watchers(d)
-      safe_close_conn(d)
+
+      teardown_watchers(d)
     end
 
     def fail_all(d, error)
@@ -520,10 +504,32 @@ module PgPipeline
       )
     end
 
+    def watcher_wait_timeout
+      scheduler = Fiber.scheduler
+      return nil if scheduler&.respond_to?(:fiber_interrupt)
+
+      WATCHER_POLL_INTERVAL
+    end
+
+    def wait_socket_readable(d, timeout)
+      loop do
+        return false unless d.running
+        return true if d.socket.wait_readable(timeout)
+      end
+    end
+
+    def wait_socket_writable(d, timeout)
+      loop do
+        return false unless d.running
+        return true if d.socket.wait_writable(timeout)
+      end
+    end
+
     def reader_watcher(d)
+      timeout = watcher_wait_timeout
+
       while d.running
-        d.socket.wait_readable
-        break unless d.running
+        break unless wait_socket_readable(d, timeout)
 
         d.events.enqueue(:readable)
         command = d.reader_rearm.dequeue
@@ -536,11 +542,13 @@ module PgPipeline
     end
 
     def writer_watcher(d)
+      timeout = watcher_wait_timeout
+
       while d.running
         command = d.writer_commands.dequeue
         break unless command == :wait_writable && d.running
+        break unless wait_socket_writable(d, timeout)
 
-        d.socket.wait_writable
         d.events.enqueue(:writable) if d.running
       end
     rescue StandardError => e
@@ -549,21 +557,87 @@ module PgPipeline
       end
     end
 
-    def stop_watchers(d)
-      reader = d.reader_task
-      writer = d.writer_task
+    def teardown_watchers(d)
+      tasks = release_watchers(d)
+      close_wait_points(d)
+      begin
+        d.socket&.close
+      rescue StandardError
+        nil
+      end
+      stop_watchers(tasks)
+      join_watchers(d, tasks)
+      d.socket = nil
+      safe_close_conn(d)
+      nil
+    end
+
+    def release_watchers(d)
+      tasks = [d.reader_task, d.writer_task].compact
       d.reader_task = nil
       d.writer_task = nil
+      tasks
+    end
+
+    def close_wait_points(d)
+      [d.reader_rearm, d.writer_commands, d.events].each do |queue|
+        queue&.close
+      rescue StandardError
+        nil
+      end
+    end
+
+    def stop_watchers(tasks)
+      Array(tasks).each do |task|
+        task.stop
+      rescue Runtime::Cancel, StandardError
+        nil
+      end
+
+      nil
+    end
 
-      [reader, writer].each do |task|
-        task&.stop
-      rescue Async::Cancel, StandardError
+    def join_watchers(d, tasks)
+      Array(tasks).each do |task|
+        task.wait(WATCHER_JOIN_TIMEOUT)
+      rescue Runtime::TimeoutError
+        d.leaked_watchers += 1
+        warn_leaked_task(d, task, WATCHER_JOIN_TIMEOUT)
+      rescue Runtime::Cancel, StandardError
         nil
       end
 
       nil
     end
 
+    def join_owner(d)
+      owner = d.owner_task
+      return if owner.nil?
+      return if Fiber.current.equal?(owner.fiber)
+
+      begin
+        owner.wait(OWNER_JOIN_TIMEOUT)
+      rescue Runtime::TimeoutError
+        d.leaked_watchers += 1
+        warn_leaked_task(d, owner, OWNER_JOIN_TIMEOUT)
+      rescue Runtime::Cancel, StandardError
+        nil
+      end
+
+      nil
+    end
+
+    def warn_leaked_task(d, task, timeout)
+      return if ENV["PG_PIPELINE_SILENCE_WARNINGS"]
+
+      warn(
+        "pg_pipeline: task #{task.name.inspect} did not exit within " \
+        "#{timeout}s and has been leaked (total #{d.leaked_watchers}). " \
+        "Closing the socket and, on schedulers without #fiber_interrupt, the " \
+        "#{WATCHER_POLL_INTERVAL}s watcher poll did not release the fiber in time."
+      )
+    end
+
     def safe_close_conn(d)
       d.conn.close unless d.conn.finished?
     rescue StandardError

data/lib/pg_pipeline/errors.rb

@@ -8,6 +8,7 @@ module PgPipeline
   class ConnectionLostError < Error; end
   class NotDispatchedError < ConnectionLostError; end
   class IndeterminateResultError < ConnectionLostError; end
+  class IndeterminateCommitError < IndeterminateResultError; end
   class ShutdownError < Error; end
   class ProtocolError < Error; end
   class RecursiveCheckoutError < Error; end

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
 
@@ -58,7 +58,8 @@ module PgPipeline
       @health_failures = 0
       @supervisor_error = nil
       @supervisor = nil
-      @task_parent = nil
+      @supervisor_wake = Runtime::Notification.new
+      @pause_failures = 0
 
       @prepared_statements = {}
       @prepared_generation = 0
@@ -70,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
@@ -174,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
 
@@ -184,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
@@ -220,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
@@ -229,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
 
@@ -284,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
@@ -292,9 +292,29 @@ module PgPipeline
         end
 
         break if @closing
+        break unless supervisor_pause
+      end
+    end
 
-        sleep(supervisor_sleep_interval)
+    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
+
+      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
@@ -305,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!
 
@@ -316,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
@@ -339,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
 
@@ -376,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
@@ -413,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
@@ -447,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)
@@ -513,7 +548,6 @@ module PgPipeline
       close_free_pinned
       @pinned_gate = nil
       @supervisor = nil
-      @task_parent = nil
       @started = false
       @closing = false
       @closed = false