database/sql: fix transaction leak

When the user context which passed in (*DB)BeginTx is canceled or
timeout, the current implementation could cause db transaction leak
in some extreme scenario.

Goroutine 1:
        Call (*DB) BeginTx begins a transaction with a userContext.
        In (*DB)BeginTx, a new goroutine (*Tx)awaitDone
        which monitor context and rollback tx if needed will be created

Goroutine 2(awaitDone):
        block on tx.ctx.Done()

Goroutine 1:
        Execute some insert or update sqls on the database

Goroutine 1:
        Commit the transaction, (*Tx)Commit set
        the atomic variable tx.done to 1

Goroutine 3(maybe global timer):
        Cancel userContext which be passed in Tx

Goroutine 1:
        (*Tx)Commit checks tx.ctx.Done().
        Due to the context has been canceled, it will return
        context.Canceled or context.DeadlineExceeded error immediately
        and abort the real COMMIT operation of transaction

Goroutine 2:
        Release with tx.ctx.Done() signal, execute (*Tx)rollback.
        However the atomic variable tx.done is 1 currently,
        it will return ErrTxDone error immediately and
        abort the real ROLLBACK operation of transaction

Fixes #22976

Change-Id: I3bc23adf25db823861d91e33d3cca6189fb1171d
Reviewed-on: https://go-review.googlesource.com/81736
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
Reviewed-by: Daniel Theophanes <kardianos@gmail.com>

1 files changed, 9 insertions, 3 deletions

diff --git a/src/database/sql/sql.go b/src/database/sql/sql.go
index 22458c0aeb..9f4fa14534 100644
--- a/src/database/sql/sql.go
+++ b/src/database/sql/sql.go
@@ -1913,14 +1913,20 @@ func (tx *Tx) closePrepared() {
 
 // Commit commits the transaction.
 func (tx *Tx) Commit() error {
-	if !atomic.CompareAndSwapInt32(&tx.done, 0, 1) {
-		return ErrTxDone
-	}
+	// Check context first to avoid transaction leak.
+	// If put it behind tx.done CompareAndSwap statement, we cant't ensure
+	// the consistency between tx.done and the real COMMIT operation.
 	select {
 	default:
 	case <-tx.ctx.Done():
+		if atomic.LoadInt32(&tx.done) == 1 {
+			return ErrTxDone
+		}
 		return tx.ctx.Err()
 	}
+	if !atomic.CompareAndSwapInt32(&tx.done, 0, 1) {
+		return ErrTxDone
+	}
 	var err error
 	withLock(tx.dc, func() {
 		err = tx.txi.Commit()