enhance: Make dataNode.import.maxConcurrentTaskNum dynamic (milvus-io…

…#37102)

Resize import execution pool when config
`dataNode.import.maxConcurrentTaskNum` update.

issue: milvus-io#37095

Signed-off-by: bigsheeper <[email protected]>

internal/datanode/importv2/pool.go

@@ -17,8 +17,14 @@
 package importv2
 
 import (
+	"context"
+	"runtime"
 	"sync"
 
+	"go.uber.org/zap"
+
+	"github.com/milvus-io/milvus/pkg/config"
+	"github.com/milvus-io/milvus/pkg/log"
 	"github.com/milvus-io/milvus/pkg/util/conc"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
 )
@@ -29,10 +35,32 @@ var (
 )
 
 func initExecPool() {
+	pt := paramtable.Get()
+	initPoolSize := paramtable.Get().DataNodeCfg.MaxConcurrentImportTaskNum.GetAsInt()
 	execPool = conc.NewPool[any](
-		paramtable.Get().DataNodeCfg.MaxConcurrentImportTaskNum.GetAsInt(),
-		conc.WithPreAlloc(true),
+		initPoolSize,
+		conc.WithPreAlloc(false), // pre alloc must be false to resize pool dynamically, use warmup to alloc worker here
+		conc.WithDisablePurge(true),
 	)
+	conc.WarmupPool(execPool, runtime.LockOSThread)
+
+	watchKey := pt.DataNodeCfg.MaxConcurrentImportTaskNum.Key
+	pt.Watch(watchKey, config.NewHandler(watchKey, resizeExecPool))
+	log.Info("init import execution pool done", zap.Int("size", initPoolSize))
+}
+
+func resizeExecPool(evt *config.Event) {
+	if evt.HasUpdated {
+		newSize := paramtable.Get().DataNodeCfg.MaxConcurrentImportTaskNum.GetAsInt()
+		log := log.Ctx(context.Background()).With(zap.Int("newSize", newSize))
+
+		err := GetExecPool().Resize(newSize)
+		if err != nil {
+			log.Warn("failed to resize pool", zap.Error(err))
+			return
+		}
+		log.Info("pool resize successfully")
+	}
 }
 
 func GetExecPool() *conc.Pool[any] {

internal/datanode/importv2/pool_test.go

@@ -0,0 +1,63 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package importv2
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/milvus-io/milvus/pkg/config"
+	"github.com/milvus-io/milvus/pkg/util/paramtable"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestResizePools(t *testing.T) {
+	paramtable.Get().Init(paramtable.NewBaseTable(paramtable.SkipRemote(true)))
+	pt := paramtable.Get()
+
+	defer func() {
+		_ = pt.Reset(pt.DataNodeCfg.MaxConcurrentImportTaskNum.Key)
+	}()
+
+	t.Run("ExecPool", func(t *testing.T) {
+		expectedCap := pt.DataNodeCfg.MaxConcurrentImportTaskNum.GetAsInt()
+		assert.Equal(t, expectedCap, GetExecPool().Cap())
+		resizeExecPool(&config.Event{
+			HasUpdated: true,
+		})
+		assert.Equal(t, expectedCap, GetExecPool().Cap())
+
+		_ = pt.Save(pt.DataNodeCfg.MaxConcurrentImportTaskNum.Key, fmt.Sprintf("%d", expectedCap*2))
+		expectedCap = pt.DataNodeCfg.MaxConcurrentImportTaskNum.GetAsInt()
+		resizeExecPool(&config.Event{
+			HasUpdated: true,
+		})
+		assert.Equal(t, expectedCap, GetExecPool().Cap())
+
+		_ = pt.Save(pt.DataNodeCfg.MaxConcurrentImportTaskNum.Key, "0")
+		resizeExecPool(&config.Event{
+			HasUpdated: true,
+		})
+		assert.Equal(t, expectedCap, GetExecPool().Cap(), "pool shall not be resized when newSize is 0")
+
+		_ = pt.Save(pt.DataNodeCfg.MaxConcurrentImportTaskNum.Key, "invalid")
+		resizeExecPool(&config.Event{
+			HasUpdated: true,
+		})
+		assert.Equal(t, expectedCap, GetExecPool().Cap())
+	})
+}