376 lines
10 KiB
V
376 lines
10 KiB
V
module pg
|
|
|
|
import sync
|
|
import time
|
|
|
|
// PoolStats reports the current state of a `DB`'s connection pool.
|
|
pub struct PoolStats {
|
|
pub:
|
|
max_open_connections int // configured limit (0 = unlimited)
|
|
open_connections int // total in-use + idle conns
|
|
in_use int // conns currently checked out
|
|
idle int // conns parked as idle
|
|
wait_count int // number of callers currently blocked on acquire
|
|
}
|
|
|
|
// PoolConfig configures pool behavior at construction time.
|
|
@[params]
|
|
pub struct PoolConfig {
|
|
pub:
|
|
max_open_conns int // 0 = unlimited
|
|
max_idle_conns int = 2 // 0 = keep no idle conns
|
|
conn_max_lifetime time.Duration // 0 = unlimited
|
|
}
|
|
|
|
// IdleSlot is the pool's internal representation of a pooled libpq handle.
|
|
// The pool tracks raw PGconn* metadata in idle/waiter channels; `acquire`
|
|
// allocates a fresh `&Conn` wrapping the slot and `release` extracts the
|
|
// metadata back. Keeping the wrapper separate from the slot means a stale
|
|
// `&Conn` kept by user code after `close()` cannot be revived even when
|
|
// the pool re-hands the same physical PGconn* to another caller.
|
|
struct IdleSlot {
|
|
handle voidptr
|
|
created_at time.Time
|
|
mut:
|
|
bad bool
|
|
}
|
|
|
|
@[heap]
|
|
struct Pool {
|
|
mut:
|
|
mu &sync.Mutex = unsafe { nil }
|
|
conninfo string
|
|
max_open int
|
|
max_idle int
|
|
max_lifetime time.Duration
|
|
idle []IdleSlot
|
|
open_count int
|
|
waiters []chan IdleSlot
|
|
closed bool
|
|
// last_ids stores the per-thread last-inserted id so DB.last_id() returns
|
|
// the value captured on the same pooled conn that ran the INSERT, instead
|
|
// of calling LASTVAL() on whatever conn the pool happens to hand out next
|
|
// (which is session-scoped and would return 0 or a wrong value).
|
|
last_ids_mu &sync.Mutex = unsafe { nil }
|
|
last_ids map[u64]i64
|
|
}
|
|
|
|
fn new_pool(conninfo string, cfg PoolConfig) &Pool {
|
|
mut p := &Pool{
|
|
mu: sync.new_mutex()
|
|
last_ids_mu: sync.new_mutex()
|
|
conninfo: conninfo
|
|
max_open: cfg.max_open_conns
|
|
max_idle: cfg.max_idle_conns
|
|
max_lifetime: cfg.conn_max_lifetime
|
|
}
|
|
if p.max_open < 0 {
|
|
p.max_open = 0
|
|
}
|
|
if p.max_idle < 0 {
|
|
p.max_idle = 0
|
|
}
|
|
return p
|
|
}
|
|
|
|
// wrap_slot builds the fresh `&Conn` that gets handed to a caller. Each
|
|
// acquire allocates a new wrapper so a previously-checked-out (and since
|
|
// released) `&Conn` reference can never be reused to reach the underlying
|
|
// PGconn*.
|
|
fn (p &Pool) wrap_slot(slot IdleSlot) &Conn {
|
|
return &Conn{
|
|
conn: slot.handle
|
|
pool: unsafe { p }
|
|
created_at: slot.created_at
|
|
bad: slot.bad
|
|
}
|
|
}
|
|
|
|
fn (mut p Pool) acquire() !&Conn {
|
|
for {
|
|
p.mu.lock()
|
|
if p.closed {
|
|
p.mu.unlock()
|
|
return error('pg: pool is closed')
|
|
}
|
|
// Pop newest idle slot (LIFO keeps the freshest connection warm)
|
|
for p.idle.len > 0 {
|
|
slot := p.idle.last()
|
|
p.idle.delete_last()
|
|
if slot_expired(slot, p.max_lifetime) || slot_bad(slot) {
|
|
physical_close_handle(slot.handle)
|
|
p.open_count--
|
|
continue
|
|
}
|
|
p.mu.unlock()
|
|
return p.wrap_slot(slot)
|
|
}
|
|
// Capacity available: open a new conn outside the lock
|
|
if p.max_open == 0 || p.open_count < p.max_open {
|
|
p.open_count++
|
|
conninfo := p.conninfo
|
|
p.mu.unlock()
|
|
slot := connect_slot(conninfo) or {
|
|
p.mu.lock()
|
|
p.open_count--
|
|
p.mu.unlock()
|
|
return err
|
|
}
|
|
// close() may have run while we were dialing outside the lock.
|
|
// Honor the close contract by tearing the fresh handle down
|
|
// instead of returning a live Conn after shutdown.
|
|
p.mu.lock()
|
|
if p.closed {
|
|
p.open_count--
|
|
p.mu.unlock()
|
|
physical_close_handle(slot.handle)
|
|
return error('pg: pool is closed')
|
|
}
|
|
p.mu.unlock()
|
|
return p.wrap_slot(slot)
|
|
}
|
|
// At capacity: wait for a release/close/cap-raise to signal us.
|
|
// Senders transfer slot ownership via this channel; a sentinel slot
|
|
// with a nil handle means "capacity changed, retry the acquire
|
|
// loop" (used by set_max_open when raising the limit).
|
|
waiter := chan IdleSlot{cap: 1}
|
|
p.waiters << waiter
|
|
p.mu.unlock()
|
|
slot := <-waiter or { return error('pg: pool was closed while waiting for connection') }
|
|
if isnil(slot.handle) {
|
|
continue
|
|
}
|
|
return p.wrap_slot(slot)
|
|
}
|
|
return error('pg: unreachable')
|
|
}
|
|
|
|
fn (mut p Pool) release(conn &Conn) {
|
|
if isnil(conn) {
|
|
return
|
|
}
|
|
mut c := unsafe { conn }
|
|
// Detach the wrapper from the physical handle before doing anything else.
|
|
// This is what makes a stale `&Conn` reference held by user code inert:
|
|
// any subsequent method call on it will see `c.conn == nil` and error.
|
|
// It is also the idempotency guard against double-close — a second
|
|
// release() finds nothing to return to the pool and is a no-op.
|
|
if isnil(c.conn) {
|
|
return
|
|
}
|
|
slot := IdleSlot{
|
|
handle: c.conn
|
|
created_at: c.created_at
|
|
bad: c.bad
|
|
}
|
|
c.conn = unsafe { nil }
|
|
c.pool = unsafe { nil }
|
|
p.mu.lock()
|
|
if p.closed {
|
|
// close() only decremented open_count for slots it found parked;
|
|
// in-use conns are accounted for here as they trickle back.
|
|
p.open_count--
|
|
p.mu.unlock()
|
|
physical_close_handle(slot.handle)
|
|
return
|
|
}
|
|
// Discard broken or expired conns
|
|
if slot_bad(slot) || slot_expired(slot, p.max_lifetime) {
|
|
physical_close_handle(slot.handle)
|
|
p.open_count--
|
|
// A capacity slot just opened; hand it to a waiter as a fresh conn
|
|
if p.waiters.len > 0 && (p.max_open == 0 || p.open_count < p.max_open) {
|
|
waiter := p.waiters[0]
|
|
p.waiters.delete(0)
|
|
p.open_count++
|
|
conninfo := p.conninfo
|
|
p.mu.unlock()
|
|
new_slot := connect_slot(conninfo) or {
|
|
p.mu.lock()
|
|
p.open_count--
|
|
// Capacity is now open but the dial just failed. Wake every
|
|
// other parked waiter too, otherwise they hang forever waiting
|
|
// for a release that will never come (e.g. max_open=1).
|
|
extras := p.waiters.clone()
|
|
p.waiters = []chan IdleSlot{}
|
|
p.mu.unlock()
|
|
waiter.close()
|
|
for w in extras {
|
|
w.close()
|
|
}
|
|
return
|
|
}
|
|
p.mu.lock()
|
|
if p.closed {
|
|
// Pool closed during the dial: drop the new conn and signal the waiter.
|
|
p.open_count--
|
|
p.mu.unlock()
|
|
physical_close_handle(new_slot.handle)
|
|
waiter.close()
|
|
return
|
|
}
|
|
waiter <- new_slot
|
|
p.mu.unlock()
|
|
return
|
|
}
|
|
p.mu.unlock()
|
|
return
|
|
}
|
|
// If a recent set_max_open() shrank the cap below open_count, this
|
|
// returning slot has to be retired even when a waiter is queued —
|
|
// otherwise the fast hand-off keeps open_count pinned above the new
|
|
// limit forever under steady traffic. Wake the waiter with a retry
|
|
// sentinel so it re-evaluates capacity under the new cap.
|
|
if p.max_open > 0 && p.open_count > p.max_open {
|
|
physical_close_handle(slot.handle)
|
|
p.open_count--
|
|
if p.waiters.len > 0 {
|
|
waiter := p.waiters[0]
|
|
p.waiters.delete(0)
|
|
waiter <- IdleSlot{
|
|
handle: unsafe { nil }
|
|
}
|
|
}
|
|
p.mu.unlock()
|
|
return
|
|
}
|
|
// Healthy conn: prefer handing it directly to a waiter. Send under the
|
|
// lock (cap:1 makes it non-blocking) so close() can't slip in and orphan
|
|
// the popped waiter with a live conn.
|
|
if p.waiters.len > 0 {
|
|
waiter := p.waiters[0]
|
|
p.waiters.delete(0)
|
|
waiter <- slot
|
|
p.mu.unlock()
|
|
return
|
|
}
|
|
// Park as idle, unless we'd exceed max_idle (0 = keep no idle conns)
|
|
// or we are over a recently-shrunk max_open and need to converge.
|
|
if p.idle.len >= p.max_idle || (p.max_open > 0 && p.open_count > p.max_open) {
|
|
physical_close_handle(slot.handle)
|
|
p.open_count--
|
|
p.mu.unlock()
|
|
return
|
|
}
|
|
p.idle << slot
|
|
p.mu.unlock()
|
|
}
|
|
|
|
fn (mut p Pool) close() {
|
|
p.mu.lock()
|
|
if p.closed {
|
|
p.mu.unlock()
|
|
return
|
|
}
|
|
p.closed = true
|
|
for slot in p.idle {
|
|
physical_close_handle(slot.handle)
|
|
}
|
|
idle_len := p.idle.len
|
|
p.idle = []IdleSlot{}
|
|
p.open_count -= idle_len
|
|
for waiter in p.waiters {
|
|
waiter.close()
|
|
}
|
|
p.waiters = []chan IdleSlot{}
|
|
p.mu.unlock()
|
|
}
|
|
|
|
fn (mut p Pool) stats() PoolStats {
|
|
p.mu.lock()
|
|
stats := PoolStats{
|
|
max_open_connections: p.max_open
|
|
open_connections: p.open_count
|
|
idle: p.idle.len
|
|
in_use: p.open_count - p.idle.len
|
|
wait_count: p.waiters.len
|
|
}
|
|
p.mu.unlock()
|
|
return stats
|
|
}
|
|
|
|
fn (mut p Pool) set_max_open(n int) {
|
|
mut nn := n
|
|
if nn < 0 {
|
|
nn = 0
|
|
}
|
|
p.mu.lock()
|
|
p.max_open = nn
|
|
// Shrink: drop idle slots until we are back under the new cap. Without
|
|
// this, acquire() pops idle slots before checking max_open, so two
|
|
// callers could still both succeed after set_max_open_conns(1) when two
|
|
// warm conns are parked. In-use conns can't be reclaimed here — they get
|
|
// discarded on release once open_count is over the cap.
|
|
if nn > 0 {
|
|
for p.idle.len > 0 && p.open_count > nn {
|
|
slot := p.idle.last()
|
|
p.idle.delete_last()
|
|
physical_close_handle(slot.handle)
|
|
p.open_count--
|
|
}
|
|
}
|
|
// Raise: wake parked waiters so they can retry the acquire loop and
|
|
// dial against the new headroom. Without this nudge they stay blocked
|
|
// until some other release happens, which may never come if all current
|
|
// conns are long-lived.
|
|
if p.waiters.len > 0 && (nn == 0 || p.open_count < nn) {
|
|
spare := if nn == 0 { p.waiters.len } else { nn - p.open_count }
|
|
mut n_wake := if spare < p.waiters.len { spare } else { p.waiters.len }
|
|
for n_wake > 0 {
|
|
waiter := p.waiters[0]
|
|
p.waiters.delete(0)
|
|
waiter <- IdleSlot{
|
|
handle: unsafe { nil }
|
|
}
|
|
n_wake--
|
|
}
|
|
}
|
|
p.mu.unlock()
|
|
}
|
|
|
|
fn (mut p Pool) set_max_idle(n int) {
|
|
mut nn := n
|
|
if nn < 0 {
|
|
nn = 0
|
|
}
|
|
p.mu.lock()
|
|
p.max_idle = nn
|
|
for p.idle.len > nn {
|
|
slot := p.idle.last()
|
|
p.idle.delete_last()
|
|
physical_close_handle(slot.handle)
|
|
p.open_count--
|
|
}
|
|
p.mu.unlock()
|
|
}
|
|
|
|
fn (mut p Pool) set_conn_max_lifetime(d time.Duration) {
|
|
p.mu.lock()
|
|
p.max_lifetime = d
|
|
p.mu.unlock()
|
|
}
|
|
|
|
// stash_last_id records `id` as the last-inserted id for the current thread.
|
|
// Called by `DB.insert` right after running INSERT on the pinned conn, so the
|
|
// next `DB.last_id()` call on the same thread returns this exact value rather
|
|
// than running LASTVAL() on a different pooled session.
|
|
fn (mut p Pool) stash_last_id(id int) {
|
|
tid := sync.thread_id()
|
|
p.last_ids_mu.lock()
|
|
p.last_ids[tid] = i64(id)
|
|
p.last_ids_mu.unlock()
|
|
}
|
|
|
|
// take_last_id returns the last-inserted id stashed by this thread's most
|
|
// recent `DB.insert`, or 0 if there is none. Consuming on read keeps the
|
|
// map bounded as threads come and go, and prevents a future thread that
|
|
// reuses this tid from observing a stale id before doing its own insert.
|
|
fn (mut p Pool) take_last_id() int {
|
|
tid := sync.thread_id()
|
|
p.last_ids_mu.lock()
|
|
id := p.last_ids[tid] or { i64(0) }
|
|
p.last_ids.delete(tid)
|
|
p.last_ids_mu.unlock()
|
|
return int(id)
|
|
}
|