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

issue: #38318

Signed-off-by: jaime <yun.zhang@zilliz.com>
This commit is contained in:
jaime 2024-12-11 15:38:49 +08:00 committed by GitHub
parent 9be106dedf
commit 37a52286b1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 12 additions and 3 deletions

View File

@ -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,7 +115,6 @@ func NewDispatcher(ctx context.Context,
err = stream.Seek(ctx, []*Pos{position}, includeCurrentMsg)
if err != nil {
stream.Close()
log.Error("seek failed", zap.Error(err))
return nil, err
}
@ -114,7 +122,7 @@ func NewDispatcher(ctx context.Context,
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

View File

@ -55,6 +55,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