Repository: jackc/pgx
Pull request #1094 (2021-10-03 19:59:15)

Eager initialize minpoolsize on connect
Currently minPoolSize is ignored on connect and the idle connections are only created once the background check runs.
This breaks developer expectations and has a big performance impact on applications that are autoscaled.
Autoscaled applications will immediatly receive lots of traffic and if the pool isn't ready the connections will have to be created the first few requests (lazy) instead of the startup of the application.

Currently our applications use this workaround:
```go
func warmupConnections(ctx context.Context, conn *pgxpool.Pool, targetConnections int) error {
	for i := 0; i < targetConnections; i++ {
		c, err := conn.Acquire(ctx)
		if err != nil {
			return err
		}
		defer c.Release()
	}
	return nil
}
```

But it would be nice if this library would enforce the min pool size at all times.

BornToGrill: Forgot to run with the -race flag. Fixed my tests.
(TestPoolBackgroundChecksMaxConnIdleTime is flakey on my machine even on master).

Patch:
diff --git a/doc.go b/doc.go
index 51b0d9f4..222f9047 100644
--- a/doc.go
+++ b/doc.go
@@ -309,7 +309,7 @@ CopyFrom can be faster than an insert with as few as 5 rows.
 Listen and Notify
 
 pgx can listen to the PostgreSQL notification system with the `Conn.WaitForNotification` method. It blocks until a
-context is received or the context is canceled.
+notification is received or the context is canceled.
 
     _, err := conn.Exec(context.Background(), "listen channelname")
     if err != nil {
diff --git a/log/zerologadapter/adapter.go b/log/zerologadapter/adapter.go
index 9bfa5788..6e8b4b94 100644
--- a/log/zerologadapter/adapter.go
+++ b/log/zerologadapter/adapter.go
@@ -9,9 +9,10 @@ import (
 )
 
 type Logger struct {
-	logger     zerolog.Logger
-	withFunc   func(context.Context, zerolog.Context) zerolog.Context
-	skipModule bool
+	logger      zerolog.Logger
+	withFunc    func(context.Context, zerolog.Context) zerolog.Context
+	fromContext bool
+	skipModule  bool
 }
 
 // option options for configuring the logger when creating a new logger.
@@ -38,13 +39,28 @@ func NewLogger(logger zerolog.Logger, options ...option) *Logger {
 	l := Logger{
 		logger: logger,
 	}
+	l.init(options)
+	return &l
+}
+
+// NewContextLogger creates logger that extracts the zerolog.Logger from the
+// context.Context by using `zerolog.Ctx`. The zerolog.DefaultContextLogger will
+// be used if no logger is associated with the context.
+func NewContextLogger(options ...option) *Logger {
+	l := Logger{
+		fromContext: true,
+	}
+	l.init(options)
+	return &l
+}
+
+func (pl *Logger) init(options []option) {
 	for _, opt := range options {
-		opt(&l)
+		opt(pl)
 	}
-	if !l.skipModule {
-		l.logger = l.logger.With().Str("module", "pgx").Logger()
+	if !pl.skipModule {
+		pl.logger = pl.logger.With().Str("module", "pgx").Logger()
 	}
-	return &l
 }
 
 func (pl *Logger) Log(ctx context.Context, level pgx.LogLevel, msg string, data map[string]interface{}) {
@@ -63,10 +79,24 @@ func (pl *Logger) Log(ctx context.Context, level pgx.LogLevel, msg string, data
 	default:
 		zlevel = zerolog.DebugLevel
 	}
-	zctx := pl.logger.With()
+
+	var zctx zerolog.Context
+	if pl.fromContext {
+		logger := zerolog.Ctx(ctx)
+		zctx = logger.With()
+	} else {
+		zctx = pl.logger.With()
+	}
 	if pl.withFunc != nil {
 		zctx = pl.withFunc(ctx, zctx)
 	}
-	pgxlog := zctx.Fields(data).Logger()
-	pgxlog.WithLevel(zlevel).Msg(msg)
+
+	pgxlog := zctx.Logger()
+	event := pgxlog.WithLevel(zlevel)
+	if event.Enabled() {
+		if pl.fromContext && !pl.skipModule {
+			event.Str("module", "pgx")
+		}
+		event.Fields(data).Msg(msg)
+	}
 }
diff --git a/pgxpool/pool.go b/pgxpool/pool.go
index f1a79c7f..86ba40c9 100644
--- a/pgxpool/pool.go
+++ b/pgxpool/pool.go
@@ -225,6 +225,12 @@ func ConnectConfig(ctx context.Context, config *Config) (*Pool, error) {
 	go p.backgroundHealthCheck()
 
 	if !config.LazyConnect {
+		if err := p.createIdleResources(ctx, int(p.minConns)); err != nil {
+			// Couldn't create resources for minpool size. Close unhealthy pool.
+			p.Close()
+			return nil, err
+		}
+
 		// Initially establish one connection
 		res, err := p.p.Acquire(ctx)
 		if err != nil {
@@ -377,6 +383,31 @@ func (p *Pool) checkMinConns() {
 	}
 }
 
+func (p *Pool) createIdleResources(parentCtx context.Context, targetResources int) error {
+	ctx, cancel := context.WithCancel(parentCtx)
+	defer cancel()
+
+	errs := make(chan error, targetResources)
+
+	for i := 0; i < targetResources; i++ {
+		go func() {
+			err := p.p.CreateResource(ctx)
+			errs <- err
+		}()
+	}
+
+	var firstError error
+	for i := 0; i < targetResources; i++ {
+		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) (*Conn, error) {
 	for {


Test patch:
diff --git a/log/zerologadapter/adapter_test.go b/log/zerologadapter/adapter_test.go
index deda6ed4..3a11cbc0 100644
--- a/log/zerologadapter/adapter_test.go
+++ b/log/zerologadapter/adapter_test.go
@@ -38,6 +38,23 @@ func TestLogger(t *testing.T) {
 		}
 	})
 
+	t.Run("from context", func(t *testing.T) {
+		var buf bytes.Buffer
+		zlogger := zerolog.New(&buf)
+		ctx := zlogger.WithContext(context.Background())
+		logger := zerologadapter.NewContextLogger()
+		logger.Log(ctx, pgx.LogLevelInfo, "hello", map[string]interface{}{"one": "two"})
+		const want = `{"level":"info","module":"pgx","one":"two","message":"hello"}
+`
+
+		got := buf.String()
+		if got != want {
+			t.Log(got)
+			t.Log(want)
+			t.Errorf("%s != %s", got, want)
+		}
+	})
+
 	var buf bytes.Buffer
 	type key string
 	var ck key
@@ -52,7 +69,8 @@ func TestLogger(t *testing.T) {
 				logWith = logWith.Str("req_id", id)
 			}
 			return logWith
-		}))
+		}),
+	)
 
 	t.Run("no request id", func(t *testing.T) {
 		buf.Reset()
@@ -77,5 +95,4 @@ func TestLogger(t *testing.T) {
 			t.Errorf("%s != %s", got, want)
 		}
 	})
