fix load segment hangs (milvus-io#18770)

Signed-off-by: Wei Liu <[email protected]>

Signed-off-by: Wei Liu <[email protected]>

Author: weiliu1031

internal/mq/msgstream/mqwrapper/id.go

@@ -24,4 +24,6 @@ type MessageID interface {
 	AtEarliestPosition() bool
 
 	LessOrEqualThan(msgID []byte) (bool, error)
+
+	Equal(msgID []byte) (bool, error)
 }

internal/mq/msgstream/mqwrapper/kafka/kafka_id.go

@@ -19,6 +19,10 @@ func (kid *kafkaID) AtEarliestPosition() bool {
 	return kid.messageID <= 0
 }
 
+func (kid *kafkaID) Equal(msgID []byte) (bool, error) {
+	return kid.messageID == DeserializeKafkaID(msgID), nil
+}
+
 func (kid *kafkaID) LessOrEqualThan(msgID []byte) (bool, error) {
 	return kid.messageID <= DeserializeKafkaID(msgID), nil
 }

internal/mq/msgstream/mqwrapper/kafka/kafka_id_test.go

@@ -43,6 +43,23 @@ func TestKafkaID_LessOrEqualThan(t *testing.T) {
 	}
 }
 
+func TestKafkaID_Equal(t *testing.T) {
+	rid1 := &kafkaID{messageID: 0}
+	rid2 := &kafkaID{messageID: 0}
+
+	{
+		ret, err := rid1.Equal(rid1.Serialize())
+		assert.Nil(t, err)
+		assert.True(t, ret)
+	}
+
+	{
+		ret, err := rid1.Equal(rid2.Serialize())
+		assert.Nil(t, err)
+		assert.False(t, ret)
+	}
+}
+
 func Test_SerializeKafkaID(t *testing.T) {
 	bin := SerializeKafkaID(10)
 	assert.NotNil(t, bin)

internal/mq/msgstream/mqwrapper/pulsar/pulsar_id.go

@@ -60,6 +60,21 @@ func (pid *pulsarID) LessOrEqualThan(msgID []byte) (bool, error) {
 	return false, nil
 }
 
+func (pid *pulsarID) Equal(msgID []byte) (bool, error) {
+	pMsgID, err := pulsar.DeserializeMessageID(msgID)
+	if err != nil {
+		return false, err
+	}
+
+	if pid.messageID.LedgerID() == pMsgID.LedgerID() &&
+		pid.messageID.EntryID() == pMsgID.EntryID() &&
+		pid.messageID.BatchIdx() == pMsgID.BatchIdx() {
+		return true, nil
+	}
+
+	return false, nil
+}
+
 // SerializePulsarMsgID returns the serialized message ID
 func SerializePulsarMsgID(messageID pulsar.MessageID) []byte {
 	return messageID.Serialize()

internal/mq/msgstream/mqwrapper/pulsar/pulsar_id_test.go

@@ -72,6 +72,30 @@ func TestLessOrEqualThan(t *testing.T) {
 	assert.False(t, ret)
 }
 
+func TestPulsarID_Equal(t *testing.T) {
+	msg1 := pulsar.EarliestMessageID()
+	pid1 := &pulsarID{
+		messageID: msg1,
+	}
+
+	msg2 := pulsar.LatestMessageID()
+	pid2 := &pulsarID{
+		messageID: msg2,
+	}
+
+	{
+		ret, err := pid1.Equal(pid1.Serialize())
+		assert.Nil(t, err)
+		assert.True(t, ret)
+	}
+
+	{
+		ret, err := pid1.Equal(pid2.Serialize())
+		assert.Nil(t, err)
+		assert.False(t, ret)
+	}
+}
+
 func Test_SerializePulsarMsgID(t *testing.T) {
 	mid := pulsar.EarliestMessageID()
 

internal/mq/msgstream/mqwrapper/rmq/rmq_id.go

@@ -44,6 +44,11 @@ func (rid *rmqID) LessOrEqualThan(msgID []byte) (bool, error) {
 	return rid.messageID < rMsgID, nil
 }
 
+func (rid *rmqID) Equal(msgID []byte) (bool, error) {
+	rMsgID := DeserializeRmqID(msgID)
+	return rid.messageID == rMsgID, nil
+}
+
 // SerializeRmqID is used to serialize a message ID to byte array
 func SerializeRmqID(messageID int64) []byte {
 	b := make([]byte, 8)

internal/mq/msgstream/mqwrapper/rmq/rmq_id_test.go

@@ -62,6 +62,29 @@ func TestLessOrEqualThan(t *testing.T) {
 	assert.False(t, ret)
 }
 
+func Test_Equal(t *testing.T) {
+	rid1 := &rmqID{
+		messageID: 0,
+	}
+
+	rid2 := &rmqID{
+		messageID: math.MaxInt64,
+	}
+
+	{
+		ret, err := rid1.Equal(rid1.Serialize())
+		assert.Nil(t, err)
+		assert.True(t, ret)
+
+	}
+
+	{
+		ret, err := rid1.Equal(rid2.Serialize())
+		assert.Nil(t, err)
+		assert.False(t, ret)
+	}
+}
+
 func Test_SerializeRmqID(t *testing.T) {
 	bin := SerializeRmqID(10)
 	assert.NotNil(t, bin)

internal/querynode/segment_loader.go

@@ -664,7 +664,8 @@ func (loader *segmentLoader) FromDmlCPLoadDelete(ctx context.Context, collection
 		return err
 	}
 
-	if lastMsgID.AtEarliestPosition() {
+	reachLatest, _ := lastMsgID.Equal(position.MsgID)
+	if reachLatest || lastMsgID.AtEarliestPosition() {
 		log.Info("there is no more delta msg", zap.Int64("Collection ID", collectionID), zap.String("channel", pChannelName))
 		return nil
 	}

internal/querynode/segment_loader_test.go

@@ -23,6 +23,10 @@ import (
 	"runtime"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+
 	"github.com/milvus-io/milvus/internal/common"
 	"github.com/milvus-io/milvus/internal/mq/msgstream"
 	"github.com/milvus-io/milvus/internal/proto/commonpb"
@@ -32,9 +36,6 @@ import (
 	"github.com/milvus-io/milvus/internal/storage"
 	"github.com/milvus-io/milvus/internal/util/concurrency"
 	"github.com/milvus-io/milvus/internal/util/funcutil"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/mock"
-	"github.com/stretchr/testify/require"
 )
 
 func TestSegmentLoader_loadSegment(t *testing.T) {
@@ -588,20 +589,31 @@ func TestSegmentLoader_testFromDmlCPLoadDelete(t *testing.T) {
 	{
 		mockMsg := &mockMsgID{}
 		mockMsg.On("AtEarliestPosition").Return(false, nil)
+		mockMsg.On("Equal", mock.AnythingOfType("string")).Return(false, nil)
 		testSeekFailWhenConsumingDeltaMsg(ctx, t, position, mockMsg)
 	}
 
 	//test no more data when get last msg successfully
 	{
 		mockMsg := &mockMsgID{}
 		mockMsg.On("AtEarliestPosition").Return(true, nil)
+		mockMsg.On("Equal", mock.AnythingOfType("string")).Return(false, nil)
+		assert.Nil(t, testConsumingDeltaMsg(ctx, t, position, true, mockMsg))
+	}
+
+	// test already reach latest position
+	{
+		mockMsg := &mockMsgID{}
+		mockMsg.On("AtEarliestPosition").Return(false, nil)
+		mockMsg.On("Equal", mock.AnythingOfType("string")).Return(true, nil)
 		assert.Nil(t, testConsumingDeltaMsg(ctx, t, position, true, mockMsg))
 	}
 
 	//test consume after seeking when get last msg successfully
 	{
 		mockMsg := &mockMsgID{}
 		mockMsg.On("AtEarliestPosition").Return(false, nil)
+		mockMsg.On("Equal", mock.AnythingOfType("string")).Return(false, nil)
 		mockMsg.On("LessOrEqualThan", mock.AnythingOfType("string")).Return(true, nil)
 		assert.Nil(t, testConsumingDeltaMsg(ctx, t, position, true, mockMsg))
 	}
@@ -610,6 +622,7 @@ func TestSegmentLoader_testFromDmlCPLoadDelete(t *testing.T) {
 	{
 		mockMsg := &mockMsgID{}
 		mockMsg.On("AtEarliestPosition").Return(false, nil)
+		mockMsg.On("Equal", mock.AnythingOfType("string")).Return(false, nil)
 		mockMsg.On("LessOrEqualThan", mock.AnythingOfType("string")).Return(true, errors.New(""))
 		assert.NotNil(t, testConsumingDeltaMsg(ctx, t, position, true, mockMsg))
 	}
@@ -618,6 +631,7 @@ func TestSegmentLoader_testFromDmlCPLoadDelete(t *testing.T) {
 	{
 		mockMsg := &mockMsgID{}
 		mockMsg.On("AtEarliestPosition").Return(false, nil)
+		mockMsg.On("Equal", mock.AnythingOfType("string")).Return(false, nil)
 		mockMsg.On("LessOrEqualThan", mock.AnythingOfType("string")).Return(true, errors.New(""))
 		assert.NotNil(t, testConsumingDeltaMsg(ctx, t, position, false, mockMsg))
 	}
@@ -690,6 +704,15 @@ func (m2 *mockMsgID) LessOrEqualThan(msgID []byte) (bool, error) {
 	return ret.(bool), nil
 }
 
+func (m2 *mockMsgID) Equal(msgID []byte) (bool, error) {
+	args := m2.Called()
+	ret := args.Get(0)
+	if args.Get(1) != nil {
+		return false, args.Get(1).(error)
+	}
+	return ret.(bool), nil
+}
+
 type LoadDeleteMsgStream struct {
 	msgstream.MsgStream
 	mock.Mock