pg_pipeline 0.2.4 → 0.3.0

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"
@@ -20,6 +19,17 @@ module PgPipeline
                   :draining, :needs_flush, :writer_armed, :request_event_pending,
                   :owner_task, :reader_task, :writer_task
 
+    attr_writer :readable_events, :results_read, :units_completed, :flush_calls,
+                :flush_incomplete, :dispatches, :leaked_watchers
+
+    def readable_events = @readable_events || 0
+    def results_read = @results_read || 0
+    def units_completed = @units_completed || 0
+    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)
       @max_in_flight = DriverOps.positive_integer!(max_in_flight, :max_in_flight)
@@ -27,30 +37,22 @@ module PgPipeline
       @conn = conn
       @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
 
-      @inflight = []
-      @dispatching = nil
-      @submitting = 0
+      @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
 
-      @accepting = false
-      @running = false
-      @draining = false
-      @needs_flush = false
-      @writer_armed = false
-      @request_event_pending = false
-
-      @socket = nil
-      @owner_task = nil
-      @reader_task = nil
-      @writer_task = nil
+      @inflight = []
     end
 
-    def start(parent: Async::Task.current) = DriverOps.start(self, parent)
+    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
@@ -63,29 +65,38 @@ module PgPipeline
         pending: @requests.size,
         in_flight: @inflight.size,
         submitting: @submitting,
-        needs_flush: @needs_flush
+        needs_flush: @needs_flush,
+        fast_sync: @caps.fast_sync?,
+        readable_events: readable_events,
+        results_read: results_read,
+        units_completed: units_completed,
+        flush_calls: flush_calls,
+        flush_incomplete: flush_incomplete,
+        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),
+        leaked_watchers: leaked_watchers
       }
     end
 
     def health_check(timeout)
       return true unless available?
 
-      probe = Request.new(sql: "SELECT 1", params: [])
-      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
+      probe = Request.build("SELECT 1", nil)
+      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)
@@ -96,19 +107,30 @@ module PgPipeline
   end
 
   module DriverOps
+    WATCHER_JOIN_TIMEOUT = 2.0
+    OWNER_JOIN_TIMEOUT = 5.0
+
     module_function
 
-    SUCCESS_STATUSES = [
-      PG::PGRES_EMPTY_QUERY,
-      PG::PGRES_COMMAND_OK,
-      PG::PGRES_TUPLES_OK
-    ].freeze
+    def ratio(numerator, denominator)
+      return 0.0 if denominator.zero?
 
-    COPY_STATUSES = [
-      PG::PGRES_COPY_IN,
-      PG::PGRES_COPY_OUT,
-      PG::PGRES_COPY_BOTH
-    ].freeze
+      (numerator.to_f / denominator).round(3)
+    end
+
+    def warn_flush_coupling_once(caps)
+      return if @flush_coupling_warned
+      return if caps.fast_sync?
+      return if ENV["PG_PIPELINE_SILENCE_WARNINGS"]
+
+      @flush_coupling_warned = true
+      warn(
+        "pg_pipeline: libpq #{caps.libpq_version} couples pipeline Sync with flush " \
+        "(no PQsendPipelineSync). Queries still pipeline and still amortise RTT, but " \
+        "one flush per unit caps local throughput; libpq >= 17 is recommended for " \
+        "maximum throughput. Set PG_PIPELINE_SILENCE_WARNINGS=1 to silence this."
+      )
+    end
 
     def positive_integer!(value, name)
       integer = Integer(value)
@@ -119,8 +141,9 @@ 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
@@ -129,15 +152,15 @@ module PgPipeline
       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
 
@@ -164,7 +187,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
 
@@ -174,7 +197,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
 
@@ -196,7 +219,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
@@ -213,6 +241,7 @@ module PgPipeline
         nil
       when :readable
         begin
+          d.readable_events += 1
           read_available(d)
           input_changed = true
         ensure
@@ -238,6 +267,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)
@@ -260,6 +290,7 @@ module PgPipeline
         request.dispatched!
         d.inflight << request
         d.dispatching = nil
+        d.dispatches += 1
         dispatched = true
       end
 
@@ -286,9 +317,6 @@ module PgPipeline
       rescue ProtocolError
         raise
       rescue StandardError => e
-        # ruby-pg prepares/encodes query parameters before calling PQsend*. A
-        # Ruby-side encoder/coercion exception therefore belongs only to this
-        # request and must not poison unrelated work already in the pipeline.
         request.reject!(e)
         return false
       end
