From 5d1304845307085b14ee92d74648a35e7ae1e3ea Mon Sep 17 00:00:00 2001
From: Aleksey Myasnikov <asmyasnikov@ydb.tech>
Date: Thu, 16 Jan 2025 15:52:32 +0300
Subject: [PATCH] processing of panics in retriers

---
 retry/sql.go      | 12 +---------
 retry/sql_test.go | 59 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/retry/sql.go b/retry/sql.go
index 29ddeb6b2..9b6aab209 100644
--- a/retry/sql.go
+++ b/retry/sql.go
@@ -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
index 4f42d06dd..432e4f167 100644
--- 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)
+				},
+			)
+		})
+	})
+}