pg_pipeline 0.3.2 → 0.3.4

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,14 +1,63 @@
1
1
  # frozen_string_literal: true
2
2
 
3
- require "pg"
3
+ require 'pg'
4
4
 
5
- require_relative "errors"
6
- require_relative "runtime"
7
- require_relative "connection_driver"
8
- require_relative "prepared_statement"
9
- require_relative "server_caps"
5
+ require_relative 'errors'
6
+ require_relative 'runtime'
7
+ require_relative 'connection_driver'
8
+ require_relative 'prepared_statement'
9
+ require_relative 'type_maps'
10
+ require_relative 'server_caps'
10
11
 
11
12
  module PgPipeline
13
+ PoolSettings = Data.define(
14
+ :pipeline_size, :pinned_size, :max_pending, :max_in_flight, :reconnect,
15
+ :reconnect_interval, :reconnect_backoff_max, :health_check, :health_interval,
16
+ :health_timeout, :cancel_pinned_on_abort
17
+ ) do
18
+ def self.build(options)
19
+ unknown = options.keys - members
20
+ raise ArgumentError, "unknown keyword: #{unknown.first}" unless unknown.empty?
21
+
22
+ new(**normalize_settings(defaults.merge(options)))
23
+ end
24
+
25
+ def self.defaults
26
+ {
27
+ pipeline_size: Pool::DEFAULT_PIPELINE_SIZE,
28
+ pinned_size: Pool::DEFAULT_PINNED_SIZE,
29
+ max_pending: ConnectionDriver::DEFAULT_MAX_PENDING,
30
+ max_in_flight: ConnectionDriver::DEFAULT_MAX_IN_FLIGHT,
31
+ reconnect: true,
32
+ reconnect_interval: Pool::DEFAULT_RECONNECT_INTERVAL,
33
+ reconnect_backoff_max: Pool::DEFAULT_RECONNECT_BACKOFF_MAX,
34
+ health_check: true,
35
+ health_interval: Pool::DEFAULT_HEALTH_INTERVAL,
36
+ health_timeout: Pool::DEFAULT_HEALTH_TIMEOUT,
37
+ cancel_pinned_on_abort: true
38
+ }
39
+ end
40
+
41
+ def self.normalize_settings(values)
42
+ {
43
+ pipeline_size: PoolOps.positive_integer!(values.fetch(:pipeline_size), :pipeline_size),
44
+ pinned_size: PoolOps.nonnegative_integer!(values.fetch(:pinned_size), :pinned_size),
45
+ max_pending: values.fetch(:max_pending),
46
+ max_in_flight: values.fetch(:max_in_flight),
47
+ reconnect: values.fetch(:reconnect),
48
+ reconnect_interval: PoolOps.finite_float!(values.fetch(:reconnect_interval), :reconnect_interval),
49
+ reconnect_backoff_max: PoolOps.finite_float!(values.fetch(:reconnect_backoff_max), :reconnect_backoff_max),
50
+ health_check: values.fetch(:health_check),
51
+ health_interval: PoolOps.finite_float!(values.fetch(:health_interval), :health_interval, allow_zero: true),
52
+ health_timeout: PoolOps.finite_float!(values.fetch(:health_timeout), :health_timeout),
53
+ cancel_pinned_on_abort: values.fetch(:cancel_pinned_on_abort)
54
+ }
55
+ end
56
+
57
+ private_class_method :defaults, :normalize_settings
58
+ end
59
+ private_constant :PoolSettings
60
+
12
61
  class Pool
13
62
  DEFAULT_PIPELINE_SIZE = 4
14
63
  DEFAULT_PINNED_SIZE = 2
@@ -23,30 +72,20 @@ module PgPipeline
23
72
 
24
73
  attr_reader :reconnects, :pipeline_size
25
74
 
