Skip to content

ratelimits: Refactor Reset() to accept a batch of Transactions #8503

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+158 −76
Open

ratelimits: Refactor Reset() to accept a batch of Transactions #8503

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions ra/ra.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1438,8 +1438,13 @@ func (ra *RegistrationAuthorityImpl) countFailedValidations(ctx context.Context,
// resetAccountPausingLimit resets bucket to maximum capacity for given account.
// There is no reason to surface errors from this function to the Subscriber.
func (ra *RegistrationAuthorityImpl) resetAccountPausingLimit(ctx context.Context, regId int64, ident identifier.ACMEIdentifier) {
bucketKey := ratelimits.NewRegIdIdentValueBucketKey(ratelimits.FailedAuthorizationsForPausingPerDomainPerAccount, regId, ident.Value)
err := ra.limiter.Reset(ctx, bucketKey)
txns, err := ra.txnBuilder.NewPausingResetTransactions(regId, ident)
if err != nil {
ra.log.Warningf("building reset transaction for regID=[%d] identifier=[%s]: %s", regId, ident.Value, err)
return
}

err = ra.limiter.BatchReset(ctx, txns)
if err != nil {
ra.log.Warningf("resetting bucket for regID=[%d] identifier=[%s]: %s", regId, ident.Value, err)
}
Expand Down
16 changes: 10 additions & 6 deletions ra/ra_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -715,7 +715,8 @@ func TestPerformValidation_FailedValidationsTriggerPauseIdentifiersRatelimit(t *
domain := randomDomain()
ident := identifier.NewDNS(domain)
authzPB := createPendingAuthorization(t, sa, registration.Id, ident, fc.Now().Add(12*time.Hour))
bucketKey := ratelimits.NewRegIdIdentValueBucketKey(ratelimits.FailedAuthorizationsForPausingPerDomainPerAccount, authzPB.RegistrationID, ident.Value)
bucketKey, err := ratelimits.BuildBucketKey(ratelimits.FailedAuthorizationsForPausingPerDomainPerAccount, authzPB.RegistrationID, ident, nil, netip.Addr{})
test.AssertNotError(t, err, "building bucket key")

// Set the stored TAT to indicate that this bucket has exhausted its quota.
err = rl.BatchSet(context.Background(), map[string]time.Time{
Expand Down Expand Up @@ -758,15 +759,17 @@ func TestPerformValidation_FailedValidationsTriggerPauseIdentifiersRatelimit(t *

// mockRLSourceWithSyncDelete is a mock ratelimits.Source that forwards all
// method calls to an inner Source, but also performs a blocking write to a
// channel when Delete is called to allow the tests to synchronize.
// channel when BatchDelete is called to allow the tests to synchronize.
type mockRLSourceWithSyncDelete struct {
ratelimits.Source
out chan<- string
}

func (rl mockRLSourceWithSyncDelete) Delete(ctx context.Context, bucketKey string) error {
err := rl.Source.Delete(ctx, bucketKey)
rl.out <- bucketKey
func (rl mockRLSourceWithSyncDelete) BatchDelete(ctx context.Context, bucketKeys []string) error {
err := rl.Source.BatchDelete(ctx, bucketKeys)
for _, bucketKey := range bucketKeys {
rl.out <- bucketKey
}
return err
}

Expand All @@ -791,7 +794,8 @@ func TestPerformValidation_FailedThenSuccessfulValidationResetsPauseIdentifiersR
domain := randomDomain()
ident := identifier.NewDNS(domain)
authzPB := createPendingAuthorization(t, sa, registration.Id, ident, fc.Now().Add(12*time.Hour))
bucketKey := ratelimits.NewRegIdIdentValueBucketKey(ratelimits.FailedAuthorizationsForPausingPerDomainPerAccount, authzPB.RegistrationID, ident.Value)
bucketKey, err := ratelimits.BuildBucketKey(ratelimits.FailedAuthorizationsForPausingPerDomainPerAccount, authzPB.RegistrationID, ident, nil, netip.Addr{})
test.AssertNotError(t, err, "building bucket key")

// Set a stored TAT so that we can tell when it's been reset.
err = rl.BatchSet(context.Background(), map[string]time.Time{
Expand Down
60 changes: 30 additions & 30 deletions ratelimits/gcra_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,16 +16,16 @@ func TestDecide(t *testing.T) {
limit.precompute()

// Begin by using 1 of our 10 requests.
d := maybeSpend(clk, Transaction{"test", limit, 1, true, true}, clk.Now())
d := maybeSpend(clk, Transaction{"test", limit, 1, true, true, false}, clk.Now())
test.Assert(t, d.allowed, "should be allowed")
test.AssertEquals(t, d.remaining, int64(9))
test.AssertEquals(t, d.retryIn, time.Duration(0))
test.AssertEquals(t, d.resetIn, time.Second)
// Transaction is set when we're allowed.
test.AssertEquals(t, d.transaction, Transaction{"test", limit, 1, true, true})
test.AssertEquals(t, d.transaction, Transaction{"test", limit, 1, true, true, false})

// Immediately use another 9 of our remaining requests.
d = maybeSpend(clk, Transaction{"test", limit, 9, true, true}, d.newTAT)
d = maybeSpend(clk, Transaction{"test", limit, 9, true, true, false}, d.newTAT)
test.Assert(t, d.allowed, "should be allowed")
test.AssertEquals(t, d.remaining, int64(0))
// We should have to wait 1 second before we can use another request but we
Expand All @@ -37,21 +37,21 @@ func TestDecide(t *testing.T) {
test.AssertEquals(t, d.newTAT, clk.Now().Add(time.Second*10))

// Let's try using just 1 more request without waiting.
d = maybeSpend(clk, Transaction{"test", limit, 1, true, true}, d.newTAT)
d = maybeSpend(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
test.Assert(t, !d.allowed, "should not be allowed")
test.AssertEquals(t, d.remaining, int64(0))
test.AssertEquals(t, d.retryIn, time.Second)
test.AssertEquals(t, d.resetIn, time.Second*10)
// Transaction is set when we're denied.
test.AssertEquals(t, d.transaction, Transaction{"test", limit, 1, true, true})
test.AssertEquals(t, d.transaction, Transaction{"test", limit, 1, true, true, false})

// Let's try being exactly as patient as we're told to be.
clk.Add(d.retryIn)
d = maybeSpend(clk, Transaction{"test", limit, 0, true, true}, d.newTAT)
d = maybeSpend(clk, Transaction{"test", limit, 0, true, true, false}, d.newTAT)
test.AssertEquals(t, d.remaining, int64(1))

// We are 1 second in the future, we should have 1 new request.
d = maybeSpend(clk, Transaction{"test", limit, 1, true, true}, d.newTAT)
d = maybeSpend(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
test.Assert(t, d.allowed, "should be allowed")
test.AssertEquals(t, d.remaining, int64(0))
test.AssertEquals(t, d.retryIn, time.Second)
Expand All @@ -61,7 +61,7 @@ func TestDecide(t *testing.T) {
clk.Add(d.resetIn)

// We should have 10 new requests. If we use 1 we should have 9 remaining.
d = maybeSpend(clk, Transaction{"test", limit, 1, true, true}, d.newTAT)
d = maybeSpend(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
test.Assert(t, d.allowed, "should be allowed")
test.AssertEquals(t, d.remaining, int64(9))
test.AssertEquals(t, d.retryIn, time.Duration(0))
Expand All @@ -72,7 +72,7 @@ func TestDecide(t *testing.T) {

// We should still have 9 remaining because we're still 1ms shy of the
// refill time.
d = maybeSpend(clk, Transaction{"test", limit, 0, true, true}, d.newTAT)
d = maybeSpend(clk, Transaction{"test", limit, 0, true, true, false}, d.newTAT)
test.Assert(t, d.allowed, "should be allowed")
test.AssertEquals(t, d.remaining, int64(9))
test.AssertEquals(t, d.retryIn, time.Duration(0))
Expand All @@ -83,15 +83,15 @@ func TestDecide(t *testing.T) {
clk.Add(20 * time.Hour)

// C'mon, big money, no whammies, no whammies, STOP!
d = maybeSpend(clk, Transaction{"test", limit, 0, true, true}, d.newTAT)
d = maybeSpend(clk, Transaction{"test", limit, 0, true, true, false}, d.newTAT)
test.Assert(t, d.allowed, "should be allowed")
test.AssertEquals(t, d.remaining, int64(10))
test.AssertEquals(t, d.retryIn, time.Duration(0))
test.AssertEquals(t, d.resetIn, time.Duration(0))

// Turns out that the most we can accrue is 10 (limit.Burst). Let's empty
// this bucket out so we can try something else.
d = maybeSpend(clk, Transaction{"test", limit, 10, true, true}, d.newTAT)
d = maybeSpend(clk, Transaction{"test", limit, 10, true, true, false}, d.newTAT)
test.Assert(t, d.allowed, "should be allowed")
test.AssertEquals(t, d.remaining, int64(0))
// We should have to wait 1 second before we can use another request but we
Expand All @@ -101,14 +101,14 @@ func TestDecide(t *testing.T) {
test.AssertEquals(t, d.resetIn, time.Second*10)

// If you spend 0 while you have 0 you should get 0.
d = maybeSpend(clk, Transaction{"test", limit, 0, true, true}, d.newTAT)
d = maybeSpend(clk, Transaction{"test", limit, 0, true, true, false}, d.newTAT)
test.Assert(t, d.allowed, "should be allowed")
test.AssertEquals(t, d.remaining, int64(0))
test.AssertEquals(t, d.retryIn, time.Duration(0))
test.AssertEquals(t, d.resetIn, time.Second*10)

// We don't play by the rules, we spend 1 when we have 0.
d = maybeSpend(clk, Transaction{"test", limit, 1, true, true}, d.newTAT)
d = maybeSpend(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
test.Assert(t, !d.allowed, "should not be allowed")
test.AssertEquals(t, d.remaining, int64(0))
test.AssertEquals(t, d.retryIn, time.Second)
Expand All @@ -118,7 +118,7 @@ func TestDecide(t *testing.T) {
clk.Add(d.retryIn)

// Our patience pays off, we should have 1 new request. Let's use it.
d = maybeSpend(clk, Transaction{"test", limit, 1, true, true}, d.newTAT)
d = maybeSpend(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
test.Assert(t, d.allowed, "should be allowed")
test.AssertEquals(t, d.remaining, int64(0))
test.AssertEquals(t, d.retryIn, time.Second)
Expand All @@ -130,7 +130,7 @@ func TestDecide(t *testing.T) {
// Attempt to spend 7 when we only have 5. We should be denied but the
// decision should reflect a retry of 2 seconds, the time it would take to
// refill from 5 to 7.
d = maybeSpend(clk, Transaction{"test", limit, 7, true, true}, d.newTAT)
d = maybeSpend(clk, Transaction{"test", limit, 7, true, true, false}, d.newTAT)
test.Assert(t, !d.allowed, "should not be allowed")
test.AssertEquals(t, d.remaining, int64(5))
test.AssertEquals(t, d.retryIn, time.Second*2)
Expand All @@ -143,28 +143,28 @@ func TestMaybeRefund(t *testing.T) {
limit.precompute()

// Begin by using 1 of our 10 requests.
d := maybeSpend(clk, Transaction{"test", limit, 1, true, true}, clk.Now())
d := maybeSpend(clk, Transaction{"test", limit, 1, true, true, false}, clk.Now())
test.Assert(t, d.allowed, "should be allowed")
test.AssertEquals(t, d.remaining, int64(9))
test.AssertEquals(t, d.retryIn, time.Duration(0))
test.AssertEquals(t, d.resetIn, time.Second)
// Transaction is set when we're refunding.
test.AssertEquals(t, d.transaction, Transaction{"test", limit, 1, true, true})
test.AssertEquals(t, d.transaction, Transaction{"test", limit, 1, true, true, false})

// Refund back to 10.
d = maybeRefund(clk, Transaction{"test", limit, 1, true, true}, d.newTAT)
d = maybeRefund(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
test.AssertEquals(t, d.remaining, int64(10))
test.AssertEquals(t, d.retryIn, time.Duration(0))
test.AssertEquals(t, d.resetIn, time.Duration(0))

// Refund 0, we should still have 10.
d = maybeRefund(clk, Transaction{"test", limit, 0, true, true}, d.newTAT)
d = maybeRefund(clk, Transaction{"test", limit, 0, true, true, false}, d.newTAT)
test.AssertEquals(t, d.remaining, int64(10))
test.AssertEquals(t, d.retryIn, time.Duration(0))
test.AssertEquals(t, d.resetIn, time.Duration(0))

// Spend 1 more of our 10 requests.
d = maybeSpend(clk, Transaction{"test", limit, 1, true, true}, d.newTAT)
d = maybeSpend(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
test.Assert(t, d.allowed, "should be allowed")
test.AssertEquals(t, d.remaining, int64(9))
test.AssertEquals(t, d.retryIn, time.Duration(0))
Expand All @@ -174,16 +174,16 @@ func TestMaybeRefund(t *testing.T) {
clk.Add(d.resetIn)

// Attempt to refund from 10 to 11.
d = maybeRefund(clk, Transaction{"test", limit, 1, true, true}, d.newTAT)
d = maybeRefund(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
test.Assert(t, !d.allowed, "should not be allowed")
test.AssertEquals(t, d.remaining, int64(10))
test.AssertEquals(t, d.retryIn, time.Duration(0))
test.AssertEquals(t, d.resetIn, time.Duration(0))
// Transaction is set when our bucket is full.
test.AssertEquals(t, d.transaction, Transaction{"test", limit, 1, true, true})
test.AssertEquals(t, d.transaction, Transaction{"test", limit, 1, true, true, false})

// Spend 10 all 10 of our requests.
d = maybeSpend(clk, Transaction{"test", limit, 10, true, true}, d.newTAT)
d = maybeSpend(clk, Transaction{"test", limit, 10, true, true, false}, d.newTAT)
test.Assert(t, d.allowed, "should be allowed")
test.AssertEquals(t, d.remaining, int64(0))
// We should have to wait 1 second before we can use another request but we
Expand All @@ -193,7 +193,7 @@ func TestMaybeRefund(t *testing.T) {
test.AssertEquals(t, d.resetIn, time.Second*10)

// Attempt a refund of 10.
d = maybeRefund(clk, Transaction{"test", limit, 10, true, true}, d.newTAT)
d = maybeRefund(clk, Transaction{"test", limit, 10, true, true, false}, d.newTAT)
test.AssertEquals(t, d.remaining, int64(10))
test.AssertEquals(t, d.retryIn, time.Duration(0))
test.AssertEquals(t, d.resetIn, time.Duration(0))
Expand All @@ -202,17 +202,17 @@ func TestMaybeRefund(t *testing.T) {
clk.Add(11 * time.Second)

// Attempt to refund to 11, then ensure it's still 10.
d = maybeRefund(clk, Transaction{"test", limit, 1, true, true}, d.newTAT)
d = maybeRefund(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
test.Assert(t, !d.allowed, "should be allowed")
test.AssertEquals(t, d.remaining, int64(10))
test.AssertEquals(t, d.retryIn, time.Duration(0))
test.AssertEquals(t, d.resetIn, time.Duration(0))
// Transaction is set when our TAT is in the past.
test.AssertEquals(t, d.transaction, Transaction{"test", limit, 1, true, true})
test.AssertEquals(t, d.transaction, Transaction{"test", limit, 1, true, true, false})

// Spend 5 of our 10 requests, then refund 1.
d = maybeSpend(clk, Transaction{"test", limit, 5, true, true}, d.newTAT)
d = maybeRefund(clk, Transaction{"test", limit, 1, true, true}, d.newTAT)
d = maybeSpend(clk, Transaction{"test", limit, 5, true, true, false}, d.newTAT)
d = maybeRefund(clk, Transaction{"test", limit, 1, true, true, false}, d.newTAT)
test.Assert(t, d.allowed, "should be allowed")
test.AssertEquals(t, d.remaining, int64(6))
test.AssertEquals(t, d.retryIn, time.Duration(0))
Expand All @@ -221,15 +221,15 @@ func TestMaybeRefund(t *testing.T) {
clk.Add(time.Millisecond * 2500)

// Ensure we have 8.5 requests.
d = maybeSpend(clk, Transaction{"test", limit, 0, true, true}, d.newTAT)
d = maybeSpend(clk, Transaction{"test", limit, 0, true, true, false}, d.newTAT)
test.Assert(t, d.allowed, "should be allowed")
test.AssertEquals(t, d.remaining, int64(8))
test.AssertEquals(t, d.retryIn, time.Duration(0))
// Check that ResetIn represents the fractional earned request.
test.AssertEquals(t, d.resetIn, time.Millisecond*1500)

// Refund 2 requests, we should only have 10, not 10.5.
d = maybeRefund(clk, Transaction{"test", limit, 2, true, true}, d.newTAT)
d = maybeRefund(clk, Transaction{"test", limit, 2, true, true, false}, d.newTAT)
test.AssertEquals(t, d.remaining, int64(10))
test.AssertEquals(t, d.retryIn, time.Duration(0))
test.AssertEquals(t, d.resetIn, time.Duration(0))
Expand Down
26 changes: 22 additions & 4 deletions ratelimits/limiter.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -434,11 +434,29 @@ func (l *Limiter) BatchRefund(ctx context.Context, txns []Transaction) (*Decisio
return batchDecision, nil
}

// Reset resets the specified bucket to its maximum capacity. The new bucket
// state is persisted to the underlying datastore before returning.
func (l *Limiter) Reset(ctx context.Context, bucketKey string) error {
// BatchReset resets the specified buckets to their maximum capacity using the
// provided reset Transactions. The new bucket state is persisted to the
// underlying datastore before returning.
func (l *Limiter) BatchReset(ctx context.Context, txns []Transaction) error {
var bucketKeys []string
for _, txn := range txns {
if txn.allowOnly() {
// Ignore allow-only transactions.
continue
}
if !txn.resetOnly() {
return fmt.Errorf("found reset-only transaction, received check=%t spend=%t reset=%t", txn.check, txn.spend, txn.reset)
}
if slices.Contains(bucketKeys, txn.bucketKey) {
return fmt.Errorf("found duplicate bucket %q in batch", txn.bucketKey)
}
bucketKeys = append(bucketKeys, txn.bucketKey)
}
if len(bucketKeys) == 0 {
return nil
}
// Remove cancellation from the request context so that transactions are not
// interrupted by a client disconnect.
ctx = context.WithoutCancel(ctx)
return l.source.Delete(ctx, bucketKey)
return l.source.BatchDelete(ctx, bucketKeys)
}
23 changes: 13 additions & 10 deletions ratelimits/limiter_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,14 @@ func setup(t *testing.T) (context.Context, map[string]*Limiter, *TransactionBuil
}, newTestTransactionBuilder(t), clk, randIP.String()
}

func mustReset(t *testing.T, l *Limiter, ctx context.Context, limit *Limit, bucketKey string) {
Copy link
Contributor

@aarongable aarongable Nov 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a nit, but: I think we tend to use must to mean "will panic if a precondition fails", rather than "will fail a test assertion". I don't think I see any other mustFoo functions in the codebase that take a testing.T as an argument; meanwhile we do have mustFoo test functions that use our must.Do helper.

Maybe this is totally fine? I think the meaning and usage is mostly clear. But it just blurs the line between this helper being a test function that readers need to care about, versus it being boilerplate that we're clearly stating can never fail.

t.Helper()
txn, err := newResetTransaction(limit, bucketKey)
test.AssertNotError(t, err, "txn should be valid")
err = l.BatchReset(ctx, []Transaction{txn})
test.AssertNotError(t, err, "should not error")
}

func TestLimiter_CheckWithLimitOverrides(t *testing.T) {
t.Parallel()
testCtx, limiters, txnBuilder, clk, testIP := setup(t)
Expand Down Expand Up @@ -153,10 +161,8 @@ func TestLimiter_CheckWithLimitOverrides(t *testing.T) {
test.AssertEquals(t, d.resetIn, time.Millisecond*50)

// Reset between tests.
err = l.Reset(testCtx, overriddenBucketKey)
test.AssertNotError(t, err, "should not error")
err = l.Reset(testCtx, normalBucketKey)
test.AssertNotError(t, err, "should not error")
mustReset(t, l, testCtx, overriddenLimit, overriddenBucketKey)
mustReset(t, l, testCtx, normalLimit, normalBucketKey)

// Spend the same bucket but in a batch with a Transaction that is
// check-only. This should succeed, but the decision should reflect
Expand Down Expand Up @@ -238,8 +244,7 @@ func TestLimiter_CheckWithLimitOverrides(t *testing.T) {
test.AssertEquals(t, d.resetIn, time.Millisecond*50)

// Reset between tests.
err = l.Reset(testCtx, overriddenBucketKey)
test.AssertNotError(t, err, "should not error")
mustReset(t, l, testCtx, overriddenLimit, overriddenBucketKey)
})
}
}
Expand Down Expand Up @@ -278,8 +283,7 @@ func TestLimiter_InitializationViaCheckAndSpend(t *testing.T) {
test.AssertEquals(t, d.retryIn, time.Duration(0))

// Reset our bucket.
err = l.Reset(testCtx, bucketKey)
test.AssertNotError(t, err, "should not error")
mustReset(t, l, testCtx, limit, bucketKey)

// Similar to above, but we'll use Spend() to actually initialize
// the bucket. Spend should return the same result as Check.
Expand Down Expand Up @@ -400,8 +404,7 @@ func TestLimiter_RefundAndReset(t *testing.T) {
test.AssertEquals(t, d.remaining, int64(0))
test.AssertEquals(t, d.resetIn, time.Second)

err = l.Reset(testCtx, bucketKey)
test.AssertNotError(t, err, "should not error")
mustReset(t, l, testCtx, limit, bucketKey)

// Attempt to spend 20 more requests, this should succeed.
d, err = l.Spend(testCtx, txn20)
Expand Down
4 changes: 2 additions & 2 deletions ratelimits/names.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -409,7 +409,7 @@ func BuildBucketKey(name Name, regId int64, singleIdent identifier.ACMEIdentifie
if err != nil {
return "", err
}
return NewRegIdIdentValueBucketKey(name, regId, coveringIdent), nil
return newRegIdIdentValueBucketKey(name, regId, coveringIdent), nil
}
if regId == 0 {
return "", makeMissingErr("regId")
Expand All @@ -429,7 +429,7 @@ func BuildBucketKey(name Name, regId int64, singleIdent identifier.ACMEIdentifie
return "", makeMissingErr("regId")
}
// Default: use 'enum:regId:identValue' bucket key format.
return NewRegIdIdentValueBucketKey(name, regId, singleIdent.Value), nil
return newRegIdIdentValueBucketKey(name, regId, singleIdent.Value), nil
}
if regId == 0 {
return "", makeMissingErr("regId")
Expand Down
Loading
Loading