Fix pulsar reader not close (#14543)

Signed-off-by: godchen0212 <qingxiang.chen@zilliz.com>

internal/msgstream/mq_msgstream.go

@@ -226,6 +226,13 @@ func (ms *mqMsgStream) Close() {
 			consumer.Close()
 		}
 	}
+
+	for _, reader := range ms.readers {
+		if reader != nil {
+			reader.Close()
+		}
+	}
+	ms.client.Close()
 }
 
 func (ms *mqMsgStream) ComputeProduceChannelIndexes(tsMsgs []TsMsg) [][]int32 {

internal/querynode/segment_loader.go

@@ -462,10 +462,14 @@ func (loader *segmentLoader) FromDmlCPLoadDelete(ctx context.Context, collection
 	if err != nil {
 		return err
 	}
+	defer stream.Close()
 	pChannelName := rootcoord.ToPhysicalChannel(position.ChannelName)
 	position.ChannelName = pChannelName
 	stream.AsReader([]string{pChannelName}, fmt.Sprintf("querynode-%d-%d", Params.QueryNodeCfg.QueryNodeID, collectionID))
-	stream.SeekReaders([]*internalpb.MsgPosition{position})
+	err = stream.SeekReaders([]*internalpb.MsgPosition{position})
+	if err != nil {
+		return err
+	}
 
 	delData := &deleteData{
 		deleteIDs:        make(map[UniqueID][]int64),
@@ -518,7 +522,6 @@ func (loader *segmentLoader) FromDmlCPLoadDelete(ctx context.Context, collection
 		go deletePk(loader.historicalReplica, delData, segmentID, &wg)
 	}
 	wg.Wait()
-	stream.Close()
 	log.Debug("from dml check point load done", zap.Any("msg id", position.GetMsgID()))
 	return nil
 }

internal/querynode/task.go

@@ -625,7 +625,7 @@ func (w *watchDeltaChannelsTask) Execute(ctx context.Context) error {
 
 	for _, info := range w.req.Infos {
 		if err := w.node.loader.FromDmlCPLoadDelete(w.ctx, collectionID, info.SeekPosition); err != nil {
-			return errors.New("watchDeltaChannelsTask failed, error = " + err.Error())
+			return errors.New("watchDeltaChannelsTask from dml cp load delete failed, error = " + err.Error())
 		}
 	}
 

internal/querynode/task_test.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"testing"
 
+	"github.com/apache/pulsar-client-go/pulsar"
 	"github.com/stretchr/testify/assert"
 
 	"github.com/milvus-io/milvus/internal/msgstream"
@@ -452,8 +453,8 @@ func TestTask_watchDeltaChannelsTask(t *testing.T) {
 				CollectionID: defaultCollectionID,
 				ChannelName:  defaultDeltaChannel,
 				SeekPosition: &internalpb.MsgPosition{
-					ChannelName: defaultDeltaChannel,
-					MsgID:       []byte{1, 2, 3},
+					ChannelName: defaultDMLChannel,
+					MsgID:       pulsar.EarliestMessageID().Serialize(),
 					MsgGroup:    defaultSubName,
 					Timestamp:   0,
 				},

internal/util/mqclient/rmq_client.go

@@ -110,6 +110,5 @@ func (rc *rmqClient) BytesToMsgID(id []byte) (MessageID, error) {
 }
 
 func (rc *rmqClient) Close() {
-	// TODO(yukun): What to do here?
-	// rc.client.Close()
+	rc.client.Close()
 }

internal/util/mqclient/rmq_consumer.go

@@ -25,6 +25,7 @@ type RmqConsumer struct {
 	closeCh    chan struct{}
 	once       sync.Once
 	skip       int32
+	wg         sync.WaitGroup
 }
 
 // Subscription returns the subscription name of this consumer
@@ -37,7 +38,9 @@ func (rc *RmqConsumer) Chan() <-chan Message {
 	if rc.msgChannel == nil {
 		rc.once.Do(func() {
 			rc.msgChannel = make(chan Message, 256)
+			rc.wg.Add(1)
 			go func() {
+				defer rc.wg.Done()
 				for { //nolint:gosimple
 					select {
 					case msg, ok := <-rc.c.Chan():
@@ -78,6 +81,6 @@ func (rc *RmqConsumer) Ack(message Message) {
 
 // Close is used to free the resources of this consumer
 func (rc *RmqConsumer) Close() {
-	rc.c.Close()
 	close(rc.closeCh)
+	rc.wg.Wait()
 }

internal/util/mqclient/rmq_producer.go

@@ -38,5 +38,6 @@ func (rp *rmqProducer) Send(ctx context.Context, message *ProducerMessage) (Mess
 
 // Close does nothing currently
 func (rp *rmqProducer) Close() {
-	// TODO(yukun): may need to destroy topic
+	//TODO: close producer. Now it has bug
+	//rp.p.Close()
 }

internal/util/rocksmq/client/rocksmq/client_impl.go

@@ -182,7 +182,6 @@ func (c *client) CreateReader(readerOptions ReaderOptions) (Reader, error) {
 
 // Close close the channel to notify rocksmq to stop operation and close rocksmq server
 func (c *client) Close() {
-	// TODO(yukun): Should call server.close() here?
 	c.closeOnce.Do(func() {
 		close(c.closeCh)
 		c.wg.Wait()

internal/util/rocksmq/server/rocksmq/rocksmq_impl.go

@@ -1002,6 +1002,20 @@ func (rmq *rocksmq) getReader(topicName, readerName string) *rocksmqReader {
 	return nil
 }
 
+func (rmq *rocksmq) getAndDeleteReader(topicName, readerName string) *rocksmqReader {
+	if vals, ok := rmq.readers.Load(topicName); ok {
+		readers := vals.([]*rocksmqReader)
+		for i, v := range vals.([]*rocksmqReader) {
+			if v.readerName == readerName {
+				readers[i] = readers[len(readers)-1]
+				rmq.readers.Store(topicName, readers[:len(readers)-1])
+				return v
+			}
+		}
+	}
+	return nil
+}
+
 // ReaderSeek seek a reader to the pointed position
 func (rmq *rocksmq) ReaderSeek(topicName string, readerName string, msgID UniqueID) error {
 	if rmq.isClosed() {
@@ -1046,10 +1060,11 @@ func (rmq *rocksmq) CloseReader(topicName string, readerName string) {
 	if rmq.isClosed() {
 		return
 	}
-	reader := rmq.getReader(topicName, readerName)
+	reader := rmq.getAndDeleteReader(topicName, readerName)
 	if reader == nil {
 		log.Warn("reader not exist", zap.String("topic", topicName), zap.String("readerName", readerName))
 		return
 	}
 	reader.Close()
+	reader = nil
 }