Module: PgPipeline::DriverOps

Defined in:
lib/pg_pipeline/connection_driver.rb

Constant Summary collapse

WATCHER_JOIN_TIMEOUT =
2.0
OWNER_JOIN_TIMEOUT =
5.0
WATCHER_POLL_INTERVAL =
0.25

Class Method Summary collapse

Class Method Details

.abort!(d, error) ⇒ Object



192
193
194
195
196
197
198
199
200
# File 'lib/pg_pipeline/connection_driver.rb', line 192

def abort!(d, error)
  return unless d.running

  d.accepting = false
  d.requests.close(not_dispatched_error(error))
  d.events.enqueue([:abort, error])
  join_owner(d)
  nil
end

.abort_timed_out_health_probe(d, probe) ⇒ Object



202
203
204
205
206
207
208
# File 'lib/pg_pipeline/connection_driver.rb', line 202

def abort_timed_out_health_probe(d, probe)
  return true if probe.settled? || !d.running
  return true unless exclusive_health_probe?(d, probe)

  abort!(d, ConnectionLostError.new("idle health check timed out"))
  false
end

.arm_writer(d) ⇒ Object



366
367
368
369
370
371
372
# File 'lib/pg_pipeline/connection_driver.rb', line 366

def arm_writer(d)
  return if d.writer_armed
  return unless d.running

  d.writer_armed = true
  d.writer_commands.enqueue(:wait_writable)
end

.clear_result(result) ⇒ Object



647
648
649
# File 'lib/pg_pipeline/connection_driver.rb', line 647

def clear_result(result)
  result.clear if result.respond_to?(:clear)
end

.close_wait_points(d) ⇒ Object



582
583
584
585
586
587
588
# File 'lib/pg_pipeline/connection_driver.rb', line 582

def close_wait_points(d)
  [d.reader_rearm, d.writer_commands, d.events].each do |queue|
    queue&.close
  rescue StandardError
    nil
  end
end

.complete_front(d, request) ⇒ Object

Raises:



435
436
437
438
439
440
441
# File 'lib/pg_pipeline/connection_driver.rb', line 435

def complete_front(d, request)
  raise ProtocolError, "sync does not match FIFO front" unless request.equal?(d.inflight.first)

  d.inflight.shift
  d.units_completed += 1
  request.finish!
end

.drain_results(d) ⇒ Object



380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
# File 'lib/pg_pipeline/connection_driver.rb', line 380

def drain_results(d)
  read = 0

  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

.drained?(d) ⇒ Boolean

Returns:

  • (Boolean)


449
450
451
# File 'lib/pg_pipeline/connection_driver.rb', line 449

def drained?(d)
  d.submitting.zero? && d.requests.empty? && d.inflight.empty? && !d.needs_flush && d.dispatching.nil?
end

.ensure_before_query_boundary!(request, status) ⇒ Object

Raises:



429
430
431
432
433
# File 'lib/pg_pipeline/connection_driver.rb', line 429

def ensure_before_query_boundary!(request, status)
  return unless request.query_boundary_seen?

  raise ProtocolError, "result status #{status} arrived after query boundary"
end

.exclusive_health_probe?(d, probe) ⇒ Boolean

Returns:

  • (Boolean)


210
211
212
213
214
215
216
217
# File 'lib/pg_pipeline/connection_driver.rb', line 210

def exclusive_health_probe?(d, probe)
  d.accepting &&
    d.dispatching.nil? &&
    d.submitting.zero? &&
    d.requests.empty? &&
    d.inflight.length == 1 &&
    d.inflight.first.equal?(probe)
end

.fail_all(d, error) ⇒ Object



474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
# File 'lib/pg_pipeline/connection_driver.rb', line 474

