Add Redis rate limiter

pkg/middlewares/ratelimiter/in_memory_limiter.go

@@ -0,0 +1,72 @@
+package ratelimiter
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/mailgun/ttlmap"
+	"github.com/rs/zerolog"
+	"golang.org/x/time/rate"
+)
+
+type inMemoryRateLimiter struct {
+	rate  rate.Limit // reqs/s
+	burst int64
+	// maxDelay is the maximum duration we're willing to wait for a bucket reservation to become effective, in nanoseconds.
+	// For now it is somewhat arbitrarily set to 1/(2*rate).
+	maxDelay time.Duration
+	// Each rate limiter for a given source is stored in the buckets ttlmap.
+	// To keep this ttlmap constrained in size,
+	// each ratelimiter is "garbage collected" when it is considered expired.
+	// It is considered expired after it hasn't been used for ttl seconds.
+	ttl     int
+	buckets *ttlmap.TtlMap // actual buckets, keyed by source.
+
+	logger *zerolog.Logger
+}
+
+func newInMemoryRateLimiter(rate rate.Limit, burst int64, maxDelay time.Duration, ttl int, logger *zerolog.Logger) (*inMemoryRateLimiter, error) {
+	buckets, err := ttlmap.NewConcurrent(maxSources)
+	if err != nil {
+		return nil, fmt.Errorf("creating ttlmap: %w", err)
+	}
+
+	return &inMemoryRateLimiter{
+		rate:     rate,
+		burst:    burst,
+		maxDelay: maxDelay,
+		ttl:      ttl,
+		logger:   logger,
+		buckets:  buckets,
+	}, nil
+}
+
+func (i *inMemoryRateLimiter) Allow(_ context.Context, source string) (*time.Duration, error) {
+	// Get bucket which contains limiter information.
+	var bucket *rate.Limiter
+	if rlSource, exists := i.buckets.Get(source); exists {
+		bucket = rlSource.(*rate.Limiter)
+	} else {
+		bucket = rate.NewLimiter(i.rate, int(i.burst))
+	}
+
+	// We Set even in the case where the source already exists,
+	// because we want to update the expiryTime everytime we get the source,
+	// as the expiryTime is supposed to reflect the activity (or lack thereof) on that source.
+	if err := i.buckets.Set(source, bucket, i.ttl); err != nil {
+		return nil, fmt.Errorf("setting buckets: %w", err)
+	}
+
+	res := bucket.Reserve()
+	if !res.OK() {
+		return nil, nil
+	}
+
+	delay := res.Delay()
+	if delay > i.maxDelay {
+		res.Cancel()
+	}
+
+	return &delay, nil
+}