26
- def initialize(connection_args,
27
- pipeline_size: DEFAULT_PIPELINE_SIZE,
28
- pinned_size: DEFAULT_PINNED_SIZE,
29
- max_pending: ConnectionDriver::DEFAULT_MAX_PENDING,
30
- max_in_flight: ConnectionDriver::DEFAULT_MAX_IN_FLIGHT,
31
- reconnect: true,
32
- reconnect_interval: DEFAULT_RECONNECT_INTERVAL,
33
- reconnect_backoff_max: DEFAULT_RECONNECT_BACKOFF_MAX,
34
- health_check: true,
35
- health_interval: DEFAULT_HEALTH_INTERVAL,
36
- health_timeout: DEFAULT_HEALTH_TIMEOUT,
37
- cancel_pinned_on_abort: true)
75
+ def initialize(connection_args, **options)
38
76
  @connection_args = connection_args
39
- @pipeline_size = PoolOps.positive_integer!(pipeline_size, :pipeline_size)
40
- @pinned_size = PoolOps.nonnegative_integer!(pinned_size, :pinned_size)
41
- @max_pending = max_pending
42
- @max_in_flight = max_in_flight
43
- @reconnect = reconnect
44
- @reconnect_interval = PoolOps.finite_float!(reconnect_interval, :reconnect_interval)
45
- @reconnect_backoff_max = PoolOps.finite_float!(reconnect_backoff_max, :reconnect_backoff_max)
46
- @health_check = health_check
47
- @health_interval = PoolOps.finite_float!(health_interval, :health_interval, allow_zero: true)
48
- @health_timeout = PoolOps.finite_float!(health_timeout, :health_timeout)
49
- @cancel_pinned_on_abort = cancel_pinned_on_abort
77
+ settings = PoolSettings.build(options)
78
+ @pipeline_size = settings.pipeline_size
79
+ @pinned_size = settings.pinned_size
80
+ @max_pending = settings.max_pending
81
+ @max_in_flight = settings.max_in_flight
82
+ @reconnect = settings.reconnect
83
+ @reconnect_interval = settings.reconnect_interval
84
+ @reconnect_backoff_max = settings.reconnect_backoff_max
85
+ @health_check = settings.health_check
86
+ @health_interval = settings.health_interval
87
+ @health_timeout = settings.health_timeout
88
+ @cancel_pinned_on_abort = settings.cancel_pinned_on_abort
50
89
 
51
90
  @drivers = []
52
91
  @driver_backoff = []
@@ -64,6 +103,7 @@ module PgPipeline
64
103
  @prepared_statements = {}
65
104
  @prepared_generation = 0
66
105
  @statement_sequence = 0
106
+ @type_maps = TypeMaps.new
67
107
 
68
108
  @pinned_free = []
69
109
  @pinned_in_use = {}
@@ -79,20 +119,10 @@ module PgPipeline
79
119
  end
80
120
 
81
121
  def start
82
- raise Error, "pool already started" if @started
83
- if @closing || @closed
84
- raise ShutdownError, "pool is closing or was closed and cannot be restarted; create a new Pool"
85
- end
86
- raise Error, "pool start requires an active Fiber scheduler" unless Fiber.scheduler
122
+ validate_start!
87
123
 
88
124
  begin
89
- @pipeline_size.times { @drivers << start_pipeline_driver }
90
- @driver_backoff = Array.new(@drivers.size, 0.0)
91
- @driver_attempts = Array.new(@drivers.size, 0)
92
- @driver_last_health = Array.new(@drivers.size, monotonic)
93
- @pinned_gate = Runtime::Semaphore.new(@pinned_size) if @pinned_size.positive?
94
- @started = true
95
- @supervisor = Runtime.spawn(name: :supervisor) { supervise } if @reconnect || @health_check
125
+ start_runtime
96
126
  rescue Exception
97
127
  cleanup_partial_start
98
128
  raise
@@ -101,6 +131,35 @@ module PgPipeline
101
131
  self
102
132
  end
103
133
 
