refactoring of retry budget sources

ydb-platform · Apr 23, 2024 · 862d8f4 · 862d8f4
1 parent 2a635d1
commit 862d8f4
Show file tree

Hide file tree

Showing 25 changed files with 112 additions and 128 deletions.
diff --git a/config/config.go b/config/config.go
@@ -12,7 +12,7 @@ import (
 	balancerConfig "github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer/config"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/config"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/retry"
+	"github.com/ydb-platform/ydb-go-sdk/v3/retry/budget"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
@@ -159,9 +159,9 @@ func WithTrace(t trace.Driver, opts ...trace.DriverComposeOption) Option { //nol
 }
 
 // Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
-func WithRetryLimiter(l retry.Budget) Option {
+func WithRetryBudget(b budget.Budget) Option {
 	return func(c *Config) {
-		config.SetRetryLimiter(&c.Common, l)
+		config.SetRetryBudget(&c.Common, b)
 	}
 }
 

diff --git a/internal/balancer/balancer.go b/internal/balancer/balancer.go
@@ -89,7 +89,7 @@ func (b *Balancer) clusterDiscovery(ctx context.Context) (err error) {
 		},
 		retry.WithIdempotent(true),
 		retry.WithTrace(b.driverConfig.TraceRetry()),
-		retry.WithBudget(b.driverConfig.RetryLimiter()),
+		retry.WithBudget(b.driverConfig.RetryBudget()),
 	)
 }
 

