test/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2024-11-30 02:48:45 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

enhance: get msg type from the msg header to reduce the Unmarshal usage (#36409)

Browse Source 
/kind improvement

Signed-off-by: SimFG <bang.fu@zilliz.com>
This commit is contained in:
SimFG 2024-09-24 10:19:20 +08:00 committed by GitHub
parent 51cdee1f84
commit ddadefcb62
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 41 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
pkg/mq/common/message.go
View File

@ -16,6 +16,14 @@
package common
import (
"fmt"
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
)
// ProducerOptions contains the options of a producer
type ProducerOptions struct {
// The topic that this Producer will publish
@ -65,3 +73,30 @@ const (
// SubscriptionPositionUnkown indicates we don't care about the consumer location, since we are doing another seek or only some meta api over that
SubscriptionPositionUnknown
)
const MsgTypeKey = "msg_type"
func GetMsgType(msg Message) (commonpb.MsgType, error) {
msgType := commonpb.MsgType_Undefined
properties := msg.Properties()
if properties != nil {
if val, ok := properties[MsgTypeKey]; ok {
msgType = commonpb.MsgType(commonpb.MsgType_value[val])
}
}
if msgType == commonpb.MsgType_Undefined {
header := commonpb.MsgHeader{}
if msg.Payload() == nil {
return msgType, fmt.Errorf("failed to unmarshal message header, payload is empty")
}
err := proto.Unmarshal(msg.Payload(), &header)
if err != nil {
return msgType, fmt.Errorf("failed to unmarshal message header, err %s", err.Error())
}
if header.Base == nil {
return msgType, fmt.Errorf("failed to unmarshal message, header is uncomplete")
}
msgType = header.Base.MsgType
}
return msgType, nil
}

17
pkg/mq/msgstream/mq_msgstream.go
View File

@ -325,7 +325,9 @@ func (ms *mqMsgStream) Produce(msgPack *MsgPack) error {
return err
}
msg := &common.ProducerMessage{Payload: m, Properties: map[string]string{}}
msg := &common.ProducerMessage{Payload: m, Properties: map[string]string{
common.MsgTypeKey: v.Msgs[i].Type().String(),
}}
InjectCtx(spanCtx, msg.Properties)
ms.producerLock.RLock()
@ -396,18 +398,11 @@ func (ms *mqMsgStream) getTsMsgFromConsumerMsg(msg common.Message) (TsMsg, error
// GetTsMsgFromConsumerMsg get TsMsg from consumer message
func GetTsMsgFromConsumerMsg(unmarshalDispatcher UnmarshalDispatcher, msg common.Message) (TsMsg, error) {
header := commonpb.MsgHeader{}
if msg.Payload() == nil {
return nil, fmt.Errorf("failed to unmarshal message header, payload is empty")
}
err := proto.Unmarshal(msg.Payload(), &header)
msgType, err := common.GetMsgType(msg)
if err != nil {
return nil, fmt.Errorf("failed to unmarshal message header, err %s", err.Error())
return nil, err
}
if header.Base == nil {
return nil, fmt.Errorf("failed to unmarshal message, header is uncomplete")
}
tsMsg, err := unmarshalDispatcher.Unmarshal(msg.Payload(), header.Base.MsgType)
tsMsg, err := unmarshalDispatcher.Unmarshal(msg.Payload(), msgType)
if err != nil {
return nil, fmt.Errorf("failed to unmarshal tsMsg, err %s", err.Error())
}