diff --git a/CHANGELOG.md b/CHANGELOG.md index feacaa99e..cda75a2b0 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/config/config.go b/config/config.go index 20dc44bb6..5cf9a0a09 100644 --- a/config/config.go +++ b/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)) diff --git a/connection.go b/connection.go index fd2822a1f..28b81788d 100644 --- a/connection.go +++ b/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) diff --git a/internal/config/config.go b/internal/config/config.go index f7a094d7e..dcf89c4b7 100644 --- a/internal/config/config.go +++ b/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...) +} diff --git a/internal/xcontext/retry.go b/internal/xcontext/retry.go new file mode 100644 index 000000000..912aa1f13 --- /dev/null +++ b/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{} +} diff --git a/metrics/retry.go b/metrics/retry.go index a0eaca54f..8d3a059ff 100644 --- a/metrics/retry.go +++ b/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 } diff --git a/metrics/traces.go b/metrics/traces.go index a601b5bcc..5fc978606 100644 --- a/metrics/traces.go +++ b/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)), ) } diff --git a/options.go b/options.go index 1f92d4ce0..80fba7603 100644 --- a/options.go +++ b/options.go @@ -249,7 +249,7 @@ 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...)) @@ -257,6 +257,21 @@ func WithTraceDriver(trace trace.Driver, opts ...trace.DriverComposeOption) Opti } } +// 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( diff --git a/retry/retry.go b/retry/retry.go index 2c2589a07..87ea3d5bd 100644 --- a/retry/retry.go +++ b/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 { diff --git a/retry/retry_test.go b/retry/retry_test.go index f467b3dcc..bb9818775 100644 --- a/retry/retry_test.go +++ b/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() diff --git a/retry/sql.go b/retry/sql.go index f6e6dd3c9..b248c12bc 100644 --- a/retry/sql.go +++ b/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)