Skip to content

Commit

Permalink
Optimistic and check-only transaction support for Check/Spend
Browse files Browse the repository at this point in the history
  • Loading branch information
beautifulentropy committed Nov 21, 2023
1 parent 600a37e commit 0e5b7d5
Show file tree
Hide file tree
Showing 5 changed files with 50 additions and 37 deletions.
20 changes: 10 additions & 10 deletions ratelimits/bucket.go
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,7 @@ func newOptimisticTransaction(b bucketId, cost int64) Transaction {
}

// NewRegistrationsPerIPAddressTransaction returns a Transaction for the
// provided IP address.
// NewRegistrationsPerIPAddress limit for the provided IP address.
func NewRegistrationsPerIPAddressTransaction(ip net.IP, cost int64) (Transaction, error) {
bucketId, err := newIPAddressBucketId(NewRegistrationsPerIPAddress, ip)
if err != nil {
Expand All @@ -157,8 +157,9 @@ func NewRegistrationsPerIPAddressTransaction(ip net.IP, cost int64) (Transaction
return newTransaction(bucketId, cost), nil
}

// NewRegistrationsPerIPv6RangeTransaction returns a Transaction for the /48
// IPv6 range containing the provided IPv6 address.
// NewRegistrationsPerIPv6RangeTransaction returns a Transaction for the
// NewRegistrationsPerIPv6Range limit for the /48 IPv6 range which contains the
// provided IPv6 address.
func NewRegistrationsPerIPv6RangeTransaction(ip net.IP, cost int64) (Transaction, error) {
bucketId, err := newIPv6RangeCIDRBucketId(NewRegistrationsPerIPv6Range, ip)
if err != nil {
Expand All @@ -168,7 +169,7 @@ func NewRegistrationsPerIPv6RangeTransaction(ip net.IP, cost int64) (Transaction
}

// NewOrdersPerAccountTransaction returns a Transaction for the provided ACME
// registration Id.
// NewOrdersPerAccount limit for the provided ACME registration Id.
func NewOrdersPerAccountTransaction(regId, cost int64) (Transaction, error) {
bucketId, err := newRegIdBucketId(NewOrdersPerAccount, regId)
if err != nil {
Expand All @@ -177,9 +178,9 @@ func NewOrdersPerAccountTransaction(regId, cost int64) (Transaction, error) {
return newTransaction(bucketId, cost), nil
}

// NewFailedAuthorizationsPerAccountCheckOnlyTransaction returns a Transaction
// for the provided ACME registration Id which, when processed as part of a
// batch call, will only check whether the bucket is over threshold.
// NewFailedAuthorizationsPerAccountCheckOnlyTransaction returns a 0 cost
// check-only Transaction for the provided ACME registration Id for the
// FailedAuthorizationsPerAccount limit.
func NewFailedAuthorizationsPerAccountCheckOnlyTransaction(regId int64) (Transaction, error) {
bucketId, err := newRegIdBucketId(FailedAuthorizationsPerAccount, regId)
if err != nil {
Expand All @@ -189,7 +190,7 @@ func NewFailedAuthorizationsPerAccountCheckOnlyTransaction(regId int64) (Transac
}

// NewFailedAuthorizationsPerAccountTransaction returns a Transaction for the
// provided ACME registration Id.
// FailedAuthorizationsPerAccount limit for the provided ACME registration Id.
func NewFailedAuthorizationsPerAccountTransaction(regId, cost int64) (Transaction, error) {
bucketId, err := newRegIdBucketId(FailedAuthorizationsPerAccount, regId)
if err != nil {
Expand All @@ -199,8 +200,7 @@ func NewFailedAuthorizationsPerAccountTransaction(regId, cost int64) (Transactio
}

// NewCertificatesPerDomainTransactions returns a slice of Transactions for the
// provided order domain names. The cost specified will be applied per eTLD+1
// name present in the orderDomains.
// CertificatesPerDomain limit for the provided order domain names.
//
// Note: when overrides to the CertificatesPerDomainPerAccount are configured
// for the subscriber, the cost:
Expand Down
2 changes: 1 addition & 1 deletion ratelimits/gcra.go
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ func maybeSpend(clk clock.Clock, rl limit, tat time.Time, cost int64) *Decision
// or greater. A cost will only be refunded up to the burst capacity of the
// limit. A partial refund is still considered successful.
func maybeRefund(clk clock.Clock, rl limit, tat time.Time, cost int64) *Decision {
if cost <= 0 || cost > rl.Burst {
if cost < 0 || cost > rl.Burst {
// The condition above is checked in the Refund method of Limiter. If
// this panic is reached, it means that the caller has introduced a bug.
panic("invalid cost for maybeRefund")
Expand Down
6 changes: 6 additions & 0 deletions ratelimits/gcra_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -152,6 +152,12 @@ func Test_maybeRefund(t *testing.T) {
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, limit, d.newTAT, 0)
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, limit, d.newTAT, 1)
test.Assert(t, d.Allowed, "should be allowed")
Expand Down
49 changes: 32 additions & 17 deletions ratelimits/limiter.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,11 +22,8 @@ const (
Denied = "denied"
)

// ErrInvalidCost indicates that the cost specified was <= 0.
var ErrInvalidCost = fmt.Errorf("invalid cost, must be > 0")

// ErrInvalidCostForCheck indicates that the check cost specified was < 0.
var ErrInvalidCostForCheck = fmt.Errorf("invalid check cost, must be >= 0")
// ErrInvalidCost indicates that the cost specified was < 0.
var ErrInvalidCost = fmt.Errorf("invalid cost, must be >= 0")

// ErrInvalidCostOverLimit indicates that the cost specified was > limit.Burst.
var ErrInvalidCostOverLimit = fmt.Errorf("invalid cost, must be <= limit.Burst")
Expand Down Expand Up @@ -127,7 +124,7 @@ type Decision struct {
// state is persisted to the underlying datastore.
func (l *Limiter) Check(ctx context.Context, txn Transaction) (*Decision, error) {
if txn.cost < 0 {
return nil, ErrInvalidCostForCheck
return nil, ErrInvalidCost
}

limit, err := l.getLimit(txn.limitName, txn.bucketKey)
Expand Down Expand Up @@ -165,7 +162,7 @@ func (l *Limiter) Check(ctx context.Context, txn Transaction) (*Decision, error)
// state is persisted to the underlying datastore, if applicable, before
// returning.
func (l *Limiter) Spend(ctx context.Context, txn Transaction) (*Decision, error) {
if txn.cost <= 0 {
if txn.cost < 0 {
return nil, ErrInvalidCost
}

Expand All @@ -190,7 +187,7 @@ func (l *Limiter) Spend(ctx context.Context, txn Transaction) (*Decision, error)
// Remove cancellation from the request context so that transactions are not
// interrupted by a client disconnect.
ctx = context.WithoutCancel(ctx)
tat, err := l.source.Get(ctx, txn.bucketKey)
currentTAT, err := l.source.Get(ctx, txn.bucketKey)
if err != nil {
if errors.Is(err, ErrBucketNotFound) {
// First request from this client.
Expand All @@ -206,7 +203,7 @@ func (l *Limiter) Spend(ctx context.Context, txn Transaction) (*Decision, error)
return nil, err
}

d := maybeSpend(l.clk, limit, tat, txn.cost)
d := maybeSpend(l.clk, limit, currentTAT, txn.cost)

if limit.isOverride {
// Calculate the current utilization of the override limit.
Expand All @@ -215,6 +212,15 @@ func (l *Limiter) Spend(ctx context.Context, txn Transaction) (*Decision, error)
}

if !d.Allowed {
if txn.optimistic {
// Denials for optimistic transactions are NOT persisted.
return maybeSpend(l.clk, limit, currentTAT, 0), nil
}
return d, nil
}

if currentTAT == d.newTAT || txn.checkOnly {
// No-op.
return d, nil
}

Expand All @@ -235,7 +241,7 @@ func (l *Limiter) prepareBatch(txns []Transaction) ([]batchTransaction, []string
var batchTxns []batchTransaction
var bucketKeys []string
for _, txn := range txns {
if txn.cost <= 0 {
if txn.cost < 0 {
return nil, nil, ErrInvalidCost
}
limit, err := l.getLimit(txn.limitName, txn.bucketKey)
Expand Down Expand Up @@ -316,17 +322,20 @@ func (l *Limiter) BatchSpend(ctx context.Context, txns []Transaction) (*Decision
newTATs := make(map[string]time.Time)

for _, txn := range batch {
tat, exists := tats[txn.bucketKey]
currentTAT, exists := tats[txn.bucketKey]
if !exists {
// First request from this client. Initialize the bucket with a TAT of
// "now", which is equivalent to a full bucket.
tat = l.clk.Now()
currentTAT = l.clk.Now()
}

// Spend the cost and update the consolidated decision.
d := maybeSpend(l.clk, txn.limit, tat, txn.cost)
d := maybeSpend(l.clk, txn.limit, currentTAT, txn.cost)
if d.Allowed && !txn.checkOnly {
// Check-only transactions are NOT persisted.
if currentTAT == d.newTAT {
// No-op.
continue
}
newTATs[txn.bucketKey] = d.newTAT
}

Expand Down Expand Up @@ -368,7 +377,7 @@ func (l *Limiter) BatchSpend(ctx context.Context, txns []Transaction) (*Decision
// requests remaining, a refund request of 7 will result in the bucket reaching
// its maximum capacity of 10, not 12.
func (l *Limiter) Refund(ctx context.Context, txn Transaction) (*Decision, error) {
if txn.cost <= 0 {
if txn.cost < 0 {
return nil, ErrInvalidCost
}

Expand Down Expand Up @@ -435,15 +444,21 @@ func (l *Limiter) BatchRefund(ctx context.Context, txns []Transaction) (*Decisio
continue
}

tat, exists := tats[txn.bucketKey]
currentTAT, exists := tats[txn.bucketKey]
if !exists {
// Ignore non-existent bucket.
continue
}

// Refund the cost and update the consolidated decision.
d := maybeRefund(l.clk, txn.limit, tat, txn.cost)
d := maybeRefund(l.clk, txn.limit, currentTAT, txn.cost)
if currentTAT == d.newTAT {
// No-op.
continue
}

if d.Allowed {
// Refund is possible.
newTATs[txn.bucketKey] = d.newTAT
}
batchDecision.consolidate(d)
Expand Down
10 changes: 1 addition & 9 deletions ratelimits/limiter_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -320,18 +320,10 @@ func Test_Limiter_RefundAndSpendCostErr(t *testing.T) {
bucketId, err := newIPAddressBucketId(NewRegistrationsPerIPAddress, net.ParseIP(testIP))
test.AssertNotError(t, err, "should not error")

// Spend a cost of 0, which should fail.
_, err = l.Spend(testCtx, newTransaction(bucketId, 0))
test.AssertErrorIs(t, err, ErrInvalidCost)

// Spend a negative cost, which should fail.
_, err = l.Spend(testCtx, newTransaction(bucketId, -1))
test.AssertErrorIs(t, err, ErrInvalidCost)

// Refund a cost of 0, which should fail.
_, err = l.Refund(testCtx, newTransaction(bucketId, 0))
test.AssertErrorIs(t, err, ErrInvalidCost)

// Refund a negative cost, which should fail.
_, err = l.Refund(testCtx, newTransaction(bucketId, -1))
test.AssertErrorIs(t, err, ErrInvalidCost)
Expand All @@ -348,7 +340,7 @@ func Test_Limiter_CheckWithBadCost(t *testing.T) {
test.AssertNotError(t, err, "should not error")

_, err = l.Check(testCtx, newTransaction(bucketId, -1))
test.AssertErrorIs(t, err, ErrInvalidCostForCheck)
test.AssertErrorIs(t, err, ErrInvalidCost)
})
}
}
Expand Down

0 comments on commit 0e5b7d5

Please sign in to comment.