def fail_all(d, error)
  uncertain = []
  uncertain << d.dispatching if d.dispatching
  uncertain.concat(d.inflight)
  queued = d.requests.drain

  uncertain.compact.uniq.each do |request|
    request.reject!(indeterminate_error(error)) unless request.settled?
  end

  queued.each do |request|
    request.reject!(not_dispatched_error(error)) unless request.settled?
  end

  d.dispatching = nil
  d.inflight.clear
end

.fatal_close(d, error) ⇒ Object



463
464
465
466
467
468
469
470
471
472
# File 'lib/pg_pipeline/connection_driver.rb', line 463

def fatal_close(d, error)
  return unless d.running || d.accepting

  d.accepting = false
  d.running = false
  d.requests.close(not_dispatched_error(error))
  fail_all(d, error)

  teardown_watchers(d)
end

.finish_graceful_close(d) ⇒ Object



453
454
455
456
457
458
459
460
461
# File 'lib/pg_pipeline/connection_driver.rb', line 453

def finish_graceful_close(d)
  d.accepting = false
  d.running = false
  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

.flush_output(d) ⇒ Object



351
352
353
354
355
356
357
358
359
360
361
362
363
364
# File 'lib/pg_pipeline/connection_driver.rb', line 351

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
  raise ConnectionLostError, "flush failed: #{e.class}: #{e.message}"
end

.graceful_close(d) ⇒ Object



182
183
184
185
186
187
188
189
190
# File 'lib/pg_pipeline/connection_driver.rb', line 182

def graceful_close(d)
  return unless d.running

  d.accepting = false
  d.requests.close(ShutdownError.new("driver is closing"))
  d.events.enqueue(:close)
  join_owner(d)
  nil
end

.indeterminate_error(error) ⇒ Object



498
499
500
501
502
503
504
505
# File 'lib/pg_pipeline/connection_driver.rb', line 498

def indeterminate_error(error)
  return error if error.is_a?(IndeterminateResultError)

  IndeterminateResultError.new(
    "#{error.message}; request was dispatched but its Sync was not observed, " \
    "so execution/commit outcome is indeterminate"
  )
end

.join_owner(d) ⇒ Object



613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
# File 'lib/pg_pipeline/connection_driver.rb', line 613

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

.join_watchers(d, tasks) ⇒ Object



600
601
602
603
604
605
606
607
608
609
610
611
# File 'lib/pg_pipeline/connection_driver.rb', line 600

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

.not_dispatched_error(error) ⇒ Object



492
493
494
495
496
# File 'lib/pg_pipeline/connection_driver.rb', line 492

def not_dispatched_error(error)
  return error if error.is_a?(NotDispatchedError)

  NotDispatchedError.new("#{error.message}; request was not dispatched")
end

.notify_requests(d) ⇒ Object



265
266
267
268
269
270
271
# File 'lib/pg_pipeline/connection_driver.rb', line 265

def notify_requests(d)
  return if d.request_event_pending
  return unless d.running

  d.request_event_pending = true
  d.events.enqueue(:requests)
end

.owner_loop(d) ⇒ Object



219
220
221
222
223
224
225
226
227
228
229
230
# File 'lib/pg_pipeline/connection_driver.rb', line 219

def owner_loop(d)
  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
  fatal_close(d, ShutdownError.new("driver owner stopped before shutdown completed")) if d.running
end

.positive_integer!(value, name) ⇒ Object



134
135
136
137
138
139
140
141
# File 'lib/pg_pipeline/connection_driver.rb', line 134

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

.process_event(d, event) ⇒ Object



232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
# File 'lib/pg_pipeline/connection_driver.rb', line 232

def process_event(d, event)
  input_changed = false
  case event
  when :requests
    d.request_event_pending = false
  when :submission_finished
    nil
  when :readable
    begin
      d.readable_events += 1
      read_available(d)
      input_changed = true
    ensure
      d.reader_rearm.enqueue(:rearm) if d.running
    end
  when :writable
    d.writer_armed = false
    flush_output(d)
  when :close
    d.draining = true
  when Array
    tag, payload = event
    if tag == :abort
      fatal_close(d, payload)
      return
    end
  end

  drain_results(d) if input_changed
  pump_requests(d) if d.running
  finish_graceful_close(d) if d.running && d.draining && drained?(d)
