Merge pull request #1184 from ydb-platform/retry-quoter

added experimental `retry/budget` package

CHANGELOG.md

@@ -1,3 +1,4 @@
+* Added experimental package `retry/budget` for limit second and subsequent retry attempts 
 * Refactored internals for enabling `containedctx` linter
 * Fixed the hanging semaphore issue on coordination session reconnect
 

config/config.go

@@ -12,6 +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/retry/budget"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
@@ -157,6 +158,13 @@ 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 WithRetryBudget(b budget.Budget) Option {
+	return func(c *Config) {
+		config.SetRetryBudget(&c.Common, b)
+	}
+}
+
 func WithTraceRetry(t *trace.Retry, opts ...trace.RetryComposeOption) Option {
 	return func(c *Config) {
 		config.SetTraceRetry(&c.Common, t, opts...)

internal/backoff/delay.go

@@ -0,0 +1,43 @@
+package backoff
+
+import (
+	"time"
+)
+
+type (
+	delayOptions struct {
+		fast Backoff
+		slow Backoff
+	}
+	delayOption func(o *delayOptions)
+)
+
+func WithFastBackoff(fast Backoff) delayOption {
+	return func(o *delayOptions) {
+		o.fast = fast
+	}
+}
+
+func WithSlowBackoff(slow Backoff) delayOption {
+	return func(o *delayOptions) {
+		o.slow = slow
+	}
+}
+
+func Delay(t Type, i int, opts ...delayOption) time.Duration {
+	optionsHolder := delayOptions{
+		fast: Fast,
+		slow: Slow,
+	}
+	for _, opt := range opts {
+		opt(&optionsHolder)
+	}
+	switch t {
+	case TypeFast:
+		return optionsHolder.fast.Delay(i)
+	case TypeSlow:
+		return optionsHolder.slow.Delay(i)
+	default:
+		return 0
+	}
+}

internal/backoff/delay_test.go

@@ -0,0 +1,92 @@
+package backoff
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
+)
+
+func TestDelay(t *testing.T) {
+	for _, tt := range []struct {
+		name string
+		act  time.Duration
+		exp  time.Duration
+	}{
+		{
+			name: xtest.CurrentFileLine(),
+			act:  Delay(TypeNoBackoff, 0),
+			exp:  0,
+		},
+		{
+			name: xtest.CurrentFileLine(),
+			act:  Delay(TypeNoBackoff, 1),
+			exp:  0,
+		},
+		{
+			name: xtest.CurrentFileLine(),
+			act:  Delay(TypeNoBackoff, 2),
+			exp:  0,
+		},
+		{
+			name: xtest.CurrentFileLine(),
+			act: Delay(TypeFast, 0, WithFastBackoff(New(
+				WithSlotDuration(fastSlot),
+				WithCeiling(6),
+				WithJitterLimit(1),
+			))),
+			exp: 5 * time.Millisecond,
+		},
+		{
+			name: xtest.CurrentFileLine(),
+			act: Delay(TypeFast, 1, WithFastBackoff(New(
+				WithSlotDuration(fastSlot),
+				WithCeiling(6),
+				WithJitterLimit(1),
+			))),
+			exp: 10 * time.Millisecond,
+		},
+		{
+			name: xtest.CurrentFileLine(),
+			act: Delay(TypeFast, 3, WithFastBackoff(New(
+				WithSlotDuration(fastSlot),
+				WithCeiling(6),
+				WithJitterLimit(1),
+			))),
+			exp: 40 * time.Millisecond,
+		},
+		{
+			name: xtest.CurrentFileLine(),
+			act: Delay(TypeSlow, 0, WithSlowBackoff(New(
+				WithSlotDuration(slowSlot),
+				WithCeiling(6),
+				WithJitterLimit(1),
+			))),
+			exp: time.Second,
+		},
+		{
+			name: xtest.CurrentFileLine(),
+			act: Delay(TypeSlow, 1, WithSlowBackoff(New(
+				WithSlotDuration(slowSlot),
+				WithCeiling(6),
+				WithJitterLimit(1),
+			))),
+			exp: 2 * time.Second,
+		},
+		{
+			name: xtest.CurrentFileLine(),
+			act: Delay(TypeSlow, 3, WithSlowBackoff(New(
+				WithSlotDuration(slowSlot),
+				WithCeiling(6),
+				WithJitterLimit(1),
+			))),
+			exp: 8 * time.Second,
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			require.Equal(t, tt.exp, tt.act)
+		})
+	}
+}

