runner: close rate limiter on Stop (#635)

redis client is never closed, even after graceful shutdown. It is possible to observe this by creating a runner, and calling Stop(), and stopping redis. Runner is expected to be stopped, and all components closed, however still logs can be seen that redis is unavailable. this patch allows runner to actually gracefully close rate limit implementations, so all components are closed gracefully on Stop. Signed-off-by: Johannes Brüderl <[email protected]>
envoyproxy · Aug 1, 2024 · f4af2db · f4af2db
1 parent 05bf226
commit f4af2db
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 9 deletions.
diff --git a/src/redis/cache_impl.go b/src/redis/cache_impl.go
@@ -1,6 +1,7 @@
 package redis
 
 import (
+	"io"
 	"math/rand"
 
 	"github.com/coocood/freecache"
@@ -12,15 +13,18 @@ import (
 	"github.com/envoyproxy/ratelimit/src/utils"
 )
 
-func NewRateLimiterCacheImplFromSettings(s settings.Settings, localCache *freecache.Cache, srv server.Server, timeSource utils.TimeSource, jitterRand *rand.Rand, expirationJitterMaxSeconds int64, statsManager stats.Manager) limiter.RateLimitCache {
+func NewRateLimiterCacheImplFromSettings(s settings.Settings, localCache *freecache.Cache, srv server.Server, timeSource utils.TimeSource, jitterRand *rand.Rand, expirationJitterMaxSeconds int64, statsManager stats.Manager) (limiter.RateLimitCache, io.Closer) {
+	closer := &utils.MultiCloser{}
 	var perSecondPool Client
 	if s.RedisPerSecond {
 		perSecondPool = NewClientImpl(srv.Scope().Scope("redis_per_second_pool"), s.RedisPerSecondTls, s.RedisPerSecondAuth, s.RedisPerSecondSocketType,
 			s.RedisPerSecondType, s.RedisPerSecondUrl, s.RedisPerSecondPoolSize, s.RedisPerSecondPipelineWindow, s.RedisPerSecondPipelineLimit, s.RedisTlsConfig, s.RedisHealthCheckActiveConnection, srv)
+		closer.Closers = append(closer.Closers, perSecondPool)
 	}
 
 	otherPool := NewClientImpl(srv.Scope().Scope("redis_pool"), s.RedisTls, s.RedisAuth, s.RedisSocketType, s.RedisType, s.RedisUrl, s.RedisPoolSize,
 		s.RedisPipelineWindow, s.RedisPipelineLimit, s.RedisTlsConfig, s.RedisHealthCheckActiveConnection, srv)
+	closer.Closers = append(closer.Closers, otherPool)
 
 	return NewFixedRateLimitCacheImpl(
 		otherPool,
@@ -33,5 +37,5 @@ func NewRateLimiterCacheImplFromSettings(s settings.Settings, localCache *freeca
 		s.CacheKeyPrefix,
 		statsManager,
 		s.StopCacheKeyIncrementWhenOverlimit,
-	)
+	), closer
 }
diff --git a/src/service_cmd/runner/runner.go b/src/service_cmd/runner/runner.go
@@ -32,10 +32,11 @@ import (
 )
 
 type Runner struct {
-	statsManager stats.Manager
-	settings     settings.Settings
-	srv          server.Server
-	mu           sync.Mutex
+	statsManager    stats.Manager
+	settings        settings.Settings
+	srv             server.Server
+	mu              sync.Mutex
+	ratelimitCloser io.Closer
 }
 
 func NewRunner(s settings.Settings) Runner {
@@ -80,7 +81,7 @@ func (runner *Runner) GetStatsStore() gostats.Store {
 	return runner.statsManager.GetStatsStore()
 }
 
-func createLimiter(srv server.Server, s settings.Settings, localCache *freecache.Cache, statsManager stats.Manager) limiter.RateLimitCache {
+func createLimiter(srv server.Server, s settings.Settings, localCache *freecache.Cache, statsManager stats.Manager) (limiter.RateLimitCache, io.Closer) {
 	switch s.BackendType {
 	case "redis", "":
 		return redis.NewRateLimiterCacheImplFromSettings(
@@ -99,7 +100,7 @@ func createLimiter(srv server.Server, s settings.Settings, localCache *freecache
 			rand.New(utils.NewLockedSource(time.Now().Unix())),
 			localCache,
 			srv.Scope(),
-			statsManager)
+			statsManager), &utils.MultiCloser{} // memcache client can't be closed
 	default:
 		logger.Fatalf("Invalid setting for BackendType: %s", s.BackendType)
 		panic("This line should not be reachable")
@@ -147,8 +148,11 @@ func (runner *Runner) Run() {
 	runner.srv = srv
 	runner.mu.Unlock()
 
+	limiter, limiterCloser := createLimiter(srv, s, localCache, runner.statsManager)
+	runner.ratelimitCloser = limiterCloser
+
 	service := ratelimit.NewService(
-		createLimiter(srv, s, localCache, runner.statsManager),
+		limiter,
 		srv.Provider(),
 		runner.statsManager,
 		srv.HealthChecker(),
@@ -184,4 +188,8 @@ func (runner *Runner) Stop() {
 	if srv != nil {
 		srv.Stop()
 	}
+
+	if runner.ratelimitCloser != nil {
+		_ = runner.ratelimitCloser.Close()
+	}
 }
diff --git a/src/utils/multi_closer.go b/src/utils/multi_closer.go
@@ -0,0 +1,18 @@
+package utils
+
+import (
+	"errors"
+	"io"
+)
+
+type MultiCloser struct {
+	Closers []io.Closer
+}
+
+func (m *MultiCloser) Close() error {
+	var e error
+	for _, closer := range m.Closers {
+		e = errors.Join(closer.Close())
+	}
+	return e
+}