diff --git a/internal/config/config.go b/internal/config/config.go
@@ -3,19 +3,18 @@ package config
 import (
 	"time"
 
-	retry2 "github.com/ydb-platform/ydb-go-sdk/v3/internal/retry"
-	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
+	"github.com/ydb-platform/ydb-go-sdk/v3/retry/budget"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
-var defaultRetryLimiter = retry.Budget(-1)
+var defaultRetryBudget = budget.New(-1)
 
 type Common struct {
 	operationTimeout     time.Duration
 	operationCancelAfter time.Duration
 	disableAutoRetry     bool
 	traceRetry           trace.Retry
-	retryLimiter         retry2.Budget
+	retryBudget          budget.Budget
 
 	panicCallback func(e interface{})
 }
@@ -53,12 +52,12 @@ func (c *Common) TraceRetry() *trace.Retry {
 	return &c.traceRetry
 }
 
-func (c *Common) RetryLimiter() retry2.Budget {
-	if c.retryLimiter == nil {
-		return defaultRetryLimiter
+func (c *Common) RetryBudget() budget.Budget {
+	if c.retryBudget == nil {
+		return defaultRetryBudget
 	}
 
-	return c.retryLimiter
+	return c.retryBudget
 }
 
 // SetOperationTimeout define the maximum amount of time a YDB server will process
@@ -95,6 +94,6 @@ func SetTraceRetry(c *Common, t *trace.Retry, opts ...trace.RetryComposeOption)
 	c.traceRetry = *c.traceRetry.Compose(t, opts...)
 }
 
-func SetRetryLimiter(c *Common, l retry2.Budget) {
-	c.retryLimiter = l
+func SetRetryBudget(c *Common, b budget.Budget) {
+	c.retryBudget = b
 }
diff --git a/internal/coordination/client.go b/internal/coordination/client.go
@@ -117,7 +117,7 @@ func (c *Client) CreateNode(ctx context.Context, path string, config coordinatio
 		retry.WithStackTrace(),
 		retry.WithIdempotent(true),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 }
 
@@ -150,7 +150,7 @@ func (c *Client) AlterNode(ctx context.Context, path string, config coordination
 		retry.WithStackTrace(),
 		retry.WithIdempotent(true),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 }
 
@@ -211,7 +211,7 @@ func (c *Client) DropNode(ctx context.Context, path string) (finalErr error) {
 		retry.WithStackTrace(),
 		retry.WithIdempotent(true),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 }
 
@@ -271,7 +271,7 @@ func (c *Client) DescribeNode(
 		retry.WithStackTrace(),
 		retry.WithIdempotent(true),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 	if err != nil {
 		return nil, nil, xerrors.WithStackTrace(err)

diff --git a/internal/query/options/retry.go b/internal/query/options/retry.go
@@ -2,8 +2,8 @@ package options
 
 import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/tx"
-	retry2 "github.com/ydb-platform/ydb-go-sdk/v3/internal/retry"
 	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
+	"github.com/ydb-platform/ydb-go-sdk/v3/retry/budget"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
@@ -96,8 +96,8 @@ func WithTrace(t *trace.Query) traceOption {
 	return traceOption{t: t}
 }
 
-func WithBudget(l retry2.Budget) retryOptionsOption {
-	return []retry.Option{retry.WithBudget(l)}
+func WithRetryBudget(b budget.Budget) retryOptionsOption {
+	return []retry.Option{retry.WithBudget(b)}
 }
 
 func ParseDoOpts(t *trace.Query, opts ...DoOption) (s *doSettings) {

diff --git a/internal/ratelimiter/client.go b/internal/ratelimiter/client.go
@@ -63,7 +63,7 @@ func (c *Client) CreateResource(
 		retry.WithStackTrace(),
 		retry.WithIdempotent(true),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 }
 
@@ -113,7 +113,7 @@ func (c *Client) AlterResource(
 		retry.WithStackTrace(),
 		retry.WithIdempotent(true),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 }
 
@@ -163,7 +163,7 @@ func (c *Client) DropResource(
 		retry.WithStackTrace(),
 		retry.WithIdempotent(true),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 }
 
@@ -209,7 +209,7 @@ func (c *Client) ListResource(
 		retry.WithIdempotent(true),
 		retry.WithStackTrace(),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 
 	return list, err
@@ -269,7 +269,7 @@ func (c *Client) DescribeResource(
 		retry.WithIdempotent(true),
 		retry.WithStackTrace(),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 
 	return
@@ -338,7 +338,7 @@ func (c *Client) AcquireResource(
 	return retry.Retry(ctx, call,
 		retry.WithStackTrace(),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 }
 

diff --git a/internal/scheme/client.go b/internal/scheme/client.go
@@ -64,7 +64,7 @@ func (c *Client) MakeDirectory(ctx context.Context, path string) (finalErr error
 		retry.WithStackTrace(),
 		retry.WithIdempotent(true),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 }
 
@@ -104,7 +104,7 @@ func (c *Client) RemoveDirectory(ctx context.Context, path string) (finalErr err
 		retry.WithStackTrace(),
 		retry.WithIdempotent(true),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 }
 
@@ -146,7 +146,7 @@ func (c *Client) ListDirectory(ctx context.Context, path string) (d scheme.Direc
 		retry.WithIdempotent(true),
 		retry.WithStackTrace(),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 
 	return d, xerrors.WithStackTrace(err)
@@ -210,7 +210,7 @@ func (c *Client) DescribePath(ctx context.Context, path string) (e scheme.Entry,
 		retry.WithIdempotent(true),
 		retry.WithStackTrace(),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 
 	return e, xerrors.WithStackTrace(err)
@@ -272,7 +272,7 @@ func (c *Client) ModifyPermissions(
 		retry.WithStackTrace(),
 		retry.WithIdempotent(true),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 }
 

diff --git a/internal/scripting/client.go b/internal/scripting/client.go
@@ -60,7 +60,7 @@ func (c *Client) Execute(
 	err = retry.Retry(ctx, call,
 		retry.WithStackTrace(),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 
 	return r, xerrors.WithStackTrace(err)
@@ -140,7 +140,7 @@ func (c *Client) Explain(
 		retry.WithStackTrace(),
 		retry.WithIdempotent(true),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 
 	return e, xerrors.WithStackTrace(err)
@@ -215,7 +215,7 @@ func (c *Client) StreamExecute(
 	err = retry.Retry(ctx, call,
 		retry.WithStackTrace(),
 		retry.WithTrace(c.config.TraceRetry()),
-		retry.WithBudget(c.config.RetryLimiter()),
+		retry.WithBudget(c.config.RetryBudget()),
 	)
 
 	return r, xerrors.WithStackTrace(err)

diff --git a/internal/table/retry.go b/internal/table/retry.go
@@ -99,7 +99,7 @@ func (c *Client) retryOptions(opts ...table.Option) *table.Options {
 		),
 		RetryOptions: []retry.Option{
 			retry.WithTrace(c.config.TraceRetry()),
-			retry.WithBudget(c.config.RetryLimiter()),
+			retry.WithBudget(c.config.RetryBudget()),
 		},
 	}
 	for _, opt := range opts {

diff --git a/internal/topic/topicclientinternal/client.go b/internal/topic/topicclientinternal/client.go
@@ -87,7 +87,7 @@ func (c *Client) Alter(ctx context.Context, path string, opts ...topicoptions.Al
 		return retry.Retry(ctx, call,
 			retry.WithIdempotent(true),
 			retry.WithTrace(c.cfg.TraceRetry()),
-			retry.WithBudget(c.cfg.RetryLimiter()),
+			retry.WithBudget(c.cfg.RetryBudget()),
 		)
 	}
 
@@ -120,7 +120,7 @@ func (c *Client) Create(
 		return retry.Retry(ctx, call,
 			retry.WithIdempotent(true),
 			retry.WithTrace(c.cfg.TraceRetry()),
-			retry.WithBudget(c.cfg.RetryLimiter()),
+			retry.WithBudget(c.cfg.RetryBudget()),
 		)
 	}
 
@@ -158,7 +158,7 @@ func (c *Client) Describe(
 		err = retry.Retry(ctx, call,
 			retry.WithIdempotent(true),
 			retry.WithTrace(c.cfg.TraceRetry()),
-			retry.WithBudget(c.cfg.RetryLimiter()),
+			retry.WithBudget(c.cfg.RetryBudget()),
 		)
 	} else {
 		err = call(ctx)
@@ -195,7 +195,7 @@ func (c *Client) Drop(ctx context.Context, path string, opts ...topicoptions.Dro
 		return retry.Retry(ctx, call,
 			retry.WithIdempotent(true),
 			retry.WithTrace(c.cfg.TraceRetry()),
-			retry.WithBudget(c.cfg.RetryLimiter()),
+			retry.WithBudget(c.cfg.RetryBudget()),
 		)
 	}
 

diff --git a/internal/xsql/conn.go b/internal/xsql/conn.go
@@ -831,7 +831,7 @@ func (c *conn) retryIdempotent(ctx context.Context, f func(ctx context.Context)
 	err := retry.Retry(ctx, f,
 		retry.WithIdempotent(true),
 		retry.WithTrace(c.connector.traceRetry),
-		retry.WithBudget(c.connector.retryLimiter),
+		retry.WithBudget(c.connector.retryBudget),
 	)
 	if err != nil {
 		return xerrors.WithStackTrace(err)

diff --git a/internal/xsql/connector.go b/internal/xsql/connector.go
@@ -11,11 +11,11 @@ import (
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/bind"
 	metaHeaders "github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/retry"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	"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/meta"
+	"github.com/ydb-platform/ydb-go-sdk/v3/retry/budget"
 	"github.com/ydb-platform/ydb-go-sdk/v3/scheme"
 	"github.com/ydb-platform/ydb-go-sdk/v3/scripting"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table"
@@ -176,18 +176,18 @@ func WithTraceRetry(t *trace.Retry) ConnectorOption {
 	return traceRetryConnectorOption{t: t}
 }
 
-type retryLimiterConnectorOption struct {
-	l retry.Budget
+type retryBudgetConnectorOption struct {
+	b budget.Budget
 }
 
-func (l retryLimiterConnectorOption) Apply(c *Connector) error {
-	c.retryLimiter = l.l
+func (l retryBudgetConnectorOption) Apply(c *Connector) error {
+	c.retryBudget = l.b
 
 	return nil
 }
 
-func WithRetryLimiter(l retry.Budget) ConnectorOption {
-	return retryLimiterConnectorOption{l: l}
+func WithretryBudget(b budget.Budget) ConnectorOption {
+	return retryBudgetConnectorOption{b: b}
 }
 
 type fakeTxConnectorOption QueryMode
@@ -263,9 +263,9 @@ type Connector struct {
 	disableServerBalancer bool
 	idleThreshold         time.Duration
 
-	trace        *trace.DatabaseSQL
-	traceRetry   *trace.Retry
-	retryLimiter retry.Budget
+	trace       *trace.DatabaseSQL
+	traceRetry  *trace.Retry
+	retryBudget budget.Budget
 }
 
 var (
@@ -369,8 +369,8 @@ func (d *driverWrapper) TraceRetry() *trace.Retry {
 	return d.c.traceRetry
 }
 
-func (d *driverWrapper) RetryLimiter() retry.Budget {
-	return d.c.retryLimiter
+func (d *driverWrapper) RetryBudget() budget.Budget {
+	return d.c.retryBudget
 }
 
 func (d *driverWrapper) Open(_ string) (driver.Conn, error) {