golang 中根据 protobuf message name 动态实例化 protobuf 消息,消息内容通过输入 json 文件指定
背景:
项目中使用 protobuf 作为 rpc 调用协议,计划用 golang 实现一个压测工具,希望能够指定 message name 和 json 动态的构建 protobuf 消息;从 json 解析到 golang protobuf message 可以用 jsonpb,这里使用 UnmarshalString,函数签名
func UnmarshalString(str string, pb proto.Message) error
str 是 json 字符串,message 是自定义的 proto messgae 接口。于是剩下需要做的是通过 message name 获取对应的 proto.Message 接口,搜索了一下,对于 golang 没有找到很好的方法,检查了 protoc 生成的 golang 消息文件,可以按以下方式根据 message name 获取到 message type,然后利用 golang reflect 包实例画消息;
解决方式:
简单来说,就是需要根据 message name 获取到 message type, 然后利用 golang 反射实例化消息结构。从 message name 获取 message type,最直观的是维护一个 map[string]reflect.Type 的字典,protoc 生成的 golang 代码已经包含这个字典,自定义的 message 会通过 proto.RegisterType 注册到 protoTypes 和 revProtoTypes 这两个结构中,并提供 MessageName 和 MessageType 用来通过 name 获取 type 或者反之, 相关代码在 proto/properties.go 中, 由此可以实现通过 message name 获取到 message type 进而实例化消息的功能。
其它包括 enum 类型,extensions 都有相应的注册/获取函数 proto.RegisterEnum, proto.RegisterExtension;
示例:
以下以一个 rpc 消息定义为例实现从消息名称实例化一个消息实例,完整代码见 https://github.com/wuyidong/parse_pb_by_name_golang
以下一个简单 protobuf 做 rpc 协议的简单例子,我们在 package rpc 中定义了协议的一般格式,由协议头(Head)和消息本身(Body)组成,Body 全部为可选字段,用于填充具体的协议,Head 为固定格式, 其中 Head.message_type 用于标识 Body 所带有的协议类型,服务端根据 Head.message_type 路由到具体的处理过程,具体的协议如 CreateAccountRequest/CreateAccountResponse 等都作为 rpc.Body 的可选字段。
rpc.proto --> rpc 消息格式
package rpc; message RPCMessage { // 消息头部 required Head head = 1; // 消息内容 required Body body = 2; }; message Head { // 请求 uuid required string session_no = 1; // 请求消息类型 required int32 message_type = 2; }; message Body { extensions 1000 to max; }; message ResponseCode { required int32 retcode = 1; // 返回码 optional string error_messgae = 2; // 返回失败时,错误信息 };
account.proto --> 账户相关操作
package rpc.account; import "rpc.proto"; enum MessageType { CREATE_ACCOUNT_REQUEST = 1001; CREATE_ACCOUNT_RESPONSE = 1002; DELETE_ACCOUNT_REQUEST = 1003; DELETE_ACCOUNT_RESPONSE = 1004; // ... }; extend rpc.Body { optional CreateAccountRequest create_account_request = 1001; optional CreateAccountResponse create_account_response = 1002; // ... }; // account 相关操作接口 message CreateAccountRequest { required string email = 1; optional string name = 2; // 不指定则为 email optional string passwd = 3; // 初始密码为 email }; message CreateAccountResponse { required ResponseCode rc = 1; };
// ...
proto 代码编译之后,rpc.account 包被命名为 rpc_account, 以 CreateAccountRequest 为例,可以看到 protoc 编译后生成了如下 golang 代码:
// CreateAccountRequest 结构体定义 type CreateAccountRequest struct { Email *string `protobuf:"bytes,1,req,name=email" json:"email,omitempty"` Name *string `protobuf:"bytes,2,opt,name=name" json:"name,omitempty"` Passwd *string `protobuf:"bytes,3,opt,name=passwd" json:"passwd,omitempty"` XXX_unrecognized []byte `json:"-"` } // proto.Message 接口指定的函数 func (m *CreateAccountRequest) Reset() { *m = CreateAccountRequest{} } func (m *CreateAccountRequest) String() string { return proto.CompactTextString(m) } func (*CreateAccountRequest) ProtoMessage() {} // extension type 定义 var E_CreateAccountRequest = &proto.ExtensionDesc{ ExtendedType: (*rpc.Body)(nil), ExtensionType: (*CreateAccountRequest)(nil), Field: 1001, Name: "rpc.account.create_account_request", Tag: "bytes,1001,opt,name=create_account_request", Filename: "account.proto", } // 注册定义结构到 proto func init() { proto.RegisterType((*CreateAccountRequest)(nil), "rpc.account.CreateAccountRequest") proto.RegisterEnum("rpc.account.MessageType", MessageType_name, MessageType_value) proto.RegisterExtension(E_CreateAccountRequest) }
其中 init 函数中三个 Register*** 函数将 CreateAccountRequest 相关信息注册到 proto 包中:
// github.com/golang/protobuf/proto/properties.go 中 func RegisterEnum(typeName string, unusedNameMap map[int32]string, valueMap map[string]int32) { if _, ok := enumValueMaps[typeName]; ok { panic("proto: duplicate enum registered: " + typeName) } enumValueMaps[typeName] = valueMap } // EnumValueMap returns the mapping from names to integers of the // enum type enumType, or a nil if not found. func EnumValueMap(enumType string) map[string]int32 { return enumValueMaps[enumType] } func RegisterType(x Message, name string) { if _, ok := protoTypes[name]; ok { // TODO: Some day, make this a panic. log.Printf("proto: duplicate proto type registered: %s", name) return } t := reflect.TypeOf(x) protoTypes[name] = t revProtoTypes[t] = name } // MessageName returns the fully-qualified proto name for the given message type. func MessageName(x Message) string { type xname interface { XXX_MessageName() string } if m, ok := x.(xname); ok { return m.XXX_MessageName() } return revProtoTypes[reflect.TypeOf(x)] } // MessageType returns the message type (pointer to struct) for a named message. func MessageType(name string) reflect.Type { return protoTypes[name] } // github.com/golang/protobuf/proto/extensions.go 中 // RegisterExtension is called from the generated code. func RegisterExtension(desc *ExtensionDesc) { st := reflect.TypeOf(desc.ExtendedType).Elem() m := extensionMaps[st] if m == nil { m = make(map[int32]*ExtensionDesc) extensionMaps[st] = m } if _, ok := m[desc.Field]; ok { panic("proto: duplicate extension registered: " + st.String() + " " + strconv.Itoa(int(desc.Field))) } m[desc.Field] = desc } // RegisteredExtensions returns a map of the registered extensions of a // protocol buffer struct, indexed by the extension number. // The argument pb should be a nil pointer to the struct type. func RegisteredExtensions(pb Message) map[int32]*ExtensionDesc { return extensionMaps[reflect.TypeOf(pb).Elem()] }
对照 Register*** 的实现,可以看到通过 E_CreateAccountRequest 类型是注册到了 extensionMaps 下,这是个两层的map, map[extendedType]map[messageField]ExtensionType,messageFlied 为 rpc.Body 的字段标识,因此我们根据 RegisteredExtensions(rpc.Body) 可以获取到 rpc.Body 下所有的 extension 消息类型,messageFlied 则和我们之前在 MessageType 中定义的对应各消息的枚举类型一致,可以通过 EnumValueMap(rpc.account.MessageType)[rpc.account.create_account_request] 取到,因此可以通过消息名称获取到消息对应的 ExtensionDesc 类型, 其中 ExtensionType 即为消息类型。对应的我们用以下代码通过给定消息名称实例化一个消息结构:
// message id -> *proto.ExtensionDesc // 记录 rpc.Body 的拓展消息 var RPCMessageBodyExtensions map[int32]*proto.ExtensionDesc func init() { RPCMessageBodyExtensions = proto.RegisteredExtensions((*rpc.Body)(nil)) } // some utils for UMessage // msgName: rpc.account.create_account_request func GetRPCMessageObjectByName(msgName string) (msg proto.Message, err error) { msgType := reflect.TypeOf(GetRPCMessageExtension(msgName).ExtensionType) if msgType == nil { err = fmt.Errorf("can't find message type") return } // msgType is pointer msg = reflect.Indirect(reflect.New(msgType.Elem())).Addr().Interface().(proto.Message) return } // msgName: rpc.account.create_account_request // namePrefix: rpc.account // name: create_account_request func GetNamePrefix(msgName string) (prefix string) { items := strings.Split(msgName, ".") prefix = strings.Join(items[0:len(items)-1], ".") return } func GetName(msgName string) (name string) { items := strings.Split(msgName, ".") name = items[len(items)-1] return } func GetRPCMessageId(msgName string) (msgId int32) { msgTypeName := GetNamePrefix(msgName) + ".MessageType" mapMsgNameId := proto.EnumValueMap(msgTypeName) msgId = mapMsgNameId[strings.ToUpper(GetName(msgName))] return } func GetRPCMessageExtension(msgName string) (extension *proto.ExtensionDesc) { msgId := GetRPCMessageId(msgName) extension = RPCMessageBodyExtensions[msgId] return }