ratelimits: Refactor Reset() to accept Transactions

Author: beautifulentropy

ra/ra.go

@@ -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.Reset(ctx, txns)
 	if err != nil {
 		ra.log.Warningf("resetting bucket for regID=[%d] identifier=[%s]: %s", regId, ident.Value, err)
 	}

ra/ra_test.go

@@ -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{
@@ -791,7 +792,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{

ratelimits/gcra_test.go

@@ -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
@@ -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)
@@ -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))
@@ -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))
@@ -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
@@ -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)
@@ -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)
@@ -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)
@@ -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))
@@ -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
@@ -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))
@@ -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))
@@ -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))

ratelimits/limiter.go

@@ -434,11 +434,35 @@ 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 {
+// Reset 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) Reset(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)
+	for _, bucketKey := range bucketKeys {
+		err := l.source.Delete(ctx, bucketKey)
+		if err != nil {
+			return fmt.Errorf("resetting bucket %q: %w", bucketKey, err)
+		}
+	}
+	return nil
 }

ratelimits/limiter_test.go

@@ -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) {
+	t.Helper()
+	txn, err := newResetTransaction(limit, bucketKey)
+	test.AssertNotError(t, err, "txn should be valid")
+	err = l.Reset(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)
@@ -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
@@ -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)
 		})
 	}
 }
@@ -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.
@@ -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)

ratelimits/names.go

@@ -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")
@@ -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")