Module: PgPipeline::DriverOps

Defined in:
lib/pg_pipeline/connection_driver.rb

Class Method Summary collapse

Class Method Details

.abort!(d, error) ⇒ Object



207
208
209
210
211
212
213
214
215
# File 'lib/pg_pipeline/connection_driver.rb', line 207

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

  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)
  nil
end

.abort_timed_out_health_probe(d, probe) ⇒ Object



217
218
219
220
221
222
223
# File 'lib/pg_pipeline/connection_driver.rb', line 217

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



378
379
380
381
382
383
# File 'lib/pg_pipeline/connection_driver.rb', line 378

def arm_writer(d)
  return if d.writer_armed

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

.clear_result(result) ⇒ Object



573
574
575
# File 'lib/pg_pipeline/connection_driver.rb', line 573

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

.complete_front(d, request) ⇒ Object

Raises:



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

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



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
428
429
430
431
432
433
434
435
436
437
438
439
# File 'lib/pg_pipeline/connection_driver.rb', line 391

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)


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

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:



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

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)


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

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



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

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



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

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)
  stop_watchers(d)
  safe_close_conn(d)
end

.finish_graceful_close(d) ⇒ Object



465
466
467
468
469
470
471
472
473
474
475
476
477
# File 'lib/pg_pipeline/connection_driver.rb', line 465

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
end

.flush_output(d) ⇒ Object



362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
# File 'lib/pg_pipeline/connection_driver.rb', line 362

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



197
198
199
200
201
202
203
204
205
# File 'lib/pg_pipeline/connection_driver.rb', line 197

def graceful_close(d)
  return unless d.running

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

.indeterminate_error(error) ⇒ Object



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

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

.not_dispatched_error(error) ⇒ Object



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

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



276
277
278
279
280
281
# File 'lib/pg_pipeline/connection_driver.rb', line 276

def notify_requests(d)
  return if d.request_event_pending

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

.owner_loop(d) ⇒ Object



234
235
236
237
238
239
240
# File 'lib/pg_pipeline/connection_driver.rb', line 234

def owner_loop(d)
  process_event(d, d.events.dequeue) while d.running
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



149
150
151
152
153
154
155
156
# File 'lib/pg_pipeline/connection_driver.rb', line 149

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



242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
# File 'lib/pg_pipeline/connection_driver.rb', line 242

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



283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
# File 'lib/pg_pipeline/connection_driver.rb', line 283

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



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

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



129
130
131
132
133
# File 'lib/pg_pipeline/connection_driver.rb', line 129

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

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

.read_available(d) ⇒ Object



385
386
387
388
389
# File 'lib/pg_pipeline/connection_driver.rb', line 385

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



523
524
525
526
527
528
529
530
531
532
533
534
535
536
# File 'lib/pg_pipeline/connection_driver.rb', line 523

def reader_watcher(d)
  while d.running
    d.socket.wait_readable
    break unless d.running

    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

.reusable_after_send_rejection?(d) ⇒ Boolean

Returns:

  • (Boolean)


354
355
356
357
358
359
360
# File 'lib/pg_pipeline/connection_driver.rb', line 354

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



567
568
569
570
571
# File 'lib/pg_pipeline/connection_driver.rb', line 567

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

.send_command(conn, request) ⇒ Object



337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
# File 'lib/pg_pipeline/connection_driver.rb', line 337

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



307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
# File 'lib/pg_pipeline/connection_driver.rb', line 307

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, parent) ⇒ Object



158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
# File 'lib/pg_pipeline/connection_driver.rb', line 158

def start(d, parent)
  raise Error, "driver already started" if d.running

  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
rescue Exception
  d.accepting = false
  d.running = false
  stop_watchers(d)
  safe_close_conn(d)
  raise
end

.stop_watchers(d) ⇒ Object



552
553
554
555
556
557
558
559
560
561
562
563
564
565
# File 'lib/pg_pipeline/connection_driver.rb', line 552

def stop_watchers(d)
  reader = d.reader_task
  writer = d.writer_task
  d.reader_task = nil
  d.writer_task = nil

  [reader, writer].each do |task|
    task&.stop
  rescue Async::Cancel, StandardError
    nil
  end

  nil
end

.submit(d, request) ⇒ Object

Raises:



180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
# File 'lib/pg_pipeline/connection_driver.rb', line 180

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

.warn_flush_coupling_once(caps) ⇒ Object



135
136
137
138
139
140
141
142
143
144
145
146
147
# File 'lib/pg_pipeline/connection_driver.rb', line 135

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

.writer_watcher(d) ⇒ Object



538
539
540
541
542
543
544
545
546
547
548
549
550
# File 'lib/pg_pipeline/connection_driver.rb', line 538

def writer_watcher(d)
  while d.running
    command = d.writer_commands.dequeue
    break unless command == :wait_writable && d.running

    d.socket.wait_writable
    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