Fix kafka consumer connection leak (#27224)

Signed-off-by: jaime <yun.zhang@zilliz.com>

Makefile

@@ -202,7 +202,7 @@ generated-proto: download-milvus-proto build-3rdparty
 	@(env bash $(PWD)/scripts/generate_proto.sh)
 
 build-cpp: generated-proto
-	@echo "Building Milvus cpp library ... ${AZURE_OPTION}"
+	@echo "Building Milvus cpp library ..."
 	@(env bash $(PWD)/scripts/core_build.sh -t ${mode} -f "$(CUSTOM_THIRDPARTY_PATH)" -n ${disk_index} -y ${use_dynamic_simd} ${AZURE_OPTION} -x ${index_engine})
 
 build-cpp-gpu: generated-proto

pkg/mq/msgstream/mqwrapper/kafka/kafka_client_test.go

@@ -35,6 +35,7 @@ func TestMain(m *testing.M) {
 
 	broker := mockCluster.BootstrapServers()
 	Params.Save("kafka.brokerList", broker)
+	log.Info("start testing kafka broker", zap.String("address", broker))
 
 	exitCode := m.Run()
 	os.Exit(exitCode)

pkg/mq/msgstream/mqwrapper/kafka/kafka_consumer.go

@@ -126,16 +126,6 @@ func (kc *Consumer) Chan() <-chan mqwrapper.Message {
 			for {
 				select {
 				case <-kc.closeCh:
-					log.Info("close consumer ", zap.String("topic", kc.topic), zap.String("groupID", kc.groupID))
-					start := time.Now()
-					err := kc.c.Close()
-					if err != nil {
-						log.Warn("failed to close ", zap.String("topic", kc.topic), zap.Error(err))
-					}
-					cost := time.Since(start).Milliseconds()
-					if cost > 200 {
-						log.Warn("close consumer costs too long time", zap.Any("topic", kc.topic), zap.String("groupID", kc.groupID), zap.Int64("time(ms)", cost))
-					}
 					if kc.msgChannel != nil {
 						close(kc.msgChannel)
 					}
@@ -258,9 +248,25 @@ func (kc *Consumer) CheckTopicValid(topic string) error {
 	return nil
 }
 
+func (kc *Consumer) closeInternal() {
+	log.Info("close consumer ", zap.String("topic", kc.topic), zap.String("groupID", kc.groupID))
+	start := time.Now()
+	err := kc.c.Close()
+	if err != nil {
+		log.Warn("failed to close ", zap.String("topic", kc.topic), zap.Error(err))
+	}
+	cost := time.Since(start).Milliseconds()
+	if cost > 200 {
+		log.Warn("close consumer costs too long time", zap.Any("topic", kc.topic), zap.String("groupID", kc.groupID), zap.Int64("time(ms)", cost))
+	}
+}
+
 func (kc *Consumer) Close() {
 	kc.closeOnce.Do(func() {
 		close(kc.closeCh)
-		kc.wg.Wait() // wait worker exist and close the client
+		// wait work goroutine exit
+		kc.wg.Wait()
+		// close the client
+		kc.closeInternal()
 	})
 }

pkg/mq/msgstream/mqwrapper/kafka/kafka_consumer_test.go

@@ -268,3 +268,31 @@ func TestKafkaConsumer_CheckPreTopicValid(t *testing.T) {
 	err = consumer.CheckTopicValid(topic)
 	assert.NoError(t, err)
 }
+
+func TestKafkaConsumer_Close(t *testing.T) {
+	rand.Seed(time.Now().UnixNano())
+	topic := fmt.Sprintf("test-topicName-%d", rand.Int())
+
+	data1 := []int{111, 222, 333}
+	data2 := []string{"111", "222", "333"}
+	testKafkaConsumerProduceData(t, topic, data1, data2)
+
+	t.Run("close after only get latest msgID", func(t *testing.T) {
+		groupID := fmt.Sprintf("test-groupid-%d", rand.Int())
+		config := createConfig(groupID)
+		consumer, err := newKafkaConsumer(config, 16, topic, groupID, mqwrapper.SubscriptionPositionEarliest)
+		assert.NoError(t, err)
+		_, err = consumer.GetLatestMsgID()
+		assert.NoError(t, err)
+		consumer.Close()
+	})
+
+	t.Run("close after only Chan method is invoked", func(t *testing.T) {
+		groupID := fmt.Sprintf("test-groupid-%d", rand.Int())
+		config := createConfig(groupID)
+		consumer, err := newKafkaConsumer(config, 16, topic, groupID, mqwrapper.SubscriptionPositionEarliest)
+		assert.NoError(t, err)
+		<-consumer.Chan()
+		consumer.Close()
+	})
+}