From b0507c7e84bc7283aff0c698b274da60f5a21ea3 Mon Sep 17 00:00:00 2001 From: George Date: Fri, 31 Mar 2023 13:11:48 +0200 Subject: [PATCH 1/3] feat: add monthly rate limit --- Makefile | 4 ++-- src/limiter/cache_key.go | 22 +++++++++++++++--- src/utils/utilities.go | 48 +++++++++++++++++++++++++++++++++++++--- 3 files changed, 66 insertions(+), 8 deletions(-) diff --git a/Makefile b/Makefile index 2f98413d..25e93dea 100644 --- a/Makefile +++ b/Makefile @@ -125,8 +125,8 @@ docker_tests: docker run $$(tty -s && echo "-it" || echo) $(INTEGRATION_IMAGE):$(VERSION) .PHONY: docker_image -docker_image: docker_tests - docker build . -t $(IMAGE):$(VERSION) +docker_image: + docker build --platform linux/amd64 . -t $(IMAGE):$(VERSION) .PHONY: docker_push docker_push: docker_image diff --git a/src/limiter/cache_key.go b/src/limiter/cache_key.go index 4aeab204..0b892b2c 100644 --- a/src/limiter/cache_key.go +++ b/src/limiter/cache_key.go @@ -7,9 +7,9 @@ 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" + logger "github.com/sirupsen/logrus" ) type CacheKeyGenerator struct { @@ -39,7 +39,7 @@ func isPerSecondLimit(unit pb.RateLimitResponse_RateLimit_Unit) bool { return unit == pb.RateLimitResponse_RateLimit_SECOND } -// Generate a cache key for a limit lookup. +// GenerateCacheKey a cache key for a limit lookup. // @param domain supplies the cache key domain. // @param descriptor supplies the descriptor to generate the key for. // @param limit supplies the rate limit to generate the key for (may be nil). @@ -71,7 +71,23 @@ func (this *CacheKeyGenerator) GenerateCacheKey( } divider := utils.UnitToDivider(limit.Limit.Unit) - b.WriteString(strconv.FormatInt((now/divider)*divider, 10)) + + // The key needs to be the same within the time unit. If we change the function, + // then we need to make sure the key is always the same within the time unit, so it + // can be picked up on the next lookup. + // This code section handles the MONTH time unit. + if limit.Limit.Unit == pb.RateLimitResponse_RateLimit_MONTH { + // get the first day of the current month as unix time + y, m, _ := utils.CurrentTime(now).Date() + first, _ := utils.MonthInterval(y, m) + + logger.Debugf("calculating cacke key for time unit: %v, the cache key is %v", limit.Limit.Unit, first.Unix()) + b.WriteString(strconv.FormatInt(first.Unix(), 10)) + } else { + // This code section handles can handle all time units except for MONTH and YEAR + logger.Debugf("calculating cache key for time unit: %v, the cache key is %v", limit.Limit.Unit, (now/divider)*divider) + b.WriteString(strconv.FormatInt((now/divider)*divider, 10)) + } return CacheKey{ Key: b.String(), diff --git a/src/utils/utilities.go b/src/utils/utilities.go index f9ecf856..feb573b5 100644 --- a/src/utils/utilities.go +++ b/src/utils/utilities.go @@ -2,9 +2,11 @@ package utils import ( "strings" + "time" pb "github.com/envoyproxy/go-control-plane/envoy/service/ratelimit/v3" "github.com/golang/protobuf/ptypes/duration" + logger "github.com/sirupsen/logrus" ) // Interface for a time source. @@ -26,15 +28,55 @@ func UnitToDivider(unit pb.RateLimitResponse_RateLimit_Unit) int64 { return 60 * 60 case pb.RateLimitResponse_RateLimit_DAY: return 60 * 60 * 24 + case pb.RateLimitResponse_RateLimit_MONTH: + // This cannot be hardcoded to 30 days, as there are months with 31 days and the TTL will expire before the end of the month + return 60 * 60 * 24 * daysOfCurrentMonth(time.Now().Unix()) //todo: get timesource instead of time.now. } - panic("should not get here") } +func daysOfCurrentMonth(unix int64) int64 { + y, m, _ := CurrentTime(unix).Date() + + return daysIn(m, y) + +} + +// daysIn returns the number of days in a month for a given year. +func daysIn(m time.Month, year int) int64 { + // This is equivalent to time.daysIn(m, year). + return int64(time.Date(year, m+1, 0, 0, 0, 0, 0, time.UTC).Day()) +} + +// CalculateReset calculates the reset time for a given unit and time source. func CalculateReset(unit *pb.RateLimitResponse_RateLimit_Unit, timeSource TimeSource) *duration.Duration { - sec := UnitToDivider(*unit) now := timeSource.UnixNow() - return &duration.Duration{Seconds: sec - now%sec} + unitInSec := UnitToDivider(*unit) + + if *unit == pb.RateLimitResponse_RateLimit_MONTH { + y, m, _ := CurrentTime(timeSource.UnixNow()).Date() + _, lastDayOfMonth := MonthInterval(y, m) + + // calculate seconds between now and the end of the month + seconds := lastDayOfMonth.Unix() - now + logger.Debugf("the reset duration for until the end of the current month is %v seconds", seconds) + + return &duration.Duration{Seconds: seconds} + } + + // This doesn't work for months or years, because the exact number of days in a month or year are not always the same. + return &duration.Duration{Seconds: unitInSec - now%unitInSec} + +} + +func CurrentTime(unixNow int64) time.Time { + return time.Unix(unixNow, 0) +} + +func MonthInterval(y int, m time.Month) (firstDay, lastDay time.Time) { + firstDay = time.Date(y, m, 1, 0, 0, 0, 0, time.UTC) + lastDay = time.Date(y, m+1, 1, 0, 0, 0, -1, time.UTC) + return firstDay, lastDay } func Max(a uint32, b uint32) uint32 { From 4fd3bbbd74b4d7a7aef0da2a6f83aa482f1657d1 Mon Sep 17 00:00:00 2001 From: George Date: Fri, 31 Mar 2023 13:43:05 +0200 Subject: [PATCH 2/3] get current unix time in utilities --- src/limiter/base_limiter.go | 2 +- src/limiter/cache_key.go | 2 +- src/memcached/cache_impl.go | 2 +- src/redis/fixed_cache_impl.go | 4 +++- src/utils/utilities.go | 6 +++--- 5 files changed, 9 insertions(+), 7 deletions(-) diff --git a/src/limiter/base_limiter.go b/src/limiter/base_limiter.go index f1a7de22..e2d9f723 100644 --- a/src/limiter/base_limiter.go +++ b/src/limiter/base_limiter.go @@ -108,7 +108,7 @@ func (this *BaseRateLimiter) GetResponseDescriptorStatus(key string, limitInfo * // 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(key), []byte{}, int(utils.UnitToDivider(limitInfo.limit.Limit.Unit))) + err := this.localCache.Set([]byte(key), []byte{}, int(utils.UnitToDivider(limitInfo.limit.Limit.Unit, this.timeSource.UnixNow()))) if err != nil { logger.Errorf("Failing to set local cache key: %s", key) } diff --git a/src/limiter/cache_key.go b/src/limiter/cache_key.go index 0b892b2c..d3b92a66 100644 --- a/src/limiter/cache_key.go +++ b/src/limiter/cache_key.go @@ -70,7 +70,7 @@ func (this *CacheKeyGenerator) GenerateCacheKey( b.WriteByte('_') } - divider := utils.UnitToDivider(limit.Limit.Unit) + divider := utils.UnitToDivider(limit.Limit.Unit, now) // The key needs to be the same within the time unit. If we change the function, // then we need to make sure the key is always the same within the time unit, so it diff --git a/src/memcached/cache_impl.go b/src/memcached/cache_impl.go index a79451df..50aa7e9a 100644 --- a/src/memcached/cache_impl.go +++ b/src/memcached/cache_impl.go @@ -157,7 +157,7 @@ func (this *rateLimitMemcacheImpl) increaseAsync(cacheKeys []limiter.CacheKey, i _, err := this.client.Increment(cacheKey.Key, hitsAddend) if err == memcache.ErrCacheMiss { - expirationSeconds := utils.UnitToDivider(limits[i].Limit.Unit) + expirationSeconds := utils.UnitToDivider(limits[i].Limit.Unit, this.timeSource.UnixNow()) if this.expirationJitterMaxSeconds > 0 { expirationSeconds += this.jitterRand.Int63n(this.expirationJitterMaxSeconds) } diff --git a/src/redis/fixed_cache_impl.go b/src/redis/fixed_cache_impl.go index a61b319f..28cd04f1 100644 --- a/src/redis/fixed_cache_impl.go +++ b/src/redis/fixed_cache_impl.go @@ -28,6 +28,7 @@ type fixedRateLimitCacheImpl struct { // limits regardless of unit. If this client is not nil, then it // is used for limits that have a SECOND unit. perSecondClient Client + timeSource utils.TimeSource baseRateLimiter *limiter.BaseRateLimiter } @@ -74,7 +75,7 @@ func (this *fixedRateLimitCacheImpl) DoLimit( logger.Debugf("looking up cache key: %s", cacheKey.Key) - expirationSeconds := utils.UnitToDivider(limits[i].Limit.Unit) + expirationSeconds := utils.UnitToDivider(limits[i].Limit.Unit, this.timeSource.UnixNow()) if this.baseRateLimiter.ExpirationJitterMaxSeconds > 0 { expirationSeconds += this.baseRateLimiter.JitterRand.Int63n(this.baseRateLimiter.ExpirationJitterMaxSeconds) } @@ -134,6 +135,7 @@ func NewFixedRateLimitCacheImpl(client Client, perSecondClient Client, timeSourc jitterRand *rand.Rand, expirationJitterMaxSeconds int64, localCache *freecache.Cache, nearLimitRatio float32, cacheKeyPrefix string, statsManager stats.Manager) limiter.RateLimitCache { return &fixedRateLimitCacheImpl{ client: client, + timeSource: timeSource, perSecondClient: perSecondClient, baseRateLimiter: limiter.NewBaseRateLimit(timeSource, jitterRand, expirationJitterMaxSeconds, localCache, nearLimitRatio, cacheKeyPrefix, statsManager), } diff --git a/src/utils/utilities.go b/src/utils/utilities.go index feb573b5..5f1cb762 100644 --- a/src/utils/utilities.go +++ b/src/utils/utilities.go @@ -18,7 +18,7 @@ type TimeSource interface { // 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 { +func UnitToDivider(unit pb.RateLimitResponse_RateLimit_Unit, now int64) int64 { switch unit { case pb.RateLimitResponse_RateLimit_SECOND: return 1 @@ -30,7 +30,7 @@ func UnitToDivider(unit pb.RateLimitResponse_RateLimit_Unit) int64 { return 60 * 60 * 24 case pb.RateLimitResponse_RateLimit_MONTH: // This cannot be hardcoded to 30 days, as there are months with 31 days and the TTL will expire before the end of the month - return 60 * 60 * 24 * daysOfCurrentMonth(time.Now().Unix()) //todo: get timesource instead of time.now. + return 60 * 60 * 24 * daysOfCurrentMonth(now) } panic("should not get here") } @@ -51,7 +51,7 @@ func daysIn(m time.Month, year int) int64 { // CalculateReset calculates the reset time for a given unit and time source. func CalculateReset(unit *pb.RateLimitResponse_RateLimit_Unit, timeSource TimeSource) *duration.Duration { now := timeSource.UnixNow() - unitInSec := UnitToDivider(*unit) + unitInSec := UnitToDivider(*unit, timeSource.UnixNow()) if *unit == pb.RateLimitResponse_RateLimit_MONTH { y, m, _ := CurrentTime(timeSource.UnixNow()).Date() From 19639bb8d5113c17aa62e61ca5809fe98997d9c0 Mon Sep 17 00:00:00 2001 From: George Date: Fri, 31 Mar 2023 13:46:23 +0200 Subject: [PATCH 3/3] document the monthly cache key generation --- src/limiter/cache_key.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/limiter/cache_key.go b/src/limiter/cache_key.go index d3b92a66..9144e2b1 100644 --- a/src/limiter/cache_key.go +++ b/src/limiter/cache_key.go @@ -75,7 +75,8 @@ func (this *CacheKeyGenerator) GenerateCacheKey( // The key needs to be the same within the time unit. If we change the function, // then we need to make sure the key is always the same within the time unit, so it // can be picked up on the next lookup. - // This code section handles the MONTH time unit. + // This code section handles the MONTH time unit and uses the unix time of the first day of the current month for + // the cache key. if limit.Limit.Unit == pb.RateLimitResponse_RateLimit_MONTH { // get the first day of the current month as unix time y, m, _ := utils.CurrentTime(now).Date()