134
+ private
135
+
136
+ def validate_start!
137
+ raise Error, 'pool already started' if @started
138
+ if @closing || @closed
139
+ raise ShutdownError,
140
+ 'pool is closing or was closed and cannot be restarted; create a new Pool'
141
+ end
142
+ raise Error, 'pool start requires an active Fiber scheduler' unless Fiber.scheduler
143
+ end
144
+
145
+ def start_runtime
146
+ @pipeline_size.times { @drivers << start_pipeline_driver }
147
+ driver_count = @drivers.size
148
+ @driver_backoff = Array.new(driver_count, 0.0)
149
+ @driver_attempts = Array.new(driver_count, 0)
150
+ @driver_last_health = Array.new(driver_count, monotonic)
151
+ pinned_size = @pinned_size
152
+ @pinned_gate = Runtime::Semaphore.new(pinned_size) if pinned_size.positive?
153
+ @started = true
154
+ start_supervisor if @reconnect || @health_check
155
+ end
156
+
157
+ def start_supervisor
158
+ @supervisor = Runtime.spawn(name: :supervisor) { supervise }
159
+ end
160
+
161
+ public
162
+
104
163
  def closing? = @closing
105
164
 
106
165
  def pipeline_driver
@@ -108,97 +167,104 @@ module PgPipeline
108
167
 
109
168
  driver = PoolOps.select_driver_into(@drivers, @rr, @rr_slot)
110
169
  @rr = @rr_slot[0]
111
- raise NotDispatchedError, "no live pipeline connections; request was not dispatched" unless driver
170
+ raise NotDispatchedError, 'no live pipeline connections; request was not dispatched' unless driver
112
171
 
113
172
  driver
114
173
  end
115
174
 
116
- def prepare_statement(client, name, sql, param_types)
175
+ def bind_type_map(request, statement)
176
+ return unless statement.typed?
177
+
178
+ maps = @type_maps
179
+ return unless maps
180
+
181
+ maps.bind(request, statement)
182
+ end
183
+
184
+ def prepare_statement(client, name, sql, param_types, typed: false)
117
185
  ensure_available!
118
186
  logical_name = PreparedStatementOps.snapshot_name(name)
119
187
  if @prepared_statements.key?(logical_name)
120
188
  raise Error, "prepared statement #{logical_name.inspect} already exists"
121
189
  end
122
190
 
191
+ statement = register_prepared_statement(client, logical_name, sql, param_types, typed)
192
+ requests = []
193
+
194
+ begin
195
+ requests = dispatch_prepare(statement)
196
+ await_prepares!(requests)
197
+ statement
198
+ rescue Exception
199
+ drop_prepared_statement(logical_name, statement)
200
+ raise
201
+ ensure
202
+ requests.each { |request| request.cancel! unless request.settled? }
203
+ end
204
+ end
205
+
206
+ def register_prepared_statement(client, logical_name, sql, param_types, typed)
123
207
  @statement_sequence += 1
124
208
  statement = PreparedStatement.new(
125
209
  client: client,
126
210
  name: logical_name,
127
211
  physical_name: "pgp_#{@statement_sequence.to_s(36)}",
128
212
  sql: sql,
129
- param_types: param_types
213
+ param_types: param_types,
214
+ typed: typed
130
215
  )
131
216
 
132
217
  @prepared_statements[logical_name] = statement
218
+ register_typed_maps(statement) if statement.typed?
133
219
  @prepared_generation += 1
134
- requests = []
220
+ statement
221
+ end
135
222
 
136
- begin
137
- drivers = @drivers.select(&:available?)
138
- if drivers.empty?
139
- raise NotDispatchedError, "no live pipeline connections; statement was not prepared"
140
- end
223
+ def register_typed_maps(statement)
224
+ @type_maps ||= TypeMaps.new
225
+ @type_maps.register(statement)
226
+ @type_maps.ensure_bundle!(@connection_args)
227
+ end
141
228
 
142
- requests = drivers.map do |driver|
143
- Request.prepare(statement).tap { |request| driver.submit(request) }
144
- end
229
+ def dispatch_prepare(statement)
230
+ drivers = @drivers.select(&:available?)
231
+ raise NotDispatchedError, 'no live pipeline connections; statement was not prepared' if drivers.empty?
145
232
 
146
- first_error = nil
147
- requests.each do |request|
148
- begin
149
- RequestOps.clear_result(request.wait)
150
- rescue StandardError => e
151
- if first_error
152
- e.clear_result! if e.respond_to?(:clear_result!)
153
- else
154
- first_error = e
155
- end
156
- end
157
- end
158
- raise first_error if first_error
233
+ drivers.map { |driver| Request.prepare(statement).tap { |r| driver.submit(r) } }
234
+ end
159
235
 
