Allow to customize the high redis memory water mark

git-hulk · git-hulk · commit 8d3282eda003 · 2023-11-17T13:37:14.000+08:00
diff --git a/config/config.go b/config/config.go
@@ -37,9 +37,17 @@ type SecondaryStorage struct {
 	Spanner *SpannerConfig
 	// max number of jobs that storage pumps per batch
 	MaxJobPumpBatchSize int64
+	// range from 0 to 1, when the redis memory usage is greater than this value,
+	// the storage won't pump jobs to redis anymore until the memory usage is lower than this value.
+	//
+	// Default is 1, means no limit
+	HighRedisMemoryWatermark float64
 }
 
 func (storage *SecondaryStorage) validate() error {
+	if storage.HighRedisMemoryWatermark < 0 || storage.HighRedisMemoryWatermark > 1 {
+		return fmt.Errorf("invalid HighRedisMemoryWatermark: %f, should be between 0 and 1", storage.HighRedisMemoryWatermark)
+	}
 	return storage.Spanner.validate()
 }
 
@@ -156,6 +164,10 @@ func MustLoad(path string) (*Config, error) {
 		if err := conf.SecondaryStorage.validate(); err != nil {
 			return nil, err
 		}
+		if conf.SecondaryStorage.HighRedisMemoryWatermark == 0 {
+			// default to 1
+			conf.SecondaryStorage.HighRedisMemoryWatermark = 1
+		}
 	}
 	return conf, nil
 }
diff --git a/storage/manager.go b/storage/manager.go
@@ -30,11 +30,10 @@ const (
 
 	addJobSuccessStatus = "success"
 	addJobFailedStatus  = "failed"
-
-	redisMemoryUsageWatermark = 0.8 //  used_memory / max_memory
 )
 
 type Manager struct {
+	cfg     *config.Config
 	wg      sync.WaitGroup
 	mu      sync.Mutex
 	pools   map[string]engine.Engine
@@ -76,6 +75,7 @@ func NewManger(cfg *config.Config) (*Manager, error) {
 		return nil, fmt.Errorf("create redis client err: %w", err)
 	}
 	return &Manager{
+		cfg:              cfg,
 		redisCli:         redisCli,
 		storage:          storage,
 		pools:            make(map[string]engine.Engine),
@@ -87,7 +87,7 @@ func NewManger(cfg *config.Config) (*Manager, error) {
 func (m *Manager) PumpFn(name string, pool engine.Engine, threshold int64) func() bool {
 	return func() bool {
 		logger := log.Get().WithField("pool", name)
-		if isHighRedisMemUsage(m.redisCli) {
+		if isHighRedisMemUsage(m.redisCli, m.cfg.SecondaryStorage.HighRedisMemoryWatermark) {
 			logger.Error("High redis usage, storage stops pumping data")
 			return false
 		}
@@ -183,7 +183,10 @@ func (m *Manager) Shutdown() {
 	m.storage.Close()
 }
 
-func isHighRedisMemUsage(cli *redis.Client) bool {
+func isHighRedisMemUsage(cli *redis.Client, redisMemoryUsageWatermark float64) bool {
+	if redisMemoryUsageWatermark == 0 || redisMemoryUsageWatermark >= 1 {
+		return false
+	}
 	memoryInfo, err := cli.Info(context.TODO(), "memory").Result()
 	if err != nil {
 		return false

Original file line number	Diff line number	Diff line change
`@@ -37,9 +37,17 @@ type SecondaryStorage struct {`
`37`	`37`	`Spanner *SpannerConfig`
`38`	`38`	`// max number of jobs that storage pumps per batch`
`39`	`39`	`MaxJobPumpBatchSize int64`
	`40`	`+ // range from 0 to 1, when the redis memory usage is greater than this value,`
	`41`	`+ // the storage won't pump jobs to redis anymore until the memory usage is lower than this value.`
	`42`	`+ //`
	`43`	`+ // Default is 1, means no limit`
	`44`	`+ HighRedisMemoryWatermark float64`
`40`	`45`	`}`
`41`	`46`
`42`	`47`	`func (storage *SecondaryStorage) validate() error {`
	`48`	`+ if storage.HighRedisMemoryWatermark < 0 \|\| storage.HighRedisMemoryWatermark > 1 {`
	`49`	`+ return fmt.Errorf("invalid HighRedisMemoryWatermark: %f, should be between 0 and 1", storage.HighRedisMemoryWatermark)`
	`50`	`+ }`
`43`	`51`	`return storage.Spanner.validate()`
`44`	`52`	`}`
`45`	`53`
`@@ -156,6 +164,10 @@ func MustLoad(path string) (*Config, error) {`
`156`	`164`	`if err := conf.SecondaryStorage.validate(); err != nil {`
`157`	`165`	`return nil, err`
`158`	`166`	`}`
	`167`	`+ if conf.SecondaryStorage.HighRedisMemoryWatermark == 0 {`
	`168`	`+ // default to 1`
	`169`	`+ conf.SecondaryStorage.HighRedisMemoryWatermark = 1`
	`170`	`+ }`
`159`	`171`	`}`
`160`	`172`	`return conf, nil`
`161`	`173`	`}`