end

.pump_requests(d) ⇒ Object



273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
# File 'lib/pg_pipeline/connection_driver.rb', line 273

def pump_requests(d)
  dispatched = false
  while d.inflight.size < d.max_in_flight && !d.requests.empty?
    request = d.requests.dequeue
    break unless request
    next if request.cancelled?

    d.dispatching = request
    unless send_unit(d, request)
      d.dispatching = nil
      next
    end

    request.dispatched!
    d.inflight << request
    d.dispatching = nil
    d.dispatches += 1
    dispatched = true
  end

  flush_output(d) if dispatched && !d.needs_flush
end

.query_error(result) ⇒ Object



443
444
445
446
447
# File 'lib/pg_pipeline/connection_driver.rb', line 443

def query_error(result)
  message = result.error_message.to_s.strip
  message = "query failed" if message.empty?
  QueryError.new(message, cause_result: result)
end

.ratio(numerator, denominator) ⇒ Object



114
115
116
117
118
# File 'lib/pg_pipeline/connection_driver.rb', line 114

def ratio(numerator, denominator)
  return 0.0 if denominator.zero?

  (numerator.to_f / denominator).round(3)
end

.read_available(d) ⇒ Object



374
375
376
377
378
# File 'lib/pg_pipeline/connection_driver.rb', line 374

def read_available(d)
  d.conn.consume_input
rescue PG::Error => e
  raise ConnectionLostError, "read failed: #{e.class}: #{e.message}"
end

.reader_watcher(d) ⇒ Object



528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
# File 'lib/pg_pipeline/connection_driver.rb', line 528

def reader_watcher(d)
  timeout = watcher_wait_timeout

  while d.running
    break unless wait_socket_readable(d, timeout)

    d.events.enqueue(:readable)
    command = d.reader_rearm.dequeue
    break unless command == :rearm && d.running
  end
rescue StandardError => e
  if d.running
    d.events.enqueue([:abort, ConnectionLostError.new("reader watcher failed: #{e.class}: #{e.message}")])
  end
end

.release_watchers(d) ⇒ Object



575
576
577
578
579
580
# File 'lib/pg_pipeline/connection_driver.rb', line 575

def release_watchers(d)
  tasks = [d.reader_task, d.writer_task].compact
  d.reader_task = nil
  d.writer_task = nil
  tasks
end

.reusable_after_send_rejection?(d) ⇒ Boolean

Returns:

  • (Boolean)


343
344
345
346
347
348
349
# File 'lib/pg_pipeline/connection_driver.rb', line 343

def reusable_after_send_rejection?(d)
  !d.conn.finished? &&
    d.conn.status == PG::CONNECTION_OK &&
    d.conn.pipeline_status != PG::PQ_PIPELINE_OFF
rescue PG::Error
  false
end

.safe_close_conn(d) ⇒ Object



641
642
643
644
645
# File 'lib/pg_pipeline/connection_driver.rb', line 641

def safe_close_conn(d)
  d.conn.close unless d.conn.finished?
rescue StandardError
  nil
end

.send_command(conn, request) ⇒ Object



326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
# File 'lib/pg_pipeline/connection_driver.rb', line 326

def send_command(conn, request)
  case request.operation
  when :query
    conn.send_query_params(request.sql, request.params)
  when :prepare
    if request.param_types.nil?
      conn.send_prepare(request.statement_name, request.sql)
    else
      conn.send_prepare(request.statement_name, request.sql, request.param_types)
    end
  when :prepared_query
    conn.send_query_prepared(request.statement_name, request.params)
  else
    raise ProtocolError, "unsupported request operation #{request.operation.inspect}"
  end
end

.send_unit(d, request) ⇒ Object



296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
# File 'lib/pg_pipeline/connection_driver.rb', line 296

