Use simpler racy approach and adopt the batched spending.

wfe2/wfe.go

@@ -619,109 +619,71 @@ func link(url, relation string) string {
 	return fmt.Sprintf("<%s>;rel=\"%s\"", url, relation)
 }
 
-// checkNewAccountLimits checks whether sufficient limit quota exists for the
-// creation of a new account from the given IP address. If so, that quota is
-// spent. If an error is encountered during the check, it is logged but not
-// returned.
-//
-// TODO(#5545): For now we're simply exercising the new rate limiter codepath.
-// This should eventually return a berrors.RateLimit error containing the retry
-// after duration among other information available in the ratelimits.Decision.
-func (wfe *WebFrontEndImpl) checkNewAccountLimits(ctx context.Context, ip net.IP) {
+func (wfe *WebFrontEndImpl) newNewAccountLimitTransactions(ip net.IP) []ratelimits.Transaction {
 	if wfe.limiter == nil && wfe.txnBuilder == nil {
 		// Limiter is disabled.
-		return
+		return nil
 	}
 
 	warn := func(err error, limit ratelimits.Name) {
-		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
-			return
-		}
 		// TODO(#5545): Once key-value rate limits are authoritative this log
 		// line should be removed in favor of returning the error.
 		wfe.log.Warningf("checking %s rate limit: %s", limit, err)
 	}
 
+	var transactions []ratelimits.Transaction
 	txn, err := wfe.txnBuilder.RegistrationsPerIPAddressTransaction(ip)
 	if err != nil {
 		warn(err, ratelimits.NewRegistrationsPerIPAddress)
-		return
+		return nil
 	}
+	transactions = append(transactions, txn)
 
-	decision, err := wfe.limiter.Spend(ctx, txn)
-	if err != nil {
-		warn(err, ratelimits.NewRegistrationsPerIPAddress)
-		return
-	}
-	if !decision.Allowed || ip.To4() != nil {
-		// This requester is being limited or the request was made from an IPv4
-		// address.
-		return
+	if ip.To4() != nil {
+		// This request was made from an IPv4 address.
+		return transactions
 	}
 
 	txn, err = wfe.txnBuilder.RegistrationsPerIPv6RangeTransaction(ip)
 	if err != nil {
 		warn(err, ratelimits.NewRegistrationsPerIPv6Range)
+		return nil
+	}
+	return append(transactions, txn)
+}
+
+// checkNewAccountLimits checks whether sufficient limit quota exists for the
+// creation of a new account. If so, that quota is spent. If an error is
+// encountered during the check, it is logged but not returned.
+//
+// TODO(#5545): For now we're simply exercising the new rate limiter codepath.
+// This should eventually return a berrors.RateLimit error containing the retry
+// after duration among other information available in the ratelimits.Decision.
+func (wfe *WebFrontEndImpl) checkNewAccountLimits(ctx context.Context, transactions []ratelimits.Transaction) {
+	if wfe.limiter == nil && wfe.txnBuilder == nil {
+		// Limiter is disabled.
 		return
 	}
 
-	_, err = wfe.limiter.Spend(ctx, txn)
+	_, err := wfe.limiter.BatchSpend(ctx, transactions)
 	if err != nil {
-		warn(err, ratelimits.NewRegistrationsPerIPv6Range)
+		wfe.log.Errf("checking newAccount limits: %s", err)
 	}
 }
 
 // refundNewAccountLimits is typically called when a new account creation fails.
 // It refunds the limit quota consumed by the request, allowing the caller to
 // retry immediately. If an error is encountered during the refund, it is logged
 // but not returned.
-func (wfe *WebFrontEndImpl) refundNewAccountLimits(ctx context.Context, limitCheck <-chan struct{}, ip net.IP) {
+func (wfe *WebFrontEndImpl) refundNewAccountLimits(ctx context.Context, transactions []ratelimits.Transaction) {
 	if wfe.limiter == nil && wfe.txnBuilder == nil {
 		// Limiter is disabled.
 		return
 	}
 
-	warn := func(err error, limit ratelimits.Name) {
-		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
-			return
-		}
-		// TODO(#5545): Once key-value rate limits are authoritative this log
-		// line should be removed in favor of returning the error.
-		wfe.log.Warningf("refunding %s rate limit: %s", limit, err)
-	}
-
-	select {
-	case <-ctx.Done():
-		return
-	case <-limitCheck:
-		// Continue with refunding the limits.
-	}
-
-	txn, err := wfe.txnBuilder.RegistrationsPerIPAddressTransaction(ip)
-	if err != nil {
-		warn(err, ratelimits.NewRegistrationsPerIPAddress)
-		return
-	}
-
-	_, err = wfe.limiter.Refund(ctx, txn)
-	if err != nil {
-		warn(err, ratelimits.NewRegistrationsPerIPAddress)
-		return
-	}
-	if ip.To4() != nil {
-		// Request was made from an IPv4 address.
-		return
-	}
-
-	txn, err = wfe.txnBuilder.RegistrationsPerIPv6RangeTransaction(ip)
+	_, err := wfe.limiter.BatchRefund(ctx, transactions)
 	if err != nil {
-		warn(err, ratelimits.NewRegistrationsPerIPv6Range)
-		return
-	}
-
-	_, err = wfe.limiter.Refund(ctx, txn)
-	if err != nil {
-		warn(err, ratelimits.NewRegistrationsPerIPv6Range)
+		wfe.log.Errf("refunding newAccount limits: %s", err)
 	}
 }
 
@@ -849,19 +811,19 @@ func (wfe *WebFrontEndImpl) NewAccount(
 
 	// TODO(#5545): Spending and Refunding can be async until these rate limits
 	// are authoritative. This saves us from adding latency to each request.
-	limitCheck := make(chan struct{})
+	// Goroutines spun out below will respect a context deadline set by the
+	// ratelimits package and cannot be prematurely canceled by the requester.
+	txns := wfe.newNewAccountLimitTransactions(ip)
+	go wfe.checkNewAccountLimits(ctx, txns)
+
 	var newRegistrationSuccessful bool
 	var errIsRateLimit bool
-
-	go func() {
-		// Close the channel on goroutine completion.
-		defer close(limitCheck)
-		wfe.checkNewAccountLimits(ctx, ip)
-	}()
-
 	defer func() {
 		if !newRegistrationSuccessful && !errIsRateLimit {
-			go wfe.refundNewAccountLimits(ctx, limitCheck, ip)
+			// This can be a little racy, but we're not going to worry about it
+			// for now. If the check hasn't completed yet, we can pretty safely
+			// assume that the refund will be similarly delayed.
+			go wfe.refundNewAccountLimits(ctx, txns)
 		}
 	}()