Added Behavior RESET_REMAINING to reset any hits recorded in the cache for the specified rate limit

* Behavior is now a flag, this should be a backward compatible change for
  anyone using GLOBAL or NO_BATCHING but will break anyone using
  DURATION_IS_GREGORIAN. Use `HasBehavior()` function to check for behavior
  flags.

Author: thrawn01

CHANGELOG

@@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 * Allow cache users to invalidate a ratelimit after a specific time
 * Changing limit and duration before expire should now work correctly
+* Added Behavior RESET_REMAINING to reset any hits recorded in the cache for the
+  specified rate limit
 
 ## Changes
 * TokenBucketItem is now provided when `OnChange()` is called instead of `RateLimitResp`
@@ -17,6 +19,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Always include reset_time in leaky bucket responses
 * Fixed subtle bug during shutdown where PeerClient passed into goroutine could
   be out of scope/changed when routine runs
+* Behavior is now a flag, this should be a backward compatible change for
+  anyone using GLOBAL or NO_BATCHING but will break anyone using
+  DURATION_IS_GREGORIAN. Use `HasBehavior()` function to check for behavior
+  flags.
 
 ## [0.7.1] - 2019-12-10
 ### Added

algorithms.go

@@ -33,12 +33,24 @@ func tokenBucket(s Store, c Cache, r *RateLimitReq) (resp *RateLimitResp, err er
 	}
 
 	if ok {
+		if HasBehavior(r.Behavior, Behavior_RESET_REMAINING) {
+			c.Remove(r.HashKey())
+			if s != nil {
+				s.Remove(r.HashKey())
+			}
+			return &RateLimitResp{
+				Status:    Status_UNDER_LIMIT,
+				Limit:     r.Limit,
+				Remaining: r.Limit,
+				ResetTime: 0,
+			}, nil
+		}
+
 		// The following semantic allows for requests of more than the limit to be rejected, but subsequent
 		// requests within the same duration that are under the limit to succeed. IE: client attempts to
 		// send 1000 emails but 100 is their limit. The request is rejected as over the limit, but since we
 		// don't store OVER_LIMIT in the cache the client can retry within the same rate limit duration with
 		// 100 emails and the request will succeed.
-
 		t, ok := item.Value.(*TokenBucketItem)
 		if !ok {
 			// Client switched algorithms; perhaps due to a migration?
@@ -190,6 +202,10 @@ func leakyBucket(s Store, c Cache, r *RateLimitReq) (resp *RateLimitResp, err er
 			return leakyBucket(s, c, r)
 		}
 
+		if HasBehavior(r.Behavior, Behavior_RESET_REMAINING) {
+			b.Remaining = r.Limit
+		}
+
 		// Update limit and duration if they changed
 		b.Limit = r.Limit
 		b.Duration = r.Duration

functional_test.go

@@ -422,4 +422,76 @@ func TestChangeLimit(t *testing.T) {
 	}
 }
 
+func TestResetRemaining(t *testing.T) {
+	client, errs := guber.DialV1Server(cluster.GetPeer())
+	require.Nil(t, errs)
+
+	tests := []struct {
+		Remaining int64
+		Algorithm guber.Algorithm
+		Behavior  guber.Behavior
+		Status    guber.Status
+		Name      string
+		Limit     int64
+	}{
+		{
+			Name:      "Should subtract 1 from remaining",
+			Algorithm: guber.Algorithm_TOKEN_BUCKET,
+			Status:    guber.Status_UNDER_LIMIT,
+			Behavior:  guber.Behavior_BATCHING,
+			Remaining: 99,
+			Limit:     100,
+		},
+		{
+			Name:      "Should subtract 2 from remaining",
+			Algorithm: guber.Algorithm_TOKEN_BUCKET,
+			Status:    guber.Status_UNDER_LIMIT,
+			Behavior:  guber.Behavior_BATCHING,
+			Remaining: 98,
+			Limit:     100,
+		},
+		{
+			Name:      "Should reset the remaining",
+			Algorithm: guber.Algorithm_TOKEN_BUCKET,
+			Status:    guber.Status_UNDER_LIMIT,
+			Behavior:  guber.Behavior_RESET_REMAINING,
+			Remaining: 100,
+			Limit:     100,
+		},
+		{
+			Name:      "Should subtract 1 from remaining after reset",
+			Algorithm: guber.Algorithm_TOKEN_BUCKET,
+			Status:    guber.Status_UNDER_LIMIT,
+			Behavior:  guber.Behavior_BATCHING,
+			Remaining: 99,
+			Limit:     100,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.Name, func(t *testing.T) {
+			resp, err := client.GetRateLimits(context.Background(), &guber.GetRateLimitsReq{
+				Requests: []*guber.RateLimitReq{
+					{
+						Name:      "test_reset_remaining",
+						UniqueKey: "account:1234",
+						Algorithm: tt.Algorithm,
+						Duration:  guber.Millisecond * 100,
+						Behavior:  tt.Behavior,
+						Limit:     tt.Limit,
+						Hits:      1,
+					},
+				},
+			})
+			require.Nil(t, err)
+
+			rl := resp.Responses[0]
+
+			assert.Equal(t, tt.Status, rl.Status)
+			assert.Equal(t, tt.Remaining, rl.Remaining)
+			assert.Equal(t, tt.Limit, rl.Limit)
+		})
+	}
+}
+
 // TODO: Add a test for sending no rate limits RateLimitReqList.RateLimits = nil

gubernator.pb.go

@@ -120,7 +120,7 @@ enum Behavior {
   // ignore any `Hit` values provided in the current request. The effect this has is dependent on
   // algorithm chosen. For instance, if used with `TOKEN_BUCKET` it will immediately expire the
   // cache value. For `LEAKY_BUCKET` it sets the `Remaining` to `Limit`.
-  RESET_HITS = 8;
+  RESET_REMAINING = 8;
 
   // TODO: Add support for LOCAL. Which would force the rate limit to be handled by the local instance
 }