From 3e87a8b3639ed19dc93d179cad05328869ebdfa7 Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Wed, 10 Apr 2019 14:56:14 -0500
Subject: [PATCH] Conn.Close takes context

---
 conn.go           | 6 +++---
 conn_test.go      | 6 +++---
 helper_test.go    | 2 +-
 pool/pool.go      | 7 ++++++-
 pool/pool_test.go | 2 +-
 replication.go    | 2 +-
 6 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/conn.go b/conn.go
index 48b2f22b..aca77dcf 100644
--- a/conn.go
+++ b/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)
diff --git a/conn_test.go b/conn_test.go
index 95fe6357..0febefd0 100644
--- a/conn_test.go
+++ b/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 {
diff --git a/helper_test.go b/helper_test.go
index fa139cf6..74d04166 100644
--- a/helper_test.go
+++ b/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)
 	}
diff --git a/pool/pool.go b/pool/pool.go
index ad6dc458..157007c7 100644
--- a/pool/pool.go
+++ b/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
diff --git a/pool/pool_test.go b/pool/pool_test.go
index f67795e2..f5e71c84 100644
--- a/pool/pool_test.go
+++ b/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())
 
diff --git a/replication.go b/replication.go
index 5493311e..cfe7583a 100644
--- a/replication.go
+++ b/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 {