fix: stream connection leak in case of error (#38321)

issue: #38318 pr: #38320 Signed-off-by: jaime <[email protected]>
milvus-io · Dec 11, 2024 · 7cee01a · 7cee01a
1 parent e758d8e
commit 7cee01a
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/pkg/mq/msgdispatcher/dispatcher.go b/pkg/mq/msgdispatcher/dispatcher.go
@@ -91,7 +91,16 @@ func NewDispatcher(ctx context.Context,
 	log := log.With(zap.String("pchannel", pchannel),
 		zap.String("subName", subName), zap.Bool("isMain", isMain))
 	log.Info("creating dispatcher...")
-	stream, err := factory.NewTtMsgStream(ctx)
+
+	var stream msgstream.MsgStream
+	var err error
+	defer func() {
+		if err != nil && stream != nil {
+			stream.Close()
+		}
+	}()
+
+	stream, err = factory.NewTtMsgStream(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -106,15 +115,14 @@ func NewDispatcher(ctx context.Context,
 
 		err = stream.Seek(ctx, []*Pos{position}, false)
 		if err != nil {
-			stream.Close()
 			log.Error("seek failed", zap.Error(err))
 			return nil, err
 		}
 		posTime := tsoutil.PhysicalTime(position.GetTimestamp())
 		log.Info("seek successfully", zap.Uint64("posTs", position.GetTimestamp()),
 			zap.Time("posTime", posTime), zap.Duration("tsLag", time.Since(posTime)))
 	} else {
-		err := stream.AsConsumer(ctx, []string{pchannel}, subName, subPos)
+		err = stream.AsConsumer(ctx, []string{pchannel}, subName, subPos)
 		if err != nil {
 			log.Error("asConsumer failed", zap.Error(err))
 			return nil, err

diff --git a/pkg/mq/msgdispatcher/dispatcher_test.go b/pkg/mq/msgdispatcher/dispatcher_test.go
@@ -56,6 +56,7 @@ func TestDispatcher(t *testing.T) {
 	t.Run("test AsConsumer fail", func(t *testing.T) {
 		ms := msgstream.NewMockMsgStream(t)
 		ms.EXPECT().AsConsumer(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("mock error"))
+		ms.EXPECT().Close().Return()
 		factory := &msgstream.MockMqFactory{
 			NewMsgStreamFunc: func(ctx context.Context) (msgstream.MsgStream, error) {
 				return ms, nil