database/sql: prevent Tx.rollback from racing Tx.close

Previously Tx.done was being set in close, but in a Tx
rollback and Commit are the real closing methods,
and Tx.close is just a helper common to both. Prior to this
change a multiple rollback statements could be called, one
would enter close and begin closing it while the other was
still in rollback breaking it. Fix that by setting done
in rollback and Commit, not in Tx.close.

Fixes #18429

Change-Id: Ie274f60c2aa6a4a5aa38e55109c05ea9d4fe0223
Reviewed-on: https://go-review.googlesource.com/34716
Reviewed-by: Brad Fitzpatrick <[email protected]>
Run-TryBot: Brad Fitzpatrick <[email protected]>
TryBot-Result: Gobot Gobot <[email protected]>

src/database/sql/sql.go

@@ -1421,10 +1421,9 @@ 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")
 
+// close returns the connection to the pool and
+// must only be called by Tx.rollback or Tx.Commit.
 func (tx *Tx) close(err error) {
-	if !atomic.CompareAndSwapInt32(&tx.done, 0, 1) {
-		panic("double close") // internal error
-	}
 	tx.db.putConn(tx.dc, err)
 	tx.cancel()
 	tx.dc = nil
@@ -1449,7 +1448,7 @@ func (tx *Tx) closePrepared() {
 
 // Commit commits the transaction.
 func (tx *Tx) Commit() error {
-	if tx.isDone() {
+	if !atomic.CompareAndSwapInt32(&tx.done, 0, 1) {
 		return ErrTxDone
 	}
 	select {
@@ -1471,7 +1470,7 @@ func (tx *Tx) Commit() error {
 // rollback aborts the transaction and optionally forces the pool to discard
 // the connection.
 func (tx *Tx) rollback(discardConn bool) error {
-	if tx.isDone() {
+	if !atomic.CompareAndSwapInt32(&tx.done, 0, 1) {
 		return ErrTxDone
 	}
 	var err error

src/database/sql/sql_test.go

@@ -2607,6 +2607,54 @@ func TestIssue6081(t *testing.T) {
 	}
 }
 
+// TestIssue18429 attempts to stress rolling back the transaction from a context
+// cancel while simultaneously calling Tx.Rollback. Rolling back from a context
+// happens concurrently so tx.rollback and tx.Commit must gaurded to not
+// be entered twice.
+//
+// The test is composed of a context that is canceled while the query is in process
+// so the internal rollback will run concurrently with the explicitly called
+// Tx.Rollback.
+func TestIssue18429(t *testing.T) {
+	db := newTestDB(t, "people")
+	defer closeDB(t, db)
+
+	ctx := context.Background()
+	sem := make(chan bool, 20)
+	var wg sync.WaitGroup
+
+	const milliWait = 30
+
+	for i := 0; i < 100; i++ {
+		sem <- true
+		wg.Add(1)
+		go func() {
+			defer func() {
+				<-sem
+				wg.Done()
+			}()
+			qwait := (time.Duration(rand.Intn(milliWait)) * time.Millisecond).String()
+
+			ctx, cancel := context.WithTimeout(ctx, time.Duration(rand.Intn(milliWait))*time.Millisecond)
+			defer cancel()
+
+			tx, err := db.BeginTx(ctx, nil)
+			if err != nil {
+				return
+			}
+			rows, err := tx.QueryContext(ctx, "WAIT|"+qwait+"|SELECT|people|name|")
+			if rows != nil {
+				rows.Close()
+			}
+			// This call will race with the context cancel rollback to complete
+			// if the rollback itself isn't guarded.
+			tx.Rollback()
+		}()
+	}
+	wg.Wait()
+	time.Sleep(milliWait * 3 * time.Millisecond)
+}
+
 func TestConcurrency(t *testing.T) {
 	doConcurrentTest(t, new(concurrentDBQueryTest))
 	doConcurrentTest(t, new(concurrentDBExecTest))