Merge pull request #886 from ydb-platform/ctx-retry-call

move retry call detect into xcontext package

internal/table/retry.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/table/config"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table"
@@ -21,21 +22,6 @@ type SessionProvider interface {
 	Put(context.Context, *session) (err error)
 }
 
-type (
-	markRetryCallKey struct{}
-)
-
-func markRetryCall(ctx context.Context) context.Context {
-	return context.WithValue(ctx, markRetryCallKey{}, true)
-}
-
-func isRetryCalledAbove(ctx context.Context) bool {
-	if _, has := ctx.Value(markRetryCallKey{}).(bool); has {
-		return true
-	}
-	return false
-}
-
 func doTx(
 	ctx context.Context,
 	c SessionProvider,
@@ -47,7 +33,7 @@ func doTx(
 		opts.Trace = &trace.Table{}
 	}
 	attempts, onIntermediate := 0, trace.TableOnDoTx(opts.Trace, &ctx,
-		opts.Label, opts.Label, opts.Idempotent, isRetryCalledAbove(ctx),
+		opts.Label, opts.Label, opts.Idempotent, xcontext.IsNestedCall(ctx),
 	)
 	defer func() {
 		onIntermediate(err)(attempts, err)
@@ -87,7 +73,7 @@ func doTx(
 						}
 					}()
 				}
-				return op(ctx, tx)
+				return op(xcontext.MarkRetryCall(ctx), tx)
 			}()
 
 			if err != nil {
@@ -116,7 +102,7 @@ func do(
 		opts.Trace = &trace.Table{}
 	}
 	attempts, onIntermediate := 0, trace.TableOnDo(opts.Trace, &ctx,
-		opts.Label, opts.Label, opts.Idempotent, isRetryCalledAbove(ctx),
+		opts.Label, opts.Label, opts.Idempotent, xcontext.IsNestedCall(ctx),
 	)
 	defer func() {
 		onIntermediate(err)(attempts, err)
@@ -137,7 +123,7 @@ func do(
 						}
 					}()
 				}
-				return op(ctx, s)
+				return op(xcontext.MarkRetryCall(ctx), s)
 			}()
 
 			if err != nil {
@@ -156,7 +142,7 @@ func retryBackoff(
 	op table.Operation,
 	opts ...retry.Option,
 ) error {
-	return retry.Retry(markRetryCall(ctx),
+	return retry.Retry(ctx,
 		func(ctx context.Context) (err error) {
 			var s *session
 
@@ -169,8 +155,7 @@ func retryBackoff(
 				_ = p.Put(ctx, s)
 			}()
 
-			err = op(ctx, s)
-			if err != nil {
+			if err = op(ctx, s); err != nil {
 				s.checkError(err)
 				return xerrors.WithStackTrace(err)
 			}

internal/xcontext/retry_call.go

@@ -0,0 +1,18 @@
+package xcontext
+
+import "context"
+
+type (
+	markRetryCallKey struct{}
+)
+
+func MarkRetryCall(ctx context.Context) context.Context {
+	return context.WithValue(ctx, markRetryCallKey{}, true)
+}
+
+func IsNestedCall(ctx context.Context) bool {
+	if _, has := ctx.Value(markRetryCallKey{}).(bool); has {
+		return true
+	}
+	return false
+}

retry/retry.go

@@ -191,21 +191,6 @@ func WithPanicCallback(panicCallback func(e interface{})) panicCallbackOption {
 	return panicCallbackOption{callback: panicCallback}
 }
 
-type (
-	markRetryCallKey struct{}
-)
-
-func markRetryCall(ctx context.Context) context.Context {
-	return context.WithValue(ctx, markRetryCallKey{}, true)
-}
-
-func isRetryCalledAbove(ctx context.Context) bool {
-	if _, has := ctx.Value(markRetryCallKey{}).(bool); has {
-		return true
-	}
-	return false
-}
-
 // Retry provide the best effort fo retrying operation
 //
 // Retry implements internal busy loop until one of the following conditions is met:
@@ -228,7 +213,9 @@ func Retry(ctx context.Context, op retryOperation, opts ...Option) (finalErr err
 			opt.ApplyRetryOption(options)
 		}
 	}
-	ctx = xcontext.WithIdempotent(ctx, options.idempotent)
+	if options.idempotent {
+		ctx = xcontext.WithIdempotent(ctx, options.idempotent)
+	}
 	defer func() {
 		if finalErr != nil && options.stackTrace {
 			finalErr = xerrors.WithStackTrace(finalErr,
@@ -242,7 +229,7 @@ func Retry(ctx context.Context, op retryOperation, opts ...Option) (finalErr err
 
 		code           = int64(0)
 		onIntermediate = trace.RetryOnRetry(options.trace, &ctx,
-			options.label, options.label, options.idempotent, isRetryCalledAbove(ctx),
+			options.label, options.label, options.idempotent, xcontext.IsNestedCall(ctx),
 		)
 	)
 	defer func() {
@@ -271,7 +258,7 @@ func Retry(ctx context.Context, op retryOperation, opts ...Option) (finalErr err
 						}
 					}()
 				}
-				return op(markRetryCall(ctx))
+				return op(ctx)
 			}()
 
 			if err == nil {

retry/sql.go

@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"fmt"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
@@ -36,7 +37,7 @@ func WithDoRetryOptions(opts ...Option) doRetryOptionsOption {
 }
 
 // Do is a retryer of database/sql Conn with fallbacks on errors
-func Do(ctx context.Context, db *sql.DB, f func(ctx context.Context, cc *sql.Conn) error, opts ...doOption) error {
+func Do(ctx context.Context, db *sql.DB, op func(ctx context.Context, cc *sql.Conn) error, opts ...doOption) error {
 	var (
 		options  = doOptions{}
 		attempts = 0
@@ -62,7 +63,7 @@ func Do(ctx context.Context, db *sql.DB, f func(ctx context.Context, cc *sql.Con
 		defer func() {
 			_ = cc.Close()
 		}()
-		if err = f(ctx, cc); err != nil {
+		if err = op(xcontext.MarkRetryCall(ctx), cc); err != nil {
 			return unwrapErrBadConn(xerrors.WithStackTrace(err))
 		}
 		return nil
@@ -117,7 +118,7 @@ func WithTxOptions(txOptions *sql.TxOptions) txOptionsOption {
 }
 
 // DoTx is a retryer of database/sql transactions with fallbacks on errors
-func DoTx(ctx context.Context, db *sql.DB, f func(context.Context, *sql.Tx) error, opts ...doTxOption) error {
+func DoTx(ctx context.Context, db *sql.DB, op func(context.Context, *sql.Tx) error, opts ...doTxOption) error {
 	var (
 		options = doTxOptions{
 			retryOptions: []Option{},
@@ -159,7 +160,7 @@ func DoTx(ctx context.Context, db *sql.DB, f func(context.Context, *sql.Tx) erro
 				xerrors.WithStackTrace(fmt.Errorf("rollback failed: %w", errRollback)),
 			)
 		}()
-		if err = f(ctx, tx); err != nil {
+		if err = op(xcontext.MarkRetryCall(ctx), tx); err != nil {
 			return unwrapErrBadConn(xerrors.WithStackTrace(err))
 		}
 		if err = tx.Commit(); err != nil {

trace/retry.go

@@ -26,7 +26,8 @@ type (
 
 		Label      string
 		Idempotent bool
-		NestedCall bool // flag when Retry called inside head Retry
+
+		NestedCall bool // a sign for detect Retry calls inside head Retry
 	}
 	RetryLoopIntermediateInfo struct {
 		Error error