def send_unit(d, request)
  begin
    send_command(d.conn, request)
  rescue PG::UnableToSend => e
    d.dispatching = nil
    request.reject!(
      NotDispatchedError.new("query was rejected before libpq accepted it: #{e.class}: #{e.message}")
    )
    return false if reusable_after_send_rejection?(d)

    raise ConnectionLostError, "dispatch failed: #{e.class}: #{e.message}"
  rescue PG::Error => e
    d.dispatching = nil
    request.reject!(
      NotDispatchedError.new("query was rejected before libpq accepted it: #{e.class}: #{e.message}")
    )
    raise ConnectionLostError, "dispatch failed: #{e.class}: #{e.message}"
  rescue ProtocolError
    raise
  rescue StandardError => e
    request.reject!(e)
    return false
  end

  d.caps.place_sync(d.conn)
  true
rescue PG::Error => e
  raise ConnectionLostError, "dispatch Sync failed: #{e.class}: #{e.message}"
end

.start(d) ⇒ Object



143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
# File 'lib/pg_pipeline/connection_driver.rb', line 143

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 = 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
  teardown_watchers(d)
  join_owner(d)
  raise
end

.stop_watchers(tasks) ⇒ Object



590
591
592
593
594
595
596
597
598
# File 'lib/pg_pipeline/connection_driver.rb', line 590

def stop_watchers(tasks)
  Array(tasks).each do |task|
    task.stop
  rescue Runtime::Cancel, StandardError
    nil
  end

  nil
end

.submit(d, request) ⇒ Object

Raises:



165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
# File 'lib/pg_pipeline/connection_driver.rb', line 165

def submit(d, request)
  raise ShutdownError, "driver is not accepting work" unless d.accepting
  raise ProtocolError, "request must be new before submit" unless request.state == :new

  d.submitting += 1
  begin
    d.requests.enqueue(request)
    request.queued!
    notify_requests(d)
  ensure
    d.submitting -= 1
    d.events.enqueue(:submission_finished) if d.draining && d.running && d.submitting.zero?
  end

  request
end

.teardown_watchers(d) ⇒ Object



560
561
562
563
564
565
566
567
568
569
570
571
572
573
# File 'lib/pg_pipeline/connection_driver.rb', line 560

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

.wait_socket_readable(d, timeout) ⇒ Object



514
515
516
517
518
519
# File 'lib/pg_pipeline/connection_driver.rb', line 514

def wait_socket_readable(d, timeout)
  loop do
    return false unless d.running
    return true if d.socket.wait_readable(timeout)
  end
end

.wait_socket_writable(d, timeout) ⇒ Object



521
522
523
524
525
526
# File 'lib/pg_pipeline/connection_driver.rb', line 521

def wait_socket_writable(d, timeout)
  loop do
    return false unless d.running
    return true if d.socket.wait_writable(timeout)
  end
end

.warn_flush_coupling_once(caps) ⇒ Object



120
121
122
123
124
125
126
127
128
129
130
131
132
# File 'lib/pg_pipeline/connection_driver.rb', line 120

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

.warn_leaked_task(d, task, timeout) ⇒ Object



630
631
632
633
634
635
636
637
638
639
# File 'lib/pg_pipeline/connection_driver.rb', line 630

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

.watcher_wait_timeoutObject



507
508
509
510
511
512
# File 'lib/pg_pipeline/connection_driver.rb', line 507

def watcher_wait_timeout
  scheduler = Fiber.scheduler
  return nil if scheduler&.respond_to?(:fiber_interrupt)

  WATCHER_POLL_INTERVAL
end

.writer_watcher(d) ⇒ Object



544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
# File 'lib/pg_pipeline/connection_driver.rb', line 544

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.events.enqueue(:writable) if d.running
  end
rescue StandardError => e
  if d.running
    d.events.enqueue([:abort, ConnectionLostError.new("writer watcher failed: #{e.class}: #{e.message}")])
  end
end