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



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

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



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

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

.accept_tuples_ok(d, request, result) ⇒ Object



431
432
433
434
435
436
437
438
439
440
# File 'lib/pg_pipeline/connection_driver.rb', line 431

def accept_tuples_ok(d, request, result)
  ensure_before_query_boundary!(request, PG::PGRES_TUPLES_OK)
  begin
    apply_type_map(d, request, result)
    request.accept_result(result)
  rescue StandardError => e
    message = e.is_a?(Error) ? e.message : "typed result mapping failed: #{e.class}: #{e.message}"
    request.record_error!(QueryError.new(message, cause_result: result), result: result)
  end
end

.apply_type_map(d, request, result) ⇒ Object

Raises:



665
666
667
668
669
670
671
672
# File 'lib/pg_pipeline/connection_driver.rb', line 665

def apply_type_map(d, request, result)
  return unless request.type_map

  maps = d.type_maps
  raise Error, "typed statement has no type-map registry" unless maps

  maps.apply!(request, result, d.conn)
end

.arm_writer(d) ⇒ Object



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

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



674
675
676
# File 'lib/pg_pipeline/connection_driver.rb', line 674

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

.close_wait_points(d) ⇒ Object



600
601
602
603
604
605
606
# File 'lib/pg_pipeline/connection_driver.rb', line 600

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:



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

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



382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
# File 'lib/pg_pipeline/connection_driver.rb', line 382

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

      handle_pipeline_result(d, request, result)
    end
  ensure
    d.results_read += read if read.positive?
  end
end

.drained?(d) ⇒ Boolean

Returns:

  • (Boolean)


467
468
469
# File 'lib/pg_pipeline/connection_driver.rb', line 467

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:



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

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)


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

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



492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
# File 'lib/pg_pipeline/connection_driver.rb', line 492

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



481
482
483
484
485
486
487
488
489
490
# File 'lib/pg_pipeline/connection_driver.rb', line 481

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



471
472
473
474
475
476
477
478
479
# File 'lib/pg_pipeline/connection_driver.rb', line 471

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



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

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



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

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

.handle_pipeline_result(d, request, result) ⇒ Object



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
428
429
# File 'lib/pg_pipeline/connection_driver.rb', line 399

def handle_pipeline_result(d, request, result)
  if result.nil?
    request.query_boundary!
    return
  end

  status = result.result_status
  case status
  when PG::PGRES_TUPLES_OK
    accept_tuples_ok(d, request, 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
    reject_pipeline_status!(result, "server response was not understood")
  when PG::PGRES_COPY_IN, PG::PGRES_COPY_OUT, PG::PGRES_COPY_BOTH
    reject_pipeline_status!(result, "COPY is not supported on the multiplexed pipeline")
  else
    reject_pipeline_status!(result, "unexpected pipeline result status #{status}")
  end
end

.indeterminate_error(error) ⇒ Object



516
517
518
519
520
521
522
523
# File 'lib/pg_pipeline/connection_driver.rb', line 516

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



631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
# File 'lib/pg_pipeline/connection_driver.rb', line 631

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



618
619
620
621
622
623
624
625
626
627
628
629
# File 'lib/pg_pipeline/connection_driver.rb', line 618

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



510
511
512
513
514
# File 'lib/pg_pipeline/connection_driver.rb', line 510

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



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

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



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

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



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

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



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
264
265
# File 'lib/pg_pipeline/connection_driver.rb', line 234

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



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

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



461
462
463
464
465
# File 'lib/pg_pipeline/connection_driver.rb', line 461

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



116
117
118
119
120
# File 'lib/pg_pipeline/connection_driver.rb', line 116

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

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

.read_available(d) ⇒ Object



376
377
378
379
380
# File 'lib/pg_pipeline/connection_driver.rb', line 376

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



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

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

.reject_pipeline_status!(result, message) ⇒ Object

Raises:



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

def reject_pipeline_status!(result, message)
  clear_result(result)
  raise ProtocolError, message
end

.release_watchers(d) ⇒ Object



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

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)


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

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



659
660
661
662
663
# File 'lib/pg_pipeline/connection_driver.rb', line 659

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

.send_command(conn, request) ⇒ Object



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

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



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
325
326
# File 'lib/pg_pipeline/connection_driver.rb', line 298

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



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

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



608
609
610
611
612
613
614
615
616
# File 'lib/pg_pipeline/connection_driver.rb', line 608

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

  nil
end

.submit(d, request) ⇒ Object

Raises:



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

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



578
579
580
581
582
583
584
585
586
587
588
589
590
591
# File 'lib/pg_pipeline/connection_driver.rb', line 578

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



532
533
534
535
536
537
# File 'lib/pg_pipeline/connection_driver.rb', line 532

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



539
540
541
542
543
544
# File 'lib/pg_pipeline/connection_driver.rb', line 539

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



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

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



648
649
650
651
652
653
654
655
656
657
# File 'lib/pg_pipeline/connection_driver.rb', line 648

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



525
526
527
528
529
530
# File 'lib/pg_pipeline/connection_driver.rb', line 525

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

  WATCHER_POLL_INTERVAL
end

.writer_watcher(d) ⇒ Object



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

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