fix: call IntoMessageProto instead of Payload when rpc (#43678)

issue: #43677

Signed-off-by: chyezh <chyezh@outlook.com>

internal/streamingcoord/client/broadcast/broadcast_impl.go

@@ -4,7 +4,6 @@ import (
 	"context"
 
 	"github.com/milvus-io/milvus/internal/util/streamingutil/service/lazygrpc"
-	"github.com/milvus-io/milvus/pkg/v2/proto/messagespb"
 	"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
 	"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
 	"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
@@ -31,10 +30,7 @@ func (c *GRPCBroadcastServiceImpl) Broadcast(ctx context.Context, msg message.Br
 		return nil, err
 	}
 	resp, err := client.Broadcast(ctx, &streamingpb.BroadcastRequest{
-		Message: &messagespb.Message{
-			Payload:    msg.Payload(),
-			Properties: msg.Properties().ToRawMap(),
-		},
+		Message: msg.IntoMessageProto(),
 	})
 	if err != nil {
 		return nil, err

internal/streamingcoord/server/broadcaster/broadcast_task.go

@@ -11,7 +11,6 @@ import (
 	"github.com/milvus-io/milvus/internal/streamingcoord/server/broadcaster/registry"
 	"github.com/milvus-io/milvus/internal/streamingcoord/server/resource"
 	"github.com/milvus-io/milvus/pkg/v2/log"
-	"github.com/milvus-io/milvus/pkg/v2/proto/messagespb"
 	"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
 	"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
 )
@@ -44,7 +43,7 @@ func newBroadcastTaskFromBroadcastMessage(msg message.BroadcastMutableMessage, m
 		mu:     sync.Mutex{},
 		header: header,
 		task: &streamingpb.BroadcastTask{
-			Message:             &messagespb.Message{Payload: msg.Payload(), Properties: msg.Properties().ToRawMap()},
+			Message:             msg.IntoMessageProto(),
 			State:               streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_PENDING,
 			AckedVchannelBitmap: make([]byte, len(header.VChannels)),
 		},

internal/streamingcoord/server/broadcaster/broadcaster_test.go

@@ -193,10 +193,11 @@ func createNewBroadcastMsg(vchannels []string, rks ...message.ResourceKey) messa
 
 func createNewBroadcastTask(broadcastID uint64, vchannels []string, rks ...message.ResourceKey) *streamingpb.BroadcastTask {
 	msg := createNewBroadcastMsg(vchannels, rks...).WithBroadcastID(broadcastID)
+	pb := msg.IntoMessageProto()
 	return &streamingpb.BroadcastTask{
 		Message: &messagespb.Message{
-			Payload:    msg.Payload(),
-			Properties: msg.Properties().ToRawMap(),
+			Payload:    pb.Payload,
+			Properties: pb.Properties,
 		},
 		State:               streamingpb.BroadcastTaskState_BROADCAST_TASK_STATE_PENDING,
 		AckedVchannelBitmap: make([]byte, len(vchannels)),
@@ -208,10 +209,11 @@ func createNewWaitAckBroadcastTaskFromMessage(
 	state streamingpb.BroadcastTaskState,
 	bitmap []byte,
 ) *streamingpb.BroadcastTask {
+	pb := msg.IntoMessageProto()
 	return &streamingpb.BroadcastTask{
 		Message: &messagespb.Message{
-			Payload:    msg.Payload(),
-			Properties: msg.Properties().ToRawMap(),
+			Payload:    pb.Payload,
+			Properties: pb.Properties,
 		},
 		State:               state,
 		AckedVchannelBitmap: bitmap,

internal/streamingnode/client/handler/consumer/consumer_test.go

@@ -182,6 +182,7 @@ func newMockedConsumerImpl(t *testing.T, ctx context.Context, h message.Handler)
 func newConsumeResponse(id message.MessageID, msg message.MutableMessage) *streamingpb.ConsumeResponse {
 	msg.WithTimeTick(tsoutil.GetCurrentTime())
 	msg.WithLastConfirmed(walimplstest.NewTestMessageID(0))
+	pb := msg.IntoMessageProto()
 	return &streamingpb.ConsumeResponse{
 		Response: &streamingpb.ConsumeResponse_Consume{
 			Consume: &streamingpb.ConsumeMessageReponse{
@@ -189,8 +190,8 @@ func newConsumeResponse(id message.MessageID, msg message.MutableMessage) *strea
 					Id: &messagespb.MessageID{
 						Id: id.Marshal(),
 					},
-					Payload:    msg.Payload(),
-					Properties: msg.Properties().ToRawMap(),
+					Payload:    pb.Payload,
+					Properties: pb.Properties,
 				},
 			},
 		},

internal/streamingnode/client/handler/producer/produce_grpc_client.go

@@ -1,7 +1,6 @@
 package producer
 
 import (
-	"github.com/milvus-io/milvus/pkg/v2/proto/messagespb"
 	"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
 	"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
 )
@@ -17,10 +16,7 @@ func (p *produceGrpcClient) SendProduceMessage(requestID int64, msg message.Muta
 		Request: &streamingpb.ProduceRequest_Produce{
 			Produce: &streamingpb.ProduceMessageRequest{
 				RequestId: requestID,
-				Message: &messagespb.Message{
-					Payload:    msg.Payload(),
-					Properties: msg.Properties().ToRawMap(),
-				},
+				Message:   msg.IntoMessageProto(),
 			},
 		},
 	})

internal/streamingnode/server/service/handler/consumer/consume_server.go

@@ -187,6 +187,7 @@ func (c *ConsumeServer) sendImmutableMessage(msg message.ImmutableMessage) (err
 		metricsGuard.Finish(err)
 	}()
 
+	pb := msg.IntoMessageProto()
 	// Send Consumed message to client and do metrics.
 	if err := c.consumeServer.SendConsumeMessage(&streamingpb.ConsumeMessageReponse{
 		ConsumerId: c.consumerID,
@@ -194,8 +195,8 @@ func (c *ConsumeServer) sendImmutableMessage(msg message.ImmutableMessage) (err
 			Id: &messagespb.MessageID{
 				Id: msg.MessageID().Marshal(),
 			},
-			Payload:    msg.Payload(),
-			Properties: msg.Properties().ToRawMap(),
+			Payload:    pb.Payload,
+			Properties: pb.Properties,
 		},
 	}); err != nil {
 		return status.NewInner("send consume message failed: %s", err.Error())

internal/streamingnode/server/service/handler/consumer/consume_server_test.go

@@ -20,6 +20,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/v2/log"
 	"github.com/milvus-io/milvus/pkg/v2/mocks/proto/mock_streamingpb"
 	"github.com/milvus-io/milvus/pkg/v2/mocks/streaming/util/mock_message"
+	"github.com/milvus-io/milvus/pkg/v2/proto/messagespb"
 	"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
 	"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
 	"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
@@ -190,10 +191,7 @@ func TestConsumerServeSendArm(t *testing.T) {
 	msg := mock_message.NewMockImmutableMessage(t)
 	msg.EXPECT().MessageID().Return(walimplstest.NewTestMessageID(1))
 	msg.EXPECT().EstimateSize().Return(0)
-	msg.EXPECT().Payload().Return([]byte{})
-	properties := mock_message.NewMockRProperties(t)
-	properties.EXPECT().ToRawMap().Return(map[string]string{})
-	msg.EXPECT().Properties().Return(properties)
+	msg.EXPECT().IntoMessageProto().Return(&messagespb.Message{})
 	scanCh <- msg
 
 	// test send txn message.

internal/streamingnode/server/wal/adaptor/old_version_message.go

@@ -40,7 +40,7 @@ func newOldVersionImmutableMessage(
 	case *msgstream.DropCollectionMsg:
 		mutableMessage, err = newV1DropCollectionMsgFromV0(ctx, pchannel, underlyingMsg)
 	case *msgstream.InsertMsg:
-		mutableMessage = newV1InsertMsgFromV0(underlyingMsg, uint64(len(msg.Payload())))
+		mutableMessage = newV1InsertMsgFromV0(underlyingMsg, uint64(msg.EstimateSize()))
 	case *msgstream.DeleteMsg:
 		mutableMessage = newV1DeleteMsgFromV0(underlyingMsg, uint64(underlyingMsg.NumRows))
 	case *msgstream.TimeTickMsg:

internal/streamingnode/server/wal/interceptors/shard/shard_interceptor.go

@@ -143,7 +143,7 @@ func (impl *shardInterceptor) handleInsertMessage(ctx context.Context, msg messa
 		if partition.BinarySize == 0 {
 			// binary size should be set at proxy with estimate, but we don't implement it right now.
 			// use payload size instead.
-			partition.BinarySize = uint64(len(msg.Payload()))
+			partition.BinarySize = uint64(msg.EstimateSize())
 		}
 		req := &shards.AssignSegmentRequest{
 			CollectionID: header.GetCollectionId(),

pkg/.mockery_pkg.yaml

@@ -15,7 +15,6 @@ packages:
       ImmutableTxnMessage:
       MutableMessage:
       BroadcastMutableMessage:
-      RProperties:
   github.com/milvus-io/milvus/pkg/v2/streaming/walimpls:
     interfaces:
       OpenerBuilderImpls:

pkg/mocks/streaming/util/mock_message/mock_BroadcastMutableMessage.go

@@ -3,7 +3,9 @@
 package mock_message
 
 import (
+	messagespb "github.com/milvus-io/milvus/pkg/v2/proto/messagespb"
 	message "github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
+
 	mock "github.com/stretchr/testify/mock"
 
 	zapcore "go.uber.org/zap/zapcore"
@@ -159,6 +161,53 @@ func (_c *MockBroadcastMutableMessage_EstimateSize_Call) RunAndReturn(run func()
 	return _c
 }
 
+// IntoMessageProto provides a mock function with no fields
+func (_m *MockBroadcastMutableMessage) IntoMessageProto() *messagespb.Message {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for IntoMessageProto")
+	}
+
+	var r0 *messagespb.Message
+	if rf, ok := ret.Get(0).(func() *messagespb.Message); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*messagespb.Message)
+		}
+	}
+
+	return r0
+}
+
+// MockBroadcastMutableMessage_IntoMessageProto_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IntoMessageProto'
+type MockBroadcastMutableMessage_IntoMessageProto_Call struct {
+	*mock.Call
+}
+
+// IntoMessageProto is a helper method to define mock.On call
+func (_e *MockBroadcastMutableMessage_Expecter) IntoMessageProto() *MockBroadcastMutableMessage_IntoMessageProto_Call {
+	return &MockBroadcastMutableMessage_IntoMessageProto_Call{Call: _e.mock.On("IntoMessageProto")}
+}
+
+func (_c *MockBroadcastMutableMessage_IntoMessageProto_Call) Run(run func()) *MockBroadcastMutableMessage_IntoMessageProto_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockBroadcastMutableMessage_IntoMessageProto_Call) Return(_a0 *messagespb.Message) *MockBroadcastMutableMessage_IntoMessageProto_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockBroadcastMutableMessage_IntoMessageProto_Call) RunAndReturn(run func() *messagespb.Message) *MockBroadcastMutableMessage_IntoMessageProto_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // IsPersisted provides a mock function with no fields
 func (_m *MockBroadcastMutableMessage) IsPersisted() bool {
 	ret := _m.Called()

pkg/mocks/streaming/util/mock_message/mock_ImmutableMessage.go

@@ -3,7 +3,9 @@
 package mock_message
 
 import (
+	messagespb "github.com/milvus-io/milvus/pkg/v2/proto/messagespb"
 	message "github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
+
 	mock "github.com/stretchr/testify/mock"
 
 	zapcore "go.uber.org/zap/zapcore"
@@ -159,6 +161,53 @@ func (_c *MockImmutableMessage_EstimateSize_Call) RunAndReturn(run func() int) *
 	return _c
 }
 
+// IntoMessageProto provides a mock function with no fields
+func (_m *MockImmutableMessage) IntoMessageProto() *messagespb.Message {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for IntoMessageProto")
+	}
+
+	var r0 *messagespb.Message
+	if rf, ok := ret.Get(0).(func() *messagespb.Message); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*messagespb.Message)
+		}
+	}
+
+	return r0
+}
+
+// MockImmutableMessage_IntoMessageProto_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IntoMessageProto'
+type MockImmutableMessage_IntoMessageProto_Call struct {
+	*mock.Call
+}
+
+// IntoMessageProto is a helper method to define mock.On call
+func (_e *MockImmutableMessage_Expecter) IntoMessageProto() *MockImmutableMessage_IntoMessageProto_Call {
+	return &MockImmutableMessage_IntoMessageProto_Call{Call: _e.mock.On("IntoMessageProto")}
+}
+
+func (_c *MockImmutableMessage_IntoMessageProto_Call) Run(run func()) *MockImmutableMessage_IntoMessageProto_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockImmutableMessage_IntoMessageProto_Call) Return(_a0 *messagespb.Message) *MockImmutableMessage_IntoMessageProto_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockImmutableMessage_IntoMessageProto_Call) RunAndReturn(run func() *messagespb.Message) *MockImmutableMessage_IntoMessageProto_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // IsPersisted provides a mock function with no fields
 func (_m *MockImmutableMessage) IsPersisted() bool {
 	ret := _m.Called()

pkg/mocks/streaming/util/mock_message/mock_ImmutableTxnMessage.go

@@ -3,7 +3,9 @@
 package mock_message
 
 import (
+	messagespb "github.com/milvus-io/milvus/pkg/v2/proto/messagespb"
 	message "github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
+
 	mock "github.com/stretchr/testify/mock"
 
 	zapcore "go.uber.org/zap/zapcore"
@@ -253,6 +255,53 @@ func (_c *MockImmutableTxnMessage_EstimateSize_Call) RunAndReturn(run func() int
 	return _c
 }
 
+// IntoMessageProto provides a mock function with no fields
+func (_m *MockImmutableTxnMessage) IntoMessageProto() *messagespb.Message {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for IntoMessageProto")
+	}
+
+	var r0 *messagespb.Message
+	if rf, ok := ret.Get(0).(func() *messagespb.Message); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*messagespb.Message)
+		}
+	}
+
+	return r0
+}
+
+// MockImmutableTxnMessage_IntoMessageProto_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IntoMessageProto'
+type MockImmutableTxnMessage_IntoMessageProto_Call struct {
+	*mock.Call
+}
+
+// IntoMessageProto is a helper method to define mock.On call
+func (_e *MockImmutableTxnMessage_Expecter) IntoMessageProto() *MockImmutableTxnMessage_IntoMessageProto_Call {
+	return &MockImmutableTxnMessage_IntoMessageProto_Call{Call: _e.mock.On("IntoMessageProto")}
+}
+
+func (_c *MockImmutableTxnMessage_IntoMessageProto_Call) Run(run func()) *MockImmutableTxnMessage_IntoMessageProto_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockImmutableTxnMessage_IntoMessageProto_Call) Return(_a0 *messagespb.Message) *MockImmutableTxnMessage_IntoMessageProto_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockImmutableTxnMessage_IntoMessageProto_Call) RunAndReturn(run func() *messagespb.Message) *MockImmutableTxnMessage_IntoMessageProto_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // IsPersisted provides a mock function with no fields
 func (_m *MockImmutableTxnMessage) IsPersisted() bool {
 	ret := _m.Called()

pkg/mocks/streaming/util/mock_message/mock_MutableMessage.go

@@ -3,7 +3,9 @@
 package mock_message
 
 import (
+	messagespb "github.com/milvus-io/milvus/pkg/v2/proto/messagespb"
 	message "github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
+
 	mock "github.com/stretchr/testify/mock"
 
 	zapcore "go.uber.org/zap/zapcore"
@@ -207,6 +209,53 @@ func (_c *MockMutableMessage_IntoImmutableMessage_Call) RunAndReturn(run func(me
 	return _c
 }
 
+// IntoMessageProto provides a mock function with no fields
+func (_m *MockMutableMessage) IntoMessageProto() *messagespb.Message {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for IntoMessageProto")
+	}
+
+	var r0 *messagespb.Message
+	if rf, ok := ret.Get(0).(func() *messagespb.Message); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*messagespb.Message)
+		}
+	}
+
+	return r0
+}
+
+// MockMutableMessage_IntoMessageProto_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IntoMessageProto'
+type MockMutableMessage_IntoMessageProto_Call struct {
+	*mock.Call
+}
+
+// IntoMessageProto is a helper method to define mock.On call
+func (_e *MockMutableMessage_Expecter) IntoMessageProto() *MockMutableMessage_IntoMessageProto_Call {
+	return &MockMutableMessage_IntoMessageProto_Call{Call: _e.mock.On("IntoMessageProto")}
+}
+
+func (_c *MockMutableMessage_IntoMessageProto_Call) Run(run func()) *MockMutableMessage_IntoMessageProto_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockMutableMessage_IntoMessageProto_Call) Return(_a0 *messagespb.Message) *MockMutableMessage_IntoMessageProto_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockMutableMessage_IntoMessageProto_Call) RunAndReturn(run func() *messagespb.Message) *MockMutableMessage_IntoMessageProto_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // IsPersisted provides a mock function with no fields
 func (_m *MockMutableMessage) IsPersisted() bool {
 	ret := _m.Called()

pkg/mocks/streaming/util/mock_message/mock_RProperties.go

@@ -1,181 +0,0 @@
-// Code generated by mockery v2.53.3. DO NOT EDIT.
-
-package mock_message
-
-import mock "github.com/stretchr/testify/mock"
-
-// MockRProperties is an autogenerated mock type for the RProperties type
-type MockRProperties struct {
-	mock.Mock
-}
-
-type MockRProperties_Expecter struct {
-	mock *mock.Mock
-}
-
-func (_m *MockRProperties) EXPECT() *MockRProperties_Expecter {
-	return &MockRProperties_Expecter{mock: &_m.Mock}
-}
-
-// Exist provides a mock function with given fields: key
-func (_m *MockRProperties) Exist(key string) bool {
-	ret := _m.Called(key)
-
-	if len(ret) == 0 {
-		panic("no return value specified for Exist")
-	}
-
-	var r0 bool
-	if rf, ok := ret.Get(0).(func(string) bool); ok {
-		r0 = rf(key)
-	} else {
-		r0 = ret.Get(0).(bool)
-	}
-
-	return r0
-}
-
-// MockRProperties_Exist_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Exist'
-type MockRProperties_Exist_Call struct {
-	*mock.Call
-}
-
-// Exist is a helper method to define mock.On call
-//   - key string
-func (_e *MockRProperties_Expecter) Exist(key interface{}) *MockRProperties_Exist_Call {
-	return &MockRProperties_Exist_Call{Call: _e.mock.On("Exist", key)}
-}
-
-func (_c *MockRProperties_Exist_Call) Run(run func(key string)) *MockRProperties_Exist_Call {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
-	})
-	return _c
-}
-
-func (_c *MockRProperties_Exist_Call) Return(_a0 bool) *MockRProperties_Exist_Call {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *MockRProperties_Exist_Call) RunAndReturn(run func(string) bool) *MockRProperties_Exist_Call {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Get provides a mock function with given fields: key
-func (_m *MockRProperties) Get(key string) (string, bool) {
-	ret := _m.Called(key)
-
-	if len(ret) == 0 {
-		panic("no return value specified for Get")
-	}
-
-	var r0 string
-	var r1 bool
-	if rf, ok := ret.Get(0).(func(string) (string, bool)); ok {
-		return rf(key)
-	}
-	if rf, ok := ret.Get(0).(func(string) string); ok {
-		r0 = rf(key)
-	} else {
-		r0 = ret.Get(0).(string)
-	}
-
-	if rf, ok := ret.Get(1).(func(string) bool); ok {
-		r1 = rf(key)
-	} else {
-		r1 = ret.Get(1).(bool)
-	}
-
-	return r0, r1
-}
-
-// MockRProperties_Get_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Get'
-type MockRProperties_Get_Call struct {
-	*mock.Call
-}
-
-// Get is a helper method to define mock.On call
-//   - key string
-func (_e *MockRProperties_Expecter) Get(key interface{}) *MockRProperties_Get_Call {
-	return &MockRProperties_Get_Call{Call: _e.mock.On("Get", key)}
-}
-
-func (_c *MockRProperties_Get_Call) Run(run func(key string)) *MockRProperties_Get_Call {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
-	})
-	return _c
-}
-
-func (_c *MockRProperties_Get_Call) Return(value string, ok bool) *MockRProperties_Get_Call {
-	_c.Call.Return(value, ok)
-	return _c
-}
-
-func (_c *MockRProperties_Get_Call) RunAndReturn(run func(string) (string, bool)) *MockRProperties_Get_Call {
-	_c.Call.Return(run)
-	return _c
-}
-
-// ToRawMap provides a mock function with no fields
-func (_m *MockRProperties) ToRawMap() map[string]string {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for ToRawMap")
-	}
-
-	var r0 map[string]string
-	if rf, ok := ret.Get(0).(func() map[string]string); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(map[string]string)
-		}
-	}
-
-	return r0
-}
-
-// MockRProperties_ToRawMap_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ToRawMap'
-type MockRProperties_ToRawMap_Call struct {
-	*mock.Call
-}
-
-// ToRawMap is a helper method to define mock.On call
-func (_e *MockRProperties_Expecter) ToRawMap() *MockRProperties_ToRawMap_Call {
-	return &MockRProperties_ToRawMap_Call{Call: _e.mock.On("ToRawMap")}
-}
-
-func (_c *MockRProperties_ToRawMap_Call) Run(run func()) *MockRProperties_ToRawMap_Call {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *MockRProperties_ToRawMap_Call) Return(_a0 map[string]string) *MockRProperties_ToRawMap_Call {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *MockRProperties_ToRawMap_Call) RunAndReturn(run func() map[string]string) *MockRProperties_ToRawMap_Call {
-	_c.Call.Return(run)
-	return _c
-}
-
-// NewMockRProperties creates a new instance of MockRProperties. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
-// The first argument is typically a *testing.T value.
-func NewMockRProperties(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *MockRProperties {
-	mock := &MockRProperties{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

pkg/streaming/util/message/message.go

@@ -3,6 +3,8 @@ package message
 import (
 	"go.uber.org/zap/zapcore"
 	"google.golang.org/protobuf/proto"
+
+	"github.com/milvus-io/milvus/pkg/v2/proto/messagespb"
 )
 
 var (
@@ -24,7 +26,10 @@ type BasicMessage interface {
 	// from 1: new version after streamingnode.
 	Version() Version
 
-	// Message payload.
+	// Payload returns the message payload.
+	// If the underlying message is encrypted, the payload will be decrypted.
+	// !!! So if the message is encrypted, additional overhead will be paid for decryption.
+	// If the underlying message is not encrypted, the payload will be returned directly.
 	Payload() []byte
 
 	// EstimateSize returns the estimated size of message.
@@ -53,6 +58,9 @@ type BasicMessage interface {
 
 	// IsPersisted returns true if the message is persisted into underlying log storage.
 	IsPersisted() bool
+
+	// IntoMessageProto converts the message to a protobuf message.
+	IntoMessageProto() *messagespb.Message
 }
 
 // MutableMessage is the mutable message interface.

pkg/streaming/util/message/message_impl.go

@@ -56,6 +56,14 @@ func (m *messageImpl) IsPersisted() bool {
 	return !m.properties.Exist(messageNotPersisteted)
 }
 
+// IntoMessageProto converts the message to a protobuf message.
+func (m *messageImpl) IntoMessageProto() *messagespb.Message {
+	return &messagespb.Message{
+		Payload:    m.payload,
+		Properties: m.properties.ToRawMap(),
+	}
+}
+
 // EstimateSize returns the estimated size of current message.
 func (m *messageImpl) EstimateSize() int {
 	if ch := m.cipherHeader(); ch != nil {

pkg/streaming/walimpls/impls/kafka/wal.go

@@ -30,7 +30,8 @@ func (w *walImpl) Append(ctx context.Context, msg message.MutableMessage) (messa
 		panic("write on a wal that is not in read-write mode")
 	}
 
-	properties := msg.Properties().ToRawMap()
+	pb := msg.IntoMessageProto()
+	properties := pb.Properties
 	headers := make([]kafka.Header, 0, len(properties))
 	for key, value := range properties {
 		header := kafka.Header{Key: key, Value: []byte(value)}
@@ -41,7 +42,7 @@ func (w *walImpl) Append(ctx context.Context, msg message.MutableMessage) (messa
 
 	if err := w.p.Produce(&kafka.Message{
 		TopicPartition: kafka.TopicPartition{Topic: &topic, Partition: 0},
-		Value:          msg.Payload(),
+		Value:          pb.Payload,
 		Headers:        headers,
 	}, ch); err != nil {
 		return nil, err

pkg/streaming/walimpls/impls/pulsar/wal.go

@@ -82,14 +82,15 @@ func (w *walImpl) Append(ctx context.Context, msg message.MutableMessage) (messa
 	if err != nil {
 		return nil, errors.Wrap(err, "get producer from future")
 	}
+	pb := msg.IntoMessageProto()
 	id, err := p.Send(ctx, &pulsar.ProducerMessage{
-		Payload:    msg.Payload(),
-		Properties: msg.Properties().ToRawMap(),
+		Payload:    pb.Payload,
+		Properties: pb.Properties,
 	})
 	if w.backlogClearHelper != nil {
 		// Observe the append traffic even if the message is not sent successfully.
 		// Because if the write is failed, the message may be already written to the pulsar topic.
-		w.backlogClearHelper.ObserveAppend(len(msg.Payload()))
+		w.backlogClearHelper.ObserveAppend(msg.EstimateSize())
 	}
 	if err != nil {
 		w.Log().RatedWarn(1, "send message to pulsar failed", zap.Error(err))

pkg/streaming/walimpls/impls/rmq/wal.go

@@ -34,10 +34,10 @@ func (w *walImpl) Append(ctx context.Context, msg message.MutableMessage) (messa
 	if w.Channel().AccessMode != types.AccessModeRW {
 		panic("write on a wal that is not in read-write mode")
 	}
-
+	pb := msg.IntoMessageProto()
 	id, err := w.p.SendForStreamingService(&common.ProducerMessage{
-		Payload:    msg.Payload(),
-		Properties: msg.Properties().ToRawMap(),
+		Payload:    pb.Payload,
+		Properties: pb.Properties,
 	})
 	if err != nil {
 		w.Log().RatedWarn(1, "send message to rmq failed", zap.Error(err))

pkg/streaming/walimpls/impls/walimplstest/message_log.go

@@ -45,10 +45,11 @@ func (l *messageLog) Append(_ context.Context, msg message.MutableMessage) (mess
 	l.cond.LockAndBroadcast()
 	defer l.cond.L.Unlock()
 	id := l.id
+	pb := msg.IntoMessageProto()
 	newEntry := entry{
 		ID:         id,
-		Payload:    msg.Payload(),
-		Properties: msg.Properties().ToRawMap(),
+		Payload:    pb.Payload,
+		Properties: pb.Properties,
 	}
 	data, err := json.Marshal(newEntry)
 	if err != nil {

pkg/streaming/walimpls/impls/wp/wal.go

@@ -31,11 +31,11 @@ func (w *walImpl) Append(ctx context.Context, msg message.MutableMessage) (messa
 	if w.Channel().AccessMode != types.AccessModeRW {
 		panic("write on a wal that is not in read-write mode")
 	}
-
+	pb := msg.IntoMessageProto()
 	r := w.p.Write(ctx,
 		&wp.WriterMessage{
-			Payload:    msg.Payload(),
-			Properties: msg.Properties().ToRawMap(),
+			Payload:    pb.Payload,
+			Properties: pb.Properties,
 		},
 	)
 	if r.Err != nil {