Sync rate limit resets for daily/monthly

pkg/ratelimit/rate_limit.go

@@ -18,9 +18,7 @@ const (
 	periodMonthly = "monthly"
 
 	// Time constants
-	dailyHours   = 24
-	monthlyHours = 30 * 24
-	bufferTime   = 60 // Additional expiration buffer in seconds
+	bufferTime = 60 // Additional expiration buffer in seconds
 
 	// Script responses
 	scriptAllowed = 1
@@ -69,29 +67,34 @@ type QuotaOptions struct {
 // CheckRateLimits checks all rate limits in a single call and returns remaining counts
 // Returns a RateLimitStatus with the allowed flag and remaining counts for all limits
 func (rl *RateLimit) CheckRateLimits(ctx context.Context, user uint64, config RateLimitConfig) (*RateLimitStatus, error) {
-	now := time.Now().UnixNano()
+	now := time.Now()
 	burstKey := formatKey(prefixLeakyBucket, user)
 	dailyKey := formatKey(fmt.Sprintf("%s:%s", prefixQuota, periodDaily), user)
 	monthlyKey := formatKey(fmt.Sprintf("%s:%s", prefixQuota, periodMonthly), user)
 
 	// Calculate parameters
 	burstDuration := time.Duration(config.BurstDurationSeconds) * time.Second
-	cutoff := now - burstDuration.Nanoseconds()
 	burstExpiration := int(burstDuration.Seconds()) + bufferTime
-	dailyExpiration := int(dailyHours * time.Hour.Seconds())
-	monthlyExpiration := int(monthlyHours * time.Hour.Seconds())
+
+	// Calculate time until next midnight for daily reset
+	nextMidnight := time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location())
+	dailyExpiration := int(nextMidnight.Sub(now).Seconds()) + bufferTime
+
+	// Calculate time until the first day of the next month for monthly reset
+	nextMonth := time.Date(now.Year(), now.Month()+1, 1, 0, 0, 0, 0, now.Location())
+	monthlyExpiration := int(nextMonth.Sub(now).Seconds()) + bufferTime
 
 	// Execute the Lua script
 	result, err := rl.redis.Eval(ctx, combinedRateLimitScript,
 		[]string{burstKey, dailyKey, monthlyKey},
-		now,                 // ARGV[1] - Current timestamp
-		cutoff,              // ARGV[2] - Cutoff timestamp
-		config.BurstLimit,   // ARGV[3] - Burst limit
-		burstExpiration,     // ARGV[4] - Burst expiration
-		config.DailyLimit,   // ARGV[5] - Daily limit
-		dailyExpiration,     // ARGV[6] - Daily expiration
-		config.MonthlyLimit, // ARGV[7] - Monthly limit
-		monthlyExpiration,   // ARGV[8] - Monthly expiration
+		now.UnixNano(),                     // ARGV[1] - Current timestamp
+		now.Add(-burstDuration).UnixNano(), // ARGV[2] - Cutoff timestamp
+		config.BurstLimit,                  // ARGV[3] - Burst limit
+		burstExpiration,                    // ARGV[4] - Burst expiration
+		config.DailyLimit,                  // ARGV[5] - Daily limit
+		dailyExpiration,                    // ARGV[6] - Daily expiration
+		config.MonthlyLimit,                // ARGV[7] - Monthly limit
+		monthlyExpiration,                  // ARGV[8] - Monthly expiration
 	).Result()
 
 	if err != nil {
@@ -121,19 +124,19 @@ func (rl *RateLimit) CheckRateLimits(ctx context.Context, user uint64, config Ra
 // GetRateLimitStatus retrieves the current usage status of daily and monthly rate limits
 // without consuming any of the limits.
 func (rl *RateLimit) GetRateLimitStatus(ctx context.Context, user uint64, config RateLimitConfig) (*RateLimitStatus, error) {
-	now := time.Now().UnixNano()
+	now := time.Now()
 	burstKey := formatKey(prefixLeakyBucket, user)
 	dailyKey := formatKey(fmt.Sprintf("%s:%s", prefixQuota, periodDaily), user)
 	monthlyKey := formatKey(fmt.Sprintf("%s:%s", prefixQuota, periodMonthly), user)
 
 	// Calculate parameters
 	burstDuration := time.Duration(config.BurstDurationSeconds) * time.Second
-	cutoff := now - burstDuration.Nanoseconds()
+	cutoff := now.Add(-burstDuration).UnixNano()
 
 	// Execute a read-only Lua script that doesn't modify any counters
 	result, err := rl.redis.Eval(ctx, readOnlyRateLimitScript,
 		[]string{burstKey, dailyKey, monthlyKey},
-		now,                 // ARGV[1] - Current timestamp
+		now.UnixNano(),      // ARGV[1] - Current timestamp
 		cutoff,              // ARGV[2] - Cutoff timestamp
 		config.BurstLimit,   // ARGV[3] - Burst limit
 		config.DailyLimit,   // ARGV[4] - Daily limit