From fb654e555632fb5c26be905c08e82bc5758d6308 Mon Sep 17 00:00:00 2001
From: Ashraf Hanafy <ashie1287@gmail.com>
Date: Mon, 1 Nov 2021 10:48:19 -0400
Subject: [PATCH 1/2] rollback on panic (need help with tests)

---
 crdb/common.go | 40 ++++++++++++++++++++++++++--------------
 go.mod         |  2 +-
 2 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/crdb/common.go b/crdb/common.go
index 68b2be9..2857027 100644
--- a/crdb/common.go
+++ b/crdb/common.go
@@ -39,15 +39,23 @@ type Tx interface {
 // fn is subject to the same restrictions as the fn passed to ExecuteTx.
 func ExecuteInTx(ctx context.Context, tx Tx, fn func() error) (err error) {
 	defer func() {
-		if err == nil {
+		r := recover()
+
+		if r == nil && err == nil {
 			// Ignore commit errors. The tx has already been committed by RELEASE.
 			_ = tx.Commit(ctx)
-		} else {
-			// We always need to execute a Rollback() so sql.DB releases the
-			// connection.
-			_ = tx.Rollback(ctx)
+			return
+		}
+
+		// We always need to execute a Rollback() so sql.DB releases the
+		// connection.
+		_ = tx.Rollback(ctx)
+
+		if r != nil {
+			panic(r)
 		}
 	}()
+
 	// Specify that we intend to retry this txn in case of CockroachDB retryable
 	// errors.
 	if err = tx.Exec(ctx, "SAVEPOINT cockroach_restart"); err != nil {
@@ -58,31 +66,35 @@ func ExecuteInTx(ctx context.Context, tx Tx, fn func() error) (err error) {
 	const maxRetries = 50
 	retryCount := 0
 	for {
-		released := false
-		err = fn()
-		if err == nil {
+		releaseFailed := false
+		if err = fn(); err == nil {
 			// RELEASE acts like COMMIT in CockroachDB. We use it since it gives us an
 			// opportunity to react to retryable errors, whereas tx.Commit() doesn't.
-			released = true
 			if err = tx.Exec(ctx, "RELEASE SAVEPOINT cockroach_restart"); err == nil {
 				return nil
 			}
+			releaseFailed = true
 		}
-		// We got an error; let's see if it's a retryable one and, if so, restart.
+
+		// We got an error (i.e. err != nil) if at this point
+
+		// Let's see if it's a retryable one and, if so, restart.
 		if !errIsRetryable(err) {
-			if released {
+			if releaseFailed {
 				err = newAmbiguousCommitError(err)
 			}
 			return err
 		}
 
-		if retryErr := tx.Exec(ctx, "ROLLBACK TO SAVEPOINT cockroach_restart"); retryErr != nil {
-			return newTxnRestartError(retryErr, err)
+		if rollbackErr := tx.Exec(ctx, "ROLLBACK TO SAVEPOINT cockroach_restart"); rollbackErr != nil {
+			err = newTxnRestartError(rollbackErr, err)
+			return err
 		}
 
 		retryCount++
 		if retryCount > maxRetries {
-			return newMaxRetriesExceededError(err, maxRetries)
+			err = newMaxRetriesExceededError(err, maxRetries)
+			return err
 		}
 	}
 }
diff --git a/go.mod b/go.mod
index 7825140..0b11d96 100644
--- a/go.mod
+++ b/go.mod
@@ -3,7 +3,7 @@ module github.com/cockroachdb/cockroach-go/v2
 go 1.13
 
 require (
-	github.com/gofrs/flock v0.8.1 // indirect
+	github.com/gofrs/flock v0.8.1
 	github.com/jackc/pgx/v4 v4.10.1
 	github.com/jmoiron/sqlx v1.3.1
 	github.com/lib/pq v1.10.0

From 142038cf907583268708e94ef484e7e20167b6a0 Mon Sep 17 00:00:00 2001
From: Ashraf Hanafy <ashie1287@gmail.com>
Date: Thu, 9 Dec 2021 12:41:36 -0500
Subject: [PATCH 2/2] Fix comment, undo unnecessary var assignment

---
 crdb/common.go | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/crdb/common.go b/crdb/common.go
index 2857027..8039923 100644
--- a/crdb/common.go
+++ b/crdb/common.go
@@ -67,7 +67,8 @@ func ExecuteInTx(ctx context.Context, tx Tx, fn func() error) (err error) {
 	retryCount := 0
 	for {
 		releaseFailed := false
-		if err = fn(); err == nil {
+		err = fn()
+		if err == nil {
 			// RELEASE acts like COMMIT in CockroachDB. We use it since it gives us an
 			// opportunity to react to retryable errors, whereas tx.Commit() doesn't.
 			if err = tx.Exec(ctx, "RELEASE SAVEPOINT cockroach_restart"); err == nil {
@@ -76,9 +77,7 @@ func ExecuteInTx(ctx context.Context, tx Tx, fn func() error) (err error) {
 			releaseFailed = true
 		}
 
-		// We got an error (i.e. err != nil) if at this point
-
-		// Let's see if it's a retryable one and, if so, restart.
+		// We got an error; let's see if it's a retryable one and, if so, restart.
 		if !errIsRetryable(err) {
 			if releaseFailed {
 				err = newAmbiguousCommitError(err)
@@ -87,14 +86,12 @@ func ExecuteInTx(ctx context.Context, tx Tx, fn func() error) (err error) {
 		}
 
 		if rollbackErr := tx.Exec(ctx, "ROLLBACK TO SAVEPOINT cockroach_restart"); rollbackErr != nil {
-			err = newTxnRestartError(rollbackErr, err)
-			return err
+			return newTxnRestartError(rollbackErr, err)
 		}
 
 		retryCount++
 		if retryCount > maxRetries {
-			err = newMaxRetriesExceededError(err, maxRetries)
-			return err
+			return newMaxRetriesExceededError(err, maxRetries)
 		}
 	}
 }