Make clean up buckets interval configurable

Author: kgantsov

cmd/limiter/main.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
+	"time"
 
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
@@ -20,6 +21,8 @@ var (
 	redisPort int    // HTTP Port
 	logLevel  string // Log level
 	logMode   string // Log mode
+
+	cleanUpBucketsInterval = 300 // Clean up buckets interval in seconds
 )
 
 func main() {
@@ -29,6 +32,7 @@ func main() {
 	flag.IntVar(&redisPort, "redis_port", 46379, "Redis Port")
 	flag.StringVar(&logMode, "log_mode", "console", "Log mode: console, stackdriver")
 	flag.StringVar(&logLevel, "log_level", "info", "Log level")
+	flag.IntVar(&cleanUpBucketsInterval, "cleanup_interval", 300, "Clean up buckets interval in seconds")
 
 	flag.Parse()
 
@@ -39,7 +43,7 @@ func main() {
 
 	zerolog.TimeFieldFormat = zerolog.TimeFormatUnix
 
-	rateLimiter := limiter.NewRateLimiter()
+	rateLimiter := limiter.NewRateLimiter(time.Duration(cleanUpBucketsInterval) * time.Second)
 	go rateLimiter.StartCleanUpFullBuckets()
 
 	app := http_server.NewApp(

pkg/limiter/benchmark_test.go

@@ -9,7 +9,7 @@ import (
 
 func BenchmarkReduce_100_1000_1_1000_1(b *testing.B) {
 	numberOfKeys := 1000
-	rl := NewRateLimiter()
+	rl := NewRateLimiter(300 * time.Second)
 
 	b.RunParallel(func(pb *testing.PB) {
 		r := rand.New(rand.NewSource(time.Now().UnixNano()))
@@ -22,7 +22,7 @@ func BenchmarkReduce_100_1000_1_1000_1(b *testing.B) {
 }
 func BenchmarkReduce_10000000_1000_10_1000_10(b *testing.B) {
 	numberOfKeys := 10000000
-	rl := NewRateLimiter()
+	rl := NewRateLimiter(300 * time.Second)
 
 	b.RunParallel(func(pb *testing.PB) {
 		r := rand.New(rand.NewSource(time.Now().UnixNano()))

pkg/limiter/limiter.go

@@ -27,10 +27,13 @@ type Shard struct {
 type RateLimiter struct {
 	shards []*Shard
 	length int64
+
+	cleanUpBucketsInterval time.Duration
 }
 
-func NewRateLimiter() *RateLimiter {
+func NewRateLimiter(cleanUpBucketsInterval time.Duration) *RateLimiter {
 	rateLimiter := new(RateLimiter)
+	rateLimiter.cleanUpBucketsInterval = cleanUpBucketsInterval
 
 	rateLimiter.shards = make([]*Shard, SHARDS)
 
@@ -133,7 +136,7 @@ func (l *RateLimiter) CleanUpFullBuckets() {
 
 func (l *RateLimiter) StartCleanUpFullBuckets() {
 	for {
-		time.Sleep(60 * time.Second)
+		time.Sleep(l.cleanUpBucketsInterval)
 		started := time.Now()
 
 		l.CleanUpFullBuckets()

pkg/limiter/limiter_test.go

@@ -10,7 +10,7 @@ import (
 
 func TestSlowRateLimiter(t *testing.T) {
 
-	rl := NewRateLimiter()
+	rl := NewRateLimiter(300 * time.Second)
 
 	val, _ := rl.Reduce("login", 5, 2, 5, 1)
 	assert.Equal(t, int64(4), val)
@@ -46,7 +46,7 @@ func TestSlowRateLimiter(t *testing.T) {
 
 func TestFastRateLimiter(t *testing.T) {
 
-	rl := NewRateLimiter()
+	rl := NewRateLimiter(300 * time.Second)
 
 	for i := 999; i >= 0; i-- {
 		val, _ := rl.Reduce("api_call", 1000, 1, 1000, 1)
@@ -66,7 +66,7 @@ func TestFastRateLimiter(t *testing.T) {
 }
 
 func TestRateLimiterWithManyKeys(t *testing.T) {
-	rl := NewRateLimiter()
+	rl := NewRateLimiter(300 * time.Second)
 
 	for i := 1000_000; i >= 0; i-- {
 		val, _ := rl.Reduce(fmt.Sprintf("api_call:%d", i), 1000, 1, 1000, 1)
@@ -76,7 +76,7 @@ func TestRateLimiterWithManyKeys(t *testing.T) {
 
 func TestRateLimiterReuseTheSameKey(t *testing.T) {
 
-	rl := NewRateLimiter()
+	rl := NewRateLimiter(300 * time.Second)
 
 	val, _ := rl.Reduce("user:123", 1000, 1, 50, 1)
 	assert.Equal(t, int64(999), val)
@@ -93,7 +93,7 @@ func TestRateLimiterReuseTheSameKey(t *testing.T) {
 }
 
 func TestRateLimiterD(t *testing.T) {
-	rl := NewRateLimiter()
+	rl := NewRateLimiter(300 * time.Second)
 
 	val, _ := rl.Reduce("api_call", 100, 1, 10, 10)
 	assert.Equal(t, int64(90), val)
@@ -107,7 +107,7 @@ func TestRateLimiterD(t *testing.T) {
 }
 
 func TestCleanUpFullBuckets(t *testing.T) {
-	rl := NewRateLimiter()
+	rl := NewRateLimiter(300 * time.Second)
 
 	assert.Equal(t, int64(0), rl.Len())
 

pkg/redis_server/redis_server_test.go

@@ -22,7 +22,7 @@ func TestServerBasic(t *testing.T) {
 	port := 56379
 
 	app := &http_server.App{
-		RateLimiter: limiter.NewRateLimiter(),
+		RateLimiter: limiter.NewRateLimiter(300 * time.Second),
 		PathMap:     make(map[string]string),
 	}