diff --git a/internal/ratelimit/ratelimit.go b/internal/ratelimit/ratelimit.go index 6aa333e..07c1620 100644 --- a/internal/ratelimit/ratelimit.go +++ b/internal/ratelimit/ratelimit.go @@ -77,10 +77,22 @@ func (s *service) Allow( ) (bool, int, time.Time, error) { redisKey := fmt.Sprintf("rate_limit:%s:%s", ip, key) - // Increment counter - count, err := s.redis.Incr(ctx, redisKey).Result() + // EXPIRE NX runs on every request, not only the first, so a counter that lost its TTL + // (e.g. an EXPIRE not yet replicated when Redis failed over) heals instead of growing forever. + var ( + incrCmd *redis.IntCmd + ttlCmd *redis.DurationCmd + ) + + _, err := s.redis.TxPipelined(ctx, func(pipe redis.Pipeliner) error { + incrCmd = pipe.Incr(ctx, redisKey) + pipe.ExpireNX(ctx, redisKey, window) + ttlCmd = pipe.TTL(ctx, redisKey) + + return nil + }) if err != nil { - s.log.WithError(err).Error("failed to increment rate limit counter in Redis") + s.log.WithError(err).Error("failed to update rate limit counter in Redis") // Handle failure based on configured mode if s.failureMode == "fail_closed" { @@ -91,19 +103,11 @@ func (s *service) Allow( return true, 0, time.Time{}, nil } - // Set expiry on first request - if count == 1 { - if expireErr := s.redis.Expire(ctx, redisKey, window).Err(); expireErr != nil { - s.log.WithError(expireErr).Warn("failed to set rate limit TTL") - } - } - - // Calculate reset time - ttl, err := s.redis.TTL(ctx, redisKey).Result() - if err != nil { - s.log.WithError(err).Warn("failed to get rate limit TTL") + count := incrCmd.Val() - ttl = window // Fallback + ttl := ttlCmd.Val() + if ttl <= 0 { + ttl = window } resetAt := time.Now().Add(ttl) diff --git a/internal/ratelimit/ratelimit_test.go b/internal/ratelimit/ratelimit_test.go index ada8d19..bea375f 100644 --- a/internal/ratelimit/ratelimit_test.go +++ b/internal/ratelimit/ratelimit_test.go @@ -426,3 +426,41 @@ func TestService_Allow_ConcurrentRequests(t *testing.T) { require.NoError(t, err) assert.Equal(t, int64(numGoroutines), count, "all requests should be counted") } + +// TestService_Allow_RestoresMissingTTL verifies that a counter left without a TTL (as after a +// Redis failover drops an unreplicated EXPIRE) gets one again instead of blocking forever. +func TestService_Allow_RestoresMissingTTL(t *testing.T) { + mr := miniredis.RunT(t) + defer mr.Close() + + client := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + defer client.Close() + + logger := logrus.New() + logger.SetOutput(io.Discard) + + svc := NewService(logger, client, "fail_open") + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + key := "rate_limit:10.42.25.1:health" + require.NoError(t, mr.Set(key, "106440")) + assert.Equal(t, time.Duration(0), mr.TTL(key), "precondition: counter has no TTL") + + allowed, _, resetAt, err := svc.Allow(ctx, "10.42.25.1", "health", 100000, time.Minute) + require.NoError(t, err) + assert.False(t, allowed, "counter is still over the limit for the current window") + assert.True(t, resetAt.After(time.Now()), "reset time should be in the future") + + ttl := mr.TTL(key) + assert.Greater(t, ttl, 50*time.Second, "TTL should be restored") + assert.LessOrEqual(t, ttl, time.Minute) + + mr.FastForward(time.Minute + time.Second) + + allowed, remaining, _, err := svc.Allow(ctx, "10.42.25.1", "health", 100000, time.Minute) + require.NoError(t, err) + assert.True(t, allowed, "counter should reset once the restored TTL expires") + assert.Equal(t, 99999, remaining) +}