Conn.Close takes context

conn.go

@@ -199,7 +199,7 @@ func connect(ctx context.Context, config *ConnConfig, connInfo *pgtype.ConnInfo)
 	if c.ConnInfo == minimalConnInfo {
 		err = c.initConnInfo()
 		if err != nil {
-			c.Close()
+			c.Close(ctx)
 			return nil, err
 		}
 	}
@@ -435,7 +435,7 @@ func (c *Conn) LocalAddr() (net.Addr, error) {
 
 // Close closes a connection. It is safe to call Close on a already closed
 // connection.
-func (c *Conn) Close() error {
+func (c *Conn) Close(ctx context.Context) error {
 	c.mux.Lock()
 	defer c.mux.Unlock()
 
@@ -444,7 +444,7 @@ func (c *Conn) Close() error {
 	}
 	c.status = connStatusClosed
 
-	err := c.pgConn.Close(context.TODO())
+	err := c.pgConn.Close(ctx)
 	c.causeOfDeath = errors.New("Closed")
 	if c.shouldLog(LogLevelInfo) {
 		c.log(LogLevelInfo, "closed connection", nil)

conn_test.go

@@ -72,7 +72,7 @@ func TestConnect(t *testing.T) {
 		t.Errorf("Did not connect as specified user (%v)", config.Config.User)
 	}
 
-	err = conn.Close()
+	err = conn.Close(context.Background())
 	if err != nil {
 		t.Fatal("Unable to close connection")
 	}
@@ -481,7 +481,7 @@ func TestFatalRxError(t *testing.T) {
 	}()
 
 	otherConn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
-	defer otherConn.Close()
+	defer otherConn.Close(context.Background())
 
 	if _, err := otherConn.Exec(context.Background(), "select pg_terminate_backend($1)", conn.PID()); err != nil {
 		t.Fatalf("Unable to kill backend PostgreSQL process: %v", err)
@@ -504,7 +504,7 @@ func TestFatalTxError(t *testing.T) {
 			defer closeConn(t, conn)
 
 			otherConn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
-			defer otherConn.Close()
+			defer otherConn.Close(context.Background())
 
 			_, err := otherConn.Exec(context.Background(), "select pg_terminate_backend($1)", conn.PID())
 			if err != nil {

helper_test.go

@@ -40,7 +40,7 @@ func mustReplicationConnect(t testing.TB, config pgx.ConnConfig) *pgx.Replicatio
 }
 
 func closeConn(t testing.TB, conn *pgx.Conn) {
-	err := conn.Close()
+	err := conn.Close(context.Background())
 	if err != nil {
 		t.Fatalf("conn.Close unexpectedly failed: %v", err)
 	}

pool/pool.go

@@ -2,6 +2,7 @@ package pool
 
 import (
 	"context"
+	"time"
 
 	"github.com/jackc/pgconn"
 	"github.com/jackc/pgx"
@@ -20,7 +21,11 @@ func Connect(ctx context.Context, connString string) (*Pool, error) {
 	maxConnections := 5 // TODO - unhard-code
 	p.p = puddle.NewPool(
 		func(ctx context.Context) (interface{}, error) { return pgx.Connect(ctx, connString) },
-		func(value interface{}) { value.(*pgx.Conn).Close() },
+		func(value interface{}) {
+			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+			value.(*pgx.Conn).Close(ctx)
+			cancel()
+		},
 		maxConnections)
 
 	// Initially establish one connection

pool/pool_test.go

@@ -160,7 +160,7 @@ func TestConnReleaseDestroysClosedConn(t *testing.T) {
 	c, err := pool.Acquire(ctx)
 	require.NoError(t, err)
 
-	c.Conn().Close()
+	c.Conn().Close(ctx)
 
 	assert.Equal(t, 1, pool.Stat().TotalConns())
 

replication.go

@@ -203,7 +203,7 @@ func (rc *ReplicationConn) SendStandbyStatus(k *StandbyStatus) (err error) {
 }
 
 func (rc *ReplicationConn) Close() error {
-	return rc.c.Close()
+	return rc.c.Close(context.TODO())
 }
 
 func (rc *ReplicationConn) IsAlive() bool {