fix: analyzer memory leak because function runner not close (#41839)

relate: #41213

---------

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

Author: aoiasd

internal/datanode/importv2/util.go

@@ -378,6 +378,8 @@ func RunBm25Function(task *ImportTask, data *storage.InsertData) error {
 			continue
 		}
 
+		defer runner.Close()
+
 		inputFieldIDs := lo.Map(runner.GetInputFields(), func(field *schemapb.FieldSchema, _ int) int64 { return field.GetFieldID() })
 		inputDatas := make([]any, 0, len(inputFieldIDs))
 		for _, inputFieldID := range inputFieldIDs {

internal/flushcommon/pipeline/flow_graph_embedding_node.go

@@ -171,6 +171,12 @@ func (eNode *embeddingNode) Operate(in []Msg) []Msg {
 	return []Msg{fgMsg}
 }
 
+func (eNode *embeddingNode) Close() {
+	for _, runner := range eNode.functionRunners {
+		runner.Close()
+	}
+}
+
 func BuildSparseFieldData(array *schemapb.SparseFloatArray) storage.FieldData {
 	return &storage.SparseFloatVectorFieldData{
 		SparseFloatArray: schemapb.SparseFloatArray{

internal/flushcommon/pipeline/flow_graph_embedding_node_test.go

@@ -73,6 +73,7 @@ func TestEmbeddingNode_BM25_Operator(t *testing.T) {
 	t.Run("normal case", func(t *testing.T) {
 		node, err := newEmbeddingNode("test-channel", metaCache)
 		assert.NoError(t, err)
+		defer node.Close()
 
 		var output []Msg
 		assert.NotPanics(t, func() {
@@ -114,6 +115,7 @@ func TestEmbeddingNode_BM25_Operator(t *testing.T) {
 	t.Run("with close msg", func(t *testing.T) {
 		node, err := newEmbeddingNode("test-channel", metaCache)
 		assert.NoError(t, err)
+		defer node.Close()
 
 		var output []Msg
 
@@ -131,6 +133,7 @@ func TestEmbeddingNode_BM25_Operator(t *testing.T) {
 	t.Run("prepare insert failed", func(t *testing.T) {
 		node, err := newEmbeddingNode("test-channel", metaCache)
 		assert.NoError(t, err)
+		defer node.Close()
 
 		assert.Panics(t, func() {
 			node.Operate([]Msg{
@@ -152,6 +155,7 @@ func TestEmbeddingNode_BM25_Operator(t *testing.T) {
 	t.Run("embedding failed", func(t *testing.T) {
 		node, err := newEmbeddingNode("test-channel", metaCache)
 		assert.NoError(t, err)
+		defer node.Close()
 
 		node.functionRunners[0].GetSchema().Type = 0
 		assert.Panics(t, func() {

internal/querynodev2/pipeline/embedding_node.go

@@ -206,6 +206,12 @@ func (eNode *embeddingNode) Operate(in Msg) Msg {
 	return nodeMsg
 }
 
+func (eNode *embeddingNode) Close() {
+	for _, functionRunner := range eNode.functionRunners {
+		functionRunner.Close()
+	}
+}
+
 func getEmbeddingFieldDatas(datas []*schemapb.FieldData, fieldIDs ...int64) ([]any, error) {
 	result := []any{}
 	for _, fieldID := range fieldIDs {

internal/querynodev2/pipeline/embedding_node_test.go

@@ -140,8 +140,9 @@ func (suite *EmbeddingNodeSuite) TestCreateEmbeddingNode() {
 		collSchema := proto.Clone(suite.collectionSchema).(*schemapb.CollectionSchema)
 		collection := segments.NewCollectionWithoutSegcoreForTest(suite.collectionID, collSchema)
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
-		_, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
+		node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
 		suite.NoError(err)
+		defer node.Close()
 	})
 }
 
@@ -151,6 +152,7 @@ func (suite *EmbeddingNodeSuite) TestOperator() {
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
 		node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
 		suite.NoError(err)
+		defer node.Close()
 
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(nil).Once()
 		suite.Panics(func() {
@@ -163,6 +165,7 @@ func (suite *EmbeddingNodeSuite) TestOperator() {
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Times(2)
 		node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
 		suite.NoError(err)
+		defer node.Close()
 
 		suite.Panics(func() {
 			node.Operate(&insertNodeMsg{
@@ -192,6 +195,7 @@ func (suite *EmbeddingNodeSuite) TestOperator() {
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Times(2)
 		node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
 		suite.NoError(err)
+		defer node.Close()
 
 		suite.NotPanics(func() {
 			output := node.Operate(&insertNodeMsg{
@@ -213,6 +217,7 @@ func (suite *EmbeddingNodeSuite) TestAddInsertData() {
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
 		node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
 		suite.NoError(err)
+		defer node.Close()
 
 		// transfer insert msg failed because rowbase data not support sparse vector
 		insertDatas := make(map[int64]*delegator.InsertData)
@@ -237,6 +242,7 @@ func (suite *EmbeddingNodeSuite) TestAddInsertData() {
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
 		node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
 		suite.NoError(err)
+		defer node.Close()
 
 		insertDatas := make(map[int64]*delegator.InsertData)
 		err = node.addInsertData(insertDatas, suite.msgs[0], collection)
@@ -250,6 +256,7 @@ func (suite *EmbeddingNodeSuite) TestBM25Embedding() {
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
 		node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
 		suite.NoError(err)
+		defer node.Close()
 
 		runner := function.NewMockFunctionRunner(suite.T())
 		runner.EXPECT().BatchRun(mock.Anything).Return(nil, errors.New("mock error"))
@@ -265,6 +272,7 @@ func (suite *EmbeddingNodeSuite) TestBM25Embedding() {
 		suite.colManager.EXPECT().Get(suite.collectionID).Return(collection).Once()
 		node, err := newEmbeddingNode(suite.collectionID, suite.channel, suite.manager, 128)
 		suite.NoError(err)
+		defer node.Close()
 
 		runner := function.NewMockFunctionRunner(suite.T())
 		runner.EXPECT().BatchRun(mock.Anything).Return([]interface{}{1}, nil)

internal/util/function/bm25_function.go

@@ -251,6 +251,10 @@ func (v *BM25FunctionRunner) GetInputFields() []*schemapb.FieldSchema {
 	return []*schemapb.FieldSchema{v.inputField}
 }
 
+func (v *BM25FunctionRunner) Close() {
+	v.tokenizer.Destroy()
+}
+
 func buildSparseFloatArray(mapdata []map[uint32]float32) *schemapb.SparseFloatArray {
 	dim := int64(0)
 	bytes := lo.Map(mapdata, func(sparseMap map[uint32]float32, _ int) []byte {

internal/util/function/function.go

@@ -30,6 +30,8 @@ type FunctionRunner interface {
 	GetSchema() *schemapb.FunctionSchema
 	GetOutputFields() []*schemapb.FieldSchema
 	GetInputFields() []*schemapb.FieldSchema
+
+	Close()
 }
 
 func NewFunctionRunner(coll *schemapb.CollectionSchema, schema *schemapb.FunctionSchema) (FunctionRunner, error) {

internal/util/function/mock_function.go

@@ -325,3 +325,9 @@ func (v *MultiAnalyzerBM25FunctionRunner) GetOutputFields() []*schemapb.FieldSch
 func (v *MultiAnalyzerBM25FunctionRunner) GetInputFields() []*schemapb.FieldSchema {
 	return v.inputFields
 }
+
+func (v *MultiAnalyzerBM25FunctionRunner) Close() {
+	for _, analyzer := range v.analyzers {
+		analyzer.Destroy()
+	}
+}

internal/util/function/multi_analyzer_bm25_function.go

@@ -24,6 +24,8 @@ type Node interface {
 	Name() string
 	MaxQueueLength() int32
 	Operate(in Msg) Msg
+
+	Close()
 }
 
 type nodeCtx struct {
@@ -56,6 +58,9 @@ func (node *BaseNode) MaxQueueLength() int32 {
 	return node.maxQueueLength
 }
 
+func (node *BaseNode) Close() {
+}
+
 func NewBaseNode(name string, maxQueryLength int32) *BaseNode {
 	return &BaseNode{
 		name:           name,

internal/util/pipeline/node.go

@@ -70,6 +70,7 @@ func (p *pipeline) Start() error {
 
 func (p *pipeline) Close() {
 	for _, node := range p.nodes {
+		node.node.Close()
 		if node.Checker != nil {
 			node.Checker.Close()
 		}