160
- statement
161
- rescue Exception
162
- if @prepared_statements.delete(logical_name)
163
- @prepared_generation += 1
236
+ def await_prepares!(requests)
237
+ first_error = nil
238
+ requests.each do |request|
239
+ RequestOps.clear_result(request.wait)
240
+ rescue StandardError => e
241
+ if first_error
242
+ e.clear_result! if e.respond_to?(:clear_result!)
243
+ else
244
+ first_error = e
164
245
  end
165
- raise
166
- ensure
167
- requests.each { |request| request.cancel! unless request.settled? }
168
246
  end
247
+ raise first_error if first_error
248
+ end
249
+
250
+ def drop_prepared_statement(logical_name, statement)
251
+ return unless @prepared_statements.delete(logical_name)
252
+
253
+ @type_maps.unregister(statement.physical_name) if statement&.typed? && @type_maps
254
+ @prepared_generation += 1
169
255
  end
170
256
 
171
- def with_pinned
257
+ def with_pinned(&block)
172
258
  ensure_available!
173
259
  raise @pinned_error if @pinned_error
174
- raise Error, "pinned pool is disabled (pinned_size=0)" if @pinned_size.zero?
260
+ raise Error, 'pinned pool is disabled (pinned_size=0)' if @pinned_size.zero?
175
261
 
176
262
  owner = Fiber.current
177
- if @pinned_owners[owner].positive?
178
- raise RecursiveCheckoutError,
179
- "nested Client#session/transaction on the same fiber is not allowed; " \
180
- "use Transaction#savepoint for nested atomicity"
181
- end
263
+ assert_not_nested_pinned!(owner)
182
264
 
183
265
  @pinned_gate.acquire do
184
- raise @pinned_error if @pinned_error
185
- raise ShutdownError, "pool is closing" if @closing
186
-
187
- conn = nil
188
- @pinned_active += 1
189
- @pinned_owners[owner] += 1
190
-
191
- begin
192
- conn = @pinned_free.pop || PoolOps.new_connection(@connection_args)
193
- raise @pinned_error if @pinned_error
194
- raise ShutdownError, "pool is closing" if @closing
195
-
196
- @pinned_in_use[conn] = owner
197
- yield conn
198
- ensure
199
- @pinned_in_use.delete(conn) if conn
200
- release_pinned(owner, conn)
201
- end
266
+ assert_pinned_open!
267
+ run_pinned(owner, &block)
202
268
  end
203
269
  end
204
270
 
@@ -230,7 +296,7 @@ module PgPipeline
230
296
  return unless @started
231
297
 
232
298
  if @pinned_owners[Fiber.current].positive?
233
- raise Error, "cannot close the pool from inside Client#session/transaction"
299
+ raise Error, 'cannot close the pool from inside Client#session/transaction'
234
300
  end
235
301
 
236
302
  @closing = true
@@ -265,9 +331,7 @@ module PgPipeline
265
331
 
266
332
  begin
267
333
  first_error = preferred_shutdown_error(first_error, stop_supervisor)
268
- if @cancel_pinned_on_abort
269
- first_error = preferred_shutdown_error(first_error, cancel_in_use_pinned)
270
- end
334
+ first_error = preferred_shutdown_error(first_error, cancel_in_use_pinned) if @cancel_pinned_on_abort
271
335
  first_error = preferred_shutdown_error(first_error, close_all_drivers(:abort!))
272
336
  ensure
273
337
  close_free_pinned
@@ -304,7 +368,7 @@ module PgPipeline
304
368
  @supervisor_error = e
305
369
  @pause_failures = (@pause_failures || 0) + 1
306
370
 
307
- unless ENV["PG_PIPELINE_SILENCE_WARNINGS"]
371
+ unless ENV['PG_PIPELINE_SILENCE_WARNINGS']
308
372
  warn("pg_pipeline: supervisor pause failed (#{e.class}: #{e.message}) " \
309
373
  "[#{@pause_failures}/#{MAX_PAUSE_FAILURES}]")
