Skip to content
This repository has been archived by the owner on Apr 19, 2024. It is now read-only.

Commit

Permalink
Added Behavior RESET_REMAINING to reset any hits recorded in the cach…
Browse files Browse the repository at this point in the history
…e 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.
  • Loading branch information
thrawn01 committed Jan 7, 2020
1 parent e6ca0e7 commit 8ab97c3
Show file tree
Hide file tree
Showing 6 changed files with 172 additions and 64 deletions.
6 changes: 6 additions & 0 deletions CHANGELOG
Original file line number Diff line number Diff line change
Expand Up @@ -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`
Expand All @@ -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
Expand Down
18 changes: 17 additions & 1 deletion algorithms.go
Original file line number Diff line number Diff line change
Expand Up @@ -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?
Expand Down Expand Up @@ -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
Expand Down
72 changes: 72 additions & 0 deletions functional_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
117 changes: 63 additions & 54 deletions gubernator.pb.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion proto/gubernator.proto
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down
Loading

0 comments on commit 8ab97c3

Please sign in to comment.