database/sql: on Tx rollback, retain connection if driver can reset session

Previously the Tx would drop the connection after rolling back from
a context cancel. Now if the driver can reset the session,
keep the connection.

Change-Id: Ie6a3124275632787629844d91a06bb2e70cc060b
Reviewed-on: https://go-review.googlesource.com/c/go/+/216241
Reviewed-by: Emmanuel Odeke <emm.odeke@gmail.com>
Run-TryBot: Emmanuel Odeke <emm.odeke@gmail.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>

diff --git a/src/database/sql/sql.go b/src/database/sql/sql.go
index 3710264dcf6097dcd25db6f51662185726a8a1d8..b63d5591f6d9d2c9decffd7cabe15efdd74e6828 100644 (file)
--- a/src/database/sql/sql.go
+++ b/src/database/sql/sql.go
@@ -1746,7 +1746,11 @@ func (db *DB) begin(ctx context.Context, opts *TxOptions, strategy connReuseStra
 // beginDC starts a transaction. The provided dc must be valid and ready to use.
 func (db *DB) beginDC(ctx context.Context, dc *driverConn, release func(error), opts *TxOptions) (tx *Tx, err error) {
        var txi driver.Tx
+       keepConnOnRollback := false
        withLock(dc, func() {
+               _, hasSessionResetter := dc.ci.(driver.SessionResetter)
+               _, hasConnectionValidator := dc.ci.(driver.Validator)
+               keepConnOnRollback = hasSessionResetter && hasConnectionValidator
                txi, err = ctxDriverBegin(ctx, opts, dc.ci)
        })
        if err != nil {
@@ -1758,12 +1762,13 @@ func (db *DB) beginDC(ctx context.Context, dc *driverConn, release func(error),
        // The cancel function in Tx will be called after done is set to true.
        ctx, cancel := context.WithCancel(ctx)
        tx = &Tx{
-               db:          db,
-               dc:          dc,
-               releaseConn: release,
-               txi:         txi,
-               cancel:      cancel,
-               ctx:         ctx,
+               db:                 db,
+               dc:                 dc,
+               releaseConn:        release,
+               txi:                txi,
+               cancel:             cancel,
+               keepConnOnRollback: keepConnOnRollback,
+               ctx:                ctx,
        }
        go tx.awaitDone()
        return tx, nil
@@ -2025,6 +2030,11 @@ type Tx struct {
        // Use atomic operations on value when checking value.
        done int32
 
+       // keepConnOnRollback is true if the driver knows
+       // how to reset the connection's session and if need be discard
+       // the connection.
+       keepConnOnRollback bool
+
        // All Stmts prepared for this transaction. These will be closed after the
        // transaction has been committed or rolled back.
        stmts struct {
@@ -2050,7 +2060,10 @@ func (tx *Tx) awaitDone() {
        // transaction is closed and the resources are released.  This
        // rollback does nothing if the transaction has already been
        // committed or rolled back.
-       tx.rollback(true)
+       // Do not discard the connection if the connection knows
+       // how to reset the session.
+       discardConnection := !tx.keepConnOnRollback
+       tx.rollback(discardConnection)
 }
 
 func (tx *Tx) isDone() bool {

diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index 41265caed8060cfc72511b474656fd17666e7e20..00a67c79ee073a8c21ac8c26ac895b9704c14d85 100644 (file)
--- a/src/database/sql/sql_test.go
+++ b/src/database/sql/sql_test.go
@@ -442,6 +442,7 @@ func TestTxContextWait(t *testing.T) {
                }
                t.Fatal(err)
        }
+       tx.keepConnOnRollback = false
 
        // This will trigger the *fakeConn.Prepare method which will take time
        // performing the query. The ctxDriverPrepare func will check the context
@@ -454,6 +455,35 @@ func TestTxContextWait(t *testing.T) {
        waitForFree(t, db, 5*time.Second, 0)
 }
 
+// TestTxContextWaitNoDiscard is the same as TestTxContextWait, but should not discard
+// the final connection.
+func TestTxContextWaitNoDiscard(t *testing.T) {
+       db := newTestDB(t, "people")
+       defer closeDB(t, db)
+
+       ctx, cancel := context.WithTimeout(context.Background(), 15*time.Millisecond)
+       defer cancel()
+
+       tx, err := db.BeginTx(ctx, nil)
+       if err != nil {
+               // Guard against the context being canceled before BeginTx completes.
+               if err == context.DeadlineExceeded {
+                       t.Skip("tx context canceled prior to first use")
+               }
+               t.Fatal(err)
+       }
+
+       // This will trigger the *fakeConn.Prepare method which will take time
+       // performing the query. The ctxDriverPrepare func will check the context
+       // after this and close the rows and return an error.
+       _, err = tx.QueryContext(ctx, "WAIT|1s|SELECT|people|age,name|")
+       if err != context.DeadlineExceeded {
+               t.Fatalf("expected QueryContext to error with context deadline exceeded but returned %v", err)
+       }
+
+       waitForFree(t, db, 5*time.Second, 1)
+}
+
 // TestUnsupportedOptions checks that the database fails when a driver that
 // doesn't implement ConnBeginTx is used with non-default options and an
 // un-cancellable context.