From 256ccb8f65cacf96e0c9eaaff67032ac998ea625 Mon Sep 17 00:00:00 2001
From: Bingyi Sun <sunbingyi1992@gmail.com>
Date: Wed, 16 Mar 2022 16:47:21 +0800
Subject: [PATCH] Fix rootcoord goroutine leak (#16071)

issue: #15658
Signed-off-by: sunby <bingyi.sun@zilliz.com>

Co-authored-by: sunby <bingyi.sun@zilliz.com>
---
 internal/mq/msgstream/mq_msgstream.go | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/internal/mq/msgstream/mq_msgstream.go b/internal/mq/msgstream/mq_msgstream.go
index b3496cba54..0959d82514 100644
--- a/internal/mq/msgstream/mq_msgstream.go
+++ b/internal/mq/msgstream/mq_msgstream.go
@@ -186,9 +186,6 @@ func (ms *mqMsgStream) Start() {
 }
 
 func (ms *mqMsgStream) Close() {
-	if !atomic.CompareAndSwapInt32(&ms.closed, 0, 1) {
-		return
-	}
 	ms.streamCancel()
 	ms.wait.Wait()
 
@@ -204,6 +201,10 @@ func (ms *mqMsgStream) Close() {
 	}
 
 	ms.client.Close()
+
+	if !atomic.CompareAndSwapInt32(&ms.closed, 0, 1) {
+		return
+	}
 	close(ms.receiveBuf)
 }