ratelimits: Correctly handle stale and concurrently initialized buckets

ratelimits/limiter.go

@@ -277,10 +277,12 @@ func (l *Limiter) BatchSpend(ctx context.Context, txns []Transaction) (*Decision
 	batchDecision := allowedDecision
 	newBuckets := make(map[string]time.Time)
 	incrBuckets := make(map[string]increment)
+	staleBuckets := make(map[string]time.Time)
 	txnOutcomes := make(map[Transaction]string)
 
 	for _, txn := range batch {
 		tat, bucketExists := tats[txn.bucketKey]
+		originalTAT := tat
 		if !bucketExists {
 			// First request from this client.
 			tat = l.clk.Now()
@@ -294,14 +296,18 @@ func (l *Limiter) BatchSpend(ctx context.Context, txns []Transaction) (*Decision
 		}
 
 		if d.allowed && (tat != d.newTAT) && txn.spend {
-			// New bucket state should be persisted.
-			if bucketExists {
+			if !bucketExists {
+				// No bucket exists, initialize a new bucket with BatchSetNotExist.
+				newBuckets[txn.bucketKey] = d.newTAT
+			} else if originalTAT.Before(l.clk.Now()) {
+				// A bucket exists, with a TAT in the past, update it with BatchSet.
+				staleBuckets[txn.bucketKey] = d.newTAT
+			} else {
+				// A bucket exists, with a TAT in the future, update it with BatchIncrement.
 				incrBuckets[txn.bucketKey] = increment{
 					cost: time.Duration(txn.cost * txn.limit.emissionInterval),
 					ttl:  time.Duration(txn.limit.burstOffset),
 				}
-			} else {
-				newBuckets[txn.bucketKey] = d.newTAT
 			}
 		}
 
@@ -319,14 +325,28 @@ func (l *Limiter) BatchSpend(ctx context.Context, txns []Transaction) (*Decision
 
 	if batchDecision.allowed {
 		if len(newBuckets) > 0 {
-			err = l.source.BatchSet(ctx, newBuckets)
+			createdBuckets, err := l.source.BatchSetNotExisting(ctx, newBuckets)
 			if err != nil {
 				return nil, fmt.Errorf("batch set for %d keys: %w", len(newBuckets), err)
 			}
+			for k, created := range createdBuckets {
+				if !created {
+					// A bucket was created by another request, fall back to
+					// BatchSet.
+					staleBuckets[k] = newBuckets[k]
+				}
+			}
+		}
+
+		if len(staleBuckets) > 0 {
+			err := l.source.BatchSet(ctx, staleBuckets)
+			if err != nil {
+				return nil, fmt.Errorf("batch set for %d keys: %w", len(staleBuckets), err)
+			}
 		}
 
 		if len(incrBuckets) > 0 {
-			err = l.source.BatchIncrement(ctx, incrBuckets)
+			err := l.source.BatchIncrement(ctx, incrBuckets)
 			if err != nil {
 				return nil, fmt.Errorf("batch increment for %d keys: %w", len(incrBuckets), err)
 			}

ratelimits/source.go

@@ -20,6 +20,11 @@ type Source interface {
 	//    the underlying storage client implementation).
 	BatchSet(ctx context.Context, bucketKeys map[string]time.Time) error
 
+	// BatchSetNotExisting attempts to set TATs for the specified bucketKeys if
+	// they do not already exist. Returns a map indicating which keys were set
+	// successfully.
+	BatchSetNotExisting(ctx context.Context, buckets map[string]time.Time) (map[string]bool, error)
+
 	// BatchIncrement updates the TATs for the specified bucketKeys, similar to
 	// BatchSet. Implementations MUST ensure non-blocking operations by either:
 	//   a) applying a deadline or timeout to the context WITHIN the method, or
@@ -79,6 +84,20 @@ func (in *inmem) BatchSet(_ context.Context, bucketKeys map[string]time.Time) er
 	return nil
 }
 
+func (in *inmem) BatchSetNotExisting(_ context.Context, bucketKeys map[string]time.Time) (map[string]bool, error) {
+	in.Lock()
+	defer in.Unlock()
+	results := make(map[string]bool, len(bucketKeys))
+	for k, v := range bucketKeys {
+		_, ok := in.m[k]
+		if !ok {
+			in.m[k] = v
+			results[k] = true
+		}
+	}
+	return results, nil
+}
+
 func (in *inmem) BatchIncrement(_ context.Context, bucketKeys map[string]increment) error {
 	in.Lock()
 	defer in.Unlock()

ratelimits/source_redis.go

@@ -108,6 +108,38 @@ func (r *RedisSource) BatchSet(ctx context.Context, buckets map[string]time.Time
 	return nil
 }
 
+// BatchSetNotExisting attempts to set TATs for the specified bucketKeys if they
+// do not already exist. Returns a map indicating which keys were set successfully.
+func (r *RedisSource) BatchSetNotExisting(ctx context.Context, buckets map[string]time.Time) (map[string]bool, error) {
+	start := r.clk.Now()
+
+	pipeline := r.client.Pipeline()
+	cmds := make(map[string]*redis.BoolCmd, len(buckets))
+	for bucketKey, tat := range buckets {
+		// Set a TTL of TAT + 10 minutes to account for clock skew.
+		ttl := tat.UTC().Sub(r.clk.Now()) + 10*time.Minute
+		cmds[bucketKey] = pipeline.SetNX(ctx, bucketKey, tat.UTC().UnixNano(), ttl)
+	}
+
+	_, err := pipeline.Exec(ctx)
+	totalLatency := r.clk.Since(start)
+	perSetLatency := totalLatency / time.Duration(len(buckets))
+	if err != nil {
+		r.observeLatency("batchsetnotexisting", totalLatency, err)
+		return nil, err
+	}
+
+	results := make(map[string]bool, len(buckets))
+	for bucketKey, cmd := range cmds {
+		success, _ := cmd.Result()
+		results[bucketKey] = success
+		r.observeLatency("batchsetnotexisting_entry", perSetLatency, nil)
+	}
+
+	r.observeLatency("batchsetnotexisting", totalLatency, nil)
+	return results, nil
+}
+
 // BatchIncrement updates TATs for the specified bucketKeys using a pipelined
 // Redis Transaction in order to reduce the number of round-trips to each Redis
 // shard.