enable allowN and allowAtMost to peek at contents (#47)

* add failing test for allowN increment of zero

* fix lua script to allow increment by zero (effectively peek)

In past versions of redis_rate, you could peek at the current state
in the redis limiter by calling allowN but not incrementing. The
move to lua script did not preserve this (unadvertised) ability.

This adds it back, the only issue was trying to SET a redis
row with an expiry of zero if the row didn't exist. Since the
caller doesn't want to increment at all, this changes the SET
to not be called if the value is absent.

* same test and fix for allowAtMost increment by zero

* add test for peeking after a write

Author: danehammer

lua.go

@@ -57,7 +57,9 @@ if remaining < 0 then
 end
 
 local reset_after = new_tat - now
-redis.call("SET", rate_limit_key, new_tat, "EX", math.ceil(reset_after))
+if reset_after > 0 then
+  redis.call("SET", rate_limit_key, new_tat, "EX", math.ceil(reset_after))
+end
 local retry_after = -1
 return {cost, remaining, tostring(retry_after), tostring(reset_after)}
 `)
@@ -121,7 +123,9 @@ local increment = emission_interval * cost
 local new_tat = tat + increment
 
 local reset_after = new_tat - now
-redis.call("SET", rate_limit_key, new_tat, "EX", math.ceil(reset_after))
+if reset_after > 0 then
+  redis.call("SET", rate_limit_key, new_tat, "EX", math.ceil(reset_after))
+end
 
 return {
   cost,

rate_test.go

@@ -68,6 +68,36 @@ func TestAllow(t *testing.T) {
 	require.InDelta(t, res.ResetAfter, 999*time.Millisecond, float64(10*time.Millisecond))
 }
 
+func TestAllowN_IncrementZero(t *testing.T) {
+	ctx := context.Background()
+	l := rateLimiter()
+	limit := redis_rate.PerSecond(10)
+
+	// Check for a row that's not there
+	res, err := l.AllowN(ctx, "test_id", limit, 0)
+	require.Nil(t, err)
+	require.Equal(t, res.Allowed, 0)
+	require.Equal(t, res.Remaining, 10)
+	require.Equal(t, res.RetryAfter, time.Duration(-1))
+	require.Equal(t, res.ResetAfter, time.Duration(0))
+
+	// Now increment it
+	res, err = l.Allow(ctx, "test_id", limit)
+	require.Nil(t, err)
+	require.Equal(t, res.Allowed, 1)
+	require.Equal(t, res.Remaining, 9)
+	require.Equal(t, res.RetryAfter, time.Duration(-1))
+	require.InDelta(t, res.ResetAfter, 100*time.Millisecond, float64(10*time.Millisecond))
+
+	// Peek again
+	res, err = l.AllowN(ctx, "test_id", limit, 0)
+	require.Nil(t, err)
+	require.Equal(t, res.Allowed, 0)
+	require.Equal(t, res.Remaining, 9)
+	require.Equal(t, res.RetryAfter, time.Duration(-1))
+	require.InDelta(t, res.ResetAfter, 100*time.Millisecond, float64(10*time.Millisecond))
+}
+
 func TestRetryAfter(t *testing.T) {
 	limit := redis_rate.Limit{
 		Rate:   1,
@@ -146,6 +176,36 @@ func TestAllowAtMost(t *testing.T) {
 	require.InDelta(t, res.ResetAfter, 999*time.Millisecond, float64(10*time.Millisecond))
 }
 
+func TestAllowAtMost_IncrementZero(t *testing.T) {
+	ctx := context.Background()
+	l := rateLimiter()
+	limit := redis_rate.PerSecond(10)
+
+	// Check for a row that isn't there
+	res, err := l.AllowAtMost(ctx, "test_id", limit, 0)
+	require.Nil(t, err)
+	require.Equal(t, res.Allowed, 0)
+	require.Equal(t, res.Remaining, 10)
+	require.Equal(t, res.RetryAfter, time.Duration(-1))
+	require.Equal(t, res.ResetAfter, time.Duration(0))
+
+	// Now increment it
+	res, err = l.Allow(ctx, "test_id", limit)
+	require.Nil(t, err)
+	require.Equal(t, res.Allowed, 1)
+	require.Equal(t, res.Remaining, 9)
+	require.Equal(t, res.RetryAfter, time.Duration(-1))
+	require.InDelta(t, res.ResetAfter, 100*time.Millisecond, float64(10*time.Millisecond))
+
+	// Peek again
+	res, err = l.AllowAtMost(ctx, "test_id", limit, 0)
+	require.Nil(t, err)
+	require.Equal(t, res.Allowed, 0)
+	require.Equal(t, res.Remaining, 9)
+	require.Equal(t, res.RetryAfter, time.Duration(-1))
+	require.InDelta(t, res.ResetAfter, 100*time.Millisecond, float64(10*time.Millisecond))
+}
+
 func BenchmarkAllow(b *testing.B) {
 	ctx := context.Background()
 	l := rateLimiter()