refactor & add base window testing

Signed-off-by: zufardhiyaulhaq <[email protected]>

README.md

@@ -190,7 +190,7 @@ For a limit of 60 requests per hour, there can only 60 requests in a single time
 Fixed window algorithm does not care when did the request arrive, all 60 can arrive at 01:01 or 01:50 and the limit will still reset at 02:00.
 
 2. Rolling window
-For a limit of 60 requests per hour. Initially it is able to take a burst of 60 requests at once, then the limit is restored by 1 each minute. Requests are allowed as long as there's still some available limit.
+For a limit of 60 requests per hour. Initially rate limiter can take a burst of 60 requests at once, then the limit is restored by 1 each minute. Requests are allowed as long as there's still some available limit.
 
 Configure rate limit algorithm with `RATE_LIMIT_ALGORITHM` environment variable.
 Use `FIXED_WINDOW` and `ROLLING_WINDOW` respectively.

src/algorithm/base_window.go

@@ -39,7 +39,7 @@ func (w *WindowImpl) GetResponseDescriptorStatus(key string, limit *config.RateL
 	isOverLimit, limitRemaining, durationUntilReset := w.algorithm.IsOverLimit(limit, int64(results), hitsAddend)
 
 	if !isOverLimit {
-		duration := w.algorithm.CalculateReset(true, limit, w.timeSource)
+		duration := w.algorithm.CalculateReset(isOverLimit, limit, w.timeSource)
 		return &pb.RateLimitResponse_DescriptorStatus{
 			Code:               pb.RateLimitResponse_OK,
 			CurrentLimit:       limit.Limit,
@@ -55,7 +55,7 @@ func (w *WindowImpl) GetResponseDescriptorStatus(key string, limit *config.RateL
 				logger.Errorf("Failing to set local cache key: %s", key)
 			}
 		}
-		duration := w.algorithm.CalculateReset(false, limit, w.timeSource)
+		duration := w.algorithm.CalculateReset(isOverLimit, limit, w.timeSource)
 		return &pb.RateLimitResponse_DescriptorStatus{
 			Code:               pb.RateLimitResponse_OVER_LIMIT,
 			CurrentLimit:       limit.Limit,
@@ -80,12 +80,6 @@ func (w *WindowImpl) GenerateCacheKeys(request *pb.RateLimitRequest,
 	return w.cacheKeyGenerator.GenerateCacheKeys(request, limits, uint32(hitsAddend), timestamp)
 }
 
-func PopulateStats(limit *config.RateLimit, nearLimit uint64, overLimit uint64, overLimitWithLocalCache uint64) {
-	limit.Stats.NearLimit.Add(nearLimit)
-	limit.Stats.OverLimit.Add(overLimit)
-	limit.Stats.OverLimitWithLocalCache.Add(overLimitWithLocalCache)
-}
-
 func (w *WindowImpl) GetExpirationSeconds() int64 {
 	return w.algorithm.GetExpirationSeconds()
 }
@@ -94,6 +88,12 @@ func (w *WindowImpl) GetResultsAfterIncrease() int64 {
 	return w.algorithm.GetResultsAfterIncrease()
 }
 
+func PopulateStats(limit *config.RateLimit, nearLimit uint64, overLimit uint64, overLimitWithLocalCache uint64) {
+	limit.Stats.NearLimit.Add(nearLimit)
+	limit.Stats.OverLimit.Add(overLimit)
+	limit.Stats.OverLimitWithLocalCache.Add(overLimitWithLocalCache)
+}
+
 func NewWindow(algorithm RatelimitAlgorithm, cacheKeyPrefix string, localCache *freecache.Cache, timeSource utils.TimeSource) *WindowImpl {
 	return &WindowImpl{
 		algorithm:         algorithm,

src/algorithm/fixed_window.go

@@ -1,9 +1,10 @@
 package algorithm
 
 import (
-	"github.com/golang/protobuf/ptypes/duration"
 	"math"
 
+	"github.com/golang/protobuf/ptypes/duration"
+
 	"github.com/coocood/freecache"
 	"github.com/envoyproxy/ratelimit/src/config"
 	"github.com/envoyproxy/ratelimit/src/utils"
@@ -53,6 +54,14 @@ func (fw *FixedWindowImpl) GetResultsAfterIncrease() int64 {
 	return 0
 }
 
+func (fw *FixedWindowImpl) CalculateSimpleReset(limit *config.RateLimit, timeSource utils.TimeSource) *duration.Duration {
+	return utils.CalculateFixedReset(limit.Limit, timeSource)
+}
+
+func (fw *FixedWindowImpl) CalculateReset(isOverLimit bool, limit *config.RateLimit, timeSource utils.TimeSource) *duration.Duration {
+	return fw.CalculateSimpleReset(limit, timeSource)
+}
+
 func NewFixedWindowAlgorithm(timeSource utils.TimeSource, localCache *freecache.Cache, nearLimitRatio float32, cacheKeyPrefix string) *FixedWindowImpl {
 	return &FixedWindowImpl{
 		timeSource:        timeSource,
@@ -61,11 +70,3 @@ func NewFixedWindowAlgorithm(timeSource utils.TimeSource, localCache *freecache.
 		nearLimitRatio:    nearLimitRatio,
 	}
 }
-
-func (c *FixedWindowImpl) CalculateSimpleReset(limit *config.RateLimit, timeSource utils.TimeSource) *duration.Duration {
-	return utils.CalculateFixedReset(limit.Limit, timeSource)
-}
-
-func (c *FixedWindowImpl) CalculateReset(isOverLimit bool, limit *config.RateLimit, timeSource utils.TimeSource) *duration.Duration {
-	return c.CalculateSimpleReset(limit, timeSource)
-}

src/algorithm/rolling_window.go

@@ -76,25 +76,25 @@ func (rw *RollingWindowImpl) GetResultsAfterIncrease() int64 {
 	return rw.newTat
 }
 
-func NewRollingWindowAlgorithm(timeSource utils.TimeSource, localCache *freecache.Cache, nearLimitRatio float32, cacheKeyPrefix string) *RollingWindowImpl {
-	return &RollingWindowImpl{
-		timeSource:        timeSource,
-		cacheKeyGenerator: utils.NewCacheKeyGenerator(cacheKeyPrefix),
-		localCache:        localCache,
-		nearLimitRatio:    nearLimitRatio,
-	}
-}
-
 func (rw *RollingWindowImpl) CalculateSimpleReset(limit *config.RateLimit, timeSource utils.TimeSource) *duration.Duration {
 	secondsToReset := utils.UnitToDivider(limit.Limit.Unit)
 	secondsToReset -= utils.NanosecondsToSeconds(timeSource.UnixNanoNow()) % secondsToReset
 	return &duration.Duration{Seconds: secondsToReset}
 }
 
 func (rw *RollingWindowImpl) CalculateReset(isOverLimit bool, limit *config.RateLimit, timeSource utils.TimeSource) *duration.Duration {
-	if isOverLimit {
+	if !isOverLimit {
 		return utils.NanosecondsToDuration(rw.newTat - rw.arrivedAt)
 	} else {
 		return utils.NanosecondsToDuration(int64(math.Ceil(float64(rw.tat - rw.arrivedAt))))
 	}
 }
+
+func NewRollingWindowAlgorithm(timeSource utils.TimeSource, localCache *freecache.Cache, nearLimitRatio float32, cacheKeyPrefix string) *RollingWindowImpl {
+	return &RollingWindowImpl{
+		timeSource:        timeSource,
+		cacheKeyGenerator: utils.NewCacheKeyGenerator(cacheKeyPrefix),
+		localCache:        localCache,
+		nearLimitRatio:    nearLimitRatio,
+	}
+}

test/algorithm/base_window_test.go

@@ -0,0 +1,157 @@
+package algorithm
+
+import (
+	"testing"
+
+	"github.com/coocood/freecache"
+	pb "github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v3"
+	"github.com/envoyproxy/ratelimit/src/algorithm"
+	"github.com/envoyproxy/ratelimit/src/config"
+	"github.com/envoyproxy/ratelimit/src/utils"
+	"github.com/envoyproxy/ratelimit/test/common"
+	mock_utils "github.com/envoyproxy/ratelimit/test/mocks/utils"
+	"github.com/golang/mock/gomock"
+	"github.com/golang/protobuf/ptypes/duration"
+	stats "github.com/lyft/gostats"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetResponseDescriptorStatus(t *testing.T) {
+	assert := assert.New(t)
+	controller := gomock.NewController(t)
+	defer controller.Finish()
+
+	timeSource := mock_utils.NewMockTimeSource(controller)
+	statsStore := stats.NewStore(stats.NewNullSink(), false)
+
+	// Fixed Window algorithm
+	fixedAlgorithm := algorithm.NewFixedWindowAlgorithm(timeSource, nil, 0.8, "")
+	baseAlgorithm := algorithm.NewWindow(fixedAlgorithm, "", nil, timeSource)
+
+	key := "key_value"
+	limit := config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_SECOND, "key_value", statsStore)
+	var results int64 = 1
+	var hitsAddend int64 = 1
+	isOverLimitWithLocalCache := false
+
+	timeSource.EXPECT().UnixNow().Return(int64(1)).MaxTimes(2)
+
+	expectedResult := &pb.RateLimitResponse_DescriptorStatus{
+		Code:               pb.RateLimitResponse_OK,
+		CurrentLimit:       limit.Limit,
+		LimitRemaining:     9,
+		DurationUntilReset: utils.CalculateFixedReset(limit.Limit, timeSource)}
+
+	actualResult := baseAlgorithm.GetResponseDescriptorStatus(key, limit, results, isOverLimitWithLocalCache, hitsAddend)
+	assert.Equal(expectedResult, actualResult)
+
+	// Rolling Window algorithm
+	rollingAlgorithm := algorithm.NewRollingWindowAlgorithm(timeSource, nil, 0.8, "")
+	baseAlgorithm = algorithm.NewWindow(rollingAlgorithm, "", nil, timeSource)
+
+	key = "key_value"
+	limit = config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_SECOND, "key_value", statsStore)
+	results = 0
+	hitsAddend = 1
+	isOverLimitWithLocalCache = false
+
+	timeSource.EXPECT().UnixNanoNow().Return(int64(1e9)).MaxTimes(1)
+
+	expectedResult = &pb.RateLimitResponse_DescriptorStatus{
+		Code:               pb.RateLimitResponse_OK,
+		CurrentLimit:       limit.Limit,
+		LimitRemaining:     9,
+		DurationUntilReset: &duration.Duration{Nanos: 1e8}}
+
+	actualResult = baseAlgorithm.GetResponseDescriptorStatus(key, limit, results, isOverLimitWithLocalCache, hitsAddend)
+	assert.Equal(expectedResult, actualResult)
+}
+
+func TestIsOverLimitWithLocalCache(t *testing.T) {
+	assert := assert.New(t)
+	controller := gomock.NewController(t)
+	defer controller.Finish()
+
+	key := "key_value"
+
+	timeSource := mock_utils.NewMockTimeSource(controller)
+
+	// Fixed Window algorithm
+	fixedLocalCache := freecache.NewCache(100)
+
+	fixedAlgorithm := algorithm.NewFixedWindowAlgorithm(timeSource, fixedLocalCache, 0.8, "")
+	baseAlgorithm := algorithm.NewWindow(fixedAlgorithm, "", fixedLocalCache, timeSource)
+
+	assert.Equal(false, baseAlgorithm.IsOverLimitWithLocalCache(key))
+
+	fixedLocalCache.Set([]byte(key), []byte{}, 1)
+	assert.Equal(true, baseAlgorithm.IsOverLimitWithLocalCache(key))
+
+	// Rolling Window algorithm
+	rollingLocalCache := freecache.NewCache(100)
+
+	rollingAlgorithm := algorithm.NewRollingWindowAlgorithm(timeSource, rollingLocalCache, 0.8, "")
+	baseAlgorithm = algorithm.NewWindow(rollingAlgorithm, "", rollingLocalCache, timeSource)
+
+	assert.Equal(false, baseAlgorithm.IsOverLimitWithLocalCache(key))
+
+	rollingLocalCache.Set([]byte(key), []byte{}, 1)
+	assert.Equal(true, baseAlgorithm.IsOverLimitWithLocalCache(key))
+}
+
+func TestGenerateCacheKeys(t *testing.T) {
+	assert := assert.New(t)
+	controller := gomock.NewController(t)
+	defer controller.Finish()
+
+	timeSource := mock_utils.NewMockTimeSource(controller)
+	statsStore := stats.NewStore(stats.NewNullSink(), false)
+
+	var hitsAddend int64 = 1
+	request := common.NewRateLimitRequest("domain", [][][2]string{{{"key", "value"}}}, 1)
+	limit := []*config.RateLimit{config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_SECOND, "key_value", statsStore)}
+
+	// Fixed Window algorithm
+	fixedAlgorithm := algorithm.NewFixedWindowAlgorithm(timeSource, nil, 0.8, "")
+	baseAlgorithm := algorithm.NewWindow(fixedAlgorithm, "", nil, timeSource)
+
+	timeSource.EXPECT().UnixNow().Return(int64(1)).MaxTimes(1)
+
+	expectedResult := []utils.CacheKey([]utils.CacheKey{{Key: "domain_key_value_1", PerSecond: true}})
+	actualResult := baseAlgorithm.GenerateCacheKeys(request, limit, hitsAddend, 1)
+	assert.Equal(expectedResult, actualResult)
+
+	// Rolling Window algorithm
+	rollingAlgorithm := algorithm.NewRollingWindowAlgorithm(timeSource, nil, 0.8, "")
+	baseAlgorithm = algorithm.NewWindow(rollingAlgorithm, "", nil, timeSource)
+
+	timeSource.EXPECT().UnixNanoNow().Return(int64(1e9)).MaxTimes(1)
+
+	expectedResult = []utils.CacheKey([]utils.CacheKey{{Key: "domain_key_value_0", PerSecond: true}})
+	actualResult = baseAlgorithm.GenerateCacheKeys(request, limit, hitsAddend, 0)
+	assert.Equal(expectedResult, actualResult)
+}
+
+func TestPopulateStats(t *testing.T) {
+	assert := assert.New(t)
+	controller := gomock.NewController(t)
+	defer controller.Finish()
+
+	statsStore := stats.NewStore(stats.NewNullSink(), false)
+	limit := config.NewRateLimit(10, pb.RateLimitResponse_RateLimit_SECOND, "key_value", statsStore)
+
+	algorithm.PopulateStats(limit, 1, 0, 0)
+	assert.Equal(uint64(1), limit.Stats.NearLimit.Value())
+	assert.Equal(uint64(0), limit.Stats.OverLimit.Value())
+	assert.Equal(uint64(0), limit.Stats.OverLimitWithLocalCache.Value())
+
+	algorithm.PopulateStats(limit, 0, 1, 0)
+	assert.Equal(uint64(1), limit.Stats.NearLimit.Value())
+	assert.Equal(uint64(1), limit.Stats.OverLimit.Value())
+	assert.Equal(uint64(0), limit.Stats.OverLimitWithLocalCache.Value())
+
+	algorithm.PopulateStats(limit, 0, 0, 1)
+	assert.Equal(uint64(1), limit.Stats.NearLimit.Value())
+	assert.Equal(uint64(1), limit.Stats.OverLimit.Value())
+	assert.Equal(uint64(1), limit.Stats.OverLimitWithLocalCache.Value())
+}