fix: the close operation of rmq consumer is not sync

Signed-off-by: chyezh <[email protected]>
milvus-io · Dec 25, 2024 · d404aef · d404aef
1 parent acc8fb7
commit d404aef
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/pkg/mq/mqimpl/rocksmq/client/client_impl.go b/pkg/mq/mqimpl/rocksmq/client/client_impl.go
@@ -124,7 +124,10 @@ func (c *client) Subscribe(options ConsumerOptions) (Consumer, error) {
 }
 
 func (c *client) consume(consumer *consumer) {
-	defer c.wg.Done()
+	defer func() {
+		close(consumer.stopCh)
+		c.wg.Done()
+	}()
 
 	if err := c.blockUntilInitDone(consumer); err != nil {
 		log.Warn("consumer init failed", zap.Error(err))

diff --git a/pkg/mq/mqimpl/rocksmq/client/consumer_impl.go b/pkg/mq/mqimpl/rocksmq/client/consumer_impl.go
@@ -28,6 +28,7 @@ type consumer struct {
 
 	startOnce sync.Once
 
+	stopCh    chan struct{}
 	msgMutex  chan struct{}
 	initCh    chan struct{}
 	messageCh chan common.Message
@@ -58,6 +59,7 @@ func newConsumer(c *client, options ConsumerOptions) (*consumer, error) {
 		client:       c,
 		consumerName: options.SubscriptionName,
 		options:      options,
+		stopCh:       make(chan struct{}),
 		msgMutex:     make(chan struct{}, 1),
 		initCh:       initCh,
 		messageCh:    messageCh,
@@ -134,6 +136,7 @@ func (c *consumer) Close() {
 	if err != nil {
 		log.Warn("Consumer close failed", zap.String("topicName", c.topic), zap.String("groupName", c.consumerName), zap.Error(err))
 	}
+	<-c.stopCh
 }
 
 func (c *consumer) GetLatestMsgID() (int64, error) {