310
374
  end
@@ -318,11 +382,13 @@ module PgPipeline
318
382
  end
319
383
 
320
384
  def supervisor_sleep_interval
385
+ reconnect_interval = @reconnect_interval
386
+ health_interval = @health_interval
321
387
  candidates = []
322
- candidates << @reconnect_interval if @reconnect
323
- candidates << @health_interval if @health_check && @health_interval.positive?
324
- interval = candidates.min || @reconnect_interval
325
- interval.positive? ? interval : @reconnect_interval
388
+ candidates << reconnect_interval if @reconnect
389
+ candidates << health_interval if @health_check && health_interval.positive?
390
+ interval = candidates.min || reconnect_interval
391
+ interval.positive? ? interval : reconnect_interval
326
392
  end
327
393
 
328
394
  def reap_and_replace
@@ -330,46 +396,61 @@ module PgPipeline
330
396
  ensure_driver_slots!
331
397
 
332
398
  @drivers.each_index do |index|
333
- driver = @drivers[index]
334
- next if driver.available?
335
- next unless driver.dead?
336
- next if now < (@driver_backoff[index] || 0.0)
399
+ replace_unavailable_driver(index, now)
400
+ end
401
+ end
337
402
 
338
- begin
339
- @drivers[index] = start_pipeline_driver
340
- @driver_backoff[index] = 0.0
341
- @driver_attempts[index] = 0
342
- @driver_last_health[index] = monotonic
343
- @reconnects += 1
344
- rescue StandardError
345
- @driver_attempts[index] = (@driver_attempts[index] || 0) + 1
346
- @driver_backoff[index] = now + next_backoff(@driver_attempts[index])
347
- end
403
+ def replace_unavailable_driver(index, now)
404
+ driver = @drivers[index]
405
+ return if driver.available? || !driver.dead? || now < (@driver_backoff[index] || 0.0)
406
+
407
+ begin
408
+ @drivers[index] = start_pipeline_driver
409
+ reset_driver_retry_state(index)
410
+ rescue StandardError
411
+ backoff_driver_retry(index, now)
348
412
  end
349
413
  end
350
414
 
415
+ def reset_driver_retry_state(index)
416
+ @driver_backoff[index] = 0.0
417
+ @driver_attempts[index] = 0
418
+ @driver_last_health[index] = monotonic
419
+ @reconnects += 1
420
+ end
421
+
422
+ def backoff_driver_retry(index, now)
423
+ attempts = (@driver_attempts[index] || 0) + 1
424
+ @driver_attempts[index] = attempts
425
+ @driver_backoff[index] = now + next_backoff(attempts)
426
+ end
427
+
351
428
  def health_probe
352
429
  now = monotonic
353
430
  ensure_driver_slots!
354
431
 
355
432
  @drivers.each_index do |index|
356
- driver = @drivers[index]
357
- next unless driver.available?
358
- next unless driver.load.zero?
359
- next if now - (@driver_last_health[index] || 0.0) < @health_interval
433
+ probe_driver_health(index, now)
434
+ end
435
+ end
360
436
 
361
- @driver_last_health[index] = now
437
+ def probe_driver_health(index, now)
438
+ driver = @drivers[index]
439
+ return unless driver.available? && driver.load.zero?
440
+ return if now - (@driver_last_health[index] || 0.0) < @health_interval
362
441
 
363
- begin
364
- next if driver.health_check(@health_timeout)
442
+ @driver_last_health[index] = now
443
+ run_driver_health_check(driver)
444
+ end
365
445
 
366
- @health_failures += 1
367
- driver.abort!
368
- rescue StandardError => e
369
- @health_failures += 1
370
- @supervisor_error = e
371
- end
372
- end
446
+ def run_driver_health_check(driver)
447
+ return if driver.health_check(@health_timeout)
448
+
449
+ @health_failures += 1
450
+ driver.abort!
451
+ rescue StandardError => e
452
+ @health_failures += 1
453
+ @supervisor_error = e
373
454
  end
374
455
 
375
456
  def ensure_driver_slots!
