Merge pull request #1255 from ydb-platform/retry-with-result

* Added `retry.RetryWithResult` helper for retry lambda and return value from lambda

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added experimental `retry.RetryWithResult` helper for retry lambda and return value from lambda
+
 ## v3.70.0
 * Fixed `config.WithDatabase` behaviour with empty database in DSN string
 * Added experimental method `query/Client.Execute` for execute query and read materialized result

retry/retry.go

@@ -256,12 +256,43 @@ func WithPanicCallback(panicCallback func(e interface{})) panicCallbackOption {
 //
 // Warning: if context without deadline or cancellation func was passed, Retry will work infinitely.
 //
+// # If you need to retry your op func on some logic errors - you must return RetryableError() from retryOperation
+//
+// Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
+func Retry(ctx context.Context, op retryOperation, opts ...Option) (finalErr error) {
+	_, err := RetryWithResult[struct{}](ctx, func(ctx context.Context) (*struct{}, error) {
+		err := op(ctx)
+		if err != nil {
+			return nil, xerrors.WithStackTrace(err)
+		}
+
+		return nil, nil //nolint:nilnil
+	}, opts...)
+	if err != nil {
+		return xerrors.WithStackTrace(err)
+	}
+
+	return nil
+}
+
+// RetryWithResult provide the best effort fo retrying operation which will return value or error
+//
+// RetryWithResult implements internal busy loop until one of the following conditions is met:
+//
+// - context was canceled or deadlined
+//
+// - retry operation returned nil as error
+//
+// Warning: if context without deadline or cancellation func was passed, RetryWithResult will work infinitely.
+//
 // If you need to retry your op func on some logic errors - you must return RetryableError() from retryOperation
 //
 //nolint:funlen
-func Retry(ctx context.Context, op retryOperation, opts ...Option) (finalErr error) {
+func RetryWithResult[T any](ctx context.Context, //nolint:revive
+	op func(context.Context) (*T, error), opts ...Option,
+) (v *T, finalErr error) {
 	options := &retryOptions{
-		call:        stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/retry.Retry"),
+		call:        stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/retry.RetryWithResult"),
 		trace:       &trace.Retry{},
 		budget:      budget.Limited(-1),
 		fastBackoff: backoff.Fast,
@@ -301,15 +332,16 @@ func Retry(ctx context.Context, op retryOperation, opts ...Option) (finalErr err
 		attempts++
 		select {
 		case <-ctx.Done():
-			return xerrors.WithStackTrace(
+			return nil, xerrors.WithStackTrace(
 				fmt.Errorf("retry failed on attempt No.%d: %w", attempts, ctx.Err()),
 			)
 
 		default:
-			err := opWithRecover(ctx, options, op)
+			var err error
+			v, err = opWithRecover(ctx, options, op)
 
 			if err == nil {
-				return nil
+				return v, nil
 			}
 
 			m := Check(err)
@@ -321,7 +353,7 @@ func Retry(ctx context.Context, op retryOperation, opts ...Option) (finalErr err
 			code = m.StatusCode()
 
 			if !m.MustRetry(options.idempotent) {
-				return xerrors.WithStackTrace(
+				return nil, xerrors.WithStackTrace(
 					fmt.Errorf("non-retryable error occurred on attempt No.%d (idempotent=%v): %w",
 						attempts, options.idempotent, err),
 				)
@@ -336,7 +368,7 @@ func Retry(ctx context.Context, op retryOperation, opts ...Option) (finalErr err
 			case <-ctx.Done():
 				t.Stop()
 
-				return xerrors.WithStackTrace(
+				return nil, xerrors.WithStackTrace(
 					xerrors.Join(
 						fmt.Errorf("attempt No.%d: %w", attempts, ctx.Err()),
 						err,
@@ -346,7 +378,7 @@ func Retry(ctx context.Context, op retryOperation, opts ...Option) (finalErr err
 				t.Stop()
 
 				if acquireErr := options.budget.Acquire(ctx); acquireErr != nil {
-					return xerrors.WithStackTrace(
+					return nil, xerrors.WithStackTrace(
 						xerrors.Join(
 							fmt.Errorf("attempt No.%d: %w", attempts, budget.ErrNoQuota),
 							acquireErr,
@@ -359,7 +391,9 @@ func Retry(ctx context.Context, op retryOperation, opts ...Option) (finalErr err
 	}
 }
 
-func opWithRecover(ctx context.Context, options *retryOptions, op retryOperation) (err error) {
+func opWithRecover[T any](ctx context.Context,
+	options *retryOptions, op func(context.Context) (*T, error),
+) (_ *T, err error) {
 	if options.panicCallback != nil {
 		defer func() {
 			if e := recover(); e != nil {

retry/retry_test.go

@@ -225,11 +225,11 @@ func TestOpWithRecover_NoPanic(t *testing.T) {
 	options := &retryOptions{
 		panicCallback: nil,
 	}
-	op := func(ctx context.Context) error {
-		return nil
+	op := func(ctx context.Context) (*struct{}, error) {
+		return nil, nil //nolint:nilnil
 	}
 
-	err := opWithRecover(ctx, options, op)
+	_, err := opWithRecover(ctx, options, op)
 
 	require.NoError(t, err)
 }
@@ -240,14 +240,66 @@ func TestOpWithRecover_WithPanic(t *testing.T) {
 	options := &retryOptions{
 		panicCallback: mockCallback.Call,
 	}
-	op := func(ctx context.Context) error {
+	op := func(ctx context.Context) (*struct{}, error) {
 		panic("test panic")
 	}
 
-	err := opWithRecover(ctx, options, op)
+	_, err := opWithRecover(ctx, options, op)
 
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "panic recovered: test panic")
 	require.True(t, mockCallback.called)
 	require.Equal(t, "test panic", mockCallback.received)
 }
+
+func TestRetryWithResult(t *testing.T) {
+	ctx := xtest.Context(t)
+	t.Run("HappyWay", func(t *testing.T) {
+		v, err := RetryWithResult(ctx, func(ctx context.Context) (*int, error) {
+			v := 123
+
+			return &v, nil
+		})
+		require.NoError(t, err)
+		require.NotNil(t, v)
+		require.EqualValues(t, 123, *v)
+	})
+	t.Run("RetryableError", func(t *testing.T) {
+		var counter int
+		v, err := RetryWithResult(ctx, func(ctx context.Context) (*int, error) {
+			counter++
+			if counter < 10 {
+				return nil, RetryableError(errors.New("test"))
+			}
+			v := counter * 123
+
+			return &v, nil
+		})
+		require.NoError(t, err)
+		require.NotNil(t, v)
+		require.EqualValues(t, 1230, *v)
+		require.EqualValues(t, 10, counter)
+	})
+	t.Run("Context", func(t *testing.T) {
+		t.Run("Cancelled", func(t *testing.T) {
+			childCtx, cancel := context.WithCancel(ctx)
+			v, err := RetryWithResult(childCtx, func(ctx context.Context) (*int, error) {
+				cancel()
+
+				return nil, ctx.Err()
+			})
+			require.ErrorIs(t, err, context.Canceled)
+			require.Nil(t, v)
+		})
+		t.Run("DeadlineExceeded", func(t *testing.T) {
+			childCtx, cancel := context.WithTimeout(ctx, 0)
+			v, err := RetryWithResult(childCtx, func(ctx context.Context) (*int, error) {
+				cancel()
+
+				return nil, ctx.Err()
+			})
+			require.ErrorIs(t, err, context.DeadlineExceeded)
+			require.Nil(t, v)
+		})
+	})
+}