From 1e1fba0588146b4c0cdb32b55c3595f70e36540d Mon Sep 17 00:00:00 2001
From: Bingyi Sun <sunbingyi1992@gmail.com>
Date: Fri, 10 May 2024 10:55:30 +0800
Subject: [PATCH] enhance: add lazy load retry configurations (#32848)

Signed-off-by: sunby <sunbingyi1992@gmail.com>
---
 pkg/util/cache/cache.go                | 11 +++++++----
 pkg/util/paramtable/component_param.go | 20 ++++++++++++++++++++
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/pkg/util/cache/cache.go b/pkg/util/cache/cache.go
index a4cb10683b..1a94d36af7 100644
--- a/pkg/util/cache/cache.go
+++ b/pkg/util/cache/cache.go
@@ -28,6 +28,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/log"
 	"github.com/milvus-io/milvus/pkg/util/lock"
 	"github.com/milvus-io/milvus/pkg/util/merr"
+	"github.com/milvus-io/milvus/pkg/util/paramtable"
 	"github.com/milvus-io/milvus/pkg/util/syncutil"
 )
 
@@ -348,18 +349,20 @@ func (c *lruCache[K, V]) getAndPin(ctx context.Context, key K) (*cacheItem[K, V]
 		}
 		timer := time.Now()
 		value, err := c.loader(ctx, key)
-		c.stats.TotalLoadTimeMs.Add(uint64(time.Since(timer).Milliseconds()))
-		// Try to evict one item if there is not enough disk space, then retry.
-		if merr.ErrServiceDiskLimitExceeded.Is(err) {
-			c.evictItems(ctx, 1)
+
+		for retryAttempt := 0; merr.ErrServiceDiskLimitExceeded.Is(err) && retryAttempt < paramtable.Get().QueryNodeCfg.LazyLoadMaxRetryTimes.GetAsInt(); retryAttempt++ {
+			// Try to evict one item if there is not enough disk space, then retry.
+			c.evictItems(ctx, paramtable.Get().QueryNodeCfg.LazyLoadMaxEvictPerRetry.GetAsInt())
 			value, err = c.loader(ctx, key)
 		}
+
 		if err != nil {
 			c.stats.LoadFailCount.Inc()
 			log.Debug("loader failed for key", zap.Any("key", key))
 			return nil, true, err
 		}
 
+		c.stats.TotalLoadTimeMs.Add(uint64(time.Since(timer).Milliseconds()))
 		c.stats.LoadSuccessCount.Inc()
 		item, err := c.setAndPin(ctx, key, value)
 		if err != nil {
diff --git a/pkg/util/paramtable/component_param.go b/pkg/util/paramtable/component_param.go
index 5f0464d4fc..ca9186b135 100644
--- a/pkg/util/paramtable/component_param.go
+++ b/pkg/util/paramtable/component_param.go
@@ -2007,6 +2007,8 @@ type queryNodeConfig struct {
 	LazyLoadWaitTimeout                  ParamItem `refreshable:"true"`
 	LazyLoadRequestResourceTimeout       ParamItem `refreshable:"true"`
 	LazyLoadRequestResourceRetryInterval ParamItem `refreshable:"true"`
+	LazyLoadMaxRetryTimes                ParamItem `refreshable:"true"`
+	LazyLoadMaxEvictPerRetry             ParamItem `refreshable:"true"`
 
 	// chunk cache
 	ReadAheadPolicy     ParamItem `refreshable:"false"`
@@ -2275,6 +2277,24 @@ func (p *queryNodeConfig) init(base *BaseTable) {
 	}
 	p.LazyLoadRequestResourceRetryInterval.Init(base.mgr)
 
+	p.LazyLoadMaxRetryTimes = ParamItem{
+		Key:          "queryNode.lazyLoadMaxRetryTimes",
+		Version:      "2.4.2",
+		DefaultValue: "1",
+		Doc:          "max retry times for lazy load, 1 by default",
+		Export:       true,
+	}
+	p.LazyLoadMaxRetryTimes.Init(base.mgr)
+
+	p.LazyLoadMaxEvictPerRetry = ParamItem{
+		Key:          "queryNode.lazyLoadMaxEvictPerRetry",
+		Version:      "2.4.2",
+		DefaultValue: "1",
+		Doc:          "max evict count for lazy load, 1 by default",
+		Export:       true,
+	}
+	p.LazyLoadMaxEvictPerRetry.Init(base.mgr)
+
 	p.ReadAheadPolicy = ParamItem{
 		Key:          "queryNode.cache.readAheadPolicy",
 		Version:      "2.3.2",