database/sql: correctly guard the query Row preventing early release

When a Tx starts a query, prevent returning the connection to the pool
until after the query finishes.

Fixes #19058

Change-Id: I2c0480d9cca9eeb173b5b3441a5aeed6f527e0ac
Reviewed-on: https://go-review.googlesource.com/40400
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

2 files changed, 13 insertions, 6 deletions

diff --git a/src/database/sql/sql.go b/src/database/sql/sql.go
index 17a0088d85..2b84cea374 100644
--- a/src/database/sql/sql.go
+++ b/src/database/sql/sql.go
@@ -1504,6 +1504,14 @@ func (tx *Tx) grabConn(ctx context.Context) (*driverConn, error) {
 	return tx.dc, nil
 }
 
+// closemuRUnlockRelease is used as a func(error) method value in
+// ExecContext and QueryContext. Unlocking in the releaseConn keeps
+// the driver conn from being returned to the connection pool until
+// the Rows has been closed.
+func (tx *Tx) closemuRUnlockRelease(error) {
+	tx.closemu.RUnlock()
+}
+
 // Closes all Stmts prepared for this transaction.
 func (tx *Tx) closePrepared() {
 	tx.stmts.Lock()
@@ -1713,13 +1721,13 @@ func (tx *Tx) Stmt(stmt *Stmt) *Stmt {
 // For example: an INSERT and UPDATE.
 func (tx *Tx) ExecContext(ctx context.Context, query string, args ...interface{}) (Result, error) {
 	tx.closemu.RLock()
-	defer tx.closemu.RUnlock()
 
 	dc, err := tx.grabConn(ctx)
 	if err != nil {
+		tx.closemu.RUnlock()
 		return nil, err
 	}
-	return tx.db.execDC(ctx, dc, func(error) {}, query, args)
+	return tx.db.execDC(ctx, dc, tx.closemuRUnlockRelease, query, args)
 }
 
 // Exec executes a query that doesn't return rows.
@@ -1731,14 +1739,14 @@ func (tx *Tx) Exec(query string, args ...interface{}) (Result, error) {
 // QueryContext executes a query that returns rows, typically a SELECT.
 func (tx *Tx) QueryContext(ctx context.Context, query string, args ...interface{}) (*Rows, error) {
 	tx.closemu.RLock()
-	defer tx.closemu.RUnlock()
 
 	dc, err := tx.grabConn(ctx)
 	if err != nil {
+		tx.closemu.RUnlock()
 		return nil, err
 	}
-	releaseConn := func(error) {}
-	return tx.db.queryDC(ctx, dc, releaseConn, query, args)
+
+	return tx.db.queryDC(ctx, dc, tx.closemuRUnlockRelease, query, args)
 }
 
 // Query executes a query that returns rows, typically a SELECT.
diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index f511aa4ac3..b5a1f850bd 100644
--- a/src/database/sql/sql_test.go
+++ b/src/database/sql/sql_test.go
@@ -2997,7 +2997,6 @@ func TestIssue18719(t *testing.T) {
 	// canceled context.
 
 	cancel()
-	waitForRowsClose(t, rows, 5*time.Second)
 }
 
 func TestConcurrency(t *testing.T) {