* Added ydb.WithTraceRetry option

CHANGELOG.md

@@ -1,3 +1,4 @@
+* Added `ydb.WithTraceRetry` option
 * Added `internal/credentials.IsAccessError(err)` helper for check access errors
 * Changed period for re-fresh static credentials token from `1/2` to `1/10` to expiration time
 

config/config.go

@@ -155,6 +155,12 @@ func WithTrace(t trace.Driver, opts ...trace.DriverComposeOption) Option { //nol
 	}
 }
 
+func WithTraceRetry(t *trace.Retry, opts ...trace.RetryComposeOption) Option {
+	return func(c *Config) {
+		config.SetTraceRetry(&c.Common, t, opts...)
+	}
+}
+
 func WithUserAgent(userAgent string) Option {
 	return func(c *Config) {
 		c.metaOptions = append(c.metaOptions, meta.WithUserAgentOption(userAgent))

connection.go

@@ -427,6 +427,7 @@ func newConnectionFromOptions(ctx context.Context, opts ...Option) (_ *Driver, e
 			WithTraceDiscovery(log.Discovery(d.logger, d.loggerDetails, d.loggerOpts...)),
 			WithTraceTopic(log.Topic(d.logger, d.loggerDetails, d.loggerOpts...)),
 			WithTraceDatabaseSQL(log.DatabaseSQL(d.logger, d.loggerDetails, d.loggerOpts...)),
+			WithTraceRetry(log.Retry(d.logger, d.loggerDetails, d.loggerOpts...)),
 		} {
 			if opt != nil {
 				err = opt(ctx, d)

internal/config/config.go

@@ -2,12 +2,15 @@ package config
 
 import (
 	"time"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
 type Common struct {
 	operationTimeout     time.Duration
 	operationCancelAfter time.Duration
 	disableAutoRetry     bool
+	traceRetry           trace.Retry
 
 	panicCallback func(e interface{})
 }
@@ -41,6 +44,10 @@ func (c *Common) OperationCancelAfter() time.Duration {
 	return c.operationCancelAfter
 }
 
+func (c *Common) TraceRetry() *trace.Retry {
+	return &c.traceRetry
+}
+
 // SetOperationTimeout define the maximum amount of time a YDB server will process
 // an operation. After timeout exceeds YDB will try to cancel operation and
 // regardless of the cancellation appropriate error will be returned to
@@ -70,3 +77,7 @@ func SetPanicCallback(c *Common, panicCallback func(e interface{})) {
 func SetAutoRetry(c *Common, autoRetry bool) {
 	c.disableAutoRetry = !autoRetry
 }
+
+func SetTraceRetry(c *Common, t *trace.Retry, opts ...trace.RetryComposeOption) {
+	c.traceRetry = *c.traceRetry.Compose(t, opts...)
+}

internal/xcontext/retry.go

@@ -0,0 +1,22 @@
+package xcontext
+
+import (
+	"context"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
+)
+
+type (
+	ctxTraceRetryKey struct{}
+)
+
+func WithTraceRetry(ctx context.Context, t *trace.Retry) context.Context {
+	return context.WithValue(ctx, ctxTraceRetryKey{}, TraceRetry(ctx).Compose(t))
+}
+
+func TraceRetry(ctx context.Context) *trace.Retry {
+	if traceRetry, ok := ctx.Value(ctxTraceRetryKey{}).(*trace.Retry); ok {
+		return traceRetry
+	}
+	return &trace.Retry{}
+}

metrics/retry.go

@@ -4,7 +4,7 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
-// Retry makes table.RetryTrace with New publishing
-func Retry(config Config) (t trace.Retry) {
+// retry makes table.RetryTrace with New publishing
+func retry(config Config) (t trace.Retry) {
 	return t
 }

metrics/traces.go

@@ -18,5 +18,6 @@ func WithTraces(config Config) ydb.Option {
 		ydb.WithTraceRatelimiter(ratelimiter(config)),
 		ydb.WithTraceDiscovery(discovery(config)),
 		ydb.WithTraceDatabaseSQL(databaseSQL(config)),
+		ydb.WithTraceRetry(retry(config)),
 	)
 }

options.go

@@ -249,14 +249,29 @@ func WithDiscoveryInterval(discoveryInterval time.Duration) Option {
 	}
 }
 
-// WithTraceDriver returns deadline which has associated Driver with it.
+// WithTraceDriver appends trace.Driver into driver traces
 func WithTraceDriver(trace trace.Driver, opts ...trace.DriverComposeOption) Option { //nolint:gocritic
 	return func(ctx context.Context, c *Driver) error {
 		c.options = append(c.options, config.WithTrace(trace, opts...))
 		return nil
 	}
 }
 
+// WithTraceRetry appends trace.Retry into retry traces
+func WithTraceRetry(t trace.Retry, opts ...trace.RetryComposeOption) Option {
+	return func(ctx context.Context, c *Driver) error {
+		c.options = append(c.options,
+			config.WithTraceRetry(&t, append(
+				[]trace.RetryComposeOption{
+					trace.WithRetryPanicCallback(c.panicCallback),
+				},
+				opts...,
+			)...),
+		)
+		return nil
+	}
+}
+
 // WithCertificate appends certificate to TLS config root certificates
 func WithCertificate(cert *x509.Certificate) Option {
 	return func(ctx context.Context, c *Driver) error {
@@ -394,7 +409,7 @@ func WithPanicCallback(panicCallback func(e interface{})) Option {
 	}
 }
 
-// WithTraceTable returns table trace option
+// WithTraceTable appends trace.Table into table traces
 func WithTraceTable(t trace.Table, opts ...trace.TableComposeOption) Option { //nolint:gocritic
 	return func(ctx context.Context, c *Driver) error {
 		c.tableOptions = append(

retry/retry.go

@@ -221,7 +221,7 @@ func Retry(ctx context.Context, op retryOperation, opts ...Option) (err error) {
 	options := &retryOptions{
 		fastBackoff: backoff.Fast,
 		slowBackoff: backoff.Slow,
-		trace:       &trace.Retry{},
+		trace:       xcontext.TraceRetry(ctx),
 	}
 	for _, opt := range opts {
 		if opt != nil {

retry/retry_test.go

@@ -121,7 +121,7 @@ func TestRetryWithCustomErrors(t *testing.T) {
 	} {
 		t.Run(tt.error.Error(), func(t *testing.T) {
 			i := 0
-			err := Retry(ctx, func(ctx context.Context) (err error) {
+			err := Retry(ctx, func(ctx context.Context) error {
 				i++
 				if i < limit {
 					return tt.error
@@ -149,7 +149,7 @@ func TestRetryTransportDeadlineExceeded(t *testing.T) {
 	} {
 		counter := 0
 		ctx, cancel := xcontext.WithTimeout(context.Background(), time.Hour)
-		err := Retry(ctx, func(ctx context.Context) (err error) {
+		err := Retry(ctx, func(ctx context.Context) error {
 			counter++
 			if !(counter < cancelCounterValue) {
 				cancel()
@@ -169,7 +169,7 @@ func TestRetryTransportCancelled(t *testing.T) {
 	} {
 		counter := 0
 		ctx, cancel := xcontext.WithCancel(context.Background())
-		err := Retry(ctx, func(ctx context.Context) (err error) {
+		err := Retry(ctx, func(ctx context.Context) error {
 			counter++
 			if !(counter < cancelCounterValue) {
 				cancel()

retry/sql.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
+	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
 type doOptions struct {
@@ -40,6 +41,13 @@ func Do(ctx context.Context, db *sql.DB, f func(ctx context.Context, cc *sql.Con
 		options  = doOptions{}
 		attempts = 0
 	)
+	if tracer, has := db.Driver().(interface {
+		TraceRetry() *trace.Retry
+	}); has {
+		options.retryOptions = append(options.retryOptions, WithTrace(
+			*tracer.TraceRetry()),
+		)
+	}
 	for _, opt := range opts {
 		if opt != nil {
 			opt.ApplyDoOption(&options)
@@ -119,6 +127,13 @@ func DoTx(ctx context.Context, db *sql.DB, f func(context.Context, *sql.Tx) erro
 		}
 		attempts = 0
 	)
+	if tracer, has := db.Driver().(interface {
+		TraceRetry() *trace.Retry
+	}); has {
+		options.retryOptions = append(options.retryOptions,
+			WithTrace(*tracer.TraceRetry()),
+		)
+	}
 	for _, opt := range opts {
 		if opt != nil {
 			opt.ApplyDoTxOption(&options)