@@ -382,13 +463,14 @@ module PgPipeline
382
463
  def next_backoff(attempts)
383
464
  exponent = [Integer(attempts) - 1, 0].max
384
465
  delay = @reconnect_interval * (2.0**exponent)
385
- delay.finite? && delay < @reconnect_backoff_max ? delay : @reconnect_backoff_max
466
+ maximum = @reconnect_backoff_max
467
+ delay.finite? && delay < maximum ? delay : maximum
386
468
  end
387
469
 
388
470
  def cancel_in_use_pinned
389
471
  cancellation = nil
390
472
 
391
- @pinned_in_use.keys.each do |conn|
473
+ @pinned_in_use.each_key do |conn|
392
474
  conn.cancel if conn.respond_to?(:cancel)
393
475
  rescue CANCEL_SIGNAL => e
394
476
  cancellation ||= e
@@ -409,7 +491,7 @@ module PgPipeline
409
491
  begin
410
492
  supervisor.wait(SUPERVISOR_JOIN_TIMEOUT)
411
493
  rescue Runtime::TimeoutError
412
- unless ENV["PG_PIPELINE_SILENCE_WARNINGS"]
494
+ unless ENV['PG_PIPELINE_SILENCE_WARNINGS']
413
495
  warn("pg_pipeline: supervisor did not exit within #{SUPERVISOR_JOIN_TIMEOUT}s and has been leaked")
414
496
  end
415
497
  rescue CANCEL_SIGNAL, StandardError
@@ -426,13 +508,11 @@ module PgPipeline
426
508
  cancellation = nil
427
509
 
428
510
  @drivers.each do |driver|
429
- begin
430
- driver.public_send(method_name)
431
- rescue CANCEL_SIGNAL => e
432
- cancellation ||= e
433
- rescue StandardError => e
434
- first_error ||= e
435
- end
511
+ driver.public_send(method_name)
512
+ rescue CANCEL_SIGNAL => e
513
+ cancellation ||= e
514
+ rescue StandardError => e
515
+ first_error ||= e
436
516
  end
437
517
 
438
518
  cancellation || first_error
@@ -446,14 +526,15 @@ module PgPipeline
446
526
  end
447
527
 
448
528
  def ensure_available!
449
- raise ShutdownError, "pool is closing or closed" if @closing || @closed
450
- raise Error, "pool not started" unless @started
529
+ raise ShutdownError, 'pool is closing or closed' if @closing || @closed
530
+ raise Error, 'pool not started' unless @started
451
531
  end
452
532
 
453
533
  def start_pipeline_driver
454
534
  conn = PoolOps.new_connection(@connection_args)
455
535
  prepare_registered_statements(conn)
456
- driver = ConnectionDriver.new(conn, max_pending: @max_pending, max_in_flight: @max_in_flight)
536
+ driver = ConnectionDriver.new(conn, max_pending: @max_pending, max_in_flight: @max_in_flight,
537
+ type_maps: @type_maps)
457
538
  driver.start
458
539
  rescue Exception
459
540
  PoolOps.safe_close(conn) if conn
@@ -466,71 +547,115 @@ module PgPipeline
466
547
  loop do
467
548
  generation = @prepared_generation || 0
468
549
  statements = (@prepared_statements || {}).values.dup
550
+ prepare_statements(conn, statements, prepared)
469
551
 
470
- statements.each do |statement|
471
- next if prepared.key?(statement.physical_name)
552
+ break if generation == (@prepared_generation || 0)
553
+ end
554
+ end
472
555
 
473
- result = if statement.param_types.nil?
474
- conn.prepare(statement.physical_name, statement.sql)
475
- else
476
- conn.prepare(statement.physical_name, statement.sql, statement.param_types)
477
- end
478
- RequestOps.clear_result(result)
479
- prepared[statement.physical_name] = true
480
- end
556
+ def prepare_statements(conn, statements, prepared)
557
+ statements.each do |statement|
558
+ prepare_statement_on_connection(conn, statement, prepared) unless prepared.key?(statement.physical_name)
559
+ end
560
+ end
481
561
 
