diff --git a/ratelimits/bucket.go b/ratelimits/bucket.go index b31fe24ddefc..83bc4780ab87 100644 --- a/ratelimits/bucket.go +++ b/ratelimits/bucket.go @@ -85,13 +85,12 @@ func newFQDNSetBucketKey(name Name, orderNames []string) (string, error) { return joinWithColon(name.EnumString(), id), nil } -// Transaction is used to represent a single rate limit Transaction. It contains -// the bucketKey identifying the limit and subscriber. A cost which MUST be -// greater than or equal to 0. Cost is variable to allow for limits such as -// CertificatesPerDomainPerAccount, where the cost is the number of domain names -// in the order. The check and spend fields are are not mutually exclusive, and -// indicate how the limiter should process the Transaction; the following are -// acceptable combinations of both: +// Transaction represents a single rate limit operation. It includes a +// bucketKey, which combines the specific rate limit enum with a unique +// identifier to form the key where the state of the "bucket" can be referenced +// or stored by the Limiter, the rate limit being enforced, a cost which MUST be +// >= 0, and check/spend fields, which indicate how the Transaction should be +// processed. The following are acceptable combinations of check/spend: // - check-and-spend: when check and spend are both true, the cost will be // checked against the bucket's capacity and spent/refunded, when possible. // - check-only: when only check is true, the cost will be checked against the @@ -121,50 +120,47 @@ func (txn Transaction) allowOnly() bool { return !txn.check && !txn.spend } -func (txn Transaction) validate() error { - if txn.allowOnly() { - return nil - } +func validateTransaction(txn Transaction) (Transaction, error) { if txn.cost < 0 { - return ErrInvalidCost + return Transaction{}, ErrInvalidCost } if txn.cost > txn.limit.Burst { - return ErrInvalidCostOverLimit + return Transaction{}, ErrInvalidCostOverLimit } - return nil + return txn, nil } -func newTransaction(limit limit, bucketKey string, cost int64) Transaction { - return Transaction{ +func newTransaction(limit limit, bucketKey string, cost int64) (Transaction, error) { + return validateTransaction(Transaction{ bucketKey: bucketKey, limit: limit, cost: cost, check: true, spend: true, - } + }) } -func newCheckOnlyTransaction(limit limit, bucketKey string, cost int64) Transaction { - return Transaction{ +func newCheckOnlyTransaction(limit limit, bucketKey string, cost int64) (Transaction, error) { + return validateTransaction(Transaction{ bucketKey: bucketKey, limit: limit, cost: cost, check: true, - } + }) } -func newSpendOnlyTransaction(limit limit, bucketKey string, cost int64) Transaction { - return Transaction{ +func newSpendOnlyTransaction(limit limit, bucketKey string, cost int64) (Transaction, error) { + return validateTransaction(Transaction{ bucketKey: bucketKey, limit: limit, cost: cost, spend: true, - } + }) } -func newAllowOnlyTransaction() Transaction { +func newAllowOnlyTransaction() (Transaction, error) { // Zero values are sufficient. - return Transaction{} + return validateTransaction(Transaction{}) } // TransactionBuilder is used to build Transactions for various rate limits. @@ -196,12 +192,11 @@ func (builder *TransactionBuilder) RegistrationsPerIPAddressTransaction(ip net.I limit, err := builder.getLimit(NewRegistrationsPerIPAddress, bucketKey) if err != nil { if errors.Is(err, errLimitDisabled) { - return newAllowOnlyTransaction(), nil + return newAllowOnlyTransaction() } return Transaction{}, err } - txn := newTransaction(limit, bucketKey, 1) - return txn, txn.validate() + return newTransaction(limit, bucketKey, 1) } // RegistrationsPerIPv6RangeTransaction returns a Transaction for the @@ -215,12 +210,11 @@ func (builder *TransactionBuilder) RegistrationsPerIPv6RangeTransaction(ip net.I limit, err := builder.getLimit(NewRegistrationsPerIPAddress, bucketKey) if err != nil { if errors.Is(err, errLimitDisabled) { - return newAllowOnlyTransaction(), nil + return newAllowOnlyTransaction() } return Transaction{}, err } - txn := newTransaction(limit, bucketKey, 1) - return txn, txn.validate() + return newTransaction(limit, bucketKey, 1) } // OrdersPerAccountTransaction returns a Transaction for the NewOrdersPerAccount @@ -233,12 +227,11 @@ func (builder *TransactionBuilder) OrdersPerAccountTransaction(regId int64) (Tra limit, err := builder.getLimit(NewRegistrationsPerIPAddress, bucketKey) if err != nil { if errors.Is(err, errLimitDisabled) { - return newAllowOnlyTransaction(), nil + return newAllowOnlyTransaction() } return Transaction{}, err } - txn := newTransaction(limit, bucketKey, 1) - return txn, txn.validate() + return newTransaction(limit, bucketKey, 1) } // FailedAuthorizationsPerAccountCheckOnlyTransaction returns a check-only @@ -252,12 +245,11 @@ func (builder *TransactionBuilder) FailedAuthorizationsPerAccountCheckOnlyTransa limit, err := builder.getLimit(NewRegistrationsPerIPAddress, bucketKey) if err != nil { if errors.Is(err, errLimitDisabled) { - return newAllowOnlyTransaction(), nil + return newAllowOnlyTransaction() } return Transaction{}, err } - txn := newCheckOnlyTransaction(limit, bucketKey, 1) - return txn, txn.validate() + return newCheckOnlyTransaction(limit, bucketKey, 1) } // FailedAuthorizationsPerAccountTransaction returns a Transaction for the @@ -270,17 +262,16 @@ func (builder *TransactionBuilder) FailedAuthorizationsPerAccountTransaction(reg limit, err := builder.getLimit(NewRegistrationsPerIPAddress, bucketKey) if err != nil { if errors.Is(err, errLimitDisabled) { - return newAllowOnlyTransaction(), nil + return newAllowOnlyTransaction() } return Transaction{}, err } - txn := newTransaction(limit, bucketKey, 1) - return txn, txn.validate() + return newTransaction(limit, bucketKey, 1) } -// CertificatesPerDomainTransactions returns a slice of Transactions for the for -// the provided order domain names. An error is returned if any of the order -// domain names are invalid. When a CertificatesPerDomainPerAccount override is +// CertificatesPerDomainTransactions returns a slice of Transactions for the +// provided order domain names. An error is returned if any of the order domain +// names are invalid. When a CertificatesPerDomainPerAccount override is // configured, two types of Transactions are returned: // - A spend-only Transaction for each per domain bucket. Spend-only transactions // will not be denied if the bucket lacks the capacity to satisfy the cost. @@ -303,7 +294,6 @@ func (builder *TransactionBuilder) CertificatesPerDomainTransactions(regId int64 } var txns []Transaction - var perAccountPerDomainCost int64 for _, name := range DomainsForRateLimiting(orderDomains) { perDomainBucketKey, err := newDomainBucketKey(CertificatesPerDomain, name) if err != nil { @@ -315,7 +305,6 @@ func (builder *TransactionBuilder) CertificatesPerDomainTransactions(regId int64 return nil, err } } - perAccountPerDomainCost += 1 if perAccountLimit.isOverride { // An override is configured for the CertificatesPerDomainPerAccount // limit. @@ -325,24 +314,21 @@ func (builder *TransactionBuilder) CertificatesPerDomainTransactions(regId int64 } // Add a check-and-spend transaction for each per account per domain // bucket. - txn := newTransaction(perAccountLimit, perAccountPerDomainKey, perAccountPerDomainCost) - err = txn.validate() + txn, err := newTransaction(perAccountLimit, perAccountPerDomainKey, 1) if err != nil { return nil, err } txns = append(txns, txn) // Add a spend-only transaction for each per domain bucket. - txn = newSpendOnlyTransaction(perDomainLimit, perDomainBucketKey, 1) - err = txn.validate() + txn, err = newSpendOnlyTransaction(perDomainLimit, perDomainBucketKey, 1) if err != nil { return nil, err } txns = append(txns, txn) } else { // Add a check-and-spend transaction for each per domain bucket. - txn := newTransaction(perDomainLimit, perDomainBucketKey, 1) - err = txn.validate() + txn, err := newTransaction(perDomainLimit, perDomainBucketKey, 1) if err != nil { return nil, err } @@ -362,10 +348,9 @@ func (builder *TransactionBuilder) CertificatesPerFQDNSetTransaction(orderNames limit, err := builder.getLimit(NewRegistrationsPerIPAddress, bucketKey) if err != nil { if errors.Is(err, errLimitDisabled) { - return newAllowOnlyTransaction(), nil + return newAllowOnlyTransaction() } return Transaction{}, err } - txn := newTransaction(limit, bucketKey, 1) - return txn, txn.validate() + return newTransaction(limit, bucketKey, 1) } diff --git a/ratelimits/limit.go b/ratelimits/limit.go index 390b428a0a10..6a6041579a8a 100644 --- a/ratelimits/limit.go +++ b/ratelimits/limit.go @@ -198,9 +198,9 @@ func newLimitRegistry(defaults, overrides string) (*limitRegistry, error) { } // getLimit returns the limit for the specified by name and bucketKey, name is -// required, bucketKey is optional. If bucketkey is left unspecified, the -// default limit for the limit specified by name is returned. If no default -// limit exists for the specified name, errLimitDisabled is returned. +// required, bucketKey is optional. If bucketkey is empty, the default for the +// limit specified by name is returned. If no default limit exists for the +// specified name, errLimitDisabled is returned. func (l *limitRegistry) getLimit(name Name, bucketKey string) (limit, error) { if !name.isValid() { // This should never happen. Callers should only be specifying the limit diff --git a/ratelimits/limiter.go b/ratelimits/limiter.go index e4b223c68108..6807bda722e2 100644 --- a/ratelimits/limiter.go +++ b/ratelimits/limiter.go @@ -152,7 +152,7 @@ func (l *Limiter) Spend(ctx context.Context, txn Transaction) (*Decision, error) } if tat == d.newTAT || txn.checkOnly() { - // No-op. + // Don't update storage return d, nil } @@ -254,10 +254,9 @@ func (l *Limiter) BatchSpend(ctx context.Context, txns []Transaction) (*Decision newTATs[txn.bucketKey] = d.newTAT } - if txn.spendOnly() { - d = allowedDecision + if !txn.spendOnly() { + batchDecision.merge(d) } - batchDecision.merge(d) } if batchDecision.Allowed { diff --git a/ratelimits/limiter_test.go b/ratelimits/limiter_test.go index 3e1e05420793..0271eb649238 100644 --- a/ratelimits/limiter_test.go +++ b/ratelimits/limiter_test.go @@ -69,12 +69,16 @@ func Test_Limiter_CheckWithLimitOverrides(t *testing.T) { test.AssertNotError(t, err, "should not error") // Attempt to spend all 40 requests, this should succeed. - d, err := l.Spend(testCtx, newTransaction(overriddenLimit, overriddenBucketKey, 40)) + overriddenTxn40, err := newTransaction(overriddenLimit, overriddenBucketKey, 40) + test.AssertNotError(t, err, "txn should be valid") + d, err := l.Spend(testCtx, overriddenTxn40) test.AssertNotError(t, err, "should not error") test.Assert(t, d.Allowed, "should be allowed") // Attempting to spend 1 more, this should fail. - d, err = l.Spend(testCtx, newTransaction(overriddenLimit, overriddenBucketKey, 1)) + overriddenTxn1, err := newTransaction(overriddenLimit, overriddenBucketKey, 1) + test.AssertNotError(t, err, "txn should be valid") + d, err = l.Spend(testCtx, overriddenTxn1) test.AssertNotError(t, err, "should not error") test.Assert(t, !d.Allowed, "should not be allowed") test.AssertEquals(t, d.Remaining, int64(0)) @@ -94,7 +98,7 @@ func Test_Limiter_CheckWithLimitOverrides(t *testing.T) { clk.Add(d.RetryIn) // We should be allowed to spend 1 more request. - d, err = l.Spend(testCtx, newTransaction(overriddenLimit, overriddenBucketKey, 1)) + d, err = l.Spend(testCtx, overriddenTxn1) test.AssertNotError(t, err, "should not error") test.Assert(t, d.Allowed, "should be allowed") test.AssertEquals(t, d.Remaining, int64(0)) @@ -105,14 +109,14 @@ func Test_Limiter_CheckWithLimitOverrides(t *testing.T) { // Quickly spend 40 requests in a row. for i := 0; i < 40; i++ { - d, err = l.Spend(testCtx, newTransaction(overriddenLimit, overriddenBucketKey, 1)) + d, err = l.Spend(testCtx, overriddenTxn1) test.AssertNotError(t, err, "should not error") test.Assert(t, d.Allowed, "should be allowed") test.AssertEquals(t, d.Remaining, int64(39-i)) } // Attempting to spend 1 more, this should fail. - d, err = l.Spend(testCtx, newTransaction(overriddenLimit, overriddenBucketKey, 1)) + d, err = l.Spend(testCtx, overriddenTxn1) test.AssertNotError(t, err, "should not error") test.Assert(t, !d.Allowed, "should not be allowed") test.AssertEquals(t, d.Remaining, int64(0)) @@ -130,7 +134,9 @@ func Test_Limiter_CheckWithLimitOverrides(t *testing.T) { // Spend the same bucket but in a batch with bucket subject to // default limits. This should succeed, but the decision should // reflect that of the default bucket. - d, err = l.BatchSpend(testCtx, []Transaction{newTransaction(overriddenLimit, overriddenBucketKey, 1), newTransaction(normalLimit, normalBucketKey, 1)}) + defaultTxn1, err := newTransaction(normalLimit, normalBucketKey, 1) + test.AssertNotError(t, err, "txn should be valid") + d, err = l.BatchSpend(testCtx, []Transaction{overriddenTxn1, defaultTxn1}) test.AssertNotError(t, err, "should not error") test.Assert(t, d.Allowed, "should be allowed") test.AssertEquals(t, d.Remaining, int64(19)) @@ -139,7 +145,7 @@ func Test_Limiter_CheckWithLimitOverrides(t *testing.T) { // Refund quota to both buckets. This should succeed, but the // decision should reflect that of the default bucket. - d, err = l.BatchRefund(testCtx, []Transaction{newTransaction(overriddenLimit, overriddenBucketKey, 1), newTransaction(normalLimit, normalBucketKey, 1)}) + d, err = l.BatchRefund(testCtx, []Transaction{overriddenTxn1, defaultTxn1}) test.AssertNotError(t, err, "should not error") test.Assert(t, d.Allowed, "should be allowed") test.AssertEquals(t, d.Remaining, int64(20)) @@ -147,7 +153,7 @@ func Test_Limiter_CheckWithLimitOverrides(t *testing.T) { test.AssertEquals(t, d.ResetIn, time.Duration(0)) // Once more. - d, err = l.BatchSpend(testCtx, []Transaction{newTransaction(overriddenLimit, overriddenBucketKey, 1), newTransaction(normalLimit, normalBucketKey, 1)}) + d, err = l.BatchSpend(testCtx, []Transaction{overriddenTxn1, defaultTxn1}) test.AssertNotError(t, err, "should not error") test.Assert(t, d.Allowed, "should be allowed") test.AssertEquals(t, d.Remaining, int64(19)) @@ -163,21 +169,27 @@ func Test_Limiter_CheckWithLimitOverrides(t *testing.T) { // Spend the same bucket but in a batch with a Transaction that is // check-only. This should succeed, but the decision should reflect // that of the default bucket. - d, err = l.BatchSpend(testCtx, []Transaction{newTransaction(overriddenLimit, overriddenBucketKey, 1), newCheckOnlyTransaction(normalLimit, normalBucketKey, 1)}) + defaultCheckOnlyTxn1, err := newCheckOnlyTransaction(normalLimit, normalBucketKey, 1) + test.AssertNotError(t, err, "txn should be valid") + d, err = l.BatchSpend(testCtx, []Transaction{overriddenTxn1, defaultCheckOnlyTxn1}) test.AssertNotError(t, err, "should not error") test.AssertEquals(t, d.Remaining, int64(19)) test.AssertEquals(t, d.RetryIn, time.Duration(0)) test.AssertEquals(t, d.ResetIn, time.Millisecond*50) // Check the remaining quota of the overridden bucket. - d, err = l.Check(testCtx, newTransaction(overriddenLimit, overriddenBucketKey, 0)) + overriddenCheckOnlyTxn0, err := newCheckOnlyTransaction(overriddenLimit, overriddenBucketKey, 0) + test.AssertNotError(t, err, "txn should be valid") + d, err = l.Check(testCtx, overriddenCheckOnlyTxn0) test.AssertNotError(t, err, "should not error") test.AssertEquals(t, d.Remaining, int64(39)) test.AssertEquals(t, d.RetryIn, time.Duration(0)) test.AssertEquals(t, d.ResetIn, time.Millisecond*25) // Check the remaining quota of the default bucket. - d, err = l.Check(testCtx, newTransaction(normalLimit, normalBucketKey, 0)) + defaultTxn0, err := newTransaction(normalLimit, normalBucketKey, 0) + test.AssertNotError(t, err, "txn should be valid") + d, err = l.Check(testCtx, defaultTxn0) test.AssertNotError(t, err, "should not error") test.AssertEquals(t, d.Remaining, int64(20)) test.AssertEquals(t, d.RetryIn, time.Duration(0)) @@ -186,21 +198,23 @@ func Test_Limiter_CheckWithLimitOverrides(t *testing.T) { // Spend the same bucket but in a batch with a Transaction that is // spend-only. This should succeed, but the decision should reflect // that of the overridden bucket. - d, err = l.BatchSpend(testCtx, []Transaction{newTransaction(overriddenLimit, overriddenBucketKey, 1), newSpendOnlyTransaction(normalLimit, normalBucketKey, 1)}) + defaultSpendOnlyTxn1, err := newSpendOnlyTransaction(normalLimit, normalBucketKey, 1) + test.AssertNotError(t, err, "txn should be valid") + d, err = l.BatchSpend(testCtx, []Transaction{overriddenTxn1, defaultSpendOnlyTxn1}) test.AssertNotError(t, err, "should not error") test.AssertEquals(t, d.Remaining, int64(38)) test.AssertEquals(t, d.RetryIn, time.Duration(0)) test.AssertEquals(t, d.ResetIn, time.Millisecond*50) // Check the remaining quota of the overridden bucket. - d, err = l.Check(testCtx, newTransaction(overriddenLimit, overriddenBucketKey, 0)) + d, err = l.Check(testCtx, overriddenCheckOnlyTxn0) test.AssertNotError(t, err, "should not error") test.AssertEquals(t, d.Remaining, int64(38)) test.AssertEquals(t, d.RetryIn, time.Duration(0)) test.AssertEquals(t, d.ResetIn, time.Millisecond*50) // Check the remaining quota of the default bucket. - d, err = l.Check(testCtx, newTransaction(normalLimit, normalBucketKey, 0)) + d, err = l.Check(testCtx, defaultTxn0) test.AssertNotError(t, err, "should not error") test.AssertEquals(t, d.Remaining, int64(19)) test.AssertEquals(t, d.RetryIn, time.Duration(0)) @@ -209,21 +223,23 @@ func Test_Limiter_CheckWithLimitOverrides(t *testing.T) { // Once more, but in now the spend-only Transaction will attempt to // spend 20 requests. The spend-only Transaction should fail, but // the decision should reflect that of the overridden bucket. - d, err = l.BatchSpend(testCtx, []Transaction{newTransaction(overriddenLimit, overriddenBucketKey, 1), newSpendOnlyTransaction(normalLimit, normalBucketKey, 20)}) + defaultSpendOnlyTxn20, err := newSpendOnlyTransaction(normalLimit, normalBucketKey, 20) + test.AssertNotError(t, err, "txn should be valid") + d, err = l.BatchSpend(testCtx, []Transaction{overriddenTxn1, defaultSpendOnlyTxn20}) test.AssertNotError(t, err, "should not error") test.AssertEquals(t, d.Remaining, int64(37)) test.AssertEquals(t, d.RetryIn, time.Duration(0)) test.AssertEquals(t, d.ResetIn, time.Millisecond*75) // Check the remaining quota of the overridden bucket. - d, err = l.Check(testCtx, newTransaction(overriddenLimit, overriddenBucketKey, 0)) + d, err = l.Check(testCtx, overriddenCheckOnlyTxn0) test.AssertNotError(t, err, "should not error") test.AssertEquals(t, d.Remaining, int64(37)) test.AssertEquals(t, d.RetryIn, time.Duration(0)) test.AssertEquals(t, d.ResetIn, time.Millisecond*75) // Check the remaining quota of the default bucket. - d, err = l.Check(testCtx, newTransaction(normalLimit, normalBucketKey, 0)) + d, err = l.Check(testCtx, defaultTxn0) test.AssertNotError(t, err, "should not error") test.AssertEquals(t, d.Remaining, int64(19)) test.AssertEquals(t, d.RetryIn, time.Duration(0)) @@ -248,7 +264,9 @@ func Test_Limiter_InitializationViaCheckAndSpend(t *testing.T) { // Check on an empty bucket should return the theoretical next state // of that bucket if the cost were spent. - d, err := l.Check(testCtx, newTransaction(limit, bucketKey, 1)) + txn1, err := newTransaction(limit, bucketKey, 1) + test.AssertNotError(t, err, "txn should be valid") + d, err := l.Check(testCtx, txn1) test.AssertNotError(t, err, "should not error") test.Assert(t, d.Allowed, "should be allowed") test.AssertEquals(t, d.Remaining, int64(19)) @@ -259,7 +277,9 @@ func Test_Limiter_InitializationViaCheckAndSpend(t *testing.T) { // However, that cost should not be spent yet, a 0 cost check should // tell us that we actually have 20 remaining. - d, err = l.Check(testCtx, newTransaction(limit, bucketKey, 0)) + txn0, err := newTransaction(limit, bucketKey, 0) + test.AssertNotError(t, err, "txn should be valid") + d, err = l.Check(testCtx, txn0) test.AssertNotError(t, err, "should not error") test.Assert(t, d.Allowed, "should be allowed") test.AssertEquals(t, d.Remaining, int64(20)) @@ -272,7 +292,7 @@ func Test_Limiter_InitializationViaCheckAndSpend(t *testing.T) { // Similar to above, but we'll use Spend() to actually initialize // the bucket. Spend should return the same result as Check. - d, err = l.Spend(testCtx, newTransaction(limit, bucketKey, 1)) + d, err = l.Spend(testCtx, txn1) test.AssertNotError(t, err, "should not error") test.Assert(t, d.Allowed, "should be allowed") test.AssertEquals(t, d.Remaining, int64(19)) @@ -283,7 +303,7 @@ func Test_Limiter_InitializationViaCheckAndSpend(t *testing.T) { // However, that cost should not be spent yet, a 0 cost check should // tell us that we actually have 19 remaining. - d, err = l.Check(testCtx, newTransaction(limit, bucketKey, 0)) + d, err = l.Check(testCtx, txn0) test.AssertNotError(t, err, "should not error") test.Assert(t, d.Allowed, "should be allowed") test.AssertEquals(t, d.Remaining, int64(19)) @@ -306,14 +326,18 @@ func Test_Limiter_DefaultLimits(t *testing.T) { test.AssertNotError(t, err, "should not error") // Attempt to spend all 20 requests, this should succeed. - d, err := l.Spend(testCtx, newTransaction(limit, bucketKey, 20)) + txn20, err := newTransaction(limit, bucketKey, 20) + test.AssertNotError(t, err, "txn should be valid") + d, err := l.Spend(testCtx, txn20) test.AssertNotError(t, err, "should not error") test.Assert(t, d.Allowed, "should be allowed") test.AssertEquals(t, d.Remaining, int64(0)) test.AssertEquals(t, d.ResetIn, time.Second) // Attempting to spend 1 more, this should fail. - d, err = l.Spend(testCtx, newTransaction(limit, bucketKey, 1)) + txn1, err := newTransaction(limit, bucketKey, 1) + test.AssertNotError(t, err, "txn should be valid") + d, err = l.Spend(testCtx, txn1) test.AssertNotError(t, err, "should not error") test.Assert(t, !d.Allowed, "should not be allowed") test.AssertEquals(t, d.Remaining, int64(0)) @@ -327,7 +351,7 @@ func Test_Limiter_DefaultLimits(t *testing.T) { clk.Add(d.RetryIn) // We should be allowed to spend 1 more request. - d, err = l.Spend(testCtx, newTransaction(limit, bucketKey, 1)) + d, err = l.Spend(testCtx, txn1) test.AssertNotError(t, err, "should not error") test.Assert(t, d.Allowed, "should be allowed") test.AssertEquals(t, d.Remaining, int64(0)) @@ -338,14 +362,14 @@ func Test_Limiter_DefaultLimits(t *testing.T) { // Quickly spend 20 requests in a row. for i := 0; i < 20; i++ { - d, err = l.Spend(testCtx, newTransaction(limit, bucketKey, 1)) + d, err = l.Spend(testCtx, txn1) test.AssertNotError(t, err, "should not error") test.Assert(t, d.Allowed, "should be allowed") test.AssertEquals(t, d.Remaining, int64(19-i)) } // Attempting to spend 1 more, this should fail. - d, err = l.Spend(testCtx, newTransaction(limit, bucketKey, 1)) + d, err = l.Spend(testCtx, txn1) test.AssertNotError(t, err, "should not error") test.Assert(t, !d.Allowed, "should not be allowed") test.AssertEquals(t, d.Remaining, int64(0)) @@ -365,19 +389,23 @@ func Test_Limiter_RefundAndReset(t *testing.T) { test.AssertNotError(t, err, "should not error") // Attempt to spend all 20 requests, this should succeed. - d, err := l.Spend(testCtx, newTransaction(limit, bucketKey, 20)) + txn20, err := newTransaction(limit, bucketKey, 20) + test.AssertNotError(t, err, "txn should be valid") + d, err := l.Spend(testCtx, txn20) test.AssertNotError(t, err, "should not error") test.Assert(t, d.Allowed, "should be allowed") test.AssertEquals(t, d.Remaining, int64(0)) test.AssertEquals(t, d.ResetIn, time.Second) // Refund 10 requests. - d, err = l.Refund(testCtx, newTransaction(limit, bucketKey, 10)) + txn10, err := newTransaction(limit, bucketKey, 10) + test.AssertNotError(t, err, "txn should be valid") + d, err = l.Refund(testCtx, txn10) test.AssertNotError(t, err, "should not error") test.AssertEquals(t, d.Remaining, int64(10)) // Spend 10 requests, this should succeed. - d, err = l.Spend(testCtx, newTransaction(limit, bucketKey, 10)) + d, err = l.Spend(testCtx, txn10) test.AssertNotError(t, err, "should not error") test.Assert(t, d.Allowed, "should be allowed") test.AssertEquals(t, d.Remaining, int64(0)) @@ -387,7 +415,7 @@ func Test_Limiter_RefundAndReset(t *testing.T) { test.AssertNotError(t, err, "should not error") // Attempt to spend 20 more requests, this should succeed. - d, err = l.Spend(testCtx, newTransaction(limit, bucketKey, 20)) + d, err = l.Spend(testCtx, txn20) test.AssertNotError(t, err, "should not error") test.Assert(t, d.Allowed, "should be allowed") test.AssertEquals(t, d.Remaining, int64(0)) @@ -397,26 +425,32 @@ func Test_Limiter_RefundAndReset(t *testing.T) { clk.Add(d.ResetIn) // Refund 1 requests above our limit, this should fail. - d, err = l.Refund(testCtx, newTransaction(limit, bucketKey, 1)) + txn1, err := newTransaction(limit, bucketKey, 1) + test.AssertNotError(t, err, "txn should be valid") + d, err = l.Refund(testCtx, txn1) test.AssertNotError(t, err, "should not error") test.Assert(t, !d.Allowed, "should not be allowed") test.AssertEquals(t, d.Remaining, int64(20)) // Spend so we can refund. - _, err = l.Spend(testCtx, newTransaction(limit, bucketKey, 1)) + _, err = l.Spend(testCtx, txn1) test.AssertNotError(t, err, "should not error") // Refund a spendOnly Transaction, which should succeed. - _, err = l.Refund(testCtx, newSpendOnlyTransaction(limit, bucketKey, 1)) + spendOnlyTxn1, err := newSpendOnlyTransaction(limit, bucketKey, 1) + test.AssertNotError(t, err, "txn should be valid") + _, err = l.Refund(testCtx, spendOnlyTxn1) test.AssertNotError(t, err, "should not error") // Spend so we can refund. - expectedDecision, err := l.Spend(testCtx, newTransaction(limit, bucketKey, 1)) + expectedDecision, err := l.Spend(testCtx, txn1) test.AssertNotError(t, err, "should not error") // Refund a checkOnly Transaction, which shouldn't error but should // return the same TAT as the previous spend. - newDecision, err := l.Refund(testCtx, newCheckOnlyTransaction(limit, bucketKey, 1)) + checkOnlyTxn1, err := newCheckOnlyTransaction(limit, bucketKey, 1) + test.AssertNotError(t, err, "txn should be valid") + newDecision, err := l.Refund(testCtx, checkOnlyTxn1) test.AssertNotError(t, err, "should not error") test.AssertEquals(t, newDecision.newTAT, expectedDecision.newTAT) })