2 files changed, 52 insertions, 14 deletions
diff --git a/src/database/sql/sql.go b/src/database/sql/sql.go
index e3580698fd..726aadb899 100644
--- a/src/database/sql/sql.go
+++ b/src/database/sql/sql.go
@@ -434,7 +434,7 @@ type DB struct {
 	maxIdleTimeClosed int64 // Total number of connections closed due to idle time.
 	maxLifetimeClosed int64 // Total number of connections closed due to max connection lifetime limit.
 
-	stop func() // stop cancels the connection opener and the session resetter.
+	stop func() // stop cancels the connection opener.
 }
 
 // connReuseStrategy determines how (*DB).conn returns database connections.
@@ -1141,7 +1141,7 @@ func (db *DB) connectionOpener(ctx context.Context) {
 
 // Open one new connection
 func (db *DB) openNewConnection(ctx context.Context) {
-	// maybeOpenNewConnctions has already executed db.numOpen++ before it sent
+	// maybeOpenNewConnections has already executed db.numOpen++ before it sent
 	// on db.openerCh. This function must execute db.numOpen-- if the
 	// connection fails or is closed before returning.
 	ci, err := db.connector.Connect(ctx)
@@ -2087,10 +2087,10 @@ func (tx *Tx) isDone() bool {
 // that has already been committed or rolled back.
 var ErrTxDone = errors.New("sql: transaction has already been committed or rolled back")
 
-// closeLocked returns the connection to the pool and
+// close returns the connection to the pool and
 // must only be called by Tx.rollback or Tx.Commit while
-// closemu is Locked and tx already canceled.
-func (tx *Tx) closeLocked(err error) {
+// tx is already canceled and won't be executed concurrently.
+func (tx *Tx) close(err error) {
 	tx.releaseConn(err)
 	tx.dc = nil
 	tx.txi = nil
@@ -2164,7 +2164,7 @@ func (tx *Tx) Commit() error {
 	// to ensure no other connection has an active query.
 	tx.cancel()
 	tx.closemu.Lock()
-	defer tx.closemu.Unlock()
+	tx.closemu.Unlock()
 
 	var err error
 	withLock(tx.dc, func() {
@@ -2173,7 +2173,7 @@ func (tx *Tx) Commit() error {
 	if err != driver.ErrBadConn {
 		tx.closePrepared()
 	}
-	tx.closeLocked(err)
+	tx.close(err)
 	return err
 }
 
@@ -2196,7 +2196,7 @@ func (tx *Tx) rollback(discardConn bool) error {
 	// to ensure no other connection has an active query.
 	tx.cancel()
 	tx.closemu.Lock()
-	defer tx.closemu.Unlock()
+	tx.closemu.Unlock()
 
 	var err error
 	withLock(tx.dc, func() {
@@ -2208,7 +2208,7 @@ func (tx *Tx) rollback(discardConn bool) error {
 	if discardConn {
 		err = driver.ErrBadConn
 	}
-	tx.closeLocked(err)
+	tx.close(err)
 	return err
 }
 
diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index 762d42f54b..c968852ade 100644
--- a/src/database/sql/sql_test.go
+++ b/src/database/sql/sql_test.go
@@ -2810,6 +2810,34 @@ func TestTxCannotCommitAfterRollback(t *testing.T) {
 	}
 }
 
+// Issue 40985 transaction statement deadlock while context cancel.
+func TestTxStmtDeadlock(t *testing.T) {
+	db := newTestDB(t, "people")
+	defer closeDB(t, db)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Millisecond)
+	defer cancel()
+	tx, err := db.BeginTx(ctx, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	stmt, err := tx.Prepare("SELECT|people|name,age|age=?")
+	if err != nil {
+		t.Fatal(err)
+	}
+	// Run number of stmt queries to reproduce deadlock from context cancel
+	for i := 0; i < 1e3; i++ {
+		// Encounter any close related errors (e.g. ErrTxDone, stmt is closed)
+		// is expected due to context cancel.
+		_, err = stmt.Query(1)
+		if err != nil {
+			break
+		}
+	}
+	_ = tx.Rollback()
+}
+
 // Issue32530 encounters an issue where a connection may
 // expire right after it comes out of a used connection pool
 // even when a new connection is requested.
@@ -2860,20 +2888,26 @@ func TestConnExpiresFreshOutOfPool(t *testing.T) {
 			waitingForConn := make(chan struct{})
 
 			go func() {
+				defer close(afterPutConn)
+
 				conn, err := db.conn(ctx, alwaysNewConn)
-				if err != nil {
-					t.Fatal(err)
+				if err == nil {
+					db.putConn(conn, err, false)
+				} else {
+					t.Errorf("db.conn: %v", err)
 				}
-				db.putConn(conn, err, false)
-				close(afterPutConn)
 			}()
 			go func() {
+				defer close(waitingForConn)
+
 				for {
+					if t.Failed() {
+						return
+					}
 					db.mu.Lock()
 					ct := len(db.connRequests)
 					db.mu.Unlock()
 					if ct > 0 {
-						close(waitingForConn)
 						return
 					}
 					time.Sleep(10 * time.Millisecond)
@@ -2882,6 +2916,10 @@ func TestConnExpiresFreshOutOfPool(t *testing.T) {
 
 			<-waitingForConn
 
+			if t.Failed() {
+				return
+			}
+
 			offsetMu.Lock()
 			if ec.expired {
 				offset = 11 * time.Second