482
- break if generation == (@prepared_generation || 0)
562
+ def prepare_statement_on_connection(conn, statement, prepared)
563
+ param_types = statement.param_types
564
+ name = statement.physical_name
565
+ sql = statement.sql
566
+ result = if param_types.nil?
567
+ conn.prepare(name, sql)
568
+ else
569
+ conn.prepare(name, sql, param_types)
570
+ end
571
+ RequestOps.clear_result(result)
572
+ prepared[name] = true
573
+ end
574
+
575
+ def assert_not_nested_pinned!(owner)
576
+ return unless @pinned_owners[owner].positive?
577
+
578
+ raise RecursiveCheckoutError,
579
+ 'nested Client#session/transaction on the same fiber is not allowed; ' \
580
+ 'use Transaction#savepoint for nested atomicity'
581
+ end
582
+
583
+ def assert_pinned_open!
584
+ raise @pinned_error if @pinned_error
585
+ raise ShutdownError, 'pool is closing' if @closing
586
+ end
587
+
588
+ def run_pinned(owner)
589
+ conn = nil
590
+ @pinned_active += 1
591
+ @pinned_owners[owner] += 1
592
+
593
+ begin
594
+ conn = take_pinned_connection
595
+ assert_pinned_open!
596
+ @pinned_in_use[conn] = owner
597
+ yield conn
598
+ ensure
599
+ @pinned_in_use.delete(conn) if conn
600
+ release_pinned(owner, conn)
483
601
  end
484
602
  end
485
603
 
604
+ def take_pinned_connection
605
+ @pinned_free.pop || PoolOps.new_connection(@connection_args)
606
+ end
607
+
486
608
  def release_pinned(owner, conn)
487
- if conn
488
- if @closing
489
- PoolOps.safe_close(conn)
490
- else
491
- recycled = recycle_pinned_connection(conn)
492
-
493
- if recycled
494
- if @closing
495
- PoolOps.safe_close(recycled)
496
- else
497
- @pinned_free << recycled
498
- end
499
- end
500
- end
501
- end
609
+ return_pinned_connection(conn) if conn
502
610
  ensure
503
611
  @pinned_active -= 1
504
- @pinned_owners[owner] -= 1
505
- @pinned_owners.delete(owner) if @pinned_owners[owner].zero?
612
+ owner_count = @pinned_owners[owner] - 1
613
+ @pinned_owners[owner] = owner_count
614
+ @pinned_owners.delete(owner) if owner_count.zero?
506
615
  @pinned_idle.signal if @pinned_active.zero?
507
616
  end
508
617
 
509
- def recycle_pinned_connection(conn)
510
- if conn.server_version < DISCARD_SEQUENCES_SERVER_VERSION
618
+ def return_pinned_connection(conn)
619
+ if @closing
511
620
  PoolOps.safe_close(conn)
512
- return PoolOps.new_connection(@connection_args)
621
+ return
622
+ end
623
+
624
+ recycled = recycle_pinned_connection(conn)
625
+ return unless recycled
626
+
627
+ if @closing
628
+ PoolOps.safe_close(recycled)
629
+ else
630
+ @pinned_free << recycled
513
631
  end
632
+ end
633
+
634
+ def recycle_pinned_connection(conn)
635
+ return replace_pinned_connection(conn) if conn.server_version < DISCARD_SEQUENCES_SERVER_VERSION
514
636
 
515
637
  PoolOps.sanitize_pinned_connection(conn)
516
638
  conn
517
639
  rescue CANCEL_SIGNAL
518
640
  PoolOps.safe_close(conn)
519
641
  raise
520
- rescue StandardError => cleanup_error
521
- PoolOps.safe_close(conn)
522
-
642
+ rescue StandardError => e
523
643
  begin
524
- PoolOps.new_connection(@connection_args)
644
+ replace_pinned_connection(conn)
525
645
  rescue StandardError => replacement_error
526
646
  @pinned_error = ConnectionLostError.new(
527
- "pinned connection reset failed (#{cleanup_error.class}: #{cleanup_error.message}) " \
647
+ "pinned connection reset failed (#{e.class}: #{e.message}) " \
528
648
  "and replacement failed (#{replacement_error.class}: #{replacement_error.message})"
529
649
  )
