Add RetryAfter and ResetAfter to AllowAtMostN

Author: vmihailenco

lua.go

@@ -45,17 +45,21 @@ local allow_at = new_tat - burst_offset
 local diff = now - allow_at
 local remaining = math.floor(diff / emission_interval + 0.5)
 
-if remaining >= 0 then
-  local reset_after = new_tat - now
-  redis.call("SET", rate_limit_key, new_tat, "EX", math.ceil(reset_after))
-  local retry_after = -1
-  return {cost, remaining, tostring(retry_after), tostring(reset_after)}
+if remaining < 0 then
+  local reset_after = tat - now
+  local retry_after = diff * -1
+  return {
+    0, -- allowed
+    0, -- remaining
+    tostring(retry_after),
+    tostring(reset_after),
+  }
 end
 
-remaining = 0
-local reset_after = tat - now
-local retry_after = diff * -1
-return {0, remaining, tostring(retry_after), tostring(reset_after)}
+local reset_after = new_tat - now
+redis.call("SET", rate_limit_key, new_tat, "EX", math.ceil(reset_after))
+local retry_after = -1
+return {cost, remaining, tostring(retry_after), tostring(reset_after)}
 `)
 
 var allowAtMost = redis.NewScript(`
@@ -96,9 +100,13 @@ local diff = now - (tat - burst_offset)
 local remaining = math.floor(diff / emission_interval + 0.5)
 
 if remaining == 0 then
+  local reset_after = tat - now
+  local retry_after = emission_interval - diff
   return {
     0, -- allowed
     0, -- remaining
+    tostring(retry_after),
+    tostring(reset_after),
   }
 end
 
@@ -118,5 +126,7 @@ redis.call("SET", rate_limit_key, new_tat, "EX", math.ceil(reset_after))
 return {
   cost,
   remaining,
+  tostring(-1),
+  tostring(reset_after),
 }
 `)

rate.go

@@ -102,7 +102,7 @@ func (l *Limiter) AllowN(
 }
 
 // AllowAtMostN reports whether at most n events may happen at time now.
-// It returns number of allowed events. RetryAfter and ResetAfter are not set.
+// It returns number of allowed events that is less than or equal to n.
 func (l *Limiter) AllowAtMostN(
 	ctx context.Context,
 	key string,
@@ -117,10 +117,22 @@ func (l *Limiter) AllowAtMostN(
 
 	values = v.([]interface{})
 
+	retryAfter, err := strconv.ParseFloat(values[2].(string), 64)
+	if err != nil {
+		return nil, err
+	}
+
+	resetAfter, err := strconv.ParseFloat(values[3].(string), 64)
+	if err != nil {
+		return nil, err
+	}
+
 	res := &Result{
-		Limit:     limit,
-		Allowed:   int(values[0].(int64)),
-		Remaining: int(values[1].(int64)),
+		Limit:      limit,
+		Allowed:    int(values[0].(int64)),
+		Remaining:  int(values[1].(int64)),
+		RetryAfter: dur(retryAfter),
+		ResetAfter: dur(resetAfter),
 	}
 	return res, nil
 }

rate_test.go

@@ -73,6 +73,8 @@ func TestAllowAtMostN(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, res.Allowed, 2)
 	assert.Equal(t, res.Remaining, 7)
+	assert.Equal(t, res.RetryAfter, time.Duration(-1))
+	assert.InDelta(t, res.ResetAfter, 300*time.Millisecond, float64(10*time.Millisecond))
 
 	res, err = l.AllowN(ctx, "test_id", limit, 0)
 	assert.Nil(t, err)
@@ -85,6 +87,8 @@ func TestAllowAtMostN(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, res.Allowed, 7)
 	assert.Equal(t, res.Remaining, 0)
+	assert.Equal(t, res.RetryAfter, time.Duration(-1))
+	assert.InDelta(t, res.ResetAfter, 999*time.Millisecond, float64(10*time.Millisecond))
 
 	res, err = l.AllowN(ctx, "test_id", limit, 0)
 	assert.Nil(t, err)
@@ -97,6 +101,8 @@ func TestAllowAtMostN(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, res.Allowed, 0)
 	assert.Equal(t, res.Remaining, 0)
+	assert.InDelta(t, res.RetryAfter, 99*time.Millisecond, float64(10*time.Millisecond))
+	assert.InDelta(t, res.ResetAfter, 999*time.Millisecond, float64(10*time.Millisecond))
 
 	res, err = l.AllowN(ctx, "test_id", limit, 1000)
 	assert.Nil(t, err)