Merge pull request #1609 from ydb-platform/retry-panics

processing of panics in retriers
ydb-platform · Jan 16, 2025 · 152f5fc · 152f5fc
2 parents f3860a6 + 9b813be
commit 152f5fc
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 12 deletions.
diff --git a/retry/sql.go b/retry/sql.go
@@ -173,7 +173,7 @@ func DoTx(ctx context.Context, db *sql.DB, op func(context.Context, *sql.Tx) err
 // DoTxWithResult is a retryer of database/sql transactions with fallbacks on errors
 //
 // Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
-func DoTxWithResult[T any](ctx context.Context, db *sql.DB, //nolint:funlen
+func DoTxWithResult[T any](ctx context.Context, db *sql.DB,
 	op func(context.Context, *sql.Tx) (T, error),
 	opts ...doTxOption,
 ) (T, error) {
@@ -211,17 +211,7 @@ func DoTxWithResult[T any](ctx context.Context, db *sql.DB, //nolint:funlen
 			return zeroValue, unwrapErrBadConn(xerrors.WithStackTrace(err))
 		}
 		defer func() {
-			if finalErr == nil {
-				return
-			}
-			errRollback := tx.Rollback()
-			if errRollback == nil {
-				return
-			}
-			finalErr = xerrors.NewWithIssues("",
-				xerrors.WithStackTrace(finalErr),
-				xerrors.WithStackTrace(fmt.Errorf("rollback failed: %w", errRollback)),
-			)
+			_ = tx.Rollback()
 		}()
 		v, err := op(xcontext.MarkRetryCall(ctx), tx)
 		if err != nil {

diff --git a/retry/sql_test.go b/retry/sql_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"database/sql/driver"
+	"errors"
 	"strconv"
 	"testing"
 	"time"
@@ -233,3 +234,61 @@ func TestDoTx(t *testing.T) {
 		})
 	}
 }
+
+func TestCleanUpResourcesOnPanicInRetryOperation(t *testing.T) {
+	panicErr := errors.New("test")
+	t.Run("Do", func(t *testing.T) {
+		m := &mockConnector{
+			t: t,
+		}
+		db := sql.OpenDB(m)
+		defer func() {
+			require.NoError(t, db.Close())
+		}()
+		require.Panics(t, func() {
+			require.Equal(t, 0, db.Stats().OpenConnections)
+			require.Equal(t, 0, db.Stats().Idle)
+			require.Equal(t, 0, db.Stats().InUse)
+			defer func() {
+				require.Equal(t, 1, db.Stats().OpenConnections)
+				require.Equal(t, 1, db.Stats().Idle)
+				require.Equal(t, 0, db.Stats().InUse)
+			}()
+			_ = Do(context.Background(), db,
+				func(ctx context.Context, cc *sql.Conn) error {
+					require.Equal(t, 1, db.Stats().OpenConnections)
+					require.Equal(t, 0, db.Stats().Idle)
+					require.Equal(t, 1, db.Stats().InUse)
+					panic(panicErr)
+				},
+			)
+		})
+	})
+	t.Run("DoTx", func(t *testing.T) {
+		m := &mockConnector{
+			t: t,
+		}
+		db := sql.OpenDB(m)
+		defer func() {
+			require.NoError(t, db.Close())
+		}()
+		require.Panics(t, func() {
+			require.Equal(t, 0, db.Stats().OpenConnections)
+			require.Equal(t, 0, db.Stats().Idle)
+			require.Equal(t, 0, db.Stats().InUse)
+			defer func() {
+				require.Equal(t, 1, db.Stats().OpenConnections)
+				require.Equal(t, 1, db.Stats().Idle)
+				require.Equal(t, 0, db.Stats().InUse)
+			}()
+			_ = DoTx(context.Background(), db,
+				func(ctx context.Context, tx *sql.Tx) error {
+					require.Equal(t, 1, db.Stats().OpenConnections)
+					require.Equal(t, 0, db.Stats().Idle)
+					require.Equal(t, 1, db.Stats().InUse)
+					panic(panicErr)
+				},
+			)
+		})
+	})
+}