@@ -327,10 +355,13 @@ module PgPipeline
     def flush_output(d)
       return unless d.running
 
+      d.flush_calls += 1
+
       if d.conn.sync_flush
         d.needs_flush = false
       else
         d.needs_flush = true
+        d.flush_incomplete += 1
         arm_writer(d)
       end
     rescue PG::Error => e
@@ -339,6 +370,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)
@@ -351,42 +383,52 @@ module PgPipeline
     end
 
     def drain_results(d)
-      while !d.inflight.empty? && !d.conn.is_busy
-        result = d.conn.sync_get_result
-        request = d.inflight.first
-        raise ProtocolError, "result without an in-flight request" unless request
+      read = 0
 
-        if result.nil?
-          request.query_boundary!
-          next
-        end
-
-        status = result.result_status
-
-        case status
-        when PG::PGRES_PIPELINE_SYNC
-          clear_result(result)
-          complete_front(d, request)
-        when PG::PGRES_PIPELINE_ABORTED
-          ensure_before_query_boundary!(request, status)
-          clear_result(result)
-          request.record_error!(PipelineAbortedError.new("pipeline unit aborted"))
-        when PG::PGRES_BAD_RESPONSE
-          clear_result(result)
-          raise ProtocolError, "server response was not understood"
-        when PG::PGRES_FATAL_ERROR
-          ensure_before_query_boundary!(request, status)
-          request.record_error!(query_error(result), result: result)
-        when *COPY_STATUSES
-          clear_result(result)
-          raise ProtocolError, "COPY is not supported on the multiplexed pipeline"
-        when *SUCCESS_STATUSES
-          ensure_before_query_boundary!(request, status)
-          request.accept_result(result)
-        else
-          clear_result(result)
-          raise ProtocolError, "unexpected pipeline result status #{status}"
+      begin
+        while !d.inflight.empty? && !d.conn.is_busy
+          result = d.conn.sync_get_result
+          read += 1
+          request = d.inflight.first
+          raise ProtocolError, "result without an in-flight request" unless request
+
+          if result.nil?
+            request.query_boundary!
+            next
+          end
+
+          status = result.result_status
+
+          case status
+          when PG::PGRES_TUPLES_OK
+            ensure_before_query_boundary!(request, status)
+            request.accept_result(result)
+          when PG::PGRES_PIPELINE_SYNC
+            clear_result(result)
+            complete_front(d, request)
+          when PG::PGRES_COMMAND_OK, PG::PGRES_EMPTY_QUERY
+            ensure_before_query_boundary!(request, status)
+            request.accept_result(result)
+          when PG::PGRES_FATAL_ERROR
+            ensure_before_query_boundary!(request, status)
+            request.record_error!(query_error(result), result: result)
+          when PG::PGRES_PIPELINE_ABORTED
+            ensure_before_query_boundary!(request, status)
+            clear_result(result)
+            request.record_error!(PipelineAbortedError.new("pipeline unit aborted"))
+          when PG::PGRES_BAD_RESPONSE
+            clear_result(result)
+            raise ProtocolError, "server response was not understood"
+          when PG::PGRES_COPY_IN, PG::PGRES_COPY_OUT, PG::PGRES_COPY_BOTH
+            clear_result(result)
+            raise ProtocolError, "COPY is not supported on the multiplexed pipeline"
+          else
+            clear_result(result)
+            raise ProtocolError, "unexpected pipeline result status #{status}"
+          end
         end
+      ensure
+        d.results_read += read if read.positive?
       end
     end
 
@@ -400,6 +442,7 @@ module PgPipeline
       raise ProtocolError, "sync does not match FIFO front" unless request.equal?(d.inflight.first)
 
       d.inflight.shift
+      d.units_completed += 1
       request.finish!
     end
 
@@ -417,14 +460,11 @@ module PgPipeline
       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)
@@ -434,8 +474,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)
@@ -500,21 +540,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
 
-      [reader, writer].each do |task|
-        task&.stop
-      rescue Async::Cancel, StandardError
+    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
 
+    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 duplexed socket_io / connection did not release the fiber " \
+        "(often parked in wait_readable/wait_writable)."
+      )
+    end
+
     def safe_close_conn(d)
       d.conn.close unless d.conn.finished?
     rescue StandardError