Module: PgPipeline::DriverOps

Defined in:
lib/pg_pipeline/connection_driver.rb

Constant Summary collapse

SUCCESS_STATUSES = 
[
  PG::PGRES_EMPTY_QUERY,
  PG::PGRES_COMMAND_OK,
  PG::PGRES_TUPLES_OK
].freeze
COPY_STATUSES = 
[
  PG::PGRES_COPY_IN,
  PG::PGRES_COPY_OUT,
  PG::PGRES_COPY_BOTH
].freeze

Class Method Summary collapse

Class Method Details

.abort!(d, error) ⇒ Object 



171
172
173
174
175
176
177
178
179
 
# File 'lib/pg_pipeline/connection_driver.rb', line 171

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 



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

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 



340
341
342
343
344
345
 
# File 'lib/pg_pipeline/connection_driver.rb', line 340

def arm_writer(d)
  return if d.writer_armed

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

.clear_result(result) ⇒ Object 



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

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

.complete_front(d, request) ⇒ Object

Raises:



399
400
401
402
403
404
 
# File 'lib/pg_pipeline/connection_driver.rb', line 399

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

  d.inflight.shift
  request.finish!
end

.drain_results(d) ⇒ Object 



353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
 
# File 'lib/pg_pipeline/connection_driver.rb', line 353

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

    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}"
    end
  end
end

.drained?(d) ⇒ Boolean

Returns:

  • (Boolean) 


412
413
414
 
# File 'lib/pg_pipeline/connection_driver.rb', line 412

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:



393
394
395
396
397
 
# File 'lib/pg_pipeline/connection_driver.rb', line 393

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) 


189
190
191
192
193
194
195
196
 
# File 'lib/pg_pipeline/connection_driver.rb', line 189

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 



441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
 
# File 'lib/pg_pipeline/connection_driver.rb', line 441

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 



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

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 



416
417
418
419
420
421
422
423
424
425
426
427
428
 
# File 'lib/pg_pipeline/connection_driver.rb', line 416

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 



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

def flush_output(d)
  return unless d.running

  if d.conn.sync_flush
    d.needs_flush = false
  else
    d.needs_flush = true
    arm_writer(d)
  end
rescue PG::Error => e
  raise ConnectionLostError, "flush failed: #{e.class}: #{e.message}"
end

.graceful_close(d) ⇒ Object 



161
162
163
164
165
166
167
168
169
 
# File 'lib/pg_pipeline/connection_driver.rb', line 161

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 



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

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 



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

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 



239
240
241
242
243
244
 
# File 'lib/pg_pipeline/connection_driver.rb', line 239

def notify_requests(d)
  return if d.request_event_pending

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

.owner_loop(d) ⇒ Object 



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

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 



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

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 



206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
 
# File 'lib/pg_pipeline/connection_driver.rb', line 206

def process_event(d, event)
  input_changed = false

  case event
  when :requests
    d.request_event_pending = false
  when :submission_finished
    nil
  when :readable
    begin
      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 



246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
 
# File 'lib/pg_pipeline/connection_driver.rb', line 246

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
    dispatched = true
  end

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

.query_error(result) ⇒ Object 



406
407
408
409
410
 
# File 'lib/pg_pipeline/connection_driver.rb', line 406

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

.read_available(d) ⇒ Object 



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

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 



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

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) 


319
320
321
322
323
324
325
 
# File 'lib/pg_pipeline/connection_driver.rb', line 319

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 



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

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

.send_command(conn, request) ⇒ Object 



302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
 
# File 'lib/pg_pipeline/connection_driver.rb', line 302

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 



269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
 
# File 'lib/pg_pipeline/connection_driver.rb', line 269

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
    # 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

  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 



122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
 
# File 'lib/pg_pipeline/connection_driver.rb', line 122

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 



503
504
505
506
507
508
509
510
511
512
513
514
515
516
 
# File 'lib/pg_pipeline/connection_driver.rb', line 503

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:



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

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

.writer_watcher(d) ⇒ Object 



489
490
491
492
493
494
495
496
497
498
499
500
501
 
# File 'lib/pg_pipeline/connection_driver.rb', line 489

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