diff --git a/pgxpool/pool.go b/pgxpool/pool.go index 0a7c6c923..e6ff92d36 100644 --- a/pgxpool/pool.go +++ b/pgxpool/pool.go @@ -619,6 +619,13 @@ func (p *Pool) Acquire(ctx context.Context) (c *Conn, err error) { return nil, err } + // Handle race condition: If the connection is assigned to this waiter at the exact + // moment the waiter's context is canceled, we must safely return it to the idle pool. + if ctx.Err() != nil { + res.Release() + return nil, ctx.Err() + } + cr := res.Value() // Destroy expired connections before doing any further work (such as diff --git a/pgxpool/repro_test.go b/pgxpool/repro_test.go new file mode 100644 index 000000000..e9dbadd01 --- /dev/null +++ b/pgxpool/repro_test.go @@ -0,0 +1,55 @@ +package pgxpool_test + +import ( + "context" + "math/rand" + "os" + "sync" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/stretchr/testify/require" +) + +func TestAcquireContextCancelRace(t *testing.T) { + t.Parallel() + + connString := os.Getenv("PGX_TEST_DATABASE") + if connString == "" { + t.Skip("PGX_TEST_DATABASE is not set") + } + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + config, err := pgxpool.ParseConfig(connString) + require.NoError(t, err) + config.MaxConns = 5 + + pool, err := pgxpool.NewWithConfig(ctx, config) + require.NoError(t, err) + defer pool.Close() + + var wg sync.WaitGroup + for i := 0; i < 1000; i++ { + wg.Add(1) + go func() { + defer wg.Done() + + timeout := time.Duration(rand.Intn(5)+1) * time.Millisecond + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + + if conn, err := pool.Acquire(ctx); err == nil { + time.Sleep(2 * time.Millisecond) + conn.Release() + } + }() + } + + wg.Wait() + time.Sleep(100 * time.Millisecond) + + require.Equal(t, int32(0), pool.Stat().AcquiredConns()) +}