530
650
  nil
531
651
  end
532
652
  end
533
653
 
654
+ def replace_pinned_connection(conn)
655
+ PoolOps.safe_close(conn)
656
+ PoolOps.new_connection(@connection_args)
657
+ end
658
+
534
659
  def wait_for_pinned_idle
535
660
  @pinned_idle.wait while @pinned_active.positive?
536
661
  end
@@ -564,14 +689,7 @@ module PgPipeline
564
689
  module PoolOps
565
690
  module_function
566
691
 
567
- def positive_integer!(value, name)
568
- integer = Integer(value)
569
- raise ArgumentError, "#{name} must be >= 1" if integer < 1
570
-
571
- integer
572
- rescue ArgumentError, TypeError
573
- raise ArgumentError, "#{name} must be an integer >= 1"
574
- end
692
+ def positive_integer!(value, name) = DriverOps.positive_integer!(value, name)
575
693
 
576
694
  def finite_float!(value, name, allow_zero: false)
577
695
  number = Float(value)
@@ -580,7 +698,7 @@ module PgPipeline
580
698
 
581
699
  number
582
700
  rescue ArgumentError, TypeError
583
- requirement = allow_zero ? "non-negative finite" : "positive finite"
701
+ requirement = allow_zero ? 'non-negative finite' : 'positive finite'
584
702
  raise ArgumentError, "#{name} must be a #{requirement} number (got #{value.inspect})"
585
703
  end
586
704
 
@@ -593,14 +711,14 @@ module PgPipeline
593
711
  raise ArgumentError, "#{name} must be an integer >= 0"
594
712
  end
595
713
 
596
- def select_driver_into(drivers, rr, slot)
714
+ def select_driver_into(drivers, round_robin, slot)
597
715
  size = drivers.length
598
716
  if size.zero?
599
- slot[0] = rr
717
+ slot[0] = round_robin
600
718
  return nil
601
719
  end
602
720
 
603
- start = rr % size
721
+ start = round_robin % size
604
722
  best = nil
605
723
  best_index = 0
606
724
  best_load = 0
@@ -621,13 +739,13 @@ module PgPipeline
621
739
  best_load = load
622
740
  end
623
741
 
624
- slot[0] = best ? (best_index + 1) % size : rr
742
+ slot[0] = best ? (best_index + 1) % size : round_robin
625
743
  best
626
744
  end
627
745
 
628
- def select_driver(drivers, rr)
629
- slot = [rr]
630
- driver = select_driver_into(drivers, rr, slot)
746
+ def select_driver(drivers, round_robin)
747
+ slot = [round_robin]
748
+ driver = select_driver_into(drivers, round_robin, slot)
631
749
  [driver, slot[0]]
632
750
  end
633
751
 
@@ -636,20 +754,21 @@ module PgPipeline
636
754
  end
637
755
 
638
756
  def sanitize_pinned_connection(conn)
639
- raise ConnectionLostError, "pinned connection is closed" if conn.finished?
640
- raise ConnectionLostError, "pinned connection is bad" unless conn.status == PG::CONNECTION_OK
757
+ raise ConnectionLostError, 'pinned connection is closed' if conn.finished?
758
+ raise ConnectionLostError, 'pinned connection is bad' unless conn.status == PG::CONNECTION_OK
641
759
 
642
- case conn.transaction_status
760
+ transaction_status = conn.transaction_status
761
+ case transaction_status
643
762
  when PG::PQTRANS_IDLE
644
763
  nil
645
764
  when PG::PQTRANS_INTRANS, PG::PQTRANS_INERROR
646
- conn.exec("ROLLBACK")
765
+ conn.exec('ROLLBACK')
647
766
  else
648
767
  raise ConnectionLostError,
649
- "pinned connection returned in unsafe transaction state #{conn.transaction_status}"
768
+ "pinned connection returned in unsafe transaction state #{transaction_status}"
650
769
  end
651
770
 
652
- conn.exec("DISCARD ALL")
771
+ conn.exec('DISCARD ALL')
653
772
  true
654
773
  end
655
774