-
 }
diff --git a/pgxpool/pool_test.go b/pgxpool/pool_test.go
index f6621df6..ac38c1a4 100644
--- a/pgxpool/pool_test.go
+++ b/pgxpool/pool_test.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -800,3 +801,179 @@ func TestIdempotentPoolClose(t *testing.T) {
 	// Close the already closed pool.
 	require.NotPanics(t, func() { pool.Close() })
 }
+
+func TestConnectCreatesMinPool(t *testing.T) {
+	t.Parallel()
+
+	config, err := pgxpool.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+
+	config.MinConns = int32(12)
+	config.MaxConns = int32(15)
+	config.LazyConnect = false
+
+	acquireAttempts := int64(0)
+	connectAttempts := int64(0)
+
+	config.BeforeAcquire = func(ctx context.Context, conn *pgx.Conn) bool {
+		atomic.AddInt64(&acquireAttempts, 1)
+		return true
+	}
+	config.BeforeConnect = func(ctx context.Context, cfg *pgx.ConnConfig) error {
+		atomic.AddInt64(&connectAttempts, 1)
+		return nil
+	}
+
+	pool, err := pgxpool.ConnectConfig(context.Background(), config)
+	require.NoError(t, err)
+	defer pool.Close()
+
+	stat := pool.Stat()
+	require.Equal(t, int32(12), stat.IdleConns())
+	require.Equal(t, int64(1), stat.AcquireCount())
+	require.Equal(t, int32(12), stat.TotalConns())
+	require.Equal(t, int64(0), acquireAttempts)
+	require.Equal(t, int64(12), connectAttempts)
+}
+func TestConnectSkipMinPoolWithLazy(t *testing.T) {
+	t.Parallel()
+
+	config, err := pgxpool.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+
+	config.MinConns = int32(12)
+	config.MaxConns = int32(15)
+	config.LazyConnect = true
+
+	acquireAttempts := int64(0)
+	connectAttempts := int64(0)
+
+	config.BeforeAcquire = func(ctx context.Context, conn *pgx.Conn) bool {
+		atomic.AddInt64(&acquireAttempts, 1)
+		return true
+	}
+	config.BeforeConnect = func(ctx context.Context, cfg *pgx.ConnConfig) error {
+		atomic.AddInt64(&connectAttempts, 1)
+		return nil
+	}
+
+	pool, err := pgxpool.ConnectConfig(context.Background(), config)
+	require.NoError(t, err)
+	defer pool.Close()
+
+	stat := pool.Stat()
+	require.Equal(t, int32(0), stat.IdleConns())
+	require.Equal(t, int64(0), stat.AcquireCount())
+	require.Equal(t, int32(0), stat.TotalConns())
+	require.Equal(t, int64(0), acquireAttempts)
+	require.Equal(t, int64(0), connectAttempts)
+}
+
+func TestConnectMinPoolZero(t *testing.T) {
+	t.Parallel()
+
+	config, err := pgxpool.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+
+	config.MinConns = int32(0)
+	config.MaxConns = int32(15)
+	config.LazyConnect = false
+
+	acquireAttempts := int64(0)
+	connectAttempts := int64(0)
+
+	config.BeforeAcquire = func(ctx context.Context, conn *pgx.Conn) bool {
+		atomic.AddInt64(&acquireAttempts, 1)
+		return true
+	}
+	config.BeforeConnect = func(ctx context.Context, cfg *pgx.ConnConfig) error {
+		atomic.AddInt64(&connectAttempts, 1)
+		return nil
+	}
+
+	pool, err := pgxpool.ConnectConfig(context.Background(), config)
+	require.NoError(t, err)
+	defer pool.Close()
+
+	stat := pool.Stat()
+	require.Equal(t, int32(1), stat.IdleConns())
+	require.Equal(t, int64(1), stat.AcquireCount())
+	require.Equal(t, int32(1), stat.TotalConns())
+	require.Equal(t, int64(0), acquireAttempts)
+	require.Equal(t, int64(1), connectAttempts)
+}
+
+func TestCreateMinPoolClosesConnectionsOnError(t *testing.T) {
+	t.Parallel()
+
+	config, err := pgxpool.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+
+	config.MinConns = int32(12)
+	config.MaxConns = int32(15)
+	config.LazyConnect = false
+
+	acquireAttempts := int64(0)
+	madeConnections := int64(0)
+	conns := make(chan *pgx.Conn, 15)
+
+	config.BeforeAcquire = func(ctx context.Context, conn *pgx.Conn) bool {
+		atomic.AddInt64(&acquireAttempts, 1)
+		return true
+	}
+	config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
+		conns <- conn
+
+		atomic.AddInt64(&madeConnections, 1)
+		mc := atomic.LoadInt64(&madeConnections)
+		if mc == 10 {
+			return errors.New("mock error")
+		}
+		return nil
+	}
+	pool, err := pgxpool.ConnectConfig(context.Background(), config)
+	require.Error(t, err)
+	require.Nil(t, pool)
+
+	close(conns)
+	for conn := range conns {
+		require.True(t, conn.IsClosed())
+	}
+
+	require.Equal(t, int64(0), acquireAttempts)
+	require.True(t, madeConnections >= 10, "Expected %d got %d", 10, madeConnections)
+}
+
+func TestCreateMinPoolReturnsFirstError(t *testing.T) {
+	t.Parallel()
+
+	config, err := pgxpool.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+
+	config.MinConns = int32(12)
+	config.MaxConns = int32(15)
+	config.LazyConnect = false
+
+	acquireAttempts := int64(0)
+	connectAttempts := int64(0)
+
+	config.BeforeAcquire = func(ctx context.Context, conn *pgx.Conn) bool {
+		atomic.AddInt64(&acquireAttempts, 1)
+		return true
+	}
+	config.BeforeConnect = func(ctx context.Context, cfg *pgx.ConnConfig) error {
+		atomic.AddInt64(&connectAttempts, 1)
+		ca := atomic.LoadInt64(&connectAttempts)
+		if ca >= 5 {
+			return fmt.Errorf("error %d", ca)
+		}
+		return nil
+	}
+
+	pool, err := pgxpool.ConnectConfig(context.Background(), config)
+	require.Nil(t, pool)
+	require.Error(t, err)
+
+	require.True(t, connectAttempts >= 5, "Expected %d got %d", 5, connectAttempts)
+	require.Equal(t, "error 5", err.Error())
+}