internal/balancer/balancer.go

@@ -89,6 +89,7 @@ func (b *Balancer) clusterDiscovery(ctx context.Context) (err error) {
 		},
 		retry.WithIdempotent(true),
 		retry.WithTrace(b.driverConfig.TraceRetry()),
+		retry.WithBudget(b.driverConfig.RetryBudget()),
 	)
 }
 

internal/config/config.go

@@ -3,14 +3,18 @@ package config
 import (
 	"time"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3/retry/budget"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
+var defaultRetryBudget = budget.Limited(-1)
+
 type Common struct {
 	operationTimeout     time.Duration
 	operationCancelAfter time.Duration
 	disableAutoRetry     bool
 	traceRetry           trace.Retry
+	retryBudget          budget.Budget
 
 	panicCallback func(e interface{})
 }
@@ -48,6 +52,14 @@ func (c *Common) TraceRetry() *trace.Retry {
 	return &c.traceRetry
 }
 
+func (c *Common) RetryBudget() budget.Budget {
+	if c.retryBudget == nil {
+		return defaultRetryBudget
+	}
+
+	return c.retryBudget
+}
+
 // 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
@@ -81,3 +93,7 @@ func SetAutoRetry(c *Common, autoRetry bool) {
 func SetTraceRetry(c *Common, t *trace.Retry, opts ...trace.RetryComposeOption) {
 	c.traceRetry = *c.traceRetry.Compose(t, opts...)
 }
+
+func SetRetryBudget(c *Common, b budget.Budget) {
+	c.retryBudget = b
+}

internal/coordination/client.go

@@ -110,9 +110,15 @@ func (c *Client) CreateNode(ctx context.Context, path string, config coordinatio
 		return createNode(ctx, c.client, request)
 	}
 
-	return retry.Retry(ctx, func(ctx context.Context) error {
-		return createNode(ctx, c.client, request)
-	}, retry.WithStackTrace(), retry.WithIdempotent(true), retry.WithTrace(c.config.TraceRetry()))
+	return retry.Retry(ctx,
+		func(ctx context.Context) error {
+			return createNode(ctx, c.client, request)
+		},
+		retry.WithStackTrace(),
+		retry.WithIdempotent(true),
+		retry.WithTrace(c.config.TraceRetry()),
+		retry.WithBudget(c.config.RetryBudget()),
+	)
 }
 
 func (c *Client) AlterNode(ctx context.Context, path string, config coordination.NodeConfig) (finalErr error) {
@@ -137,9 +143,15 @@ func (c *Client) AlterNode(ctx context.Context, path string, config coordination
 		return xerrors.WithStackTrace(call(ctx))
 	}
 
-	return retry.Retry(ctx, func(ctx context.Context) (err error) {
-		return alterNode(ctx, c.client, request)
-	}, retry.WithStackTrace(), retry.WithIdempotent(true), retry.WithTrace(c.config.TraceRetry()))
+	return retry.Retry(ctx,
+		func(ctx context.Context) (err error) {
+			return alterNode(ctx, c.client, request)
+		},
+		retry.WithStackTrace(),
+		retry.WithIdempotent(true),
+		retry.WithTrace(c.config.TraceRetry()),
+		retry.WithBudget(c.config.RetryBudget()),
+	)
 }
 
 func alterNodeRequest(
@@ -192,9 +204,15 @@ func (c *Client) DropNode(ctx context.Context, path string) (finalErr error) {
 		return xerrors.WithStackTrace(call(ctx))
 	}
 
-	return retry.Retry(ctx, func(ctx context.Context) (err error) {
-		return dropNode(ctx, c.client, request)
-	}, retry.WithStackTrace(), retry.WithIdempotent(true), retry.WithTrace(c.config.TraceRetry()))
+	return retry.Retry(ctx,
+		func(ctx context.Context) (err error) {
+			return dropNode(ctx, c.client, request)
+		},
+		retry.WithStackTrace(),
+		retry.WithIdempotent(true),
+		retry.WithTrace(c.config.TraceRetry()),
+		retry.WithBudget(c.config.RetryBudget()),
+	)
 }
 
 func dropNodeRequest(path string, operationParams *Ydb_Operations.OperationParams) *Ydb_Coordination.DropNodeRequest {
@@ -241,14 +259,20 @@ func (c *Client) DescribeNode(
 		return describeNode(ctx, c.client, request)
 	}
 
-	err := retry.Retry(ctx, func(ctx context.Context) (err error) {
-		entry, config, err = describeNode(ctx, c.client, request)
-		if err != nil {
-			return xerrors.WithStackTrace(err)
-		}
+	err := retry.Retry(ctx,
+		func(ctx context.Context) (err error) {
+			entry, config, err = describeNode(ctx, c.client, request)
+			if err != nil {
+				return xerrors.WithStackTrace(err)
+			}
 
-		return nil
-	}, retry.WithStackTrace(), retry.WithIdempotent(true), retry.WithTrace(c.config.TraceRetry()))
+			return nil
+		},
+		retry.WithStackTrace(),
+		retry.WithIdempotent(true),
+		retry.WithTrace(c.config.TraceRetry()),
+		retry.WithBudget(c.config.RetryBudget()),
+	)
 	if err != nil {
 		return nil, nil, xerrors.WithStackTrace(err)
 	}

internal/query/options/retry.go

@@ -3,16 +3,15 @@ package options
 import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/tx"
 	"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 (
-	_ DoOption = idempotentOption{}
-	_ DoOption = labelOption("")
+	_ DoOption = retryOptionsOption(nil)
 	_ DoOption = traceOption{}
 
-	_ DoTxOption = idempotentOption{}
-	_ DoTxOption = labelOption("")
+	_ DoTxOption = retryOptionsOption(nil)
 	_ DoTxOption = traceOption{}
 	_ DoTxOption = doTxSettingsOption{}
 )
@@ -36,9 +35,8 @@ type (
 		txSettings tx.Settings
 	}
 
-	idempotentOption struct{}
-	labelOption      string
-	traceOption      struct {
+	retryOptionsOption []retry.Option
+	traceOption        struct {
 		t *trace.Query
 	}
 	doTxSettingsOption struct {
@@ -62,14 +60,6 @@ func (s *doTxSettings) TxSettings() tx.Settings {
 	return s.txSettings
 }
 
-func (opt idempotentOption) applyDoTxOption(s *doTxSettings) {
-	s.doOpts = append(s.doOpts, opt)
-}
-
-func (idempotentOption) applyDoOption(s *doSettings) {
-	s.retryOpts = append(s.retryOpts, retry.WithIdempotent(true))
-}
-
 func (opt traceOption) applyDoOption(s *doSettings) {
 	s.trace = s.trace.Compose(opt.t)
 }
@@ -78,12 +68,12 @@ func (opt traceOption) applyDoTxOption(s *doTxSettings) {
 	s.doOpts = append(s.doOpts, opt)
 }
 
-func (opt labelOption) applyDoOption(s *doSettings) {
-	s.retryOpts = append(s.retryOpts, retry.WithLabel(string(opt)))
+func (opts retryOptionsOption) applyDoOption(s *doSettings) {
+	s.retryOpts = append(s.retryOpts, opts...)
 }
 
-func (opt labelOption) applyDoTxOption(s *doTxSettings) {
-	s.doOpts = append(s.doOpts, opt)
+func (opts retryOptionsOption) applyDoTxOption(s *doTxSettings) {
+	s.doOpts = append(s.doOpts, opts)
 }
 
 func (opt doTxSettingsOption) applyDoTxOption(opts *doTxSettings) {
@@ -94,18 +84,22 @@ func WithTxSettings(txSettings tx.Settings) doTxSettingsOption {
 	return doTxSettingsOption{txSettings: txSettings}
 }
 
-func WithIdempotent() idempotentOption {
-	return idempotentOption{}
+func WithIdempotent() retryOptionsOption {
+	return []retry.Option{retry.WithIdempotent(true)}
 }
 
-func WithLabel(lbl string) labelOption {
-	return labelOption(lbl)
+func WithLabel(lbl string) retryOptionsOption {
+	return []retry.Option{retry.WithLabel(lbl)}
 }
 
 func WithTrace(t *trace.Query) traceOption {
 	return traceOption{t: t}
 }
 
+func WithRetryBudget(b budget.Budget) retryOptionsOption {
+	return []retry.Option{retry.WithBudget(b)}
+}
+
 func ParseDoOpts(t *trace.Query, opts ...DoOption) (s *doSettings) {
 	s = &doSettings{
 		trace: t,