`BeforeConnect` receives `context.Background()` (instead of the pool context) from the healthcheck path

[MartinodF]

Describe the bug
Long story short, we use BeforeConnect to implement IAM Authentication with AWS / Aurora, and I traced back some panics to a missing logger in the context passed to BeforeConnect. I would expect the hook to receive a useful context (for logging, tracing, fetching credentials, ...), but while that happens for Pool.Acquire, the background goroutine refilling the pool to satisfy the minimum size calls it with a bare context.Background().

To Reproduce
Steps to reproduce the behavior:

package main

import (
	"context"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

type key int

const myKey key = 0

func main() {
	ctx := context.WithValue(context.Background(), myKey, "value")

	poolConfig, _ := pgxpool.ParseConfig("postgres://user:password@127.0.0.1:1/mydb")

	// Fail fast
	poolConfig.ConnConfig.ConnectTimeout = 200 * time.Millisecond

	// Prepare connections in the background
	poolConfig.MinConns = 1
	poolConfig.HealthCheckPeriod = 300 * time.Millisecond

	var calls int
	poolConfig.BeforeConnect = func(bcCtx context.Context, cc *pgx.ConnConfig) error {
		calls++
		if v := bcCtx.Value(myKey); v == nil {
			fmt.Printf("call %d: BeforeConnect ctx MISSING value\n", calls)
		} else {
			fmt.Printf("call %d: BeforeConnect ctx has value=%q\n", calls, v)
		}
		return nil
	}

	pool, err := pgxpool.NewWithConfig(ctx, poolConfig)
	if err != nil {
		panic(err)
	}
	defer pool.Close()

	// Let the background healthcheck fire a few times.
	time.Sleep(2 * time.Second)
}

Expected behavior

call 1: ctx has value="value"
call 2: ctx has value="value"
call 3: ctx has value="value"

Actual behavior

call 1: ctx has value="value"
call 2: ctx MISSING value
call 3: ctx MISSING value

Call 1 comes from the initial idle creation goroutine in NewWithConfig. The subsequent ones come from backgroundHealthCheck + checkMinConns.

Impact
Anyone relying on context.Value propagation in BeforeConnect (e.g. injecting a request-scoped logger or AWS credential cache) sees inconsistent errors: the hook works on Acquire, it works on initial pool warm-up, but it randomly breaks on background refills. In our case it caused a nil-pointer panic while trying to log a message about the authentication.

Either re-use the context passed to NewWithConfig, expose a new Config.HealthCheckContext, or at the very least document the behavior to avoid unpleasant surprises. I can have a go at the first option if there's agreement.

Version

• Go: go1.26.2 darwin/arm64
• pgx: v5.9.1
• puddle: v2.2.2

[MartinodF]

(I realized this is basically a follow-up to #1484 from 2023, just covering the last edge case where we call checkMinConns)

[abrightwell]

Why not set the needed values on the ctx within BeforeConnect?

Given that context.Context should be request-scoped, reuse of the same ctx passed to NewWithConfig would seem incorrect. As it is used for the initial warmup, etc, which likely has some timeout or cancellation associated with it that is specific to that purpose. Passing application-scoped values per connection request seems like what BeforeConnect should be used to solve?

Perhaps documenting this behavior is warranted, but I'm not sure as this seems like a fundamental expectation for the use of context.Context.

[MartinodF]

Why not set the needed values on the ctx within BeforeConnect?

Because you set the callback once, when initializing the pool, rather then when obtaining a connection? I am already passing a request-scoped (which is irrelevant, more below) context to Acquire & friends, but pgx doesn't pass it down to BeforeConnect in this particular instance because the operation is a background one instead.

Given that context.Context should be request-scoped, reuse of the same ctx passed to NewWithConfig would seem incorrect. As it is used for the initial warmup, etc, which likely has some timeout or cancellation associated with it that is specific to that purpose. Passing application-scoped values per connection request seems like what BeforeConnect should be used to solve?

Request-scoped is not the relevant part in this case, as this is a background maintenance operation that isn't tied to a particular request. An application-scoped context (which in my case would contain a reference to my logger, for example) is much preferable to a random context.Background?

#2546 is more or less what I would expect (see the WithCancel), with the other reasonable option being a separate config function which is called to obtain a context when necessary

[abrightwell]

Because you set the callback once, when initializing the pool, rather then when obtaining a connection?

Since BeforeConnect is called every time a new connection is created, whether at pool warmup or refill, can you share more about your use case and why you can't reference your dependencies directly in the callback?

For instance:

logger := slog.Default()

config.BeforeConnect = func(ctx context.Context, cc *pgx.ConnConfig) error {
    logger.InfoContext(ctx, "establishing new connection")
    return nil
}

config.BeforeAcquire = func(ctx context.Context, c *pgx.Conn) bool {
    logger.InfoContext(ctx, "acquiring connection")
    return true
}

An application-scoped context (which in my case would contain a reference to my logger, for example) is much preferable to a random context.Background?

Sure, I see your point, something is better than nothing. But an application-scoped context still doesn't feel like the right solution here. Perhaps a dedicated callback for building a context per connection operation would be a cleaner approach?

For instance:

BaseContext func(context.Context) context.Context // defaults to context.Background()

package logger

type loggerKey struct{}

func WithLogger(ctx context.Context) context.Context {
    return context.WithValue(ctx, loggerKey{}, slog.Default())
}

func Logger(ctx context.Context) *slog.Logger {
    return ctx.Value(loggerKey{}).(*slog.Logger)
}

// Called before every connection operation to allow enriching the context.
config.BaseContext = func(ctx context.Context) context.Context {
    return logger.WithLogger(ctx)
}

config.BeforeConnect = func(ctx context.Context, cc *pgx.ConnConfig) error {
    logger.Logger(ctx).InfoContext(ctx, "establishing new connection")
    return nil
}

config.BeforeAcquire = func(ctx context.Context, c *pgx.Conn) bool {
    logger.Logger(ctx).InfoContext(ctx, "acquiring connection")
    return true
}

Thoughts?

[jackc]

I pretty much share @abrightwell opinions here.

• It does feel wrong to keep using the ctx passed to the pool on creation
• BeforeConnect could have dependencies injected via closure
• If that is insufficient (which I am not yet convinced of), then a function that creates a base / background ctx is better than always reusing the original ctx the pool was created with.