From e8857f04a18e8650975d787b46fe280cf49bb9d4 Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Mon, 7 Feb 2022 10:44:39 -0600
Subject: [PATCH] Make BatchResults.Close safe to be called multiple times

https://github.com/jackc/pgx/issues/1138
https://github.com/jackc/pgx/issues/938
---
 batch.go             | 33 ++++++++++++++++++++++------
 pgxpool/pool_test.go | 52 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 7 deletions(-)

diff --git a/batch.go b/batch.go
index f0479ea6..7f86ad5c 100644
--- a/batch.go
+++ b/batch.go
@@ -3,6 +3,7 @@ package pgx
 import (
 	"context"
 	"errors"
+	"fmt"
 
 	"github.com/jackc/pgconn"
 )
@@ -46,17 +47,18 @@ type BatchResults interface {
 
 	// Close closes the batch operation. This must be called before the underlying connection can be used again. Any error
 	// that occurred during a batch operation may have made it impossible to resyncronize the connection with the server.
-	// In this case the underlying connection will have been closed.
+	// In this case the underlying connection will have been closed. Close is safe to call multiple times.
 	Close() error
 }
 
 type batchResults struct {
-	ctx  context.Context
-	conn *Conn
-	mrr  *pgconn.MultiResultReader
-	err  error
-	b    *Batch
-	ix   int
+	ctx    context.Context
+	conn   *Conn
+	mrr    *pgconn.MultiResultReader
+	err    error
+	b      *Batch
+	ix     int
+	closed bool
 }
 
 // Exec reads the results from the next query in the batch as if the query has been sent with Exec.
@@ -64,6 +66,9 @@ func (br *batchResults) Exec() (pgconn.CommandTag, error) {
 	if br.err != nil {
 		return nil, br.err
 	}
+	if br.closed {
+		return nil, fmt.Errorf("batch already closed")
+	}
 
 	query, arguments, _ := br.nextQueryAndArgs()
 
@@ -114,6 +119,11 @@ func (br *batchResults) Query() (Rows, error) {
 		return &connRows{err: br.err, closed: true}, br.err
 	}
 
+	if br.closed {
+		alreadyClosedErr := fmt.Errorf("batch already closed")
+		return &connRows{err: alreadyClosedErr, closed: true}, alreadyClosedErr
+	}
+
 	rows := br.conn.getRows(br.ctx, query, arguments)
 
 	if !br.mrr.NextResult() {
@@ -140,6 +150,10 @@ func (br *batchResults) Query() (Rows, error) {
 
 // QueryFunc reads the results from the next query in the batch as if the query has been sent with Conn.QueryFunc.
 func (br *batchResults) QueryFunc(scans []interface{}, f func(QueryFuncRow) error) (pgconn.CommandTag, error) {
+	if br.closed {
+		return nil, fmt.Errorf("batch already closed")
+	}
+
 	rows, err := br.Query()
 	if err != nil {
 		return nil, err
@@ -179,6 +193,11 @@ func (br *batchResults) Close() error {
 		return br.err
 	}
 
+	if br.closed {
+		return nil
+	}
+	br.closed = true
+
 	// log any queries that haven't yet been logged by Exec or Query
 	for {
 		query, args, ok := br.nextQueryAndArgs()
diff --git a/pgxpool/pool_test.go b/pgxpool/pool_test.go
index 54b688a1..9a813ae2 100644
--- a/pgxpool/pool_test.go
+++ b/pgxpool/pool_test.go
@@ -979,3 +979,55 @@ func TestCreateMinPoolReturnsFirstError(t *testing.T) {
 	require.True(t, connectAttempts >= 5, "Expected %d got %d", 5, connectAttempts)
 	require.ErrorIs(t, err, mockErr)
 }
+
+func TestPoolSendBatchBatchCloseTwice(t *testing.T) {
+	t.Parallel()
+
+	pool, err := pgxpool.Connect(context.Background(), os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+	defer pool.Close()
+
+	errChan := make(chan error)
+	testCount := 5000
+
+	for i := 0; i < testCount; i++ {
+		go func() {
+			batch := &pgx.Batch{}
+			batch.Queue("select 1")
+			batch.Queue("select 2")
+
+			br := pool.SendBatch(context.Background(), batch)
+			defer br.Close()
+
+			var err error
+			var n int32
+			err = br.QueryRow().Scan(&n)
+			if err != nil {
+				errChan <- err
+				return
+			}
+			if n != 1 {
+				errChan <- fmt.Errorf("expected 1 got %v", n)
+				return
+			}
+
+			err = br.QueryRow().Scan(&n)
+			if err != nil {
+				errChan <- err
+				return
+			}
+			if n != 2 {
+				errChan <- fmt.Errorf("expected 2 got %v", n)
+				return
+			}
+
+			err = br.Close()
+			errChan <- err
+		}()
+	}
+
+	for i := 0; i < testCount; i++ {
+		err := <-errChan
+		assert.NoError(t, err)
+	}
+}