feat(bun): generate native Go array slices for PostgreSQL array columns
Bun's pgdialect scans/appends native slices directly, so array columns
(text[], integer[], uuid[], ...) always generate as plain []string,
[]int32, etc. with an explicit "array" bun tag, regardless of --types
(sqltypes/stdlib/baselib). The SqlXxxArray wrapper types are no longer
used for Bun array columns (gorm is unaffected and keeps using them).
Adds --array-nullable pointer_slice to represent nullable array columns
as *[]T instead of []T, so callers can distinguish SQL NULL (nil) from
'{}' (pointer to an empty slice). Verified end-to-end against a live
PostgreSQL instance for NULL/{}/populated arrays in every --types mode.
Closes #13
This commit is contained in:
Hein
+52
@@ -0,0 +1,52 @@
|
||||
package pgxpool
|
||||
|
||||
import (
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
)
|
||||
|
||||
type errBatchResults struct {
|
||||
err error
|
||||
}
|
||||
|
||||
func (br errBatchResults) Exec() (pgconn.CommandTag, error) {
|
||||
return pgconn.CommandTag{}, br.err
|
||||
}
|
||||
|
||||
func (br errBatchResults) Query() (pgx.Rows, error) {
|
||||
return errRows{err: br.err}, br.err
|
||||
}
|
||||
|
||||
func (br errBatchResults) QueryRow() pgx.Row {
|
||||
return errRow{err: br.err}
|
||||
}
|
||||
|
||||
func (br errBatchResults) Close() error {
|
||||
return br.err
|
||||
}
|
||||
|
||||
type poolBatchResults struct {
|
||||
br pgx.BatchResults
|
||||
c *Conn
|
||||
}
|
||||
|
||||
func (br *poolBatchResults) Exec() (pgconn.CommandTag, error) {
|
||||
return br.br.Exec()
|
||||
}
|
||||
|
||||
func (br *poolBatchResults) Query() (pgx.Rows, error) {
|
||||
return br.br.Query()
|
||||
}
|
||||
|
||||
func (br *poolBatchResults) QueryRow() pgx.Row {
|
||||
return br.br.QueryRow()
|
||||
}
|
||||
|
||||
func (br *poolBatchResults) Close() error {
|
||||
err := br.br.Close()
|
||||
if br.c != nil {
|
||||
br.c.Release()
|
||||
br.c = nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
+134
@@ -0,0 +1,134 @@
|
||||
package pgxpool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
"github.com/jackc/puddle/v2"
|
||||
)
|
||||
|
||||
// Conn is an acquired *pgx.Conn from a Pool.
|
||||
type Conn struct {
|
||||
res *puddle.Resource[*connResource]
|
||||
p *Pool
|
||||
}
|
||||
|
||||
// Release returns c to the pool it was acquired from. Once Release has been called, other methods must not be called.
|
||||
// However, it is safe to call Release multiple times. Subsequent calls after the first will be ignored.
|
||||
func (c *Conn) Release() {
|
||||
if c.res == nil {
|
||||
return
|
||||
}
|
||||
|
||||
conn := c.Conn()
|
||||
res := c.res
|
||||
c.res = nil
|
||||
|
||||
if c.p.releaseTracer != nil {
|
||||
c.p.releaseTracer.TraceRelease(c.p, TraceReleaseData{Conn: conn})
|
||||
}
|
||||
|
||||
if conn.IsClosed() || conn.PgConn().IsBusy() || conn.PgConn().TxStatus() != 'I' {
|
||||
res.Destroy()
|
||||
// Signal to the health check to run since we just destroyed a connections
|
||||
// and we might be below minConns now
|
||||
c.p.triggerHealthCheck()
|
||||
return
|
||||
}
|
||||
|
||||
// If the pool is consistently being used, we might never get to check the
|
||||
// lifetime of a connection since we only check idle connections in checkConnsHealth
|
||||
// so we also check the lifetime here and force a health check
|
||||
if c.p.isExpired(res) {
|
||||
atomic.AddInt64(&c.p.lifetimeDestroyCount, 1)
|
||||
res.Destroy()
|
||||
// Signal to the health check to run since we just destroyed a connections
|
||||
// and we might be below minConns now
|
||||
c.p.triggerHealthCheck()
|
||||
return
|
||||
}
|
||||
|
||||
if c.p.afterRelease == nil {
|
||||
res.Release()
|
||||
return
|
||||
}
|
||||
|
||||
go func() {
|
||||
if c.p.afterRelease(conn) {
|
||||
res.Release()
|
||||
} else {
|
||||
res.Destroy()
|
||||
// Signal to the health check to run since we just destroyed a connections
|
||||
// and we might be below minConns now
|
||||
c.p.triggerHealthCheck()
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// Hijack assumes ownership of the connection from the pool. Caller is responsible for closing the connection. Hijack
|
||||
// will panic if called on an already released or hijacked connection.
|
||||
func (c *Conn) Hijack() *pgx.Conn {
|
||||
if c.res == nil {
|
||||
panic("cannot hijack already released or hijacked connection")
|
||||
}
|
||||
|
||||
conn := c.Conn()
|
||||
res := c.res
|
||||
c.res = nil
|
||||
|
||||
res.Hijack()
|
||||
|
||||
return conn
|
||||
}
|
||||
|
||||
func (c *Conn) Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) {
|
||||
return c.Conn().Exec(ctx, sql, arguments...)
|
||||
}
|
||||
|
||||
func (c *Conn) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) {
|
||||
return c.Conn().Query(ctx, sql, args...)
|
||||
}
|
||||
|
||||
func (c *Conn) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row {
|
||||
return c.Conn().QueryRow(ctx, sql, args...)
|
||||
}
|
||||
|
||||
func (c *Conn) SendBatch(ctx context.Context, b *pgx.Batch) pgx.BatchResults {
|
||||
return c.Conn().SendBatch(ctx, b)
|
||||
}
|
||||
|
||||
func (c *Conn) CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNames []string, rowSrc pgx.CopyFromSource) (int64, error) {
|
||||
return c.Conn().CopyFrom(ctx, tableName, columnNames, rowSrc)
|
||||
}
|
||||
|
||||
// Begin starts a transaction block from the *Conn without explicitly setting a transaction mode (see BeginTx with TxOptions if transaction mode is required).
|
||||
func (c *Conn) Begin(ctx context.Context) (pgx.Tx, error) {
|
||||
return c.Conn().Begin(ctx)
|
||||
}
|
||||
|
||||
// BeginTx starts a transaction block from the *Conn with txOptions determining the transaction mode.
|
||||
func (c *Conn) BeginTx(ctx context.Context, txOptions pgx.TxOptions) (pgx.Tx, error) {
|
||||
return c.Conn().BeginTx(ctx, txOptions)
|
||||
}
|
||||
|
||||
func (c *Conn) Ping(ctx context.Context) error {
|
||||
return c.Conn().Ping(ctx)
|
||||
}
|
||||
|
||||
func (c *Conn) Conn() *pgx.Conn {
|
||||
return c.connResource().conn
|
||||
}
|
||||
|
||||
func (c *Conn) connResource() *connResource {
|
||||
return c.res.Value()
|
||||
}
|
||||
|
||||
func (c *Conn) getPoolRow(r pgx.Row) *poolRow {
|
||||
return c.connResource().getPoolRow(c, r)
|
||||
}
|
||||
|
||||
func (c *Conn) getPoolRows(r pgx.Rows) *poolRows {
|
||||
return c.connResource().getPoolRows(c, r)
|
||||
}
|
||||
+27
@@ -0,0 +1,27 @@
|
||||
// Package pgxpool is a concurrency-safe connection pool for pgx.
|
||||
/*
|
||||
pgxpool implements a nearly identical interface to pgx connections.
|
||||
|
||||
Creating a Pool
|
||||
|
||||
The primary way of creating a pool is with [pgxpool.New]:
|
||||
|
||||
pool, err := pgxpool.New(context.Background(), os.Getenv("DATABASE_URL"))
|
||||
|
||||
The database connection string can be in URL or keyword/value format. PostgreSQL settings, pgx settings, and pool settings can be
|
||||
specified here. In addition, a config struct can be created by [ParseConfig].
|
||||
|
||||
config, err := pgxpool.ParseConfig(os.Getenv("DATABASE_URL"))
|
||||
if err != nil {
|
||||
// ...
|
||||
}
|
||||
config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
|
||||
// do something with every new connection
|
||||
}
|
||||
|
||||
pool, err := pgxpool.NewWithConfig(context.Background(), config)
|
||||
|
||||
A pool returns without waiting for any connections to be established. Acquire a connection immediately after creating
|
||||
the pool to check if a connection can successfully be established.
|
||||
*/
|
||||
package pgxpool
|
||||
+832
@@ -0,0 +1,832 @@
|
||||
package pgxpool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"math/rand/v2"
|
||||
"runtime"
|
||||
"strconv"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
"github.com/jackc/puddle/v2"
|
||||
)
|
||||
|
||||
var (
|
||||
defaultMaxConns = int32(4)
|
||||
defaultMinConns = int32(0)
|
||||
defaultMinIdleConns = int32(0)
|
||||
defaultMaxConnLifetime = time.Hour
|
||||
defaultMaxConnIdleTime = time.Minute * 30
|
||||
defaultHealthCheckPeriod = time.Minute
|
||||
)
|
||||
|
||||
type connResource struct {
|
||||
conn *pgx.Conn
|
||||
conns []Conn
|
||||
poolRows []poolRow
|
||||
poolRowss []poolRows
|
||||
maxAgeTime time.Time
|
||||
}
|
||||
|
||||
func (cr *connResource) getConn(p *Pool, res *puddle.Resource[*connResource]) *Conn {
|
||||
if len(cr.conns) == 0 {
|
||||
cr.conns = make([]Conn, 128)
|
||||
}
|
||||
|
||||
c := &cr.conns[len(cr.conns)-1]
|
||||
cr.conns = cr.conns[0 : len(cr.conns)-1]
|
||||
|
||||
c.res = res
|
||||
c.p = p
|
||||
|
||||
return c
|
||||
}
|
||||
|
||||
func (cr *connResource) getPoolRow(c *Conn, r pgx.Row) *poolRow {
|
||||
if len(cr.poolRows) == 0 {
|
||||
cr.poolRows = make([]poolRow, 128)
|
||||
}
|
||||
|
||||
pr := &cr.poolRows[len(cr.poolRows)-1]
|
||||
cr.poolRows = cr.poolRows[0 : len(cr.poolRows)-1]
|
||||
|
||||
pr.c = c
|
||||
pr.r = r
|
||||
|
||||
return pr
|
||||
}
|
||||
|
||||
func (cr *connResource) getPoolRows(c *Conn, r pgx.Rows) *poolRows {
|
||||
if len(cr.poolRowss) == 0 {
|
||||
cr.poolRowss = make([]poolRows, 128)
|
||||
}
|
||||
|
||||
pr := &cr.poolRowss[len(cr.poolRowss)-1]
|
||||
cr.poolRowss = cr.poolRowss[0 : len(cr.poolRowss)-1]
|
||||
|
||||
pr.c = c
|
||||
pr.r = r
|
||||
|
||||
return pr
|
||||
}
|
||||
|
||||
// Pool allows for connection reuse.
|
||||
type Pool struct {
|
||||
// 64 bit fields accessed with atomics must be at beginning of struct to guarantee alignment for certain 32-bit
|
||||
// architectures. See BUGS section of https://pkg.go.dev/sync/atomic and https://github.com/jackc/pgx/issues/1288.
|
||||
newConnsCount int64
|
||||
lifetimeDestroyCount int64
|
||||
idleDestroyCount int64
|
||||
|
||||
p *puddle.Pool[*connResource]
|
||||
config *Config
|
||||
beforeConnect func(context.Context, *pgx.ConnConfig) error
|
||||
afterConnect func(context.Context, *pgx.Conn) error
|
||||
prepareConn func(context.Context, *pgx.Conn) (bool, error)
|
||||
afterRelease func(*pgx.Conn) bool
|
||||
beforeClose func(*pgx.Conn)
|
||||
shouldPing func(context.Context, ShouldPingParams) bool
|
||||
minConns int32
|
||||
minIdleConns int32
|
||||
maxConns int32
|
||||
maxConnLifetime time.Duration
|
||||
maxConnLifetimeJitter time.Duration
|
||||
maxConnIdleTime time.Duration
|
||||
healthCheckPeriod time.Duration
|
||||
pingTimeout time.Duration
|
||||
|
||||
healthCheckMu sync.Mutex
|
||||
healthCheckTimer *time.Timer
|
||||
|
||||
healthCheckChan chan struct{}
|
||||
|
||||
acquireTracer AcquireTracer
|
||||
releaseTracer ReleaseTracer
|
||||
|
||||
closeOnce sync.Once
|
||||
closeChan chan struct{}
|
||||
}
|
||||
|
||||
// ShouldPingParams are the parameters passed to ShouldPing.
|
||||
type ShouldPingParams struct {
|
||||
Conn *pgx.Conn
|
||||
IdleDuration time.Duration
|
||||
}
|
||||
|
||||
// Config is the configuration struct for creating a pool. It must be created by [ParseConfig] and then it can be
|
||||
// modified.
|
||||
type Config struct {
|
||||
ConnConfig *pgx.ConnConfig
|
||||
|
||||
// BeforeConnect is called before a new connection is made. It is passed a copy of the underlying [pgx.ConnConfig] and
|
||||
// will not impact any existing open connections.
|
||||
BeforeConnect func(context.Context, *pgx.ConnConfig) error
|
||||
|
||||
// AfterConnect is called after a connection is established, but before it is added to the pool.
|
||||
AfterConnect func(context.Context, *pgx.Conn) error
|
||||
|
||||
// BeforeAcquire is called before a connection is acquired from the pool. It must return true to allow the
|
||||
// acquisition or false to indicate that the connection should be destroyed and a different connection should be
|
||||
// acquired.
|
||||
//
|
||||
// Deprecated: Use PrepareConn instead. If both PrepareConn and BeforeAcquire are set, PrepareConn will take
|
||||
// precedence, ignoring BeforeAcquire.
|
||||
BeforeAcquire func(context.Context, *pgx.Conn) bool
|
||||
|
||||
// PrepareConn is called before a connection is acquired from the pool. If this function returns true, the connection
|
||||
// is considered valid, otherwise the connection is destroyed. If the function returns a non-nil error, the instigating
|
||||
// query will fail with the returned error.
|
||||
//
|
||||
// Specifically, this means that:
|
||||
//
|
||||
// - If it returns true and a nil error, the query proceeds as normal.
|
||||
// - If it returns true and an error, the connection will be returned to the pool, and the instigating query will fail with the returned error.
|
||||
// - If it returns false, and an error, the connection will be destroyed, and the query will fail with the returned error.
|
||||
// - If it returns false and a nil error, the connection will be destroyed, and the instigating query will be retried on a new connection.
|
||||
PrepareConn func(context.Context, *pgx.Conn) (bool, error)
|
||||
|
||||
// AfterRelease is called after a connection is released, but before it is returned to the pool. It must return true to
|
||||
// return the connection to the pool or false to destroy the connection.
|
||||
AfterRelease func(*pgx.Conn) bool
|
||||
|
||||
// BeforeClose is called right before a connection is closed and removed from the pool.
|
||||
BeforeClose func(*pgx.Conn)
|
||||
|
||||
// ShouldPing is called after a connection is acquired from the pool. If it returns true, the connection is pinged to check for liveness.
|
||||
// If this func is not set, the default behavior is to ping connections that have been idle for at least 1 second.
|
||||
ShouldPing func(context.Context, ShouldPingParams) bool
|
||||
|
||||
// MaxConnLifetime is the duration since creation after which a connection will be automatically closed.
|
||||
MaxConnLifetime time.Duration
|
||||
|
||||
// MaxConnLifetimeJitter is the duration after MaxConnLifetime to randomly decide to close a connection.
|
||||
// This helps prevent all connections from being closed at the exact same time, starving the pool.
|
||||
MaxConnLifetimeJitter time.Duration
|
||||
|
||||
// MaxConnIdleTime is the duration after which an idle connection will be automatically closed by the health check.
|
||||
MaxConnIdleTime time.Duration
|
||||
|
||||
// PingTimeout is the maximum amount of time to wait for a connection to pong before considering it as unhealthy and
|
||||
// destroying it. If zero, the default is no timeout.
|
||||
PingTimeout time.Duration
|
||||
|
||||
// MaxConns is the maximum size of the pool. The default is the greater of 4 or runtime.NumCPU().
|
||||
MaxConns int32
|
||||
|
||||
// MinConns is the minimum size of the pool. After connection closes, the pool might dip below MinConns. A low
|
||||
// number of MinConns might mean the pool is empty after MaxConnLifetime until the health check has a chance
|
||||
// to create new connections.
|
||||
MinConns int32
|
||||
|
||||
// MinIdleConns is the minimum number of idle connections in the pool. You can increase this to ensure that
|
||||
// there are always idle connections available. This can help reduce tail latencies during request processing,
|
||||
// as you can avoid the latency of establishing a new connection while handling requests. It is superior
|
||||
// to MinConns for this purpose.
|
||||
// Similar to MinConns, the pool might temporarily dip below MinIdleConns after connection closes.
|
||||
MinIdleConns int32
|
||||
|
||||
// HealthCheckPeriod is the duration between checks of the health of idle connections.
|
||||
HealthCheckPeriod time.Duration
|
||||
|
||||
createdByParseConfig bool // Used to enforce created by ParseConfig rule.
|
||||
}
|
||||
|
||||
// Copy returns a deep copy of the config that is safe to use and modify.
|
||||
// The only exception is the tls.Config:
|
||||
// according to the tls.Config docs it must not be modified after creation.
|
||||
func (c *Config) Copy() *Config {
|
||||
newConfig := new(Config)
|
||||
*newConfig = *c
|
||||
newConfig.ConnConfig = c.ConnConfig.Copy()
|
||||
return newConfig
|
||||
}
|
||||
|
||||
// ConnString returns the connection string as parsed by pgxpool.ParseConfig into pgxpool.Config.
|
||||
func (c *Config) ConnString() string { return c.ConnConfig.ConnString() }
|
||||
|
||||
// New creates a new Pool. See [ParseConfig] for information on connString format.
|
||||
func New(ctx context.Context, connString string) (*Pool, error) {
|
||||
config, err := ParseConfig(connString)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return NewWithConfig(ctx, config)
|
||||
}
|
||||
|
||||
// NewWithConfig creates a new [Pool]. config must have been created by [ParseConfig].
|
||||
func NewWithConfig(ctx context.Context, config *Config) (*Pool, error) {
|
||||
// Default values are set in ParseConfig. Enforce initial creation by ParseConfig rather than setting defaults from
|
||||
// zero values.
|
||||
if !config.createdByParseConfig {
|
||||
panic("config must be created by ParseConfig")
|
||||
}
|
||||
|
||||
prepareConn := config.PrepareConn
|
||||
if prepareConn == nil && config.BeforeAcquire != nil {
|
||||
prepareConn = func(ctx context.Context, conn *pgx.Conn) (bool, error) {
|
||||
return config.BeforeAcquire(ctx, conn), nil
|
||||
}
|
||||
}
|
||||
|
||||
p := &Pool{
|
||||
config: config,
|
||||
beforeConnect: config.BeforeConnect,
|
||||
afterConnect: config.AfterConnect,
|
||||
prepareConn: prepareConn,
|
||||
afterRelease: config.AfterRelease,
|
||||
beforeClose: config.BeforeClose,
|
||||
minConns: config.MinConns,
|
||||
minIdleConns: config.MinIdleConns,
|
||||
maxConns: config.MaxConns,
|
||||
maxConnLifetime: config.MaxConnLifetime,
|
||||
maxConnLifetimeJitter: config.MaxConnLifetimeJitter,
|
||||
maxConnIdleTime: config.MaxConnIdleTime,
|
||||
pingTimeout: config.PingTimeout,
|
||||
healthCheckPeriod: config.HealthCheckPeriod,
|
||||
healthCheckChan: make(chan struct{}, 1),
|
||||
closeChan: make(chan struct{}),
|
||||
}
|
||||
|
||||
if t, ok := config.ConnConfig.Tracer.(AcquireTracer); ok {
|
||||
p.acquireTracer = t
|
||||
}
|
||||
|
||||
if t, ok := config.ConnConfig.Tracer.(ReleaseTracer); ok {
|
||||
p.releaseTracer = t
|
||||
}
|
||||
|
||||
if config.ShouldPing != nil {
|
||||
p.shouldPing = config.ShouldPing
|
||||
} else {
|
||||
p.shouldPing = func(ctx context.Context, params ShouldPingParams) bool {
|
||||
return params.IdleDuration > time.Second
|
||||
}
|
||||
}
|
||||
|
||||
var err error
|
||||
p.p, err = puddle.NewPool(
|
||||
&puddle.Config[*connResource]{
|
||||
Constructor: func(ctx context.Context) (*connResource, error) {
|
||||
atomic.AddInt64(&p.newConnsCount, 1)
|
||||
connConfig := p.config.ConnConfig.Copy()
|
||||
|
||||
// Connection will continue in background even if Acquire is canceled. Ensure that a connect won't hang forever.
|
||||
if connConfig.ConnectTimeout <= 0 {
|
||||
connConfig.ConnectTimeout = 2 * time.Minute
|
||||
}
|
||||
|
||||
if p.beforeConnect != nil {
|
||||
if err := p.beforeConnect(ctx, connConfig); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
conn, err := pgx.ConnectConfig(ctx, connConfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if p.afterConnect != nil {
|
||||
err = p.afterConnect(ctx, conn)
|
||||
if err != nil {
|
||||
conn.Close(ctx)
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
jitterSecs := rand.Float64() * config.MaxConnLifetimeJitter.Seconds()
|
||||
maxAgeTime := time.Now().Add(config.MaxConnLifetime).Add(time.Duration(jitterSecs) * time.Second)
|
||||
|
||||
cr := &connResource{
|
||||
conn: conn,
|
||||
conns: make([]Conn, 64),
|
||||
poolRows: make([]poolRow, 64),
|
||||
poolRowss: make([]poolRows, 64),
|
||||
maxAgeTime: maxAgeTime,
|
||||
}
|
||||
|
||||
return cr, nil
|
||||
},
|
||||
Destructor: func(value *connResource) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||
conn := value.conn
|
||||
if p.beforeClose != nil {
|
||||
p.beforeClose(conn)
|
||||
}
|
||||
conn.Close(ctx)
|
||||
select {
|
||||
case <-conn.PgConn().CleanupDone():
|
||||
case <-ctx.Done():
|
||||
}
|
||||
cancel()
|
||||
},
|
||||
MaxSize: config.MaxConns,
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
go func() {
|
||||
targetIdleResources := max(int(p.minConns), int(p.minIdleConns))
|
||||
p.createIdleResources(ctx, targetIdleResources)
|
||||
p.backgroundHealthCheck()
|
||||
}()
|
||||
|
||||
return p, nil
|
||||
}
|
||||
|
||||
// ParseConfig builds a Config from connString. It parses connString with the same behavior as [pgx.ParseConfig] with the
|
||||
// addition of the following variables:
|
||||
//
|
||||
// - pool_max_conns: integer greater than 0 (default 4)
|
||||
// - pool_min_conns: integer 0 or greater (default 0)
|
||||
// - pool_max_conn_lifetime: duration string (default 1 hour)
|
||||
// - pool_max_conn_idle_time: duration string (default 30 minutes)
|
||||
// - pool_health_check_period: duration string (default 1 minute)
|
||||
// - pool_max_conn_lifetime_jitter: duration string (default 0)
|
||||
//
|
||||
// See Config for definitions of these arguments.
|
||||
//
|
||||
// # Example Keyword/Value
|
||||
// user=jack password=secret host=pg.example.com port=5432 dbname=mydb sslmode=verify-ca pool_max_conns=10 pool_max_conn_lifetime=1h30m
|
||||
//
|
||||
// # Example URL
|
||||
// postgres://jack:secret@pg.example.com:5432/mydb?sslmode=verify-ca&pool_max_conns=10&pool_max_conn_lifetime=1h30m
|
||||
func ParseConfig(connString string) (*Config, error) {
|
||||
connConfig, err := pgx.ParseConfig(connString)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
config := &Config{
|
||||
ConnConfig: connConfig,
|
||||
createdByParseConfig: true,
|
||||
}
|
||||
|
||||
if s, ok := config.ConnConfig.Config.RuntimeParams["pool_max_conns"]; ok {
|
||||
delete(connConfig.Config.RuntimeParams, "pool_max_conns")
|
||||
n, err := strconv.ParseInt(s, 10, 32)
|
||||
if err != nil {
|
||||
return nil, pgconn.NewParseConfigError(connString, "cannot parse pool_max_conns", err)
|
||||
}
|
||||
if n < 1 {
|
||||
return nil, pgconn.NewParseConfigError(connString, "pool_max_conns too small", err)
|
||||
}
|
||||
config.MaxConns = int32(n)
|
||||
} else {
|
||||
config.MaxConns = defaultMaxConns
|
||||
if numCPU := int32(runtime.NumCPU()); numCPU > config.MaxConns {
|
||||
config.MaxConns = numCPU
|
||||
}
|
||||
}
|
||||
|
||||
if s, ok := config.ConnConfig.Config.RuntimeParams["pool_min_conns"]; ok {
|
||||
delete(connConfig.Config.RuntimeParams, "pool_min_conns")
|
||||
n, err := strconv.ParseInt(s, 10, 32)
|
||||
if err != nil {
|
||||
return nil, pgconn.NewParseConfigError(connString, "cannot parse pool_min_conns", err)
|
||||
}
|
||||
config.MinConns = int32(n)
|
||||
} else {
|
||||
config.MinConns = defaultMinConns
|
||||
}
|
||||
|
||||
if s, ok := config.ConnConfig.Config.RuntimeParams["pool_min_idle_conns"]; ok {
|
||||
delete(connConfig.Config.RuntimeParams, "pool_min_idle_conns")
|
||||
n, err := strconv.ParseInt(s, 10, 32)
|
||||
if err != nil {
|
||||
return nil, pgconn.NewParseConfigError(connString, "cannot parse pool_min_idle_conns", err)
|
||||
}
|
||||
config.MinIdleConns = int32(n)
|
||||
} else {
|
||||
config.MinIdleConns = defaultMinIdleConns
|
||||
}
|
||||
|
||||
if s, ok := config.ConnConfig.Config.RuntimeParams["pool_max_conn_lifetime"]; ok {
|
||||
delete(connConfig.Config.RuntimeParams, "pool_max_conn_lifetime")
|
||||
d, err := time.ParseDuration(s)
|
||||
if err != nil {
|
||||
return nil, pgconn.NewParseConfigError(connString, "cannot parse pool_max_conn_lifetime", err)
|
||||
}
|
||||
config.MaxConnLifetime = d
|
||||
} else {
|
||||
config.MaxConnLifetime = defaultMaxConnLifetime
|
||||
}
|
||||
|
||||
if s, ok := config.ConnConfig.Config.RuntimeParams["pool_max_conn_idle_time"]; ok {
|
||||
delete(connConfig.Config.RuntimeParams, "pool_max_conn_idle_time")
|
||||
d, err := time.ParseDuration(s)
|
||||
if err != nil {
|
||||
return nil, pgconn.NewParseConfigError(connString, "cannot parse pool_max_conn_idle_time", err)
|
||||
}
|
||||
config.MaxConnIdleTime = d
|
||||
} else {
|
||||
config.MaxConnIdleTime = defaultMaxConnIdleTime
|
||||
}
|
||||
|
||||
if s, ok := config.ConnConfig.Config.RuntimeParams["pool_health_check_period"]; ok {
|
||||
delete(connConfig.Config.RuntimeParams, "pool_health_check_period")
|
||||
d, err := time.ParseDuration(s)
|
||||
if err != nil {
|
||||
return nil, pgconn.NewParseConfigError(connString, "cannot parse pool_health_check_period", err)
|
||||
}
|
||||
config.HealthCheckPeriod = d
|
||||
} else {
|
||||
config.HealthCheckPeriod = defaultHealthCheckPeriod
|
||||
}
|
||||
|
||||
if s, ok := config.ConnConfig.Config.RuntimeParams["pool_max_conn_lifetime_jitter"]; ok {
|
||||
delete(connConfig.Config.RuntimeParams, "pool_max_conn_lifetime_jitter")
|
||||
d, err := time.ParseDuration(s)
|
||||
if err != nil {
|
||||
return nil, pgconn.NewParseConfigError(connString, "cannot parse pool_max_conn_lifetime_jitter", err)
|
||||
}
|
||||
config.MaxConnLifetimeJitter = d
|
||||
}
|
||||
|
||||
return config, nil
|
||||
}
|
||||
|
||||
// Close closes all connections in the pool and rejects future [Pool.Acquire] calls. Blocks until all connections are returned
|
||||
// to pool and closed.
|
||||
func (p *Pool) Close() {
|
||||
p.closeOnce.Do(func() {
|
||||
close(p.closeChan)
|
||||
p.p.Close()
|
||||
})
|
||||
}
|
||||
|
||||
func (p *Pool) isExpired(res *puddle.Resource[*connResource]) bool {
|
||||
return time.Now().After(res.Value().maxAgeTime)
|
||||
}
|
||||
|
||||
func (p *Pool) triggerHealthCheck() {
|
||||
const healthCheckDelay = 500 * time.Millisecond
|
||||
|
||||
p.healthCheckMu.Lock()
|
||||
defer p.healthCheckMu.Unlock()
|
||||
|
||||
if p.healthCheckTimer == nil {
|
||||
// Destroy is asynchronous so we give it time to actually remove itself from
|
||||
// the pool otherwise we might try to check the pool size too soon
|
||||
p.healthCheckTimer = time.AfterFunc(healthCheckDelay, func() {
|
||||
select {
|
||||
case <-p.closeChan:
|
||||
case p.healthCheckChan <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
p.healthCheckTimer.Reset(healthCheckDelay)
|
||||
}
|
||||
|
||||
func (p *Pool) backgroundHealthCheck() {
|
||||
ticker := time.NewTicker(p.healthCheckPeriod)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-p.closeChan:
|
||||
return
|
||||
case <-p.healthCheckChan:
|
||||
p.checkHealth()
|
||||
case <-ticker.C:
|
||||
p.checkHealth()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Pool) checkHealth() {
|
||||
for {
|
||||
// If checkMinConns failed we don't destroy any connections since we couldn't
|
||||
// even get to minConns
|
||||
if err := p.checkMinConns(); err != nil {
|
||||
// Should we log this error somewhere?
|
||||
break
|
||||
}
|
||||
if !p.checkConnsHealth() {
|
||||
// Since we didn't destroy any connections we can stop looping
|
||||
break
|
||||
}
|
||||
// Technically Destroy is asynchronous but 500ms should be enough for it to
|
||||
// remove it from the underlying pool
|
||||
select {
|
||||
case <-p.closeChan:
|
||||
return
|
||||
case <-time.After(500 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// checkConnsHealth will check all idle connections, destroy a connection if
|
||||
// it's idle or too old, and returns true if any were destroyed
|
||||
func (p *Pool) checkConnsHealth() bool {
|
||||
var destroyed bool
|
||||
totalConns := p.Stat().TotalConns()
|
||||
resources := p.p.AcquireAllIdle()
|
||||
for _, res := range resources {
|
||||
// We're okay going under minConns if the lifetime is up
|
||||
if p.isExpired(res) && totalConns >= p.minConns {
|
||||
atomic.AddInt64(&p.lifetimeDestroyCount, 1)
|
||||
res.Destroy()
|
||||
destroyed = true
|
||||
// Since Destroy is async we manually decrement totalConns.
|
||||
totalConns--
|
||||
} else if res.IdleDuration() > p.maxConnIdleTime && totalConns > p.minConns {
|
||||
atomic.AddInt64(&p.idleDestroyCount, 1)
|
||||
res.Destroy()
|
||||
destroyed = true
|
||||
// Since Destroy is async we manually decrement totalConns.
|
||||
totalConns--
|
||||
} else {
|
||||
res.ReleaseUnused()
|
||||
}
|
||||
}
|
||||
return destroyed
|
||||
}
|
||||
|
||||
func (p *Pool) checkMinConns() error {
|
||||
// TotalConns can include ones that are being destroyed but we should have
|
||||
// sleep(500ms) around all of the destroys to help prevent that from throwing
|
||||
// off this check
|
||||
|
||||
// Create the number of connections needed to get to both minConns and minIdleConns
|
||||
stat := p.Stat()
|
||||
toCreate := max(p.minConns-stat.TotalConns(), p.minIdleConns-stat.IdleConns())
|
||||
if toCreate > 0 {
|
||||
return p.createIdleResources(context.Background(), int(toCreate))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *Pool) createIdleResources(parentCtx context.Context, targetResources int) error {
|
||||
ctx, cancel := context.WithCancel(parentCtx)
|
||||
defer cancel()
|
||||
|
||||
errs := make(chan error, targetResources)
|
||||
|
||||
for range targetResources {
|
||||
go func() {
|
||||
err := p.p.CreateResource(ctx)
|
||||
// Ignore ErrNotAvailable since it means that the pool has become full since we started creating resource.
|
||||
if err == puddle.ErrNotAvailable {
|
||||
err = nil
|
||||
}
|
||||
errs <- err
|
||||
}()
|
||||
}
|
||||
|
||||
var firstError error
|
||||
for range targetResources {
|
||||
err := <-errs
|
||||
if err != nil && firstError == nil {
|
||||
cancel()
|
||||
firstError = err
|
||||
}
|
||||
}
|
||||
|
||||
return firstError
|
||||
}
|
||||
|
||||
// Acquire returns a connection ([Conn]) from the [Pool].
|
||||
func (p *Pool) Acquire(ctx context.Context) (c *Conn, err error) {
|
||||
if p.acquireTracer != nil {
|
||||
ctx = p.acquireTracer.TraceAcquireStart(ctx, p, TraceAcquireStartData{})
|
||||
defer func() {
|
||||
var conn *pgx.Conn
|
||||
if c != nil {
|
||||
conn = c.Conn()
|
||||
}
|
||||
p.acquireTracer.TraceAcquireEnd(ctx, p, TraceAcquireEndData{Conn: conn, Err: err})
|
||||
}()
|
||||
}
|
||||
|
||||
// Try to acquire from the connection pool up to maxConns + 1 times, so that
|
||||
// any that fatal errors would empty the pool and still at least try 1 fresh
|
||||
// connection.
|
||||
for range int(p.maxConns) + 1 {
|
||||
res, err := p.p.Acquire(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
cr := res.Value()
|
||||
|
||||
shouldPingParams := ShouldPingParams{Conn: cr.conn, IdleDuration: res.IdleDuration()}
|
||||
if p.shouldPing(ctx, shouldPingParams) {
|
||||
err := func() error {
|
||||
pingCtx := ctx
|
||||
if p.pingTimeout > 0 {
|
||||
var cancel context.CancelFunc
|
||||
pingCtx, cancel = context.WithTimeout(ctx, p.pingTimeout)
|
||||
defer cancel()
|
||||
}
|
||||
return cr.conn.Ping(pingCtx)
|
||||
}()
|
||||
if err != nil {
|
||||
res.Destroy()
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
if p.prepareConn != nil {
|
||||
ok, err := p.prepareConn(ctx, cr.conn)
|
||||
if !ok {
|
||||
res.Destroy()
|
||||
}
|
||||
if err != nil {
|
||||
if ok {
|
||||
res.Release()
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
return cr.getConn(p, res), nil
|
||||
}
|
||||
return nil, errors.New("pgxpool: too many failed attempts acquiring connection; likely bug in PrepareConn, BeforeAcquire, or ShouldPing hook")
|
||||
}
|
||||
|
||||
// AcquireFunc acquires a [Conn] and calls f with that [Conn]. ctx will only affect the [Pool.Acquire]. It has no effect on the
|
||||
// call of f. The return value is either an error acquiring the [Conn] or the return value of f. The [Conn] is
|
||||
// automatically released after the call of f.
|
||||
func (p *Pool) AcquireFunc(ctx context.Context, f func(*Conn) error) error {
|
||||
conn, err := p.Acquire(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer conn.Release()
|
||||
|
||||
return f(conn)
|
||||
}
|
||||
|
||||
// AcquireAllIdle atomically acquires all currently idle connections. Its intended use is for health check and
|
||||
// keep-alive functionality. It does not update pool statistics.
|
||||
func (p *Pool) AcquireAllIdle(ctx context.Context) []*Conn {
|
||||
resources := p.p.AcquireAllIdle()
|
||||
conns := make([]*Conn, 0, len(resources))
|
||||
for _, res := range resources {
|
||||
cr := res.Value()
|
||||
if p.prepareConn != nil {
|
||||
ok, err := p.prepareConn(ctx, cr.conn)
|
||||
if !ok || err != nil {
|
||||
res.Destroy()
|
||||
continue
|
||||
}
|
||||
}
|
||||
conns = append(conns, cr.getConn(p, res))
|
||||
}
|
||||
|
||||
return conns
|
||||
}
|
||||
|
||||
// Reset closes all connections, but leaves the pool open. It is intended for use when an error is detected that would
|
||||
// disrupt all connections (such as a network interruption or a server state change).
|
||||
//
|
||||
// It is safe to reset a pool while connections are checked out. Those connections will be closed when they are returned
|
||||
// to the pool.
|
||||
func (p *Pool) Reset() {
|
||||
p.p.Reset()
|
||||
}
|
||||
|
||||
// Config returns a copy of config that was used to initialize this [Pool].
|
||||
func (p *Pool) Config() *Config { return p.config.Copy() }
|
||||
|
||||
// Stat returns a pgxpool.Stat struct with a snapshot of Pool statistics.
|
||||
func (p *Pool) Stat() *Stat {
|
||||
return &Stat{
|
||||
s: p.p.Stat(),
|
||||
newConnsCount: atomic.LoadInt64(&p.newConnsCount),
|
||||
lifetimeDestroyCount: atomic.LoadInt64(&p.lifetimeDestroyCount),
|
||||
idleDestroyCount: atomic.LoadInt64(&p.idleDestroyCount),
|
||||
}
|
||||
}
|
||||
|
||||
// Exec acquires a connection from the [Pool] and executes the given SQL.
|
||||
// SQL can be either a prepared statement name or an SQL string.
|
||||
// Arguments should be referenced positionally from the SQL string as $1, $2, etc.
|
||||
// The acquired connection is returned to the pool when the [Pool.Exec] function returns.
|
||||
func (p *Pool) Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) {
|
||||
c, err := p.Acquire(ctx)
|
||||
if err != nil {
|
||||
return pgconn.CommandTag{}, err
|
||||
}
|
||||
defer c.Release()
|
||||
|
||||
return c.Exec(ctx, sql, arguments...)
|
||||
}
|
||||
|
||||
// Query acquires a connection and executes a query that returns [pgx.Rows].
|
||||
// Arguments should be referenced positionally from the SQL string as $1, $2, etc.
|
||||
// See [pgx.Rows] documentation to close the returned [pgx.Rows] and return the acquired connection to the [Pool].
|
||||
//
|
||||
// If there is an error, the returned [pgx.Rows] will be returned in an error state.
|
||||
// If preferred, ignore the error returned from [Pool.Query] and handle errors using the returned [pgx.Rows].
|
||||
//
|
||||
// For extra control over how the query is executed, the types [pgx.QueryExecMode], [pgx.QueryResultFormats], and
|
||||
// [pgx.QueryResultFormatsByOID] may be used as the first args to control exactly how the query is executed. This is rarely
|
||||
// needed. See the documentation for those types for details.
|
||||
func (p *Pool) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) {
|
||||
c, err := p.Acquire(ctx)
|
||||
if err != nil {
|
||||
return errRows{err: err}, err
|
||||
}
|
||||
|
||||
rows, err := c.Query(ctx, sql, args...)
|
||||
if err != nil {
|
||||
c.Release()
|
||||
return errRows{err: err}, err
|
||||
}
|
||||
|
||||
return c.getPoolRows(rows), nil
|
||||
}
|
||||
|
||||
// QueryRow acquires a connection and executes a query that is expected
|
||||
// to return at most one row ([pgx.Row]). Errors are deferred until [pgx.Row]'s
|
||||
// Scan method is called. If the query selects no rows, [pgx.Row]'s Scan will
|
||||
// return [pgx.ErrNoRows]. Otherwise, [pgx.Row]'s Scan scans the first selected row
|
||||
// and discards the rest. The acquired connection is returned to the [Pool] when
|
||||
// [pgx.Row]'s Scan method is called.
|
||||
//
|
||||
// Arguments should be referenced positionally from the SQL string as $1, $2, etc.
|
||||
//
|
||||
// For extra control over how the query is executed, the types [pgx.QueryExecMode], [pgx.QueryResultFormats], and
|
||||
// [pgx.QueryResultFormatsByOID] may be used as the first args to control exactly how the query is executed. This is rarely
|
||||
// needed. See the documentation for those types for details.
|
||||
func (p *Pool) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row {
|
||||
c, err := p.Acquire(ctx)
|
||||
if err != nil {
|
||||
return errRow{err: err}
|
||||
}
|
||||
|
||||
row := c.QueryRow(ctx, sql, args...)
|
||||
return c.getPoolRow(row)
|
||||
}
|
||||
|
||||
func (p *Pool) SendBatch(ctx context.Context, b *pgx.Batch) pgx.BatchResults {
|
||||
c, err := p.Acquire(ctx)
|
||||
if err != nil {
|
||||
return errBatchResults{err: err}
|
||||
}
|
||||
|
||||
br := c.SendBatch(ctx, b)
|
||||
return &poolBatchResults{br: br, c: c}
|
||||
}
|
||||
|
||||
// Begin acquires a connection from the [Pool] and starts a transaction. Unlike [database/sql], the context only affects the begin command. i.e. there is no
|
||||
// auto-rollback on context cancellation. Begin initiates a transaction block without explicitly setting a transaction mode for the block (see [Pool.BeginTx] with [pgx.TxOptions] if transaction mode is required).
|
||||
// [*Tx] is returned, which implements the [pgx.Tx] interface.
|
||||
// [Tx.Commit] or [Tx.Rollback] must be called on the returned transaction to finalize the transaction block.
|
||||
func (p *Pool) Begin(ctx context.Context) (pgx.Tx, error) {
|
||||
return p.BeginTx(ctx, pgx.TxOptions{})
|
||||
}
|
||||
|
||||
// BeginTx acquires a connection from the [Pool] and starts a transaction with [pgx.TxOptions] determining the transaction mode.
|
||||
// Unlike [database/sql], the context only affects the begin command. i.e. there is no auto-rollback on context cancellation.
|
||||
// [*Tx] is returned, which implements the [pgx.Tx] interface.
|
||||
// [Tx.Commit] or [Tx.Rollback] must be called on the returned transaction to finalize the transaction block.
|
||||
func (p *Pool) BeginTx(ctx context.Context, txOptions pgx.TxOptions) (pgx.Tx, error) {
|
||||
c, err := p.Acquire(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
t, err := c.BeginTx(ctx, txOptions)
|
||||
if err != nil {
|
||||
c.Release()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &Tx{t: t, c: c}, nil
|
||||
}
|
||||
|
||||
func (p *Pool) CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNames []string, rowSrc pgx.CopyFromSource) (int64, error) {
|
||||
c, err := p.Acquire(ctx)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer c.Release()
|
||||
|
||||
return c.Conn().CopyFrom(ctx, tableName, columnNames, rowSrc)
|
||||
}
|
||||
|
||||
// Ping acquires a connection from the [Pool] and executes an empty sql statement against it.
|
||||
// If the sql returns without error, the database [Pool.Ping] is considered successful, otherwise, the error is returned.
|
||||
func (p *Pool) Ping(ctx context.Context) error {
|
||||
c, err := p.Acquire(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer c.Release()
|
||||
return c.Ping(ctx)
|
||||
}
|
||||
+116
@@ -0,0 +1,116 @@
|
||||
package pgxpool
|
||||
|
||||
import (
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
)
|
||||
|
||||
type errRows struct {
|
||||
err error
|
||||
}
|
||||
|
||||
func (errRows) Close() {}
|
||||
func (e errRows) Err() error { return e.err }
|
||||
func (errRows) CommandTag() pgconn.CommandTag { return pgconn.CommandTag{} }
|
||||
func (errRows) FieldDescriptions() []pgconn.FieldDescription { return nil }
|
||||
func (errRows) Next() bool { return false }
|
||||
func (e errRows) Scan(dest ...any) error { return e.err }
|
||||
func (e errRows) Values() ([]any, error) { return nil, e.err }
|
||||
func (e errRows) RawValues() [][]byte { return nil }
|
||||
func (e errRows) Conn() *pgx.Conn { return nil }
|
||||
|
||||
type errRow struct {
|
||||
err error
|
||||
}
|
||||
|
||||
func (e errRow) Scan(dest ...any) error { return e.err }
|
||||
|
||||
type poolRows struct {
|
||||
r pgx.Rows
|
||||
c *Conn
|
||||
err error
|
||||
}
|
||||
|
||||
func (rows *poolRows) Close() {
|
||||
rows.r.Close()
|
||||
if rows.c != nil {
|
||||
rows.c.Release()
|
||||
rows.c = nil
|
||||
}
|
||||
}
|
||||
|
||||
func (rows *poolRows) Err() error {
|
||||
if rows.err != nil {
|
||||
return rows.err
|
||||
}
|
||||
return rows.r.Err()
|
||||
}
|
||||
|
||||
func (rows *poolRows) CommandTag() pgconn.CommandTag {
|
||||
return rows.r.CommandTag()
|
||||
}
|
||||
|
||||
func (rows *poolRows) FieldDescriptions() []pgconn.FieldDescription {
|
||||
return rows.r.FieldDescriptions()
|
||||
}
|
||||
|
||||
func (rows *poolRows) Next() bool {
|
||||
if rows.err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
n := rows.r.Next()
|
||||
if !n {
|
||||
rows.Close()
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func (rows *poolRows) Scan(dest ...any) error {
|
||||
err := rows.r.Scan(dest...)
|
||||
if err != nil {
|
||||
rows.Close()
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func (rows *poolRows) Values() ([]any, error) {
|
||||
values, err := rows.r.Values()
|
||||
if err != nil {
|
||||
rows.Close()
|
||||
}
|
||||
return values, err
|
||||
}
|
||||
|
||||
func (rows *poolRows) RawValues() [][]byte {
|
||||
return rows.r.RawValues()
|
||||
}
|
||||
|
||||
func (rows *poolRows) Conn() *pgx.Conn {
|
||||
return rows.r.Conn()
|
||||
}
|
||||
|
||||
type poolRow struct {
|
||||
r pgx.Row
|
||||
c *Conn
|
||||
err error
|
||||
}
|
||||
|
||||
func (row *poolRow) Scan(dest ...any) error {
|
||||
if row.err != nil {
|
||||
return row.err
|
||||
}
|
||||
|
||||
panicked := true
|
||||
defer func() {
|
||||
if panicked && row.c != nil {
|
||||
row.c.Release()
|
||||
}
|
||||
}()
|
||||
err := row.r.Scan(dest...)
|
||||
panicked = false
|
||||
if row.c != nil {
|
||||
row.c.Release()
|
||||
}
|
||||
return err
|
||||
}
|
||||
+91
@@ -0,0 +1,91 @@
|
||||
package pgxpool
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/jackc/puddle/v2"
|
||||
)
|
||||
|
||||
// Stat is a snapshot of Pool statistics.
|
||||
type Stat struct {
|
||||
s *puddle.Stat
|
||||
newConnsCount int64
|
||||
lifetimeDestroyCount int64
|
||||
idleDestroyCount int64
|
||||
}
|
||||
|
||||
// AcquireCount returns the cumulative count of successful acquires from the pool.
|
||||
func (s *Stat) AcquireCount() int64 {
|
||||
return s.s.AcquireCount()
|
||||
}
|
||||
|
||||
// AcquireDuration returns the total duration of all successful acquires from
|
||||
// the pool.
|
||||
func (s *Stat) AcquireDuration() time.Duration {
|
||||
return s.s.AcquireDuration()
|
||||
}
|
||||
|
||||
// AcquiredConns returns the number of currently acquired connections in the pool.
|
||||
func (s *Stat) AcquiredConns() int32 {
|
||||
return s.s.AcquiredResources()
|
||||
}
|
||||
|
||||
// CanceledAcquireCount returns the cumulative count of acquires from the pool
|
||||
// that were canceled by a context.
|
||||
func (s *Stat) CanceledAcquireCount() int64 {
|
||||
return s.s.CanceledAcquireCount()
|
||||
}
|
||||
|
||||
// ConstructingConns returns the number of conns with construction in progress in
|
||||
// the pool.
|
||||
func (s *Stat) ConstructingConns() int32 {
|
||||
return s.s.ConstructingResources()
|
||||
}
|
||||
|
||||
// EmptyAcquireCount returns the cumulative count of successful acquires from the pool
|
||||
// that waited for a resource to be released or constructed because the pool was
|
||||
// empty.
|
||||
func (s *Stat) EmptyAcquireCount() int64 {
|
||||
return s.s.EmptyAcquireCount()
|
||||
}
|
||||
|
||||
// IdleConns returns the number of currently idle conns in the pool.
|
||||
func (s *Stat) IdleConns() int32 {
|
||||
return s.s.IdleResources()
|
||||
}
|
||||
|
||||
// MaxConns returns the maximum size of the pool.
|
||||
func (s *Stat) MaxConns() int32 {
|
||||
return s.s.MaxResources()
|
||||
}
|
||||
|
||||
// TotalConns returns the total number of resources currently in the pool.
|
||||
// The value is the sum of ConstructingConns, AcquiredConns, and
|
||||
// IdleConns.
|
||||
func (s *Stat) TotalConns() int32 {
|
||||
return s.s.TotalResources()
|
||||
}
|
||||
|
||||
// NewConnsCount returns the cumulative count of new connections opened.
|
||||
func (s *Stat) NewConnsCount() int64 {
|
||||
return s.newConnsCount
|
||||
}
|
||||
|
||||
// MaxLifetimeDestroyCount returns the cumulative count of connections destroyed
|
||||
// because they exceeded MaxConnLifetime.
|
||||
func (s *Stat) MaxLifetimeDestroyCount() int64 {
|
||||
return s.lifetimeDestroyCount
|
||||
}
|
||||
|
||||
// MaxIdleDestroyCount returns the cumulative count of connections destroyed because
|
||||
// they exceeded MaxConnIdleTime.
|
||||
func (s *Stat) MaxIdleDestroyCount() int64 {
|
||||
return s.idleDestroyCount
|
||||
}
|
||||
|
||||
// EmptyAcquireWaitTime returns the cumulative time waited for successful acquires
|
||||
// from the pool for a resource to be released or constructed because the pool was
|
||||
// empty.
|
||||
func (s *Stat) EmptyAcquireWaitTime() time.Duration {
|
||||
return s.s.EmptyAcquireWaitTime()
|
||||
}
|
||||
+33
@@ -0,0 +1,33 @@
|
||||
package pgxpool
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
// AcquireTracer traces Acquire.
|
||||
type AcquireTracer interface {
|
||||
// TraceAcquireStart is called at the beginning of Acquire.
|
||||
// The returned context is used for the rest of the call and will be passed to the TraceAcquireEnd.
|
||||
TraceAcquireStart(ctx context.Context, pool *Pool, data TraceAcquireStartData) context.Context
|
||||
// TraceAcquireEnd is called when a connection has been acquired.
|
||||
TraceAcquireEnd(ctx context.Context, pool *Pool, data TraceAcquireEndData)
|
||||
}
|
||||
|
||||
type TraceAcquireStartData struct{}
|
||||
|
||||
type TraceAcquireEndData struct {
|
||||
Conn *pgx.Conn
|
||||
Err error
|
||||
}
|
||||
|
||||
// ReleaseTracer traces Release.
|
||||
type ReleaseTracer interface {
|
||||
// TraceRelease is called at the beginning of Release.
|
||||
TraceRelease(pool *Pool, data TraceReleaseData)
|
||||
}
|
||||
|
||||
type TraceReleaseData struct {
|
||||
Conn *pgx.Conn
|
||||
}
|
||||
+83
@@ -0,0 +1,83 @@
|
||||
package pgxpool
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
)
|
||||
|
||||
// Tx represents a database transaction acquired from a Pool.
|
||||
type Tx struct {
|
||||
t pgx.Tx
|
||||
c *Conn
|
||||
}
|
||||
|
||||
// Begin starts a pseudo nested transaction implemented with a savepoint.
|
||||
func (tx *Tx) Begin(ctx context.Context) (pgx.Tx, error) {
|
||||
return tx.t.Begin(ctx)
|
||||
}
|
||||
|
||||
// Commit commits the transaction and returns the associated connection back to the Pool. Commit will return an error
|
||||
// where errors.Is(ErrTxClosed) is true if the Tx is already closed, but is otherwise safe to call multiple times. If
|
||||
// the commit fails with a rollback status (e.g. the transaction was already in a broken state) then ErrTxCommitRollback
|
||||
// will be returned.
|
||||
func (tx *Tx) Commit(ctx context.Context) error {
|
||||
err := tx.t.Commit(ctx)
|
||||
if tx.c != nil {
|
||||
tx.c.Release()
|
||||
tx.c = nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// Rollback rolls back the transaction and returns the associated connection back to the Pool. Rollback will return
|
||||
// where an error where errors.Is(ErrTxClosed) is true if the Tx is already closed, but is otherwise safe to call
|
||||
// multiple times. Hence, defer tx.Rollback() is safe even if tx.Commit() will be called first in a non-error condition.
|
||||
func (tx *Tx) Rollback(ctx context.Context) error {
|
||||
err := tx.t.Rollback(ctx)
|
||||
if tx.c != nil {
|
||||
tx.c.Release()
|
||||
tx.c = nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func (tx *Tx) CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNames []string, rowSrc pgx.CopyFromSource) (int64, error) {
|
||||
return tx.t.CopyFrom(ctx, tableName, columnNames, rowSrc)
|
||||
}
|
||||
|
||||
func (tx *Tx) SendBatch(ctx context.Context, b *pgx.Batch) pgx.BatchResults {
|
||||
return tx.t.SendBatch(ctx, b)
|
||||
}
|
||||
|
||||
func (tx *Tx) LargeObjects() pgx.LargeObjects {
|
||||
return tx.t.LargeObjects()
|
||||
}
|
||||
|
||||
// Prepare creates a prepared statement with name and sql. If the name is empty,
|
||||
// an anonymous prepared statement will be used. sql can contain placeholders
|
||||
// for bound parameters. These placeholders are referenced positionally as $1, $2, etc.
|
||||
//
|
||||
// Prepare is idempotent; i.e. it is safe to call Prepare multiple times with the same
|
||||
// name and sql arguments. This allows a code path to Prepare and Query/Exec without
|
||||
// needing to first check whether the statement has already been prepared.
|
||||
func (tx *Tx) Prepare(ctx context.Context, name, sql string) (*pgconn.StatementDescription, error) {
|
||||
return tx.t.Prepare(ctx, name, sql)
|
||||
}
|
||||
|
||||
func (tx *Tx) Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) {
|
||||
return tx.t.Exec(ctx, sql, arguments...)
|
||||
}
|
||||
|
||||
func (tx *Tx) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) {
|
||||
return tx.t.Query(ctx, sql, args...)
|
||||
}
|
||||
|
||||
func (tx *Tx) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row {
|
||||
return tx.t.QueryRow(ctx, sql, args...)
|
||||
}
|
||||
|
||||
func (tx *Tx) Conn() *pgx.Conn {
|
||||
return tx.t.Conn()
|
||||
}
|
||||
+909
@@ -0,0 +1,909 @@
|
||||
// Package stdlib is the compatibility layer from pgx to database/sql.
|
||||
//
|
||||
// A database/sql connection can be established through sql.Open.
|
||||
//
|
||||
// db, err := sql.Open("pgx", "postgres://pgx_md5:secret@localhost:5432/pgx_test?sslmode=disable")
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
//
|
||||
// Or from a keyword/value string.
|
||||
//
|
||||
// db, err := sql.Open("pgx", "user=postgres password=secret host=localhost port=5432 database=pgx_test sslmode=disable")
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
//
|
||||
// Or from a *pgxpool.Pool.
|
||||
//
|
||||
// pool, err := pgxpool.New(context.Background(), os.Getenv("DATABASE_URL"))
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
//
|
||||
// db := stdlib.OpenDBFromPool(pool)
|
||||
//
|
||||
// Or a pgx.ConnConfig can be used to set configuration not accessible via connection string. In this case the
|
||||
// pgx.ConnConfig must first be registered with the driver. This registration returns a connection string which is used
|
||||
// with sql.Open.
|
||||
//
|
||||
// connConfig, _ := pgx.ParseConfig(os.Getenv("DATABASE_URL"))
|
||||
// connConfig.Tracer = &tracelog.TraceLog{Logger: myLogger, LogLevel: tracelog.LogLevelInfo}
|
||||
// connStr := stdlib.RegisterConnConfig(connConfig)
|
||||
// db, _ := sql.Open("pgx", connStr)
|
||||
//
|
||||
// pgx uses standard PostgreSQL positional parameters in queries. e.g. $1, $2. It does not support named parameters.
|
||||
//
|
||||
// db.QueryRow("select * from users where id=$1", userID)
|
||||
//
|
||||
// (*sql.Conn) Raw() can be used to get a *pgx.Conn from the standard database/sql.DB connection pool. This allows
|
||||
// operations that use pgx specific functionality.
|
||||
//
|
||||
// // Given db is a *sql.DB
|
||||
// conn, err := db.Conn(context.Background())
|
||||
// if err != nil {
|
||||
// // handle error from acquiring connection from DB pool
|
||||
// }
|
||||
//
|
||||
// err = conn.Raw(func(driverConn any) error {
|
||||
// conn := driverConn.(*stdlib.Conn).Conn() // conn is a *pgx.Conn
|
||||
// // Do pgx specific stuff with conn
|
||||
// conn.CopyFrom(...)
|
||||
// return nil
|
||||
// })
|
||||
// if err != nil {
|
||||
// // handle error that occurred while using *pgx.Conn
|
||||
// }
|
||||
//
|
||||
// # PostgreSQL Specific Data Types
|
||||
//
|
||||
// The pgtype package provides support for PostgreSQL specific types. *pgtype.Map.SQLScanner is an adapter that makes
|
||||
// these types usable as a sql.Scanner.
|
||||
//
|
||||
// m := pgtype.NewMap()
|
||||
// var a []int64
|
||||
// err := db.QueryRow("select '{1,2,3}'::bigint[]").Scan(m.SQLScanner(&a))
|
||||
package stdlib
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"database/sql/driver"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"math/rand/v2"
|
||||
"reflect"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
"github.com/jackc/pgx/v5/pgtype"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
// Only intrinsic types should be binary format with database/sql.
|
||||
var databaseSQLResultFormats pgx.QueryResultFormatsByOID
|
||||
|
||||
var pgxDriver *Driver
|
||||
|
||||
func init() {
|
||||
pgxDriver = &Driver{
|
||||
configs: make(map[string]*pgx.ConnConfig),
|
||||
}
|
||||
|
||||
// if pgx driver was already registered by different pgx major version then we
|
||||
// skip registration under the default name.
|
||||
if !slices.Contains(sql.Drivers(), "pgx") {
|
||||
sql.Register("pgx", pgxDriver)
|
||||
}
|
||||
sql.Register("pgx/v5", pgxDriver)
|
||||
|
||||
databaseSQLResultFormats = pgx.QueryResultFormatsByOID{
|
||||
pgtype.BoolOID: 1,
|
||||
pgtype.ByteaOID: 1,
|
||||
pgtype.CIDOID: 1,
|
||||
pgtype.DateOID: 1,
|
||||
pgtype.Float4OID: 1,
|
||||
pgtype.Float8OID: 1,
|
||||
pgtype.Int2OID: 1,
|
||||
pgtype.Int4OID: 1,
|
||||
pgtype.Int8OID: 1,
|
||||
pgtype.OIDOID: 1,
|
||||
pgtype.TimestampOID: 1,
|
||||
pgtype.TimestamptzOID: 1,
|
||||
pgtype.XIDOID: 1,
|
||||
}
|
||||
}
|
||||
|
||||
// OptionOpenDB options for configuring the driver when opening a new db pool.
|
||||
type OptionOpenDB func(*connector)
|
||||
|
||||
// ShouldPingParams are passed to OptionShouldPing to decide whether to ping before reusing a connection.
|
||||
type ShouldPingParams struct {
|
||||
// Conn is the underlying pgx connection.
|
||||
Conn *pgx.Conn
|
||||
// IdleDuration is how long it has been since ResetSession last ran.
|
||||
IdleDuration time.Duration
|
||||
}
|
||||
|
||||
// OptionShouldPing controls whether stdlib should issue a liveness ping before reusing a connection.
|
||||
// If the function returns true, stdlib will ping.
|
||||
// If it returns false, stdlib will skip the ping.
|
||||
// If not provided, default is ping only when IdleDuration > 1s.
|
||||
func OptionShouldPing(f func(context.Context, ShouldPingParams) bool) OptionOpenDB {
|
||||
return func(dc *connector) { dc.ShouldPing = f }
|
||||
}
|
||||
|
||||
// OptionBeforeConnect provides a callback for before connect. It is passed a shallow copy of the ConnConfig that will
|
||||
// be used to connect, so only its immediate members should be modified. Used only if db is opened with *pgx.ConnConfig.
|
||||
func OptionBeforeConnect(bc func(context.Context, *pgx.ConnConfig) error) OptionOpenDB {
|
||||
return func(dc *connector) {
|
||||
dc.BeforeConnect = bc
|
||||
}
|
||||
}
|
||||
|
||||
// OptionAfterConnect provides a callback for after connect. Used only if db is opened with *pgx.ConnConfig.
|
||||
func OptionAfterConnect(ac func(context.Context, *pgx.Conn) error) OptionOpenDB {
|
||||
return func(dc *connector) {
|
||||
dc.AfterConnect = ac
|
||||
}
|
||||
}
|
||||
|
||||
// OptionResetSession provides a callback that can be used to add custom logic prior to executing a query on the
|
||||
// connection if the connection has been used before.
|
||||
// If ResetSessionFunc returns ErrBadConn error the connection will be discarded.
|
||||
func OptionResetSession(rs func(context.Context, *pgx.Conn) error) OptionOpenDB {
|
||||
return func(dc *connector) {
|
||||
dc.ResetSession = rs
|
||||
}
|
||||
}
|
||||
|
||||
// RandomizeHostOrderFunc is a BeforeConnect hook that randomizes the host order in the provided connConfig, so that a
|
||||
// new host becomes primary each time. This is useful to distribute connections for multi-master databases like
|
||||
// CockroachDB. If you use this you likely should set https://golang.org/pkg/database/sql/#DB.SetConnMaxLifetime as well
|
||||
// to ensure that connections are periodically rebalanced across your nodes.
|
||||
func RandomizeHostOrderFunc(ctx context.Context, connConfig *pgx.ConnConfig) error {
|
||||
if len(connConfig.Fallbacks) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
newFallbacks := append([]*pgconn.FallbackConfig{{
|
||||
Host: connConfig.Host,
|
||||
Port: connConfig.Port,
|
||||
TLSConfig: connConfig.TLSConfig,
|
||||
}}, connConfig.Fallbacks...)
|
||||
|
||||
rand.Shuffle(len(newFallbacks), func(i, j int) {
|
||||
newFallbacks[i], newFallbacks[j] = newFallbacks[j], newFallbacks[i]
|
||||
})
|
||||
|
||||
// Use the one that sorted last as the primary and keep the rest as the fallbacks
|
||||
newPrimary := newFallbacks[len(newFallbacks)-1]
|
||||
connConfig.Host = newPrimary.Host
|
||||
connConfig.Port = newPrimary.Port
|
||||
connConfig.TLSConfig = newPrimary.TLSConfig
|
||||
connConfig.Fallbacks = newFallbacks[:len(newFallbacks)-1]
|
||||
return nil
|
||||
}
|
||||
|
||||
func GetConnector(config pgx.ConnConfig, opts ...OptionOpenDB) driver.Connector {
|
||||
c := connector{
|
||||
ConnConfig: config,
|
||||
BeforeConnect: func(context.Context, *pgx.ConnConfig) error { return nil }, // noop before connect by default
|
||||
AfterConnect: func(context.Context, *pgx.Conn) error { return nil }, // noop after connect by default
|
||||
ResetSession: func(context.Context, *pgx.Conn) error { return nil }, // noop reset session by default
|
||||
driver: pgxDriver,
|
||||
}
|
||||
|
||||
for _, opt := range opts {
|
||||
opt(&c)
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
// GetPoolConnector creates a new driver.Connector from the given *pgxpool.Pool. By using this be sure to set the
|
||||
// maximum idle connections of the *sql.DB created with this connector to zero since they must be managed from the
|
||||
// *pgxpool.Pool. This is required to avoid acquiring all the connections from the pgxpool and starving any direct
|
||||
// users of the pgxpool.
|
||||
func GetPoolConnector(pool *pgxpool.Pool, opts ...OptionOpenDB) driver.Connector {
|
||||
c := connector{
|
||||
pool: pool,
|
||||
ResetSession: func(context.Context, *pgx.Conn) error { return nil }, // noop reset session by default
|
||||
driver: pgxDriver,
|
||||
}
|
||||
|
||||
for _, opt := range opts {
|
||||
opt(&c)
|
||||
}
|
||||
|
||||
return c
|
||||
}
|
||||
|
||||
func OpenDB(config pgx.ConnConfig, opts ...OptionOpenDB) *sql.DB {
|
||||
c := GetConnector(config, opts...)
|
||||
return sql.OpenDB(c)
|
||||
}
|
||||
|
||||
// OpenDBFromPool creates a new *sql.DB from the given *pgxpool.Pool. Note that this method automatically sets the
|
||||
// maximum number of idle connections in *sql.DB to zero, since they must be managed from the *pgxpool.Pool. This is
|
||||
// required to avoid acquiring all the connections from the pgxpool and starving any direct users of the pgxpool. Note
|
||||
// that closing the returned *sql.DB will not close the *pgxpool.Pool.
|
||||
func OpenDBFromPool(pool *pgxpool.Pool, opts ...OptionOpenDB) *sql.DB {
|
||||
c := GetPoolConnector(pool, opts...)
|
||||
db := sql.OpenDB(c)
|
||||
db.SetMaxIdleConns(0)
|
||||
return db
|
||||
}
|
||||
|
||||
type connector struct {
|
||||
pgx.ConnConfig
|
||||
pool *pgxpool.Pool
|
||||
BeforeConnect func(context.Context, *pgx.ConnConfig) error // function to call before creation of every new connection
|
||||
AfterConnect func(context.Context, *pgx.Conn) error // function to call after creation of every new connection
|
||||
ResetSession func(context.Context, *pgx.Conn) error // function is called before a connection is reused
|
||||
ShouldPing func(context.Context, ShouldPingParams) bool // function to decide if stdlib should ping before reusing a connection
|
||||
driver *Driver
|
||||
}
|
||||
|
||||
// Connect implement driver.Connector interface
|
||||
func (c connector) Connect(ctx context.Context) (driver.Conn, error) {
|
||||
var (
|
||||
connConfig pgx.ConnConfig
|
||||
conn *pgx.Conn
|
||||
close func(context.Context) error
|
||||
err error
|
||||
)
|
||||
|
||||
if c.pool == nil {
|
||||
// Create a shallow copy of the config, so that BeforeConnect can safely modify it
|
||||
connConfig = c.ConnConfig
|
||||
|
||||
if err = c.BeforeConnect(ctx, &connConfig); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if conn, err = pgx.ConnectConfig(ctx, &connConfig); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err = c.AfterConnect(ctx, conn); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
close = conn.Close
|
||||
} else {
|
||||
var pconn *pgxpool.Conn
|
||||
|
||||
pconn, err = c.pool.Acquire(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
conn = pconn.Conn()
|
||||
|
||||
close = func(_ context.Context) error {
|
||||
pconn.Release()
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
return &Conn{
|
||||
conn: conn,
|
||||
close: close,
|
||||
driver: c.driver,
|
||||
connConfig: connConfig,
|
||||
resetSessionFunc: c.ResetSession,
|
||||
shouldPing: c.ShouldPing,
|
||||
psRefCounts: make(map[*pgconn.StatementDescription]int),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Driver implement driver.Connector interface
|
||||
func (c connector) Driver() driver.Driver {
|
||||
return c.driver
|
||||
}
|
||||
|
||||
// GetDefaultDriver returns the driver initialized in the init function
|
||||
// and used when the pgx driver is registered.
|
||||
func GetDefaultDriver() driver.Driver {
|
||||
return pgxDriver
|
||||
}
|
||||
|
||||
type Driver struct {
|
||||
configMutex sync.Mutex
|
||||
configs map[string]*pgx.ConnConfig
|
||||
sequence int
|
||||
}
|
||||
|
||||
func (d *Driver) Open(name string) (driver.Conn, error) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) // Ensure eventual timeout
|
||||
defer cancel()
|
||||
|
||||
connector, err := d.OpenConnector(name)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return connector.Connect(ctx)
|
||||
}
|
||||
|
||||
func (d *Driver) OpenConnector(name string) (driver.Connector, error) {
|
||||
return &driverConnector{driver: d, name: name}, nil
|
||||
}
|
||||
|
||||
func (d *Driver) registerConnConfig(c *pgx.ConnConfig) string {
|
||||
d.configMutex.Lock()
|
||||
connStr := fmt.Sprintf("registeredConnConfig%d", d.sequence)
|
||||
d.sequence++
|
||||
d.configs[connStr] = c
|
||||
d.configMutex.Unlock()
|
||||
return connStr
|
||||
}
|
||||
|
||||
func (d *Driver) unregisterConnConfig(connStr string) {
|
||||
d.configMutex.Lock()
|
||||
delete(d.configs, connStr)
|
||||
d.configMutex.Unlock()
|
||||
}
|
||||
|
||||
type driverConnector struct {
|
||||
driver *Driver
|
||||
name string
|
||||
}
|
||||
|
||||
func (dc *driverConnector) Connect(ctx context.Context) (driver.Conn, error) {
|
||||
var connConfig *pgx.ConnConfig
|
||||
|
||||
dc.driver.configMutex.Lock()
|
||||
connConfig = dc.driver.configs[dc.name]
|
||||
dc.driver.configMutex.Unlock()
|
||||
|
||||
if connConfig == nil {
|
||||
var err error
|
||||
connConfig, err = pgx.ParseConfig(dc.name)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
conn, err := pgx.ConnectConfig(ctx, connConfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
c := &Conn{
|
||||
conn: conn,
|
||||
close: conn.Close,
|
||||
driver: dc.driver,
|
||||
connConfig: *connConfig,
|
||||
resetSessionFunc: func(context.Context, *pgx.Conn) error { return nil },
|
||||
psRefCounts: make(map[*pgconn.StatementDescription]int),
|
||||
}
|
||||
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (dc *driverConnector) Driver() driver.Driver {
|
||||
return dc.driver
|
||||
}
|
||||
|
||||
// RegisterConnConfig registers a ConnConfig and returns the connection string to use with Open.
|
||||
func RegisterConnConfig(c *pgx.ConnConfig) string {
|
||||
return pgxDriver.registerConnConfig(c)
|
||||
}
|
||||
|
||||
// UnregisterConnConfig removes the ConnConfig registration for connStr.
|
||||
func UnregisterConnConfig(connStr string) {
|
||||
pgxDriver.unregisterConnConfig(connStr)
|
||||
}
|
||||
|
||||
type Conn struct {
|
||||
conn *pgx.Conn
|
||||
close func(context.Context) error
|
||||
driver *Driver
|
||||
connConfig pgx.ConnConfig
|
||||
resetSessionFunc func(context.Context, *pgx.Conn) error // Function is called before a connection is reused
|
||||
shouldPing func(context.Context, ShouldPingParams) bool // Function to decide if stdlib should ping before reusing a connection
|
||||
lastResetSessionTime time.Time
|
||||
|
||||
// psRefCounts contains reference counts for prepared statements. Prepare uses the underlying pgx logic to generate
|
||||
// deterministic statement names from the statement text. If this query has already been prepared then the existing
|
||||
// *pgconn.StatementDescription will be returned. However, this means that if Close is called on the returned Stmt
|
||||
// then the underlying prepared statement will be closed even when the underlying prepared statement is still in use
|
||||
// by another database/sql Stmt. To prevent this psRefCounts keeps track of how many database/sql statements are using
|
||||
// the same underlying statement and only closes the underlying statement when the reference count reaches 0.
|
||||
psRefCounts map[*pgconn.StatementDescription]int
|
||||
}
|
||||
|
||||
// Conn returns the underlying *pgx.Conn
|
||||
func (c *Conn) Conn() *pgx.Conn {
|
||||
return c.conn
|
||||
}
|
||||
|
||||
func (c *Conn) Prepare(query string) (driver.Stmt, error) {
|
||||
return c.PrepareContext(context.Background(), query)
|
||||
}
|
||||
|
||||
func (c *Conn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
|
||||
if c.conn.IsClosed() {
|
||||
return nil, driver.ErrBadConn
|
||||
}
|
||||
|
||||
sd, err := c.conn.Prepare(ctx, query, query)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.psRefCounts[sd]++
|
||||
|
||||
return &Stmt{sd: sd, conn: c}, nil
|
||||
}
|
||||
|
||||
func (c *Conn) Close() error {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
|
||||
defer cancel()
|
||||
return c.close(ctx)
|
||||
}
|
||||
|
||||
func (c *Conn) Begin() (driver.Tx, error) {
|
||||
return c.BeginTx(context.Background(), driver.TxOptions{})
|
||||
}
|
||||
|
||||
func (c *Conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
|
||||
if c.conn.IsClosed() {
|
||||
return nil, driver.ErrBadConn
|
||||
}
|
||||
|
||||
var pgxOpts pgx.TxOptions
|
||||
switch sql.IsolationLevel(opts.Isolation) {
|
||||
case sql.LevelDefault:
|
||||
case sql.LevelReadUncommitted:
|
||||
pgxOpts.IsoLevel = pgx.ReadUncommitted
|
||||
case sql.LevelReadCommitted:
|
||||
pgxOpts.IsoLevel = pgx.ReadCommitted
|
||||
case sql.LevelRepeatableRead, sql.LevelSnapshot:
|
||||
pgxOpts.IsoLevel = pgx.RepeatableRead
|
||||
case sql.LevelSerializable:
|
||||
pgxOpts.IsoLevel = pgx.Serializable
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported isolation: %v", opts.Isolation)
|
||||
}
|
||||
|
||||
if opts.ReadOnly {
|
||||
pgxOpts.AccessMode = pgx.ReadOnly
|
||||
}
|
||||
|
||||
tx, err := c.conn.BeginTx(ctx, pgxOpts)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return wrapTx{ctx: ctx, tx: tx}, nil
|
||||
}
|
||||
|
||||
func (c *Conn) ExecContext(ctx context.Context, query string, argsV []driver.NamedValue) (driver.Result, error) {
|
||||
if c.conn.IsClosed() {
|
||||
return nil, driver.ErrBadConn
|
||||
}
|
||||
|
||||
args := make([]any, len(argsV))
|
||||
convertNamedArguments(args, argsV)
|
||||
|
||||
commandTag, err := c.conn.Exec(ctx, query, args...)
|
||||
// if we got a network error before we had a chance to send the query, retry
|
||||
if err != nil {
|
||||
if pgconn.SafeToRetry(err) {
|
||||
return nil, driver.ErrBadConn
|
||||
}
|
||||
}
|
||||
return driver.RowsAffected(commandTag.RowsAffected()), err
|
||||
}
|
||||
|
||||
func (c *Conn) QueryContext(ctx context.Context, query string, argsV []driver.NamedValue) (driver.Rows, error) {
|
||||
if c.conn.IsClosed() {
|
||||
return nil, driver.ErrBadConn
|
||||
}
|
||||
|
||||
args := make([]any, 1+len(argsV))
|
||||
args[0] = databaseSQLResultFormats
|
||||
convertNamedArguments(args[1:], argsV)
|
||||
|
||||
rows, err := c.conn.Query(ctx, query, args...)
|
||||
if err != nil {
|
||||
if pgconn.SafeToRetry(err) {
|
||||
return nil, driver.ErrBadConn
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Preload first row because otherwise we won't know what columns are available when database/sql asks.
|
||||
more := rows.Next()
|
||||
if err = rows.Err(); err != nil {
|
||||
rows.Close()
|
||||
return nil, err
|
||||
}
|
||||
return &Rows{conn: c, rows: rows, skipNext: true, skipNextMore: more}, nil
|
||||
}
|
||||
|
||||
func (c *Conn) Ping(ctx context.Context) error {
|
||||
if c.conn.IsClosed() {
|
||||
return driver.ErrBadConn
|
||||
}
|
||||
|
||||
err := c.conn.Ping(ctx)
|
||||
if err != nil {
|
||||
// A Ping failure implies some sort of fatal state. The connection is almost certainly already closed by the
|
||||
// failure, but manually close it just to be sure.
|
||||
c.Close()
|
||||
return driver.ErrBadConn
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Conn) CheckNamedValue(*driver.NamedValue) error {
|
||||
// Underlying pgx supports sql.Scanner and driver.Valuer interfaces natively. So everything can be passed through directly.
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Conn) ResetSession(ctx context.Context) error {
|
||||
if c.conn.IsClosed() {
|
||||
return driver.ErrBadConn
|
||||
}
|
||||
|
||||
// Discard connection if it has an open transaction. This can happen if the
|
||||
// application did not properly commit or rollback a transaction.
|
||||
if c.conn.PgConn().TxStatus() != 'I' {
|
||||
return driver.ErrBadConn
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
idle := now.Sub(c.lastResetSessionTime)
|
||||
|
||||
doPing := idle > time.Second // default behavior: ping only if idle > 1s
|
||||
|
||||
if c.shouldPing != nil {
|
||||
doPing = c.shouldPing(ctx, ShouldPingParams{
|
||||
Conn: c.conn,
|
||||
IdleDuration: idle,
|
||||
})
|
||||
}
|
||||
|
||||
if doPing {
|
||||
if err := c.conn.PgConn().Ping(ctx); err != nil {
|
||||
return driver.ErrBadConn
|
||||
}
|
||||
}
|
||||
|
||||
c.lastResetSessionTime = now
|
||||
|
||||
return c.resetSessionFunc(ctx, c.conn)
|
||||
}
|
||||
|
||||
type Stmt struct {
|
||||
sd *pgconn.StatementDescription
|
||||
conn *Conn
|
||||
}
|
||||
|
||||
func (s *Stmt) Close() error {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
|
||||
defer cancel()
|
||||
|
||||
refCount := s.conn.psRefCounts[s.sd]
|
||||
if refCount == 1 {
|
||||
delete(s.conn.psRefCounts, s.sd)
|
||||
} else {
|
||||
s.conn.psRefCounts[s.sd]--
|
||||
return nil
|
||||
}
|
||||
|
||||
return s.conn.conn.Deallocate(ctx, s.sd.SQL)
|
||||
}
|
||||
|
||||
func (s *Stmt) NumInput() int {
|
||||
return len(s.sd.ParamOIDs)
|
||||
}
|
||||
|
||||
func (s *Stmt) Exec(argsV []driver.Value) (driver.Result, error) {
|
||||
return nil, errors.New("Stmt.Exec deprecated and not implemented")
|
||||
}
|
||||
|
||||
func (s *Stmt) ExecContext(ctx context.Context, argsV []driver.NamedValue) (driver.Result, error) {
|
||||
return s.conn.ExecContext(ctx, s.sd.SQL, argsV)
|
||||
}
|
||||
|
||||
func (s *Stmt) Query(argsV []driver.Value) (driver.Rows, error) {
|
||||
return nil, errors.New("Stmt.Query deprecated and not implemented")
|
||||
}
|
||||
|
||||
func (s *Stmt) QueryContext(ctx context.Context, argsV []driver.NamedValue) (driver.Rows, error) {
|
||||
return s.conn.QueryContext(ctx, s.sd.SQL, argsV)
|
||||
}
|
||||
|
||||
type rowValueFunc func(src []byte) (driver.Value, error)
|
||||
|
||||
type Rows struct {
|
||||
conn *Conn
|
||||
rows pgx.Rows
|
||||
valueFuncs []rowValueFunc
|
||||
skipNext bool
|
||||
skipNextMore bool
|
||||
|
||||
columnNames []string
|
||||
}
|
||||
|
||||
func (r *Rows) Columns() []string {
|
||||
if r.columnNames == nil {
|
||||
fields := r.rows.FieldDescriptions()
|
||||
r.columnNames = make([]string, len(fields))
|
||||
for i, fd := range fields {
|
||||
r.columnNames[i] = string(fd.Name)
|
||||
}
|
||||
}
|
||||
|
||||
return r.columnNames
|
||||
}
|
||||
|
||||
// ColumnTypeDatabaseTypeName returns the database system type name. If the name is unknown the OID is returned.
|
||||
func (r *Rows) ColumnTypeDatabaseTypeName(index int) string {
|
||||
if dt, ok := r.conn.conn.TypeMap().TypeForOID(r.rows.FieldDescriptions()[index].DataTypeOID); ok {
|
||||
return strings.ToUpper(dt.Name)
|
||||
}
|
||||
|
||||
return strconv.FormatInt(int64(r.rows.FieldDescriptions()[index].DataTypeOID), 10)
|
||||
}
|
||||
|
||||
const varHeaderSize = 4
|
||||
|
||||
// ColumnTypeLength returns the length of the column type if the column is a
|
||||
// variable length type. If the column is not a variable length type ok
|
||||
// should return false.
|
||||
func (r *Rows) ColumnTypeLength(index int) (int64, bool) {
|
||||
fd := r.rows.FieldDescriptions()[index]
|
||||
|
||||
switch fd.DataTypeOID {
|
||||
case pgtype.TextOID, pgtype.ByteaOID:
|
||||
return math.MaxInt64, true
|
||||
case pgtype.VarcharOID, pgtype.BPCharOID:
|
||||
return int64(fd.TypeModifier - varHeaderSize), true
|
||||
case pgtype.VarbitOID:
|
||||
return int64(fd.TypeModifier), true
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
|
||||
// ColumnTypePrecisionScale should return the precision and scale for decimal
|
||||
// types. If not applicable, ok should be false.
|
||||
func (r *Rows) ColumnTypePrecisionScale(index int) (precision, scale int64, ok bool) {
|
||||
fd := r.rows.FieldDescriptions()[index]
|
||||
|
||||
switch fd.DataTypeOID {
|
||||
case pgtype.NumericOID:
|
||||
mod := fd.TypeModifier - varHeaderSize
|
||||
precision = int64((mod >> 16) & 0xffff)
|
||||
scale = int64(mod & 0xffff)
|
||||
return precision, scale, true
|
||||
default:
|
||||
return 0, 0, false
|
||||
}
|
||||
}
|
||||
|
||||
// ColumnTypeScanType returns the value type that can be used to scan types into.
|
||||
func (r *Rows) ColumnTypeScanType(index int) reflect.Type {
|
||||
fd := r.rows.FieldDescriptions()[index]
|
||||
|
||||
switch fd.DataTypeOID {
|
||||
case pgtype.Float8OID:
|
||||
return reflect.TypeFor[float64]()
|
||||
case pgtype.Float4OID:
|
||||
return reflect.TypeFor[float32]()
|
||||
case pgtype.Int8OID:
|
||||
return reflect.TypeFor[int64]()
|
||||
case pgtype.Int4OID:
|
||||
return reflect.TypeFor[int32]()
|
||||
case pgtype.Int2OID:
|
||||
return reflect.TypeFor[int16]()
|
||||
case pgtype.BoolOID:
|
||||
return reflect.TypeFor[bool]()
|
||||
case pgtype.NumericOID:
|
||||
return reflect.TypeFor[float64]()
|
||||
case pgtype.DateOID, pgtype.TimestampOID, pgtype.TimestamptzOID:
|
||||
return reflect.TypeFor[time.Time]()
|
||||
case pgtype.ByteaOID:
|
||||
return reflect.TypeFor[[]byte]()
|
||||
default:
|
||||
return reflect.TypeFor[string]()
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Rows) Close() error {
|
||||
r.rows.Close()
|
||||
return r.rows.Err()
|
||||
}
|
||||
|
||||
func (r *Rows) Next(dest []driver.Value) error {
|
||||
m := r.conn.conn.TypeMap()
|
||||
fieldDescriptions := r.rows.FieldDescriptions()
|
||||
|
||||
if r.valueFuncs == nil {
|
||||
r.valueFuncs = make([]rowValueFunc, len(fieldDescriptions))
|
||||
|
||||
for i, fd := range fieldDescriptions {
|
||||
dataTypeOID := fd.DataTypeOID
|
||||
format := fd.Format
|
||||
|
||||
switch fd.DataTypeOID {
|
||||
case pgtype.BoolOID:
|
||||
var d bool
|
||||
scanPlan := m.PlanScan(dataTypeOID, format, &d)
|
||||
r.valueFuncs[i] = func(src []byte) (driver.Value, error) {
|
||||
err := scanPlan.Scan(src, &d)
|
||||
return d, err
|
||||
}
|
||||
case pgtype.ByteaOID:
|
||||
var d []byte
|
||||
scanPlan := m.PlanScan(dataTypeOID, format, &d)
|
||||
r.valueFuncs[i] = func(src []byte) (driver.Value, error) {
|
||||
err := scanPlan.Scan(src, &d)
|
||||
return d, err
|
||||
}
|
||||
case pgtype.CIDOID, pgtype.OIDOID, pgtype.XIDOID:
|
||||
var d pgtype.Uint32
|
||||
scanPlan := m.PlanScan(dataTypeOID, format, &d)
|
||||
r.valueFuncs[i] = func(src []byte) (driver.Value, error) {
|
||||
err := scanPlan.Scan(src, &d)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return d.Value()
|
||||
}
|
||||
case pgtype.DateOID:
|
||||
var d pgtype.Date
|
||||
scanPlan := m.PlanScan(dataTypeOID, format, &d)
|
||||
r.valueFuncs[i] = func(src []byte) (driver.Value, error) {
|
||||
err := scanPlan.Scan(src, &d)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return d.Value()
|
||||
}
|
||||
case pgtype.Float4OID:
|
||||
var d float32
|
||||
scanPlan := m.PlanScan(dataTypeOID, format, &d)
|
||||
r.valueFuncs[i] = func(src []byte) (driver.Value, error) {
|
||||
err := scanPlan.Scan(src, &d)
|
||||
return float64(d), err
|
||||
}
|
||||
case pgtype.Float8OID:
|
||||
var d float64
|
||||
scanPlan := m.PlanScan(dataTypeOID, format, &d)
|
||||
r.valueFuncs[i] = func(src []byte) (driver.Value, error) {
|
||||
err := scanPlan.Scan(src, &d)
|
||||
return d, err
|
||||
}
|
||||
case pgtype.Int2OID:
|
||||
var d int16
|
||||
scanPlan := m.PlanScan(dataTypeOID, format, &d)
|
||||
r.valueFuncs[i] = func(src []byte) (driver.Value, error) {
|
||||
err := scanPlan.Scan(src, &d)
|
||||
return int64(d), err
|
||||
}
|
||||
case pgtype.Int4OID:
|
||||
var d int32
|
||||
scanPlan := m.PlanScan(dataTypeOID, format, &d)
|
||||
r.valueFuncs[i] = func(src []byte) (driver.Value, error) {
|
||||
err := scanPlan.Scan(src, &d)
|
||||
return int64(d), err
|
||||
}
|
||||
case pgtype.Int8OID:
|
||||
var d int64
|
||||
scanPlan := m.PlanScan(dataTypeOID, format, &d)
|
||||
r.valueFuncs[i] = func(src []byte) (driver.Value, error) {
|
||||
err := scanPlan.Scan(src, &d)
|
||||
return d, err
|
||||
}
|
||||
case pgtype.JSONOID, pgtype.JSONBOID:
|
||||
var d []byte
|
||||
scanPlan := m.PlanScan(dataTypeOID, format, &d)
|
||||
r.valueFuncs[i] = func(src []byte) (driver.Value, error) {
|
||||
err := scanPlan.Scan(src, &d)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return d, nil
|
||||
}
|
||||
case pgtype.TimestampOID:
|
||||
var d pgtype.Timestamp
|
||||
scanPlan := m.PlanScan(dataTypeOID, format, &d)
|
||||
r.valueFuncs[i] = func(src []byte) (driver.Value, error) {
|
||||
err := scanPlan.Scan(src, &d)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return d.Value()
|
||||
}
|
||||
case pgtype.TimestamptzOID:
|
||||
var d pgtype.Timestamptz
|
||||
scanPlan := m.PlanScan(dataTypeOID, format, &d)
|
||||
r.valueFuncs[i] = func(src []byte) (driver.Value, error) {
|
||||
err := scanPlan.Scan(src, &d)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return d.Value()
|
||||
}
|
||||
case pgtype.XMLOID:
|
||||
var d []byte
|
||||
scanPlan := m.PlanScan(dataTypeOID, format, &d)
|
||||
r.valueFuncs[i] = func(src []byte) (driver.Value, error) {
|
||||
err := scanPlan.Scan(src, &d)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return d, nil
|
||||
}
|
||||
default:
|
||||
var d string
|
||||
scanPlan := m.PlanScan(dataTypeOID, format, &d)
|
||||
r.valueFuncs[i] = func(src []byte) (driver.Value, error) {
|
||||
err := scanPlan.Scan(src, &d)
|
||||
return d, err
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var more bool
|
||||
if r.skipNext {
|
||||
more = r.skipNextMore
|
||||
r.skipNext = false
|
||||
} else {
|
||||
more = r.rows.Next()
|
||||
}
|
||||
|
||||
if !more {
|
||||
if r.rows.Err() == nil {
|
||||
return io.EOF
|
||||
} else {
|
||||
return r.rows.Err()
|
||||
}
|
||||
}
|
||||
|
||||
for i, rv := range r.rows.RawValues() {
|
||||
if rv != nil {
|
||||
var err error
|
||||
dest[i], err = r.valueFuncs[i](rv)
|
||||
if err != nil {
|
||||
return fmt.Errorf("convert field %d failed: %w", i, err)
|
||||
}
|
||||
} else {
|
||||
dest[i] = nil
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func convertNamedArguments(args []any, argsV []driver.NamedValue) {
|
||||
for i, v := range argsV {
|
||||
if v.Value != nil {
|
||||
args[i] = v.Value.(any)
|
||||
} else {
|
||||
args[i] = nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type wrapTx struct {
|
||||
ctx context.Context
|
||||
tx pgx.Tx
|
||||
}
|
||||
|
||||
func (wtx wrapTx) Commit() error { return wtx.tx.Commit(wtx.ctx) }
|
||||
|
||||
func (wtx wrapTx) Rollback() error { return wtx.tx.Rollback(wtx.ctx) }
|
||||
+79
@@ -0,0 +1,79 @@
|
||||
# 2.2.2 (September 10, 2024)
|
||||
|
||||
* Add empty acquire time to stats (Maxim Ivanov)
|
||||
* Stop importing nanotime from runtime via linkname (maypok86)
|
||||
|
||||
# 2.2.1 (July 15, 2023)
|
||||
|
||||
* Fix: CreateResource cannot overflow pool. This changes documented behavior of CreateResource. Previously,
|
||||
CreateResource could create a resource even if the pool was full. This could cause the pool to overflow. While this
|
||||
was documented, it was documenting incorrect behavior. CreateResource now returns an error if the pool is full.
|
||||
|
||||
# 2.2.0 (February 11, 2023)
|
||||
|
||||
* Use Go 1.19 atomics and drop go.uber.org/atomic dependency
|
||||
|
||||
# 2.1.2 (November 12, 2022)
|
||||
|
||||
* Restore support to Go 1.18 via go.uber.org/atomic
|
||||
|
||||
# 2.1.1 (November 11, 2022)
|
||||
|
||||
* Fix create resource concurrently with Stat call race
|
||||
|
||||
# 2.1.0 (October 28, 2022)
|
||||
|
||||
* Concurrency control is now implemented with a semaphore. This simplifies some internal logic, resolves a few error conditions (including a deadlock), and improves performance. (Jan Dubsky)
|
||||
* Go 1.19 is now required for the improved atomic support.
|
||||
|
||||
# 2.0.1 (October 28, 2022)
|
||||
|
||||
* Fix race condition when Close is called concurrently with multiple constructors
|
||||
|
||||
# 2.0.0 (September 17, 2022)
|
||||
|
||||
* Use generics instead of interface{} (Столяров Владимир Алексеевич)
|
||||
* Add Reset
|
||||
* Do not cancel resource construction when Acquire is canceled
|
||||
* NewPool takes Config
|
||||
|
||||
# 1.3.0 (August 27, 2022)
|
||||
|
||||
* Acquire creates resources in background to allow creation to continue after Acquire is canceled (James Hartig)
|
||||
|
||||
# 1.2.1 (December 2, 2021)
|
||||
|
||||
* TryAcquire now does not block when background constructing resource
|
||||
|
||||
# 1.2.0 (November 20, 2021)
|
||||
|
||||
* Add TryAcquire (A. Jensen)
|
||||
* Fix: remove memory leak / unintentionally pinned memory when shrinking slices (Alexander Staubo)
|
||||
* Fix: Do not leave pool locked after panic from nil context
|
||||
|
||||
# 1.1.4 (September 11, 2021)
|
||||
|
||||
* Fix: Deadlock in CreateResource if pool was closed during resource acquisition (Dmitriy Matrenichev)
|
||||
|
||||
# 1.1.3 (December 3, 2020)
|
||||
|
||||
* Fix: Failed resource creation could cause concurrent Acquire to hang. (Evgeny Vanslov)
|
||||
|
||||
# 1.1.2 (September 26, 2020)
|
||||
|
||||
* Fix: Resource.Destroy no longer removes itself from the pool before its destructor has completed.
|
||||
* Fix: Prevent crash when pool is closed while resource is being created.
|
||||
|
||||
# 1.1.1 (April 2, 2020)
|
||||
|
||||
* Pool.Close can be safely called multiple times
|
||||
* AcquireAllIDle immediately returns nil if pool is closed
|
||||
* CreateResource checks if pool is closed before taking any action
|
||||
* Fix potential race condition when CreateResource and Close are called concurrently. CreateResource now checks if pool is closed before adding newly created resource to pool.
|
||||
|
||||
# 1.1.0 (February 5, 2020)
|
||||
|
||||
* Use runtime.nanotime for faster tracking of acquire time and last usage time.
|
||||
* Track resource idle time to enable client health check logic. (Patrick Ellul)
|
||||
* Add CreateResource to construct a new resource without acquiring it. (Patrick Ellul)
|
||||
* Fix deadlock race when acquire is cancelled. (Michael Tharp)
|
||||
+22
@@ -0,0 +1,22 @@
|
||||
Copyright (c) 2018 Jack Christensen
|
||||
|
||||
MIT License
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining
|
||||
a copy of this software and associated documentation files (the
|
||||
"Software"), to deal in the Software without restriction, including
|
||||
without limitation the rights to use, copy, modify, merge, publish,
|
||||
distribute, sublicense, and/or sell copies of the Software, and to
|
||||
permit persons to whom the Software is furnished to do so, subject to
|
||||
the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be
|
||||
included in all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
|
||||
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
|
||||
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
|
||||
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
|
||||
LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
|
||||
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
|
||||
WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
||||
+80
@@ -0,0 +1,80 @@
|
||||
[](https://pkg.go.dev/github.com/jackc/puddle/v2)
|
||||
|
||||
|
||||
# Puddle
|
||||
|
||||
Puddle is a tiny generic resource pool library for Go that uses the standard
|
||||
context library to signal cancellation of acquires. It is designed to contain
|
||||
the minimum functionality required for a resource pool. It can be used directly
|
||||
or it can be used as the base for a domain specific resource pool. For example,
|
||||
a database connection pool may use puddle internally and implement health checks
|
||||
and keep-alive behavior without needing to implement any concurrent code of its
|
||||
own.
|
||||
|
||||
## Features
|
||||
|
||||
* Acquire cancellation via context standard library
|
||||
* Statistics API for monitoring pool pressure
|
||||
* No dependencies outside of standard library and golang.org/x/sync
|
||||
* High performance
|
||||
* 100% test coverage of reachable code
|
||||
|
||||
## Example Usage
|
||||
|
||||
```go
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"net"
|
||||
|
||||
"github.com/jackc/puddle/v2"
|
||||
)
|
||||
|
||||
func main() {
|
||||
constructor := func(context.Context) (net.Conn, error) {
|
||||
return net.Dial("tcp", "127.0.0.1:8080")
|
||||
}
|
||||
destructor := func(value net.Conn) {
|
||||
value.Close()
|
||||
}
|
||||
maxPoolSize := int32(10)
|
||||
|
||||
pool, err := puddle.NewPool(&puddle.Config[net.Conn]{Constructor: constructor, Destructor: destructor, MaxSize: maxPoolSize})
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
||||
// Acquire resource from the pool.
|
||||
res, err := pool.Acquire(context.Background())
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
||||
// Use resource.
|
||||
_, err = res.Value().Write([]byte{1})
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
||||
// Release when done.
|
||||
res.Release()
|
||||
}
|
||||
```
|
||||
|
||||
## Status
|
||||
|
||||
Puddle is stable and feature complete.
|
||||
|
||||
* Bug reports and fixes are welcome.
|
||||
* New features will usually not be accepted if they can be feasibly implemented in a wrapper.
|
||||
* Performance optimizations will usually not be accepted unless the performance issue rises to the level of a bug.
|
||||
|
||||
## Supported Go Versions
|
||||
|
||||
puddle supports the same versions of Go that are supported by the Go project. For [Go](https://golang.org/doc/devel/release.html#policy) that is the two most recent major releases. This means puddle supports Go 1.19 and higher.
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
+24
@@ -0,0 +1,24 @@
|
||||
package puddle
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
)
|
||||
|
||||
// valueCancelCtx combines two contexts into one. One context is used for values and the other is used for cancellation.
|
||||
type valueCancelCtx struct {
|
||||
valueCtx context.Context
|
||||
cancelCtx context.Context
|
||||
}
|
||||
|
||||
func (ctx *valueCancelCtx) Deadline() (time.Time, bool) { return ctx.cancelCtx.Deadline() }
|
||||
func (ctx *valueCancelCtx) Done() <-chan struct{} { return ctx.cancelCtx.Done() }
|
||||
func (ctx *valueCancelCtx) Err() error { return ctx.cancelCtx.Err() }
|
||||
func (ctx *valueCancelCtx) Value(key any) any { return ctx.valueCtx.Value(key) }
|
||||
|
||||
func newValueCancelCtx(valueCtx, cancelContext context.Context) context.Context {
|
||||
return &valueCancelCtx{
|
||||
valueCtx: valueCtx,
|
||||
cancelCtx: cancelContext,
|
||||
}
|
||||
}
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
// Package puddle is a generic resource pool with type-parametrized api.
|
||||
/*
|
||||
|
||||
Puddle is a tiny generic resource pool library for Go that uses the standard
|
||||
context library to signal cancellation of acquires. It is designed to contain
|
||||
the minimum functionality a resource pool needs that cannot be implemented
|
||||
without concurrency concerns. For example, a database connection pool may use
|
||||
puddle internally and implement health checks and keep-alive behavior without
|
||||
needing to implement any concurrent code of its own.
|
||||
*/
|
||||
package puddle
|
||||
+85
@@ -0,0 +1,85 @@
|
||||
package genstack
|
||||
|
||||
// GenStack implements a generational stack.
|
||||
//
|
||||
// GenStack works as common stack except for the fact that all elements in the
|
||||
// older generation are guaranteed to be popped before any element in the newer
|
||||
// generation. New elements are always pushed to the current (newest)
|
||||
// generation.
|
||||
//
|
||||
// We could also say that GenStack behaves as a stack in case of a single
|
||||
// generation, but it behaves as a queue of individual generation stacks.
|
||||
type GenStack[T any] struct {
|
||||
// We can represent arbitrary number of generations using 2 stacks. The
|
||||
// new stack stores all new pushes and the old stack serves all reads.
|
||||
// Old stack can represent multiple generations. If old == new, then all
|
||||
// elements pushed in previous (not current) generations have already
|
||||
// been popped.
|
||||
|
||||
old *stack[T]
|
||||
new *stack[T]
|
||||
}
|
||||
|
||||
// NewGenStack creates a new empty GenStack.
|
||||
func NewGenStack[T any]() *GenStack[T] {
|
||||
s := &stack[T]{}
|
||||
return &GenStack[T]{
|
||||
old: s,
|
||||
new: s,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *GenStack[T]) Pop() (T, bool) {
|
||||
// Pushes always append to the new stack, so if the old once becomes
|
||||
// empty, it will remail empty forever.
|
||||
if s.old.len() == 0 && s.old != s.new {
|
||||
s.old = s.new
|
||||
}
|
||||
|
||||
if s.old.len() == 0 {
|
||||
var zero T
|
||||
return zero, false
|
||||
}
|
||||
|
||||
return s.old.pop(), true
|
||||
}
|
||||
|
||||
// Push pushes a new element at the top of the stack.
|
||||
func (s *GenStack[T]) Push(v T) { s.new.push(v) }
|
||||
|
||||
// NextGen starts a new stack generation.
|
||||
func (s *GenStack[T]) NextGen() {
|
||||
if s.old == s.new {
|
||||
s.new = &stack[T]{}
|
||||
return
|
||||
}
|
||||
|
||||
// We need to pop from the old stack to the top of the new stack. Let's
|
||||
// have an example:
|
||||
//
|
||||
// Old: <bottom> 4 3 2 1
|
||||
// New: <bottom> 8 7 6 5
|
||||
// PopOrder: 1 2 3 4 5 6 7 8
|
||||
//
|
||||
//
|
||||
// To preserve pop order, we have to take all elements from the old
|
||||
// stack and push them to the top of new stack:
|
||||
//
|
||||
// New: 8 7 6 5 4 3 2 1
|
||||
//
|
||||
s.new.push(s.old.takeAll()...)
|
||||
|
||||
// We have the old stack allocated and empty, so why not to reuse it as
|
||||
// new new stack.
|
||||
s.old, s.new = s.new, s.old
|
||||
}
|
||||
|
||||
// Len returns number of elements in the stack.
|
||||
func (s *GenStack[T]) Len() int {
|
||||
l := s.old.len()
|
||||
if s.old != s.new {
|
||||
l += s.new.len()
|
||||
}
|
||||
|
||||
return l
|
||||
}
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
package genstack
|
||||
|
||||
// stack is a wrapper around an array implementing a stack.
|
||||
//
|
||||
// We cannot use slice to represent the stack because append might change the
|
||||
// pointer value of the slice. That would be an issue in GenStack
|
||||
// implementation.
|
||||
type stack[T any] struct {
|
||||
arr []T
|
||||
}
|
||||
|
||||
// push pushes a new element at the top of a stack.
|
||||
func (s *stack[T]) push(vs ...T) { s.arr = append(s.arr, vs...) }
|
||||
|
||||
// pop pops the stack top-most element.
|
||||
//
|
||||
// If stack length is zero, this method panics.
|
||||
func (s *stack[T]) pop() T {
|
||||
idx := s.len() - 1
|
||||
val := s.arr[idx]
|
||||
|
||||
// Avoid memory leak
|
||||
var zero T
|
||||
s.arr[idx] = zero
|
||||
|
||||
s.arr = s.arr[:idx]
|
||||
return val
|
||||
}
|
||||
|
||||
// takeAll returns all elements in the stack in order as they are stored - i.e.
|
||||
// the top-most stack element is the last one.
|
||||
func (s *stack[T]) takeAll() []T {
|
||||
arr := s.arr
|
||||
s.arr = nil
|
||||
return arr
|
||||
}
|
||||
|
||||
// len returns number of elements in the stack.
|
||||
func (s *stack[T]) len() int { return len(s.arr) }
|
||||
+32
@@ -0,0 +1,32 @@
|
||||
package puddle
|
||||
|
||||
import "unsafe"
|
||||
|
||||
type ints interface {
|
||||
int | int8 | int16 | int32 | int64 | uint | uint8 | uint16 | uint32 | uint64
|
||||
}
|
||||
|
||||
// log2Int returns log2 of an integer. This function panics if val < 0. For val
|
||||
// == 0, returns 0.
|
||||
func log2Int[T ints](val T) uint8 {
|
||||
if val <= 0 {
|
||||
panic("log2 of non-positive number does not exist")
|
||||
}
|
||||
|
||||
return log2IntRange(val, 0, uint8(8*unsafe.Sizeof(val)))
|
||||
}
|
||||
|
||||
func log2IntRange[T ints](val T, begin, end uint8) uint8 {
|
||||
length := end - begin
|
||||
if length == 1 {
|
||||
return begin
|
||||
}
|
||||
|
||||
delim := begin + length/2
|
||||
mask := T(1) << delim
|
||||
if mask > val {
|
||||
return log2IntRange(val, begin, delim)
|
||||
} else {
|
||||
return log2IntRange(val, delim, end)
|
||||
}
|
||||
}
|
||||
+16
@@ -0,0 +1,16 @@
|
||||
package puddle
|
||||
|
||||
import "time"
|
||||
|
||||
// nanotime returns the time in nanoseconds since process start.
|
||||
//
|
||||
// This approach, described at
|
||||
// https://github.com/golang/go/issues/61765#issuecomment-1672090302,
|
||||
// is fast, monotonic, and portable, and avoids the previous
|
||||
// dependence on runtime.nanotime using the (unsafe) linkname hack.
|
||||
// In particular, time.Since does less work than time.Now.
|
||||
func nanotime() int64 {
|
||||
return time.Since(globalStart).Nanoseconds()
|
||||
}
|
||||
|
||||
var globalStart = time.Now()
|
||||
+710
@@ -0,0 +1,710 @@
|
||||
package puddle
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/puddle/v2/internal/genstack"
|
||||
"golang.org/x/sync/semaphore"
|
||||
)
|
||||
|
||||
const (
|
||||
resourceStatusConstructing = 0
|
||||
resourceStatusIdle = iota
|
||||
resourceStatusAcquired = iota
|
||||
resourceStatusHijacked = iota
|
||||
)
|
||||
|
||||
// ErrClosedPool occurs on an attempt to acquire a connection from a closed pool
|
||||
// or a pool that is closed while the acquire is waiting.
|
||||
var ErrClosedPool = errors.New("closed pool")
|
||||
|
||||
// ErrNotAvailable occurs on an attempt to acquire a resource from a pool
|
||||
// that is at maximum capacity and has no available resources.
|
||||
var ErrNotAvailable = errors.New("resource not available")
|
||||
|
||||
// Constructor is a function called by the pool to construct a resource.
|
||||
type Constructor[T any] func(ctx context.Context) (res T, err error)
|
||||
|
||||
// Destructor is a function called by the pool to destroy a resource.
|
||||
type Destructor[T any] func(res T)
|
||||
|
||||
// Resource is the resource handle returned by acquiring from the pool.
|
||||
type Resource[T any] struct {
|
||||
value T
|
||||
pool *Pool[T]
|
||||
creationTime time.Time
|
||||
lastUsedNano int64
|
||||
poolResetCount int
|
||||
status byte
|
||||
}
|
||||
|
||||
// Value returns the resource value.
|
||||
func (res *Resource[T]) Value() T {
|
||||
if !(res.status == resourceStatusAcquired || res.status == resourceStatusHijacked) {
|
||||
panic("tried to access resource that is not acquired or hijacked")
|
||||
}
|
||||
return res.value
|
||||
}
|
||||
|
||||
// Release returns the resource to the pool. res must not be subsequently used.
|
||||
func (res *Resource[T]) Release() {
|
||||
if res.status != resourceStatusAcquired {
|
||||
panic("tried to release resource that is not acquired")
|
||||
}
|
||||
res.pool.releaseAcquiredResource(res, nanotime())
|
||||
}
|
||||
|
||||
// ReleaseUnused returns the resource to the pool without updating when it was last used used. i.e. LastUsedNanotime
|
||||
// will not change. res must not be subsequently used.
|
||||
func (res *Resource[T]) ReleaseUnused() {
|
||||
if res.status != resourceStatusAcquired {
|
||||
panic("tried to release resource that is not acquired")
|
||||
}
|
||||
res.pool.releaseAcquiredResource(res, res.lastUsedNano)
|
||||
}
|
||||
|
||||
// Destroy returns the resource to the pool for destruction. res must not be
|
||||
// subsequently used.
|
||||
func (res *Resource[T]) Destroy() {
|
||||
if res.status != resourceStatusAcquired {
|
||||
panic("tried to destroy resource that is not acquired")
|
||||
}
|
||||
go res.pool.destroyAcquiredResource(res)
|
||||
}
|
||||
|
||||
// Hijack assumes ownership of the resource from the pool. Caller is responsible
|
||||
// for cleanup of resource value.
|
||||
func (res *Resource[T]) Hijack() {
|
||||
if res.status != resourceStatusAcquired {
|
||||
panic("tried to hijack resource that is not acquired")
|
||||
}
|
||||
res.pool.hijackAcquiredResource(res)
|
||||
}
|
||||
|
||||
// CreationTime returns when the resource was created by the pool.
|
||||
func (res *Resource[T]) CreationTime() time.Time {
|
||||
if !(res.status == resourceStatusAcquired || res.status == resourceStatusHijacked) {
|
||||
panic("tried to access resource that is not acquired or hijacked")
|
||||
}
|
||||
return res.creationTime
|
||||
}
|
||||
|
||||
// LastUsedNanotime returns when Release was last called on the resource measured in nanoseconds from an arbitrary time
|
||||
// (a monotonic time). Returns creation time if Release has never been called. This is only useful to compare with
|
||||
// other calls to LastUsedNanotime. In almost all cases, IdleDuration should be used instead.
|
||||
func (res *Resource[T]) LastUsedNanotime() int64 {
|
||||
if !(res.status == resourceStatusAcquired || res.status == resourceStatusHijacked) {
|
||||
panic("tried to access resource that is not acquired or hijacked")
|
||||
}
|
||||
|
||||
return res.lastUsedNano
|
||||
}
|
||||
|
||||
// IdleDuration returns the duration since Release was last called on the resource. This is equivalent to subtracting
|
||||
// LastUsedNanotime to the current nanotime.
|
||||
func (res *Resource[T]) IdleDuration() time.Duration {
|
||||
if !(res.status == resourceStatusAcquired || res.status == resourceStatusHijacked) {
|
||||
panic("tried to access resource that is not acquired or hijacked")
|
||||
}
|
||||
|
||||
return time.Duration(nanotime() - res.lastUsedNano)
|
||||
}
|
||||
|
||||
// Pool is a concurrency-safe resource pool.
|
||||
type Pool[T any] struct {
|
||||
// mux is the pool internal lock. Any modification of shared state of
|
||||
// the pool (but Acquires of acquireSem) must be performed only by
|
||||
// holder of the lock. Long running operations are not allowed when mux
|
||||
// is held.
|
||||
mux sync.Mutex
|
||||
// acquireSem provides an allowance to acquire a resource.
|
||||
//
|
||||
// Releases are allowed only when caller holds mux. Acquires have to
|
||||
// happen before mux is locked (doesn't apply to semaphore.TryAcquire in
|
||||
// AcquireAllIdle).
|
||||
acquireSem *semaphore.Weighted
|
||||
destructWG sync.WaitGroup
|
||||
|
||||
allResources resList[T]
|
||||
idleResources *genstack.GenStack[*Resource[T]]
|
||||
|
||||
constructor Constructor[T]
|
||||
destructor Destructor[T]
|
||||
maxSize int32
|
||||
|
||||
acquireCount int64
|
||||
acquireDuration time.Duration
|
||||
emptyAcquireCount int64
|
||||
emptyAcquireWaitTime time.Duration
|
||||
canceledAcquireCount atomic.Int64
|
||||
|
||||
resetCount int
|
||||
|
||||
baseAcquireCtx context.Context
|
||||
cancelBaseAcquireCtx context.CancelFunc
|
||||
closed bool
|
||||
}
|
||||
|
||||
type Config[T any] struct {
|
||||
Constructor Constructor[T]
|
||||
Destructor Destructor[T]
|
||||
MaxSize int32
|
||||
}
|
||||
|
||||
// NewPool creates a new pool. Returns an error iff MaxSize is less than 1.
|
||||
func NewPool[T any](config *Config[T]) (*Pool[T], error) {
|
||||
if config.MaxSize < 1 {
|
||||
return nil, errors.New("MaxSize must be >= 1")
|
||||
}
|
||||
|
||||
baseAcquireCtx, cancelBaseAcquireCtx := context.WithCancel(context.Background())
|
||||
|
||||
return &Pool[T]{
|
||||
acquireSem: semaphore.NewWeighted(int64(config.MaxSize)),
|
||||
idleResources: genstack.NewGenStack[*Resource[T]](),
|
||||
maxSize: config.MaxSize,
|
||||
constructor: config.Constructor,
|
||||
destructor: config.Destructor,
|
||||
baseAcquireCtx: baseAcquireCtx,
|
||||
cancelBaseAcquireCtx: cancelBaseAcquireCtx,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Close destroys all resources in the pool and rejects future Acquire calls.
|
||||
// Blocks until all resources are returned to pool and destroyed.
|
||||
func (p *Pool[T]) Close() {
|
||||
defer p.destructWG.Wait()
|
||||
|
||||
p.mux.Lock()
|
||||
defer p.mux.Unlock()
|
||||
|
||||
if p.closed {
|
||||
return
|
||||
}
|
||||
p.closed = true
|
||||
p.cancelBaseAcquireCtx()
|
||||
|
||||
for res, ok := p.idleResources.Pop(); ok; res, ok = p.idleResources.Pop() {
|
||||
p.allResources.remove(res)
|
||||
go p.destructResourceValue(res.value)
|
||||
}
|
||||
}
|
||||
|
||||
// Stat is a snapshot of Pool statistics.
|
||||
type Stat struct {
|
||||
constructingResources int32
|
||||
acquiredResources int32
|
||||
idleResources int32
|
||||
maxResources int32
|
||||
acquireCount int64
|
||||
acquireDuration time.Duration
|
||||
emptyAcquireCount int64
|
||||
emptyAcquireWaitTime time.Duration
|
||||
canceledAcquireCount int64
|
||||
}
|
||||
|
||||
// TotalResources returns the total number of resources currently in the pool.
|
||||
// The value is the sum of ConstructingResources, AcquiredResources, and
|
||||
// IdleResources.
|
||||
func (s *Stat) TotalResources() int32 {
|
||||
return s.constructingResources + s.acquiredResources + s.idleResources
|
||||
}
|
||||
|
||||
// ConstructingResources returns the number of resources with construction in progress in
|
||||
// the pool.
|
||||
func (s *Stat) ConstructingResources() int32 {
|
||||
return s.constructingResources
|
||||
}
|
||||
|
||||
// AcquiredResources returns the number of currently acquired resources in the pool.
|
||||
func (s *Stat) AcquiredResources() int32 {
|
||||
return s.acquiredResources
|
||||
}
|
||||
|
||||
// IdleResources returns the number of currently idle resources in the pool.
|
||||
func (s *Stat) IdleResources() int32 {
|
||||
return s.idleResources
|
||||
}
|
||||
|
||||
// MaxResources returns the maximum size of the pool.
|
||||
func (s *Stat) MaxResources() int32 {
|
||||
return s.maxResources
|
||||
}
|
||||
|
||||
// AcquireCount returns the cumulative count of successful acquires from the pool.
|
||||
func (s *Stat) AcquireCount() int64 {
|
||||
return s.acquireCount
|
||||
}
|
||||
|
||||
// AcquireDuration returns the total duration of all successful acquires from
|
||||
// the pool.
|
||||
func (s *Stat) AcquireDuration() time.Duration {
|
||||
return s.acquireDuration
|
||||
}
|
||||
|
||||
// EmptyAcquireCount returns the cumulative count of successful acquires from the pool
|
||||
// that waited for a resource to be released or constructed because the pool was
|
||||
// empty.
|
||||
func (s *Stat) EmptyAcquireCount() int64 {
|
||||
return s.emptyAcquireCount
|
||||
}
|
||||
|
||||
// EmptyAcquireWaitTime returns the cumulative time waited for successful acquires
|
||||
// from the pool for a resource to be released or constructed because the pool was
|
||||
// empty.
|
||||
func (s *Stat) EmptyAcquireWaitTime() time.Duration {
|
||||
return s.emptyAcquireWaitTime
|
||||
}
|
||||
|
||||
// CanceledAcquireCount returns the cumulative count of acquires from the pool
|
||||
// that were canceled by a context.
|
||||
func (s *Stat) CanceledAcquireCount() int64 {
|
||||
return s.canceledAcquireCount
|
||||
}
|
||||
|
||||
// Stat returns the current pool statistics.
|
||||
func (p *Pool[T]) Stat() *Stat {
|
||||
p.mux.Lock()
|
||||
defer p.mux.Unlock()
|
||||
|
||||
s := &Stat{
|
||||
maxResources: p.maxSize,
|
||||
acquireCount: p.acquireCount,
|
||||
emptyAcquireCount: p.emptyAcquireCount,
|
||||
emptyAcquireWaitTime: p.emptyAcquireWaitTime,
|
||||
canceledAcquireCount: p.canceledAcquireCount.Load(),
|
||||
acquireDuration: p.acquireDuration,
|
||||
}
|
||||
|
||||
for _, res := range p.allResources {
|
||||
switch res.status {
|
||||
case resourceStatusConstructing:
|
||||
s.constructingResources += 1
|
||||
case resourceStatusIdle:
|
||||
s.idleResources += 1
|
||||
case resourceStatusAcquired:
|
||||
s.acquiredResources += 1
|
||||
}
|
||||
}
|
||||
|
||||
return s
|
||||
}
|
||||
|
||||
// tryAcquireIdleResource checks if there is any idle resource. If there is
|
||||
// some, this method removes it from idle list and returns it. If the idle pool
|
||||
// is empty, this method returns nil and doesn't modify the idleResources slice.
|
||||
//
|
||||
// WARNING: Caller of this method must hold the pool mutex!
|
||||
func (p *Pool[T]) tryAcquireIdleResource() *Resource[T] {
|
||||
res, ok := p.idleResources.Pop()
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
res.status = resourceStatusAcquired
|
||||
return res
|
||||
}
|
||||
|
||||
// createNewResource creates a new resource and inserts it into list of pool
|
||||
// resources.
|
||||
//
|
||||
// WARNING: Caller of this method must hold the pool mutex!
|
||||
func (p *Pool[T]) createNewResource() *Resource[T] {
|
||||
res := &Resource[T]{
|
||||
pool: p,
|
||||
creationTime: time.Now(),
|
||||
lastUsedNano: nanotime(),
|
||||
poolResetCount: p.resetCount,
|
||||
status: resourceStatusConstructing,
|
||||
}
|
||||
|
||||
p.allResources.append(res)
|
||||
p.destructWG.Add(1)
|
||||
|
||||
return res
|
||||
}
|
||||
|
||||
// Acquire gets a resource from the pool. If no resources are available and the pool is not at maximum capacity it will
|
||||
// create a new resource. If the pool is at maximum capacity it will block until a resource is available. ctx can be
|
||||
// used to cancel the Acquire.
|
||||
//
|
||||
// If Acquire creates a new resource the resource constructor function will receive a context that delegates Value() to
|
||||
// ctx. Canceling ctx will cause Acquire to return immediately but it will not cancel the resource creation. This avoids
|
||||
// the problem of it being impossible to create resources when the time to create a resource is greater than any one
|
||||
// caller of Acquire is willing to wait.
|
||||
func (p *Pool[T]) Acquire(ctx context.Context) (_ *Resource[T], err error) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
p.canceledAcquireCount.Add(1)
|
||||
return nil, ctx.Err()
|
||||
default:
|
||||
}
|
||||
|
||||
return p.acquire(ctx)
|
||||
}
|
||||
|
||||
// acquire is a continuation of Acquire function that doesn't check context
|
||||
// validity.
|
||||
//
|
||||
// This function exists solely only for benchmarking purposes.
|
||||
func (p *Pool[T]) acquire(ctx context.Context) (*Resource[T], error) {
|
||||
startNano := nanotime()
|
||||
|
||||
var waitedForLock bool
|
||||
if !p.acquireSem.TryAcquire(1) {
|
||||
waitedForLock = true
|
||||
err := p.acquireSem.Acquire(ctx, 1)
|
||||
if err != nil {
|
||||
p.canceledAcquireCount.Add(1)
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
p.mux.Lock()
|
||||
if p.closed {
|
||||
p.acquireSem.Release(1)
|
||||
p.mux.Unlock()
|
||||
return nil, ErrClosedPool
|
||||
}
|
||||
|
||||
// If a resource is available in the pool.
|
||||
if res := p.tryAcquireIdleResource(); res != nil {
|
||||
waitTime := time.Duration(nanotime() - startNano)
|
||||
if waitedForLock {
|
||||
p.emptyAcquireCount += 1
|
||||
p.emptyAcquireWaitTime += waitTime
|
||||
}
|
||||
p.acquireCount += 1
|
||||
p.acquireDuration += waitTime
|
||||
p.mux.Unlock()
|
||||
return res, nil
|
||||
}
|
||||
|
||||
if len(p.allResources) >= int(p.maxSize) {
|
||||
// Unreachable code.
|
||||
panic("bug: semaphore allowed more acquires than pool allows")
|
||||
}
|
||||
|
||||
// The resource is not idle, but there is enough space to create one.
|
||||
res := p.createNewResource()
|
||||
p.mux.Unlock()
|
||||
|
||||
res, err := p.initResourceValue(ctx, res)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
p.mux.Lock()
|
||||
defer p.mux.Unlock()
|
||||
|
||||
p.emptyAcquireCount += 1
|
||||
p.acquireCount += 1
|
||||
waitTime := time.Duration(nanotime() - startNano)
|
||||
p.acquireDuration += waitTime
|
||||
p.emptyAcquireWaitTime += waitTime
|
||||
|
||||
return res, nil
|
||||
}
|
||||
|
||||
func (p *Pool[T]) initResourceValue(ctx context.Context, res *Resource[T]) (*Resource[T], error) {
|
||||
// Create the resource in a goroutine to immediately return from Acquire
|
||||
// if ctx is canceled without also canceling the constructor.
|
||||
//
|
||||
// See:
|
||||
// - https://github.com/jackc/pgx/issues/1287
|
||||
// - https://github.com/jackc/pgx/issues/1259
|
||||
constructErrChan := make(chan error)
|
||||
go func() {
|
||||
constructorCtx := newValueCancelCtx(ctx, p.baseAcquireCtx)
|
||||
value, err := p.constructor(constructorCtx)
|
||||
if err != nil {
|
||||
p.mux.Lock()
|
||||
p.allResources.remove(res)
|
||||
p.destructWG.Done()
|
||||
|
||||
// The resource won't be acquired because its
|
||||
// construction failed. We have to allow someone else to
|
||||
// take that resouce.
|
||||
p.acquireSem.Release(1)
|
||||
p.mux.Unlock()
|
||||
|
||||
select {
|
||||
case constructErrChan <- err:
|
||||
case <-ctx.Done():
|
||||
// The caller is cancelled, so no-one awaits the
|
||||
// error. This branch avoid goroutine leak.
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// The resource is already in p.allResources where it might be read. So we need to acquire the lock to update its
|
||||
// status.
|
||||
p.mux.Lock()
|
||||
res.value = value
|
||||
res.status = resourceStatusAcquired
|
||||
p.mux.Unlock()
|
||||
|
||||
// This select works because the channel is unbuffered.
|
||||
select {
|
||||
case constructErrChan <- nil:
|
||||
case <-ctx.Done():
|
||||
p.releaseAcquiredResource(res, res.lastUsedNano)
|
||||
}
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
p.canceledAcquireCount.Add(1)
|
||||
return nil, ctx.Err()
|
||||
case err := <-constructErrChan:
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
}
|
||||
|
||||
// TryAcquire gets a resource from the pool if one is immediately available. If not, it returns ErrNotAvailable. If no
|
||||
// resources are available but the pool has room to grow, a resource will be created in the background. ctx is only
|
||||
// used to cancel the background creation.
|
||||
func (p *Pool[T]) TryAcquire(ctx context.Context) (*Resource[T], error) {
|
||||
if !p.acquireSem.TryAcquire(1) {
|
||||
return nil, ErrNotAvailable
|
||||
}
|
||||
|
||||
p.mux.Lock()
|
||||
defer p.mux.Unlock()
|
||||
|
||||
if p.closed {
|
||||
p.acquireSem.Release(1)
|
||||
return nil, ErrClosedPool
|
||||
}
|
||||
|
||||
// If a resource is available now
|
||||
if res := p.tryAcquireIdleResource(); res != nil {
|
||||
p.acquireCount += 1
|
||||
return res, nil
|
||||
}
|
||||
|
||||
if len(p.allResources) >= int(p.maxSize) {
|
||||
// Unreachable code.
|
||||
panic("bug: semaphore allowed more acquires than pool allows")
|
||||
}
|
||||
|
||||
res := p.createNewResource()
|
||||
go func() {
|
||||
value, err := p.constructor(ctx)
|
||||
|
||||
p.mux.Lock()
|
||||
defer p.mux.Unlock()
|
||||
// We have to create the resource and only then release the
|
||||
// semaphore - For the time being there is no resource that
|
||||
// someone could acquire.
|
||||
defer p.acquireSem.Release(1)
|
||||
|
||||
if err != nil {
|
||||
p.allResources.remove(res)
|
||||
p.destructWG.Done()
|
||||
return
|
||||
}
|
||||
|
||||
res.value = value
|
||||
res.status = resourceStatusIdle
|
||||
p.idleResources.Push(res)
|
||||
}()
|
||||
|
||||
return nil, ErrNotAvailable
|
||||
}
|
||||
|
||||
// acquireSemAll tries to acquire num free tokens from sem. This function is
|
||||
// guaranteed to acquire at least the lowest number of tokens that has been
|
||||
// available in the semaphore during runtime of this function.
|
||||
//
|
||||
// For the time being, semaphore doesn't allow to acquire all tokens atomically
|
||||
// (see https://github.com/golang/sync/pull/19). We simulate this by trying all
|
||||
// powers of 2 that are less or equal to num.
|
||||
//
|
||||
// For example, let's immagine we have 19 free tokens in the semaphore which in
|
||||
// total has 24 tokens (i.e. the maxSize of the pool is 24 resources). Then if
|
||||
// num is 24, the log2Uint(24) is 4 and we try to acquire 16, 8, 4, 2 and 1
|
||||
// tokens. Out of those, the acquire of 16, 2 and 1 tokens will succeed.
|
||||
//
|
||||
// Naturally, Acquires and Releases of the semaphore might take place
|
||||
// concurrently. For this reason, it's not guaranteed that absolutely all free
|
||||
// tokens in the semaphore will be acquired. But it's guaranteed that at least
|
||||
// the minimal number of tokens that has been present over the whole process
|
||||
// will be acquired. This is sufficient for the use-case we have in this
|
||||
// package.
|
||||
//
|
||||
// TODO: Replace this with acquireSem.TryAcquireAll() if it gets to
|
||||
// upstream. https://github.com/golang/sync/pull/19
|
||||
func acquireSemAll(sem *semaphore.Weighted, num int) int {
|
||||
if sem.TryAcquire(int64(num)) {
|
||||
return num
|
||||
}
|
||||
|
||||
var acquired int
|
||||
for i := int(log2Int(num)); i >= 0; i-- {
|
||||
val := 1 << i
|
||||
if sem.TryAcquire(int64(val)) {
|
||||
acquired += val
|
||||
}
|
||||
}
|
||||
|
||||
return acquired
|
||||
}
|
||||
|
||||
// AcquireAllIdle acquires all currently idle resources. Its intended use is for
|
||||
// health check and keep-alive functionality. It does not update pool
|
||||
// statistics.
|
||||
func (p *Pool[T]) AcquireAllIdle() []*Resource[T] {
|
||||
p.mux.Lock()
|
||||
defer p.mux.Unlock()
|
||||
|
||||
if p.closed {
|
||||
return nil
|
||||
}
|
||||
|
||||
numIdle := p.idleResources.Len()
|
||||
if numIdle == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// In acquireSemAll we use only TryAcquire and not Acquire. Because
|
||||
// TryAcquire cannot block, the fact that we hold mutex locked and try
|
||||
// to acquire semaphore cannot result in dead-lock.
|
||||
//
|
||||
// Because the mutex is locked, no parallel Release can run. This
|
||||
// implies that the number of tokens can only decrease because some
|
||||
// Acquire/TryAcquire call can consume the semaphore token. Consequently
|
||||
// acquired is always less or equal to numIdle. Moreover if acquired <
|
||||
// numIdle, then there are some parallel Acquire/TryAcquire calls that
|
||||
// will take the remaining idle connections.
|
||||
acquired := acquireSemAll(p.acquireSem, numIdle)
|
||||
|
||||
idle := make([]*Resource[T], acquired)
|
||||
for i := range idle {
|
||||
res, _ := p.idleResources.Pop()
|
||||
res.status = resourceStatusAcquired
|
||||
idle[i] = res
|
||||
}
|
||||
|
||||
// We have to bump the generation to ensure that Acquire/TryAcquire
|
||||
// calls running in parallel (those which caused acquired < numIdle)
|
||||
// will consume old connections and not freshly released connections
|
||||
// instead.
|
||||
p.idleResources.NextGen()
|
||||
|
||||
return idle
|
||||
}
|
||||
|
||||
// CreateResource constructs a new resource without acquiring it. It goes straight in the IdlePool. If the pool is full
|
||||
// it returns an error. It can be useful to maintain warm resources under little load.
|
||||
func (p *Pool[T]) CreateResource(ctx context.Context) error {
|
||||
if !p.acquireSem.TryAcquire(1) {
|
||||
return ErrNotAvailable
|
||||
}
|
||||
|
||||
p.mux.Lock()
|
||||
if p.closed {
|
||||
p.acquireSem.Release(1)
|
||||
p.mux.Unlock()
|
||||
return ErrClosedPool
|
||||
}
|
||||
|
||||
if len(p.allResources) >= int(p.maxSize) {
|
||||
p.acquireSem.Release(1)
|
||||
p.mux.Unlock()
|
||||
return ErrNotAvailable
|
||||
}
|
||||
|
||||
res := p.createNewResource()
|
||||
p.mux.Unlock()
|
||||
|
||||
value, err := p.constructor(ctx)
|
||||
p.mux.Lock()
|
||||
defer p.mux.Unlock()
|
||||
defer p.acquireSem.Release(1)
|
||||
if err != nil {
|
||||
p.allResources.remove(res)
|
||||
p.destructWG.Done()
|
||||
return err
|
||||
}
|
||||
|
||||
res.value = value
|
||||
res.status = resourceStatusIdle
|
||||
|
||||
// If closed while constructing resource then destroy it and return an error
|
||||
if p.closed {
|
||||
go p.destructResourceValue(res.value)
|
||||
return ErrClosedPool
|
||||
}
|
||||
|
||||
p.idleResources.Push(res)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Reset destroys all resources, but leaves the pool open. It is intended for use when an error is detected that would
|
||||
// disrupt all resources (such as a network interruption or a server state change).
|
||||
//
|
||||
// It is safe to reset a pool while resources are checked out. Those resources will be destroyed when they are returned
|
||||
// to the pool.
|
||||
func (p *Pool[T]) Reset() {
|
||||
p.mux.Lock()
|
||||
defer p.mux.Unlock()
|
||||
|
||||
p.resetCount++
|
||||
|
||||
for res, ok := p.idleResources.Pop(); ok; res, ok = p.idleResources.Pop() {
|
||||
p.allResources.remove(res)
|
||||
go p.destructResourceValue(res.value)
|
||||
}
|
||||
}
|
||||
|
||||
// releaseAcquiredResource returns res to the the pool.
|
||||
func (p *Pool[T]) releaseAcquiredResource(res *Resource[T], lastUsedNano int64) {
|
||||
p.mux.Lock()
|
||||
defer p.mux.Unlock()
|
||||
defer p.acquireSem.Release(1)
|
||||
|
||||
if p.closed || res.poolResetCount != p.resetCount {
|
||||
p.allResources.remove(res)
|
||||
go p.destructResourceValue(res.value)
|
||||
} else {
|
||||
res.lastUsedNano = lastUsedNano
|
||||
res.status = resourceStatusIdle
|
||||
p.idleResources.Push(res)
|
||||
}
|
||||
}
|
||||
|
||||
// Remove removes res from the pool and closes it. If res is not part of the
|
||||
// pool Remove will panic.
|
||||
func (p *Pool[T]) destroyAcquiredResource(res *Resource[T]) {
|
||||
p.destructResourceValue(res.value)
|
||||
|
||||
p.mux.Lock()
|
||||
defer p.mux.Unlock()
|
||||
defer p.acquireSem.Release(1)
|
||||
|
||||
p.allResources.remove(res)
|
||||
}
|
||||
|
||||
func (p *Pool[T]) hijackAcquiredResource(res *Resource[T]) {
|
||||
p.mux.Lock()
|
||||
defer p.mux.Unlock()
|
||||
defer p.acquireSem.Release(1)
|
||||
|
||||
p.allResources.remove(res)
|
||||
res.status = resourceStatusHijacked
|
||||
p.destructWG.Done() // not responsible for destructing hijacked resources
|
||||
}
|
||||
|
||||
func (p *Pool[T]) destructResourceValue(value T) {
|
||||
p.destructor(value)
|
||||
p.destructWG.Done()
|
||||
}
|
||||
+28
@@ -0,0 +1,28 @@
|
||||
package puddle
|
||||
|
||||
type resList[T any] []*Resource[T]
|
||||
|
||||
func (l *resList[T]) append(val *Resource[T]) { *l = append(*l, val) }
|
||||
|
||||
func (l *resList[T]) popBack() *Resource[T] {
|
||||
idx := len(*l) - 1
|
||||
val := (*l)[idx]
|
||||
(*l)[idx] = nil // Avoid memory leak
|
||||
*l = (*l)[:idx]
|
||||
|
||||
return val
|
||||
}
|
||||
|
||||
func (l *resList[T]) remove(val *Resource[T]) {
|
||||
for i, elem := range *l {
|
||||
if elem == val {
|
||||
lastIdx := len(*l) - 1
|
||||
(*l)[i] = (*l)[lastIdx]
|
||||
(*l)[lastIdx] = nil // Avoid memory leak
|
||||
(*l) = (*l)[:lastIdx]
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
panic("BUG: removeResource could not find res in slice")
|
||||
}
|
||||
Reference in New Issue
Block a user