renamed budget.New() => budget.Limited()

added `budget.Percent()`

Author: asmyasnikov

internal/config/config.go

@@ -7,7 +7,7 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
-var defaultRetryBudget = budget.New(-1)
+var defaultRetryBudget = budget.Limited(-1)
 
 type Common struct {
 	operationTimeout     time.Duration

retry/budget/budget.go

@@ -2,36 +2,43 @@ package budget
 
 import (
 	"context"
+	"fmt"
 	"time"
 
 	"github.com/jonboulle/clockwork"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xrand"
 )
 
 type (
 	// Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
 	Budget interface {
+		// Acquire called from retryer on second and subsequence attempts
 		Acquire(ctx context.Context) error
 	}
-	budget struct {
+	fixedBudget struct {
 		clock  clockwork.Clock
 		ticker clockwork.Ticker
 		quota  chan struct{}
 		done   chan struct{}
 	}
-	option func(q *budget)
+	fixedBudgetOption func(q *fixedBudget)
+	percentBudget     struct {
+		percent int
+		rand    xrand.Rand
+	}
 )
 
-func withBudgetClock(clock clockwork.Clock) option {
-	return func(q *budget) {
+func withFixedBudgetClock(clock clockwork.Clock) fixedBudgetOption {
+	return func(q *fixedBudget) {
 		q.clock = clock
 	}
 }
 
 // Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
-func New(attemptsPerSecond int, opts ...option) *budget {
-	q := &budget{
+func Limited(attemptsPerSecond int, opts ...fixedBudgetOption) *fixedBudget {
+	q := &fixedBudget{
 		clock: clockwork.NewRealClock(),
 		done:  make(chan struct{}),
 	}
@@ -68,15 +75,15 @@ func New(attemptsPerSecond int, opts ...option) *budget {
 }
 
 // Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
-func (q *budget) Stop() {
+func (q *fixedBudget) Stop() {
 	if q.ticker != nil {
 		q.ticker.Stop()
 	}
 	close(q.done)
 }
 
 // Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
-func (q *budget) Acquire(ctx context.Context) error {
+func (q *fixedBudget) Acquire(ctx context.Context) error {
 	if err := ctx.Err(); err != nil {
 		return xerrors.WithStackTrace(err)
 	}
@@ -89,3 +96,22 @@ func (q *budget) Acquire(ctx context.Context) error {
 		return xerrors.WithStackTrace(ctx.Err())
 	}
 }
+
+// Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
+func Percent(percent int) *percentBudget {
+	if percent > 100 || percent < 0 {
+		panic(fmt.Sprintf("wrong percent value: %d", percent))
+	}
+	return &percentBudget{
+		percent: percent,
+		rand:    xrand.New(xrand.WithLock()),
+	}
+}
+
+func (b *percentBudget) Acquire(ctx context.Context) error {
+	if b.rand.Int(100) < b.percent {
+		return nil
+	}
+
+	return ErrNoQuota
+}

retry/budget/budget_test.go

@@ -15,18 +15,18 @@ import (
 func TestUnlimitedBudget(t *testing.T) {
 	xtest.TestManyTimes(t, func(t testing.TB) {
 		ctx, cancel := xcontext.WithCancel(xtest.Context(t))
-		q := New(-1)
+		q := Limited(-1)
 		require.NoError(t, q.Acquire(ctx))
 		cancel()
 		require.ErrorIs(t, q.Acquire(ctx), context.Canceled)
 	})
 }
 
-func TestBudget(t *testing.T) {
+func TestLimited(t *testing.T) {
 	xtest.TestManyTimes(t, func(t testing.TB) {
 		ctx, cancel := xcontext.WithCancel(xtest.Context(t))
 		clock := clockwork.NewFakeClock()
-		q := New(1, withBudgetClock(clock))
+		q := Limited(1, withFixedBudgetClock(clock))
 		defer q.Stop()
 		require.NoError(t, q.Acquire(ctx))
 		acquireCh := make(chan struct{})
@@ -51,3 +51,22 @@ func TestBudget(t *testing.T) {
 		require.ErrorIs(t, q.Acquire(ctx), context.Canceled)
 	})
 }
+
+func TestPercent(t *testing.T) {
+	xtest.TestManyTimes(t, func(t testing.TB) {
+		var (
+			total   = 1000000
+			percent = 0.25
+			ctx     = xtest.Context(t)
+			b       = Percent(int(percent * 100))
+			success int
+		)
+		for i := 0; i < total; i++ {
+			if b.Acquire(ctx) == nil {
+				success++
+			}
+		}
+		require.GreaterOrEqual(t, success, int(float64(total)*(percent-0.1*percent)))
+		require.LessOrEqual(t, success, int(float64(total)*(percent+0.1*percent)))
+	}, xtest.StopAfter(5*time.Second))
+}

retry/retry.go

@@ -261,7 +261,7 @@ func Retry(ctx context.Context, op retryOperation, opts ...Option) (finalErr err
 	options := &retryOptions{
 		call:        stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/retry.Retry"),
 		trace:       &trace.Retry{},
-		budget:      budget.New(-1),
+		budget:      budget.Limited(-1),
 		fastBackoff: backoff.Fast,
 		slowBackoff: backoff.Slow,
 	}

tests/integration/retry_budget_test.go

@@ -31,7 +31,7 @@ func (n noQuota) Acquire(ctx context.Context) error {
 func TestRetryBudget(t *testing.T) {
 	ctx := xtest.Context(t)
 
-	defaultLimiter := budget.New(1)
+	defaultLimiter := budget.Limited(1)
 	defer defaultLimiter.Stop()
 
 	nativeDriver, err := ydb.Open(ctx, os.Getenv("YDB_CONNECTION_STRING"),

tests/slo/database/sql/storage.go

@@ -86,7 +86,7 @@ func NewStorage(ctx context.Context, cfg *config.Config, poolSize int) (s *Stora
 	ctx, cancel := context.WithTimeout(ctx, time.Minute*5) //nolint:gomnd
 	defer cancel()
 
-	retryBudget := budget.New(int(float64(poolSize) * 0.1)) //nolint:gomnd
+	retryBudget := budget.Limited(int(float64(poolSize) * 0.1)) //nolint:gomnd
 
 	s = &Storage{
 		cfg: cfg,

tests/slo/native/query/storage.go

@@ -73,7 +73,7 @@ func NewStorage(ctx context.Context, cfg *config.Config, poolSize int) (*Storage
 	ctx, cancel := context.WithTimeout(ctx, time.Minute*5) //nolint:gomnd
 	defer cancel()
 
-	retryBudget := budget.New(int(float64(poolSize) * 0.1)) //nolint:gomnd
+	retryBudget := budget.Limited(int(float64(poolSize) * 0.1)) //nolint:gomnd
 
 	db, err := ydb.Open(ctx,
 		cfg.Endpoint+cfg.DB,

tests/slo/native/table/storage.go

@@ -74,7 +74,7 @@ func NewStorage(ctx context.Context, cfg *config.Config, poolSize int) (*Storage
 	ctx, cancel := context.WithTimeout(ctx, time.Minute*5) //nolint:gomnd
 	defer cancel()
 
-	retryBudget := budget.New(int(float64(poolSize) * 0.1)) //nolint:gomnd
+	retryBudget := budget.Limited(int(float64(poolSize) * 0.1)) //nolint:gomnd
 
 	db, err := ydb.Open(
 		ctx,