Simplify & expand Wait methods, remove FIFO

limiter.go

@@ -6,7 +6,6 @@ type Limiter[TInput any, TKey comparable] struct {
 	limits     []Limit
 	limitFuncs []LimitFunc[TInput]
 	buckets    bucketMap[TKey]
-	waiters    syncMap[TKey, *waiter]
 }
 
 // KeyFunc is a function that takes an input and returns a bucket key.

limiter_wait.go

@@ -2,8 +2,6 @@ package rate
 
 import (
 	"context"
-	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/clipperhouse/ntime"
@@ -30,12 +28,13 @@ import (
 //
 //	ctx := context.WithTimeout(ctx, limit.DurationPerToken())
 //
-// Wait offers best-effort FIFO ordering of requests. Under sustained
-// contention (when multiple requests wait longer than ~1ms), the Go runtime's
-// mutex starvation mode ensures strict FIFO ordering. Under light load,
-// ordering may be less strict but performance is optimized.
-func (r *Limiter[TInput, TKey]) Wait(ctx context.Context, input TInput) bool {
-	return r.waitN(ctx, input, ntime.Now(), 1)
+// The returned error will be non-nil if the context is cancelled.
+//
+// Wait makes no ordering guarantees. Multiple concurrent calls may
+// acquire tokens in any order.
+func (r *Limiter[TInput, TKey]) Wait(ctx context.Context, input TInput) (bool, error) {
+	allow, _, err := r.waitNWithDetails(ctx, input, ntime.Now(), 1)
+	return allow, err
 }
 
 // WaitN will poll [Limiter.AllowN] for a period of time,
@@ -59,114 +58,201 @@ func (r *Limiter[TInput, TKey]) Wait(ctx context.Context, input TInput) bool {
 //
 //	ctx := context.WithTimeout(ctx, limit.DurationPerToken())
 //
-// WaitN offers best-effort FIFO ordering of requests. Under sustained
-// contention (when multiple requests wait longer than ~1ms), the Go runtime's
-// mutex starvation mode ensures strict FIFO ordering. Under light load,
-// ordering may be less strict but performance is optimized.
-func (r *Limiter[TInput, TKey]) WaitN(ctx context.Context, input TInput, n int64) bool {
-	return r.waitN(ctx, input, ntime.Now(), n)
+// The returned error will be non-nil if the context is cancelled.
+//
+// WaitN makes no ordering guarantees. Multiple concurrent calls may
+// acquire tokens in any order.
+func (r *Limiter[TInput, TKey]) WaitN(ctx context.Context, input TInput, n int64) (bool, error) {
+	allow, _, err := r.waitNWithDetails(ctx, input, ntime.Now(), n)
+	return allow, err
 }
 
-func (r *Limiter[TInput, TKey]) waitN(ctx context.Context, input TInput, executionTime ntime.Time, n int64) bool {
-	return r.waitNWithCancellation(
-		input,
-		executionTime,
-		n,
-		ctx.Deadline,
-		ctx.Done,
-	)
+// WaitWithDetails will poll [Allow] for a period of time,
+// until it is cancelled by the passed context. It has the
+// effect of adding latency to requests instead of refusing
+// them immediately. Consider it graceful degradation.
+//
+// WaitWithDetails will return true if a token becomes available prior to
+// the context cancellation, and will consume a token. It will
+// return false if not, and therefore not consume a token. It will
+// also return the details of the request.
+//
+// Take care to create an appropriate context. You almost certainly
+// want [context.WithTimeout] or [context.WithDeadline].
+//
+// You should be conservative, as WaitWithDetails will introduce
+// backpressure on your upstream systems -- connections
+// may be held open longer, requests may queue in memory.
+//
+// A good starting place will be to timeout after waiting
+// for one token. For example:
+//
+//	ctx := context.WithTimeout(ctx, limit.DurationPerToken())
+//
+// The returned error will be non-nil if the context is cancelled.
+//
+// WaitWithDetails makes no ordering guarantees. Multiple concurrent calls may
+// acquire tokens in any order.
+func (r *Limiter[TInput, TKey]) WaitWithDetails(ctx context.Context, input TInput) (bool, Details[TInput, TKey], error) {
+	return r.WaitNWithDetails(ctx, input, 1)
 }
 
-func (r *Limiter[TInput, TKey]) waitWithCancellation(
-	input TInput,
-	startTime ntime.Time,
-	deadline func() (time.Time, bool),
-	done func() <-chan struct{},
-) bool {
-	return r.waitNWithCancellation(input, startTime, 1, deadline, done)
+// WaitNWithDetails will poll [Limiter.AllowN] for a period of time,
+// until it is cancelled by the passed context. It has the
+// effect of adding latency to requests instead of refusing
+// them immediately. Consider it graceful degradation.
+//
+// WaitNWithDetails will return true if `n` tokens become available prior to
+// the context cancellation, and will consume `n` tokens. If not,
+// it will return false, and therefore consume no tokens. It will
+// also return the details of the request.
+//
+// Take care to create an appropriate context. You almost certainly
+// want [context.WithTimeout] or [context.WithDeadline].
+//
+// You should be conservative, as WaitWithDetails will introduce
+// backpressure on your upstream systems -- connections
+// may be held open longer, requests may queue in memory.
+//
+// A good starting place will be to timeout after waiting
+// for one token. For example:
+//
+//	ctx := context.WithTimeout(ctx, limit.DurationPerToken())
+//
+// The returned error will be non-nil if the context is cancelled.
+//
+// WaitWithDetails makes no ordering guarantees. Multiple concurrent calls may
+// acquire tokens in any order.
+func (r *Limiter[TInput, TKey]) WaitNWithDetails(ctx context.Context, input TInput, n int64) (bool, Details[TInput, TKey], error) {
+	return r.waitNWithDetails(ctx, input, ntime.Now(), n)
 }
 
-// waitWithCancellation is a more testable version of wait that accepts
-// deadline and done functions instead of a context, allowing for deterministic testing.
-func (r *Limiter[TInput, TKey]) waitNWithCancellation(
+// waitNWithDetails is the internal implementation that accepts a context.
+// It is designed to be testable by accepting any context implementation,
+// including test contexts that provide deterministic behavior.
+func (r *Limiter[TInput, TKey]) waitNWithDetails(
+	ctx context.Context,
 	input TInput,
 	startTime ntime.Time,
 	n int64,
-	deadline func() (time.Time, bool),
-	done func() <-chan struct{},
-) bool {
-	if r.allowN(input, startTime, n) {
-		return true
-	}
-
-	// currentTime is an approximation of the real clock moving forward
-	// it's imprecise because it depends on time.After below.
-	// For testing purposes, we want startTime (execution time) to
-	// be a parameter.
+) (bool, Details[TInput, TKey], error) {
+	// currentTime is an approximation of the real clock moving forward.
+	// It's imprecise because it depends on time.After below.
 	currentTime := startTime
 
-	userKey := r.keyFunc(input)
-	waiter := r.getWaiter(userKey)
-
-	// Ensure cleanup happens when this waiter exits
-	defer func() {
-		// If no more waiters for this key, remove the entry to prevent memory leak,
-		// reference-counted
-		if waiter.decrement() == 0 {
-			r.waiters.delete(userKey)
-		}
-	}()
-
 	for {
-		// The goroutine at the front of the queue gets to try for a token first.
-		waiter.mu.Lock()
 		allow, details := r.allowNWithDetails(input, currentTime, n)
-		waiter.mu.Unlock()
 		if allow {
-			return true
+			return allow, details, nil
 		}
 
 		retryAfter := details.RetryAfter()
 
-		// if we can't possibly get a token, fail fast
-		if deadline, ok := deadline(); ok {
-			if deadline.Before(currentTime.Add(retryAfter).ToTime()) {
-				return false
-			}
-		}
-
 		select {
-		case <-done():
-			return false
+		case <-ctx.Done():
+			// Need to get updated details, since this cancellation
+			// event might have been a while after the last call.
+			// We'll choose the semantics of "cancellation always
+			// means deny".
+			_, details := r.peekNWithDetails(input, currentTime, n)
+			return false, details, ctx.Err()
 		case <-time.After(retryAfter):
 			currentTime = currentTime.Add(retryAfter)
 		}
 	}
 }
 
-// waiter represents a reservation queue for a specific key with reference counting
-type waiter struct {
-	mu    sync.Mutex
-	count int64 // number of active waiters for this key
+// WaitWithDebug will poll [Limiter.Allow] for a period of time,
+// until it is cancelled by the passed context. It has the
+// effect of adding latency to requests instead of refusing
+// them immediately. Consider it graceful degradation.
+//
+// WaitWithDebug will return true if a token becomes available prior to
+// the context cancellation, and will consume a token. It will
+// return false if not, and therefore not consume a token. It will
+// also return the debugs of the request.
+//
+// Take care to create an appropriate context. You almost certainly
+// want [context.WithTimeout] or [context.WithDeadline].
+//
+// The returned error will be non-nil if the context is cancelled.
+//
+// WaitNWithDebug makes no ordering guarantees. Multiple concurrent calls may
+// acquire tokens in any order.
+func (r *Limiter[TInput, TKey]) WaitWithDebug(
+	ctx context.Context,
+	input TInput,
+) (bool, []Debug[TInput, TKey], error) {
+	return r.WaitNWithDebug(ctx, input, 1)
 }
 
-// increment atomically increments the waiter count and returns the new count
-func (w *waiter) increment() int64 {
-	return atomic.AddInt64(&w.count, 1)
+// WaitNWithDebug will poll [Limiter.AllowN] for a period of time,
+// until it is cancelled by the passed context. It has the
+// effect of adding latency to requests instead of refusing
+// them immediately. Consider it graceful degradation.
+//
+// WaitNWithDebug will return true if `n` tokens become available prior to
+// the context cancellation, and will consume `n` tokens. If not,
+// it will return false, and therefore consume no tokens. It will
+// also return the debugs of the request.
+//
+// Take care to create an appropriate context. You almost certainly
+// want [context.WithTimeout] or [context.WithDeadline].
+//
+// The returned error will be non-nil if the context is cancelled.
+//
+// WaitNWithDebug makes no ordering guarantees. Multiple concurrent calls may
+// acquire tokens in any order.
+func (r *Limiter[TInput, TKey]) WaitNWithDebug(
+	ctx context.Context,
+	input TInput,
+	n int64,
+) (bool, []Debug[TInput, TKey], error) {
+	return r.waitNWithDebug(ctx, input, ntime.Now(), n)
 }
 
-// decrement atomically decrements the waiter count and returns the new count
-func (w *waiter) decrement() int64 {
-	return atomic.AddInt64(&w.count, -1)
-}
+// waitNWithDetails is the internal implementation that accepts a context.
+// It is designed to be testable by accepting any context implementation,
+// including test contexts that provide deterministic behavior.
+func (r *Limiter[TInput, TKey]) waitNWithDebug(
+	ctx context.Context,
+	input TInput,
+	startTime ntime.Time,
+	n int64,
+) (bool, []Debug[TInput, TKey], error) {
+	// currentTime is an approximation of the real clock moving forward.
+	// It's imprecise because it depends on time.After below.
+	currentTime := startTime
 
-func getWaiter() *waiter {
-	return &waiter{}
-}
+	for {
+		allow, debugs := r.allowNWithDebug(input, currentTime, n)
+		if allow {
+			return allow, debugs, nil
+		}
+
+		// len(debugs) will be greater than 0 here;
+		// if there were zero, allow was true
+
+		// Get the max retryAfter
+		retryAfter := debugs[0].RetryAfter()
+		for i := 1; i < len(debugs); i++ {
+			debug := debugs[i]
+			ra := debug.RetryAfter()
+			if ra > retryAfter {
+				retryAfter = ra
+			}
+		}
 
-// getWaiter atomically gets or creates a waiter entry and increments its reference count
-func (r *Limiter[TInput, TKey]) getWaiter(key TKey) *waiter {
-	waiter := r.waiters.loadOrStore(key, getWaiter)
-	waiter.increment()
-	return waiter
+		select {
+		case <-ctx.Done():
+			// Need to get updated debugs, since this cancellation
+			// event might have been a while after the last call.
+			// We'll choose the semantics of "cancellation always
+			// means deny".
+			_, debugs := r.peekNWithDebug(input, currentTime, n)
+			return false, debugs, ctx.Err()
+		case <-time.After(retryAfter):
+			currentTime = currentTime.Add(retryAfter)
+		}
+	}
 }