Add timeout for reader next (#12308)

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

internal/datanode/flow_graph_dmstream_input_node_test.go

@@ -104,6 +104,9 @@ func (mtm *mockTtMsgStream) SeekReaders(msgPositions []*internalpb.MsgPosition)
 func (mtm *mockTtMsgStream) Next(ctx context.Context, channelName string) (msgstream.TsMsg, error) {
 	return nil, nil
 }
+func (mtm *mockTtMsgStream) HasNext(channelName string) bool {
+	return true
+}
 
 func TestNewDmInputNode(t *testing.T) {
 	ctx := context.Background()

internal/msgstream/mq_msgstream.go

@@ -575,27 +575,29 @@ func (ms *mqMsgStream) Next(ctx context.Context, channelName string) (TsMsg, err
 	if !ok {
 		return nil, fmt.Errorf("reader for channel %s is not exist", channelName)
 	}
-	if reader.HasNext() {
-		msg, err := reader.Next(ctx)
-		if err != nil {
-			return nil, err
-		}
-		tsMsg, err := ms.getTsMsgFromConsumerMsg(msg)
-		if err != nil {
-			log.Error("Failed to getTsMsgFromConsumerMsg", zap.Error(err))
-			return nil, errors.New("Failed to getTsMsgFromConsumerMsg")
-		}
-		pos := tsMsg.Position()
-		tsMsg.SetPosition(&MsgPosition{
-			ChannelName: pos.ChannelName,
-			MsgID:       pos.MsgID,
-			Timestamp:   tsMsg.BeginTs(),
-		})
-		return tsMsg, nil
+	msg, err := reader.Next(ctx)
+	if err != nil {
+		return nil, err
 	}
-	log.Debug("All data has been read, there is no more data", zap.String("channel", channelName))
-	return nil, nil
-
+	tsMsg, err := ms.getTsMsgFromConsumerMsg(msg)
+	if err != nil {
+		log.Error("Failed to getTsMsgFromConsumerMsg", zap.Error(err))
+		return nil, errors.New("Failed to getTsMsgFromConsumerMsg")
+	}
+	pos := tsMsg.Position()
+	tsMsg.SetPosition(&MsgPosition{
+		ChannelName: pos.ChannelName,
+		MsgID:       pos.MsgID,
+		Timestamp:   tsMsg.BeginTs(),
+	})
+	return tsMsg, nil
+}
+func (ms *mqMsgStream) HasNext(channelName string) bool {
+	reader, ok := ms.readers[channelName]
+	if !ok {
+		return false
+	}
+	return reader.HasNext()
 }
 
 // Seek reset the subscription associated with this consumer to a specific position, the seek position is exclusive

internal/msgstream/mq_msgstream_test.go

@@ -1296,6 +1296,8 @@ func TestStream_MqMsgStream_Reader(t *testing.T) {
 	defer readStream.Close()
 	var seekPosition *internalpb.MsgPosition
 	for i := 0; i < n; i++ {
+		hasNext := readStream.HasNext(c)
+		assert.True(t, hasNext)
 		result, err := readStream.Next(ctx, c)
 		assert.Nil(t, err)
 		assert.Equal(t, result.ID(), int64(i))
@@ -1303,8 +1305,12 @@ func TestStream_MqMsgStream_Reader(t *testing.T) {
 			seekPosition = result.Position()
 		}
 	}
-	result, err := readStream.Next(ctx, c)
-	assert.Nil(t, err)
+	hasNext := readStream.HasNext(c)
+	assert.False(t, hasNext)
+	timeoutCtx1, cancel := context.WithTimeout(ctx, 3*time.Second)
+	defer cancel()
+	result, err := readStream.Next(timeoutCtx1, c)
+	assert.NotNil(t, err)
 	assert.Nil(t, result)
 
 	readStream2 := getPulsarReader(pulsarAddress, readerChannels)
@@ -1312,12 +1318,18 @@ func TestStream_MqMsgStream_Reader(t *testing.T) {
 	readStream2.SeekReaders([]*internalpb.MsgPosition{seekPosition})
 
 	for i := p; i < 10; i++ {
+		hasNext := readStream2.HasNext(c)
+		assert.True(t, hasNext)
 		result, err := readStream2.Next(ctx, c)
 		assert.Nil(t, err)
 		assert.Equal(t, result.ID(), int64(i))
 	}
-	result2, err := readStream2.Next(ctx, c)
-	assert.Nil(t, err)
+	hasNext = readStream2.HasNext(c)
+	assert.False(t, hasNext)
+	timeoutCtx2, cancel := context.WithTimeout(ctx, 3*time.Second)
+	defer cancel()
+	result2, err := readStream2.Next(timeoutCtx2, c)
+	assert.NotNil(t, err)
 	assert.Nil(t, result2)
 
 }

internal/msgstream/msgstream.go

@@ -69,6 +69,7 @@ type MsgStream interface {
 	BroadcastMark(*MsgPack) (map[string][]MessageID, error)
 	Consume() *MsgPack
 	Next(ctx context.Context, channelName string) (TsMsg, error)
+	HasNext(channelName string) bool
 	Seek(offset []*MsgPosition) error
 	SeekReaders(msgPositions []*internalpb.MsgPosition) error
 }

internal/proxy/mock_test.go

@@ -288,6 +288,10 @@ func (ms *simpleMockMsgStream) Next(ctx context.Context, channelName string) (ms
 	return nil, nil
 }
 
+func (ms *simpleMockMsgStream) HasNext(channelName string) bool {
+	return true
+}
+
 func (ms *simpleMockMsgStream) AsConsumerWithPosition(channels []string, subName string, position mqclient.SubscriptionInitialPosition) {
 }
 

internal/querynode/segment_loader.go

@@ -18,6 +18,7 @@ import (
 	"path"
 	"strconv"
 	"sync"
+	"time"
 
 	"go.uber.org/zap"
 
@@ -37,6 +38,8 @@ import (
 	"github.com/milvus-io/milvus/internal/util/funcutil"
 )
 
+const timeoutForEachRead = 10 * time.Second
+
 // segmentLoader is only responsible for loading the field data from binlog
 type segmentLoader struct {
 	historicalReplica ReplicaInterface
@@ -458,24 +461,30 @@ func (loader *segmentLoader) FromDmlCPLoadDelete(ctx context.Context, collection
 		deleteOffset:     make(map[UniqueID]int64),
 	}
 	log.Debug("start read msg from stream reader")
-	for {
+	for stream.HasNext(pChannelName) {
+		ctx, cancel := context.WithTimeout(ctx, timeoutForEachRead)
 		tsMsg, err := stream.Next(ctx, pChannelName)
 		if err != nil {
+			cancel()
 			return err
 		}
 		if tsMsg == nil {
-			break
+			cancel()
+			continue
 		}
 
 		if tsMsg.Type() == commonpb.MsgType_Delete {
 			dmsg := tsMsg.(*msgstream.DeleteMsg)
 			if dmsg.CollectionID != collectionID {
+				cancel()
 				continue
 			}
 			log.Debug("delete pk", zap.Any("pk", dmsg.PrimaryKeys))
 			processDeleteMessages(loader.historicalReplica, dmsg, delData)
 		}
+		cancel()
 	}
+	log.Debug("All data has been read, there is no more data", zap.String("channel", pChannelName))
 	for segmentID, pks := range delData.deleteIDs {
 		segment, err := loader.historicalReplica.getSegmentByID(segmentID)
 		if err != nil {