diff --git a/go.mod b/go.mod index 52db99472..5a82cf7e0 100644 --- a/go.mod +++ b/go.mod @@ -5,9 +5,8 @@ go 1.14 require ( github.com/alicebob/miniredis/v2 v2.11.4 github.com/cespare/xxhash v1.1.0 // indirect - github.com/cncf/udpa/go v0.0.0-20200629203442-efcf912fb354 // indirect github.com/coocood/freecache v1.1.0 - github.com/envoyproxy/go-control-plane v0.9.6 + github.com/envoyproxy/go-control-plane v0.9.7 github.com/fsnotify/fsnotify v1.4.7 // indirect github.com/golang/mock v1.4.1 github.com/golang/protobuf v1.4.2 diff --git a/go.sum b/go.sum index a077d1811..b39e91b79 100644 --- a/go.sum +++ b/go.sum @@ -13,8 +13,6 @@ github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWR github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI= github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU= github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw= -github.com/cncf/udpa/go v0.0.0-20200313221541-5f7e5dd04533 h1:8wZizuKuZVu5COB7EsBYxBQz8nRcXXn5d4Gt91eJLvU= -github.com/cncf/udpa/go v0.0.0-20200313221541-5f7e5dd04533/go.mod h1:WmhPx2Nbnhtbo57+VJT5O0JRkEi1Wbu0z5j0R8u5Hbk= github.com/cncf/udpa/go v0.0.0-20200629203442-efcf912fb354 h1:9kRtNpqLHbZVO/NNxhHp2ymxFxsHOe3x2efJGn//Tas= github.com/cncf/udpa/go v0.0.0-20200629203442-efcf912fb354/go.mod h1:WmhPx2Nbnhtbo57+VJT5O0JRkEi1Wbu0z5j0R8u5Hbk= github.com/coocood/freecache v1.1.0 h1:ENiHOsWdj1BrrlPwblhbn4GdAsMymK3pZORJ+bJGAjA= @@ -24,8 +22,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= -github.com/envoyproxy/go-control-plane v0.9.6 h1:GgblEiDzxf5ajlAZY4aC8xp7DwkrGfauFNMGdB2bBv0= -github.com/envoyproxy/go-control-plane v0.9.6/go.mod h1:GFqM7v0B62MraO4PWRedIbhThr/Rf7ev6aHOOPXeaDA= +github.com/envoyproxy/go-control-plane v0.9.7 h1:EARl0OvqMoxq/UMgMSCLnXzkaXbxzskluEBlMQCJPms= +github.com/envoyproxy/go-control-plane v0.9.7/go.mod h1:cwu0lG7PUMfa9snN8LXBig5ynNVH9qI8YYLbd1fK2po= github.com/envoyproxy/protoc-gen-validate v0.1.0 h1:EQciDnbrYxy13PgWoY8AqoxGiPrpgBZ1R8UNe3ddc+A= github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I= diff --git a/src/limiter/cache_key.go b/src/limiter/cache_key.go index a06087056..a2746f6b3 100644 --- a/src/limiter/cache_key.go +++ b/src/limiter/cache_key.go @@ -8,6 +8,7 @@ import ( pb_struct "github.com/envoyproxy/go-control-plane/envoy/extensions/common/ratelimit/v3" pb "github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v3" "github.com/envoyproxy/ratelimit/src/config" + "github.com/envoyproxy/ratelimit/src/utils" ) type CacheKeyGenerator struct { @@ -33,24 +34,6 @@ func isPerSecondLimit(unit pb.RateLimitResponse_RateLimit_Unit) bool { return unit == pb.RateLimitResponse_RateLimit_SECOND } -// Convert a rate limit into a time divider. -// @param unit supplies the unit to convert. -// @return the divider to use in time computations. -func UnitToDivider(unit pb.RateLimitResponse_RateLimit_Unit) int64 { - switch unit { - case pb.RateLimitResponse_RateLimit_SECOND: - return 1 - case pb.RateLimitResponse_RateLimit_MINUTE: - return 60 - case pb.RateLimitResponse_RateLimit_HOUR: - return 60 * 60 - case pb.RateLimitResponse_RateLimit_DAY: - return 60 * 60 * 24 - } - - panic("should not get here") -} - // Generate a cache key for a limit lookup. // @param domain supplies the cache key domain. // @param descriptor supplies the descriptor to generate the key for. @@ -81,7 +64,7 @@ func (this *CacheKeyGenerator) GenerateCacheKey( b.WriteByte('_') } - divider := UnitToDivider(limit.Limit.Unit) + divider := utils.UnitToDivider(limit.Limit.Unit) b.WriteString(strconv.FormatInt((now/divider)*divider, 10)) return CacheKey{ diff --git a/src/redis/cache_impl.go b/src/redis/cache_impl.go index c76ef067c..45186e87d 100644 --- a/src/redis/cache_impl.go +++ b/src/redis/cache_impl.go @@ -11,6 +11,8 @@ import ( "github.com/envoyproxy/ratelimit/src/limiter" "github.com/envoyproxy/ratelimit/src/server" "github.com/envoyproxy/ratelimit/src/settings" + "github.com/envoyproxy/ratelimit/src/utils" + "github.com/golang/protobuf/ptypes/duration" logger "github.com/sirupsen/logrus" "golang.org/x/net/context" ) @@ -89,7 +91,7 @@ func (this *rateLimitCacheImpl) DoLimit( logger.Debugf("looking up cache key: %s", cacheKey.Key) - expirationSeconds := limiter.UnitToDivider(limits[i].Limit.Unit) + expirationSeconds := utils.UnitToDivider(limits[i].Limit.Unit) if this.expirationJitterMaxSeconds > 0 { expirationSeconds += this.jitterRand.Int63n(this.expirationJitterMaxSeconds) } @@ -132,9 +134,10 @@ func (this *rateLimitCacheImpl) DoLimit( if isOverLimitWithLocalCache[i] { responseDescriptorStatuses[i] = &pb.RateLimitResponse_DescriptorStatus{ - Code: pb.RateLimitResponse_OVER_LIMIT, - CurrentLimit: limits[i].Limit, - LimitRemaining: 0, + Code: pb.RateLimitResponse_OVER_LIMIT, + CurrentLimit: limits[i].Limit, + LimitRemaining: 0, + DurationUntilReset: CalculateReset(limits[i].Limit, this.timeSource), } limits[i].Stats.OverLimit.Add(uint64(hitsAddend)) limits[i].Stats.OverLimitWithLocalCache.Add(uint64(hitsAddend)) @@ -152,9 +155,10 @@ func (this *rateLimitCacheImpl) DoLimit( if limitAfterIncrease > overLimitThreshold { responseDescriptorStatuses[i] = &pb.RateLimitResponse_DescriptorStatus{ - Code: pb.RateLimitResponse_OVER_LIMIT, - CurrentLimit: limits[i].Limit, - LimitRemaining: 0, + Code: pb.RateLimitResponse_OVER_LIMIT, + CurrentLimit: limits[i].Limit, + LimitRemaining: 0, + DurationUntilReset: CalculateReset(limits[i].Limit, this.timeSource), } // Increase over limit statistics. Because we support += behavior for increasing the limit, we need to @@ -179,7 +183,7 @@ func (this *rateLimitCacheImpl) DoLimit( // similar to mongo_1h, mongo_2h, etc. In the hour 1 (0h0m - 0h59m), the cache key is mongo_1h, we start // to get ratelimited in the 50th minute, the ttl of local_cache will be set as 1 hour(0h50m-1h49m). // In the time of 1h1m, since the cache key becomes different (mongo_2h), it won't get ratelimited. - err := this.localCache.Set([]byte(cacheKey.Key), []byte{}, int(limiter.UnitToDivider(limits[i].Limit.Unit))) + err := this.localCache.Set([]byte(cacheKey.Key), []byte{}, int(utils.UnitToDivider(limits[i].Limit.Unit))) if err != nil { logger.Errorf("Failing to set local cache key: %s", cacheKey.Key) } @@ -187,9 +191,10 @@ func (this *rateLimitCacheImpl) DoLimit( } else { responseDescriptorStatuses[i] = &pb.RateLimitResponse_DescriptorStatus{ - Code: pb.RateLimitResponse_OK, - CurrentLimit: limits[i].Limit, - LimitRemaining: overLimitThreshold - limitAfterIncrease, + Code: pb.RateLimitResponse_OK, + CurrentLimit: limits[i].Limit, + LimitRemaining: overLimitThreshold - limitAfterIncrease, + DurationUntilReset: CalculateReset(limits[i].Limit, this.timeSource), } // The limit is OK but we additionally want to know if we are near the limit. @@ -210,6 +215,12 @@ func (this *rateLimitCacheImpl) DoLimit( return responseDescriptorStatuses } +func CalculateReset(currentLimit *pb.RateLimitResponse_RateLimit, timeSource limiter.TimeSource) *duration.Duration { + sec := utils.UnitToDivider(currentLimit.Unit) + now := timeSource.UnixNow() + return &duration.Duration{Seconds: sec - now%sec} +} + func NewRateLimitCacheImpl(client Client, perSecondClient Client, timeSource limiter.TimeSource, jitterRand *rand.Rand, expirationJitterMaxSeconds int64, localCache *freecache.Cache) limiter.RateLimitCache { return &rateLimitCacheImpl{ client: client, diff --git a/src/utils/utilities.go b/src/utils/utilities.go new file mode 100644 index 000000000..b28619dd5 --- /dev/null +++ b/src/utils/utilities.go @@ -0,0 +1,23 @@ +package utils + +import ( + pb "github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v3" +) + +// Convert a rate limit into a time divider. +// @param unit supplies the unit to convert. +// @return the divider to use in time computations. +func UnitToDivider(unit pb.RateLimitResponse_RateLimit_Unit) int64 { + switch unit { + case pb.RateLimitResponse_RateLimit_SECOND: + return 1 + case pb.RateLimitResponse_RateLimit_MINUTE: + return 60 + case pb.RateLimitResponse_RateLimit_HOUR: + return 60 * 60 + case pb.RateLimitResponse_RateLimit_DAY: + return 60 * 60 * 24 + } + + panic("should not get here") +} diff --git a/test/integration/integration_test.go b/test/integration/integration_test.go index 0a37c5cb0..91b560120 100644 --- a/test/integration/integration_test.go +++ b/test/integration/integration_test.go @@ -3,12 +3,9 @@ package integration_test import ( - "bytes" "fmt" "io" - "io/ioutil" "math/rand" - "net/http" "os" "strconv" "testing" @@ -17,7 +14,9 @@ import ( pb_legacy "github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v2" pb "github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v3" "github.com/envoyproxy/ratelimit/src/service_cmd/runner" + "github.com/envoyproxy/ratelimit/src/utils" "github.com/envoyproxy/ratelimit/test/common" + "github.com/golang/protobuf/ptypes/duration" "github.com/stretchr/testify/assert" "golang.org/x/net/context" "google.golang.org/grpc" @@ -27,10 +26,14 @@ func newDescriptorStatus( status pb.RateLimitResponse_Code, requestsPerUnit uint32, unit pb.RateLimitResponse_RateLimit_Unit, limitRemaining uint32) *pb.RateLimitResponse_DescriptorStatus { + limit := &pb.RateLimitResponse_RateLimit{RequestsPerUnit: requestsPerUnit, Unit: unit} + sec := utils.UnitToDivider(unit) + now := time.Now().Unix() return &pb.RateLimitResponse_DescriptorStatus{ Code: status, - CurrentLimit: &pb.RateLimitResponse_RateLimit{RequestsPerUnit: requestsPerUnit, Unit: unit}, + CurrentLimit: limit, LimitRemaining: limitRemaining, + DurationUntilReset: &duration.Duration{Seconds: sec - now%sec}, } } @@ -509,34 +512,6 @@ func TestBasicConfigLegacy(t *testing.T) { response) assert.NoError(err) - json_body := []byte(`{ - "domain": "basic", - "descriptors": [ - { - "entries": [ - { - "key": "one_per_minute" - } - ] - } - ] - }`) - http_resp, _ := http.Post("http://localhost:8082/json", "application/json", bytes.NewBuffer(json_body)) - assert.Equal(http_resp.StatusCode, 200) - body, _ := ioutil.ReadAll(http_resp.Body) - http_resp.Body.Close() - assert.Equal(`{"overallCode":"OK","statuses":[{"code":"OK","currentLimit":{"requestsPerUnit":1,"unit":"MINUTE"}}]}`, string(body)) - - http_resp, _ = http.Post("http://localhost:8082/json", "application/json", bytes.NewBuffer(json_body)) - assert.Equal(http_resp.StatusCode, 429) - body, _ = ioutil.ReadAll(http_resp.Body) - http_resp.Body.Close() - assert.Equal(`{"overallCode":"OVER_LIMIT","statuses":[{"code":"OVER_LIMIT","currentLimit":{"requestsPerUnit":1,"unit":"MINUTE"}}]}`, string(body)) - - invalid_json := []byte(`{"unclosed quote: []}`) - http_resp, _ = http.Post("http://localhost:8082/json", "application/json", bytes.NewBuffer(invalid_json)) - assert.Equal(http_resp.StatusCode, 400) - response, err = c.ShouldRateLimit( context.Background(), common.NewRateLimitRequestLegacy("basic_legacy", [][][2]string{{{"key1", "foo"}}}, 1)) diff --git a/test/redis/cache_impl_test.go b/test/redis/cache_impl_test.go index 6b9a4d631..ce300b619 100644 --- a/test/redis/cache_impl_test.go +++ b/test/redis/cache_impl_test.go @@ -49,7 +49,7 @@ func testRedis(usePerSecondRedis bool) func(*testing.T) { } statsStore := stats.NewStore(stats.NewNullSink(), false) - timeSource.EXPECT().UnixNow().Return(int64(1234)) + timeSource.EXPECT().UnixNow().Return(int64(1234)).MaxTimes(3) var clientUsed *mock_redis.MockClient if usePerSecondRedis { clientUsed = perSecondClient @@ -65,14 +65,14 @@ func testRedis(usePerSecondRedis bool) func(*testing.T) { limits := []*config.RateLimit{config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_SECOND, "key_value", statsStore)} assert.Equal( - []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 5}}, + []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 5, DurationUntilReset: redis.CalculateReset(limits[0].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(1), limits[0].Stats.TotalHits.Value()) assert.Equal(uint64(0), limits[0].Stats.OverLimit.Value()) assert.Equal(uint64(0), limits[0].Stats.NearLimit.Value()) clientUsed = client - timeSource.EXPECT().UnixNow().Return(int64(1234)) + timeSource.EXPECT().UnixNow().Return(int64(1234)).MaxTimes(3) clientUsed.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key2_value2_subkey2_subvalue2_1200", uint32(1)).SetArg(1, uint32(11)).DoAndReturn(pipeAppend) clientUsed.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "EXPIRE", "domain_key2_value2_subkey2_subvalue2_1200", int64(60)).DoAndReturn(pipeAppend) @@ -89,14 +89,14 @@ func testRedis(usePerSecondRedis bool) func(*testing.T) { config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_MINUTE, "key2_value2_subkey2_subvalue2", statsStore)} assert.Equal( []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OK, CurrentLimit: nil, LimitRemaining: 0}, - {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0}}, + {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0, DurationUntilReset: redis.CalculateReset(limits[1].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(1), limits[1].Stats.TotalHits.Value()) assert.Equal(uint64(1), limits[1].Stats.OverLimit.Value()) assert.Equal(uint64(0), limits[1].Stats.NearLimit.Value()) clientUsed = client - timeSource.EXPECT().UnixNow().Return(int64(1000000)) + timeSource.EXPECT().UnixNow().Return(int64(1000000)).MaxTimes(5) clientUsed.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key3_value3_997200", uint32(1)).SetArg(1, uint32(11)).DoAndReturn(pipeAppend) clientUsed.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "EXPIRE", "domain_key3_value3_997200", int64(3600)).DoAndReturn(pipeAppend) @@ -116,8 +116,8 @@ func testRedis(usePerSecondRedis bool) func(*testing.T) { config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_DAY, "key3_value3_subkey3_subvalue3", statsStore)} assert.Equal( []*pb.RateLimitResponse_DescriptorStatus{ - {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0}, - {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0}}, + {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0, DurationUntilReset: redis.CalculateReset(limits[0].Limit, timeSource)}, + {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[1].Limit, LimitRemaining: 0, DurationUntilReset: redis.CalculateReset(limits[1].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(1), limits[0].Stats.TotalHits.Value()) assert.Equal(uint64(1), limits[0].Stats.OverLimit.Value()) @@ -178,7 +178,7 @@ func TestOverLimitWithLocalCache(t *testing.T) { localCacheStats := limiter.NewLocalCacheStats(localCache, statsStore.Scope("localcache")) // Test Near Limit Stats. Under Near Limit Ratio - timeSource.EXPECT().UnixNow().Return(int64(1000000)) + timeSource.EXPECT().UnixNow().Return(int64(1000000)).MaxTimes(3) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key4_value4_997200", uint32(1)).SetArg(1, uint32(11)).DoAndReturn(pipeAppend) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "EXPIRE", "domain_key4_value4_997200", int64(3600)).DoAndReturn(pipeAppend) @@ -191,7 +191,7 @@ func TestOverLimitWithLocalCache(t *testing.T) { assert.Equal( []*pb.RateLimitResponse_DescriptorStatus{ - {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 4}}, + {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 4, DurationUntilReset: redis.CalculateReset(limits[0].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(1), limits[0].Stats.TotalHits.Value()) assert.Equal(uint64(0), limits[0].Stats.OverLimit.Value()) @@ -202,7 +202,7 @@ func TestOverLimitWithLocalCache(t *testing.T) { testLocalCacheStats(localCacheStats, statsStore, sink, 0, 1, 1, 0, 0) // Test Near Limit Stats. At Near Limit Ratio, still OK - timeSource.EXPECT().UnixNow().Return(int64(1000000)) + timeSource.EXPECT().UnixNow().Return(int64(1000000)).MaxTimes(3) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key4_value4_997200", uint32(1)).SetArg(1, uint32(13)).DoAndReturn(pipeAppend) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "EXPIRE", "domain_key4_value4_997200", int64(3600)).DoAndReturn(pipeAppend) @@ -210,7 +210,7 @@ func TestOverLimitWithLocalCache(t *testing.T) { assert.Equal( []*pb.RateLimitResponse_DescriptorStatus{ - {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 2}}, + {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 2, DurationUntilReset: redis.CalculateReset(limits[0].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(2), limits[0].Stats.TotalHits.Value()) assert.Equal(uint64(0), limits[0].Stats.OverLimit.Value()) @@ -221,7 +221,7 @@ func TestOverLimitWithLocalCache(t *testing.T) { testLocalCacheStats(localCacheStats, statsStore, sink, 0, 2, 2, 0, 0) // Test Over limit stats - timeSource.EXPECT().UnixNow().Return(int64(1000000)) + timeSource.EXPECT().UnixNow().Return(int64(1000000)).MaxTimes(3) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key4_value4_997200", uint32(1)).SetArg(1, uint32(16)).DoAndReturn(pipeAppend) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "EXPIRE", "domain_key4_value4_997200", int64(3600)).DoAndReturn(pipeAppend) @@ -229,7 +229,7 @@ func TestOverLimitWithLocalCache(t *testing.T) { assert.Equal( []*pb.RateLimitResponse_DescriptorStatus{ - {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0}}, + {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0, DurationUntilReset: redis.CalculateReset(limits[0].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(3), limits[0].Stats.TotalHits.Value()) assert.Equal(uint64(1), limits[0].Stats.OverLimit.Value()) @@ -240,13 +240,13 @@ func TestOverLimitWithLocalCache(t *testing.T) { testLocalCacheStats(localCacheStats, statsStore, sink, 0, 2, 3, 0, 1) // Test Over limit stats with local cache - timeSource.EXPECT().UnixNow().Return(int64(1000000)) + timeSource.EXPECT().UnixNow().Return(int64(1000000)).MaxTimes(3) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key4_value4_997200", uint32(1)).Times(0) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "EXPIRE", "domain_key4_value4_997200", int64(3600)).Times(0) assert.Equal( []*pb.RateLimitResponse_DescriptorStatus{ - {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0}}, + {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0, DurationUntilReset: redis.CalculateReset(limits[0].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(4), limits[0].Stats.TotalHits.Value()) assert.Equal(uint64(2), limits[0].Stats.OverLimit.Value()) @@ -268,7 +268,7 @@ func TestNearLimit(t *testing.T) { statsStore := stats.NewStore(stats.NewNullSink(), false) // Test Near Limit Stats. Under Near Limit Ratio - timeSource.EXPECT().UnixNow().Return(int64(1000000)) + timeSource.EXPECT().UnixNow().Return(int64(1000000)).MaxTimes(3) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key4_value4_997200", uint32(1)).SetArg(1, uint32(11)).DoAndReturn(pipeAppend) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "EXPIRE", "domain_key4_value4_997200", int64(3600)).DoAndReturn(pipeAppend) @@ -281,14 +281,14 @@ func TestNearLimit(t *testing.T) { assert.Equal( []*pb.RateLimitResponse_DescriptorStatus{ - {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 4}}, + {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 4, DurationUntilReset: redis.CalculateReset(limits[0].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(1), limits[0].Stats.TotalHits.Value()) assert.Equal(uint64(0), limits[0].Stats.OverLimit.Value()) assert.Equal(uint64(0), limits[0].Stats.NearLimit.Value()) // Test Near Limit Stats. At Near Limit Ratio, still OK - timeSource.EXPECT().UnixNow().Return(int64(1000000)) + timeSource.EXPECT().UnixNow().Return(int64(1000000)).MaxTimes(3) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key4_value4_997200", uint32(1)).SetArg(1, uint32(13)).DoAndReturn(pipeAppend) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "EXPIRE", "domain_key4_value4_997200", int64(3600)).DoAndReturn(pipeAppend) @@ -296,7 +296,7 @@ func TestNearLimit(t *testing.T) { assert.Equal( []*pb.RateLimitResponse_DescriptorStatus{ - {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 2}}, + {Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 2, DurationUntilReset: redis.CalculateReset(limits[0].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(2), limits[0].Stats.TotalHits.Value()) assert.Equal(uint64(0), limits[0].Stats.OverLimit.Value()) @@ -304,7 +304,7 @@ func TestNearLimit(t *testing.T) { // Test Near Limit Stats. We went OVER_LIMIT, but the near_limit counter only increases // when we are near limit, not after we have passed the limit. - timeSource.EXPECT().UnixNow().Return(int64(1000000)) + timeSource.EXPECT().UnixNow().Return(int64(1000000)).MaxTimes(3) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key4_value4_997200", uint32(1)).SetArg(1, uint32(16)).DoAndReturn(pipeAppend) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "EXPIRE", "domain_key4_value4_997200", int64(3600)).DoAndReturn(pipeAppend) @@ -312,7 +312,7 @@ func TestNearLimit(t *testing.T) { assert.Equal( []*pb.RateLimitResponse_DescriptorStatus{ - {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0}}, + {Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0, DurationUntilReset: redis.CalculateReset(limits[0].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(3), limits[0].Stats.TotalHits.Value()) assert.Equal(uint64(1), limits[0].Stats.OverLimit.Value()) @@ -320,7 +320,7 @@ func TestNearLimit(t *testing.T) { // Now test hitsAddend that is greater than 1 // All of it under limit, under near limit - timeSource.EXPECT().UnixNow().Return(int64(1234)) + timeSource.EXPECT().UnixNow().Return(int64(1234)).MaxTimes(3) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key5_value5_1234", uint32(3)).SetArg(1, uint32(5)).DoAndReturn(pipeAppend) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "EXPIRE", "domain_key5_value5_1234", int64(1)).DoAndReturn(pipeAppend) client.EXPECT().PipeDo(gomock.Any()).Return(nil) @@ -329,14 +329,14 @@ func TestNearLimit(t *testing.T) { limits = []*config.RateLimit{config.NewRateLimit(20, pb.RateLimitResponse_RateLimit_SECOND, "key5_value5", statsStore)} assert.Equal( - []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 15}}, + []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 15, DurationUntilReset: redis.CalculateReset(limits[0].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(3), limits[0].Stats.TotalHits.Value()) assert.Equal(uint64(0), limits[0].Stats.OverLimit.Value()) assert.Equal(uint64(0), limits[0].Stats.NearLimit.Value()) // All of it under limit, some over near limit - timeSource.EXPECT().UnixNow().Return(int64(1234)) + timeSource.EXPECT().UnixNow().Return(int64(1234)).MaxTimes(3) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key6_value6_1234", uint32(2)).SetArg(1, uint32(7)).DoAndReturn(pipeAppend) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "EXPIRE", "domain_key6_value6_1234", int64(1)).DoAndReturn(pipeAppend) client.EXPECT().PipeDo(gomock.Any()).Return(nil) @@ -345,14 +345,14 @@ func TestNearLimit(t *testing.T) { limits = []*config.RateLimit{config.NewRateLimit(8, pb.RateLimitResponse_RateLimit_SECOND, "key6_value6", statsStore)} assert.Equal( - []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 1}}, + []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 1, DurationUntilReset: redis.CalculateReset(limits[0].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(2), limits[0].Stats.TotalHits.Value()) assert.Equal(uint64(0), limits[0].Stats.OverLimit.Value()) assert.Equal(uint64(1), limits[0].Stats.NearLimit.Value()) // All of it under limit, all of it over near limit - timeSource.EXPECT().UnixNow().Return(int64(1234)) + timeSource.EXPECT().UnixNow().Return(int64(1234)).MaxTimes(3) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key7_value7_1234", uint32(3)).SetArg(1, uint32(19)).DoAndReturn(pipeAppend) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "EXPIRE", "domain_key7_value7_1234", int64(1)).DoAndReturn(pipeAppend) client.EXPECT().PipeDo(gomock.Any()).Return(nil) @@ -361,14 +361,14 @@ func TestNearLimit(t *testing.T) { limits = []*config.RateLimit{config.NewRateLimit(20, pb.RateLimitResponse_RateLimit_SECOND, "key7_value7", statsStore)} assert.Equal( - []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 1}}, + []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 1, DurationUntilReset: redis.CalculateReset(limits[0].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(3), limits[0].Stats.TotalHits.Value()) assert.Equal(uint64(0), limits[0].Stats.OverLimit.Value()) assert.Equal(uint64(3), limits[0].Stats.NearLimit.Value()) // Some of it over limit, all of it over near limit - timeSource.EXPECT().UnixNow().Return(int64(1234)) + timeSource.EXPECT().UnixNow().Return(int64(1234)).MaxTimes(3) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key8_value8_1234", uint32(3)).SetArg(1, uint32(22)).DoAndReturn(pipeAppend) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "EXPIRE", "domain_key8_value8_1234", int64(1)).DoAndReturn(pipeAppend) client.EXPECT().PipeDo(gomock.Any()).Return(nil) @@ -377,14 +377,14 @@ func TestNearLimit(t *testing.T) { limits = []*config.RateLimit{config.NewRateLimit(20, pb.RateLimitResponse_RateLimit_SECOND, "key8_value8", statsStore)} assert.Equal( - []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0}}, + []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0, DurationUntilReset: redis.CalculateReset(limits[0].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(3), limits[0].Stats.TotalHits.Value()) assert.Equal(uint64(2), limits[0].Stats.OverLimit.Value()) assert.Equal(uint64(1), limits[0].Stats.NearLimit.Value()) // Some of it in all three places - timeSource.EXPECT().UnixNow().Return(int64(1234)) + timeSource.EXPECT().UnixNow().Return(int64(1234)).MaxTimes(3) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key9_value9_1234", uint32(7)).SetArg(1, uint32(22)).DoAndReturn(pipeAppend) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "EXPIRE", "domain_key9_value9_1234", int64(1)).DoAndReturn(pipeAppend) client.EXPECT().PipeDo(gomock.Any()).Return(nil) @@ -393,14 +393,14 @@ func TestNearLimit(t *testing.T) { limits = []*config.RateLimit{config.NewRateLimit(20, pb.RateLimitResponse_RateLimit_SECOND, "key9_value9", statsStore)} assert.Equal( - []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0}}, + []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0, DurationUntilReset: redis.CalculateReset(limits[0].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(7), limits[0].Stats.TotalHits.Value()) assert.Equal(uint64(2), limits[0].Stats.OverLimit.Value()) assert.Equal(uint64(4), limits[0].Stats.NearLimit.Value()) // all of it over limit - timeSource.EXPECT().UnixNow().Return(int64(1234)) + timeSource.EXPECT().UnixNow().Return(int64(1234)).MaxTimes(3) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key10_value10_1234", uint32(3)).SetArg(1, uint32(30)).DoAndReturn(pipeAppend) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "EXPIRE", "domain_key10_value10_1234", int64(1)).DoAndReturn(pipeAppend) client.EXPECT().PipeDo(gomock.Any()).Return(nil) @@ -409,7 +409,7 @@ func TestNearLimit(t *testing.T) { limits = []*config.RateLimit{config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_SECOND, "key10_value10", statsStore)} assert.Equal( - []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0}}, + []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OVER_LIMIT, CurrentLimit: limits[0].Limit, LimitRemaining: 0, DurationUntilReset: redis.CalculateReset(limits[0].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(3), limits[0].Stats.TotalHits.Value()) assert.Equal(uint64(3), limits[0].Stats.OverLimit.Value()) @@ -427,7 +427,7 @@ func TestRedisWithJitter(t *testing.T) { cache := redis.NewRateLimitCacheImpl(client, nil, timeSource, rand.New(jitterSource), 3600, nil) statsStore := stats.NewStore(stats.NewNullSink(), false) - timeSource.EXPECT().UnixNow().Return(int64(1234)) + timeSource.EXPECT().UnixNow().Return(int64(1234)).MaxTimes(3) jitterSource.EXPECT().Int63().Return(int64(100)) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "INCRBY", "domain_key_value_1234", uint32(1)).SetArg(1, uint32(5)).DoAndReturn(pipeAppend) client.EXPECT().PipeAppend(gomock.Any(), gomock.Any(), "EXPIRE", "domain_key_value_1234", int64(101)).DoAndReturn(pipeAppend) @@ -437,7 +437,7 @@ func TestRedisWithJitter(t *testing.T) { limits := []*config.RateLimit{config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_SECOND, "key_value", statsStore)} assert.Equal( - []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 5}}, + []*pb.RateLimitResponse_DescriptorStatus{{Code: pb.RateLimitResponse_OK, CurrentLimit: limits[0].Limit, LimitRemaining: 5, DurationUntilReset: redis.CalculateReset(limits[0].Limit, timeSource)}}, cache.DoLimit(nil, request, limits)) assert.Equal(uint64(1), limits[0].Stats.TotalHits.Value()) assert.Equal(uint64(0), limits[0].Stats.OverLimit.Value())