rust下根据protobuf的消息名创建对象实例

在C++里面, 我们可以根据一个消息的名称, 动态的创建一个实例

google::protobuf::Descriptor* desc =
    google::protobuf::DescriptorPool::generated_pool()
        ->FindMessageTypeByName("mypkg.MyType");
google::protobuf::Message* message =
    google::protobuf::MessageFactory::generated_factory()
        ->GetPrototype(desc)->New();

这个在protobuf里面是集成进去了, 在其他语言也有类似的东西. 

 

通过这个, 我们就让轻松实现编解码库, 而不需去构造一个映射表.

 

但是, 但是在rust里面, 是没有这种东西的. 比较难的地方是rust全局变量必须要实现Send trait, 否则是不能被共享的, 这样做确实安全, 但是对于我们实现MessageFactory就变得困难.

好在rust有thread_local和build.rs, 我们可以通过build.rs在编译proto文件的时候去遍历, 把每个消息添加到一个thread_local的hash map里面去, 从而曲线救国.

 

//首先proto message的信息
#[derive(Debug)]
struct ProtoMessageInfo {
    file: PathBuf,
    file_name: String,
    messages : Vec<String>,
}

//接下来生成一个factory.rs
fn generate_factory_file(path:&str, v: &Vec<ProtoMessageInfo>) {
    let mut template  = File::open((path.to_string() + "/factory.rs.template").as_str()).unwrap();
    let mut contents = Vec::new();
    template.read_to_end(&mut contents);

    let mut mod_file = File::create((path.to_string() + "/mod.rs").as_str()).unwrap();
    let mut factory_file = File::create((path.to_string() + "/factory.rs").as_str()).unwrap();

    mod_file.write(b"pub mod factory;\n");

    factory_file.write_all(&contents[..]);
    factory_file.write(b"\n\n");

    for item in v.iter() {
        factory_file.write_fmt(format_args!("use crate::proto::{};\n", item.file_name));
        mod_file.write_fmt(format_args!("pub mod {};\n", item.file_name));
    }

    factory_file.write(b"\npub fn init_descriptors() {");

    for file in v.iter() {
        for msg in file.messages.iter() {
            factory_file.write_fmt(format_args!("\n    register_message::<{}::{}>();", file.file_name, msg));
        }
    }

    factory_file.write(b"\n}\n");
}

fn get_proto_list(v: &Vec<ProtoMessageInfo>) -> Vec<&str> {
    let mut r = Vec::new();

    for f in v.iter() {
        r.push(f.file.to_str().unwrap());
    }

    r
}

 

比如我们有一个t.proto文件

syntax="proto3";
package t111;
option optimize_for = SPEED;

message RpcRequestHandShake
{
int32 server_id     = 1;
int64 start_time     = 2;
}

 

在build.rs内添加这样的代码:

fn main() {
    let proto_path = "src/proto";

    let v = get_protos_info(proto_path);
    let protos = get_proto_list(&v);

    protoc_rust::run(protoc_rust::Args {
        out_dir: proto_path,
        input: &protos,
        includes: &[proto_path],
        customize: Customize {
          ..Default::default()
        },
    }).expect("protoc");

    generate_factory_file(proto_path, &v);
}

 

然后就可以看到生成的factory.rs:

use protobuf::reflect::MessageDescriptor;
use protobuf::Message;

use std::cell::RefCell;
use std::collections::HashMap;


thread_local! {
    pub static GLOBAL_MAP : RefCell<HashMap<String, &'static MessageDescriptor>> = RefCell::new(HashMap::new());
}

pub fn register_message<M: Message>() {
    GLOBAL_MAP.with(|x| {
        let mut m = x.borrow_mut();
        let name = M::descriptor_static().full_name().to_string();
        if !m.contains_key(&name) {
            m.insert(name, M::descriptor_static());
        }
    })
}

pub fn get_descriptor(full_name: String) -> Option<&'static MessageDescriptor> {
    GLOBAL_MAP.with(move |x| {
        {
            let m = x.borrow_mut();
            if m.len() == 0 {
                drop(m);
                init_descriptors()
            }
        }
        {
            let m = x.borrow_mut();
            match m.get(&full_name) {
                Some(r) => Some(*r),
                None => None,
            }
        }
    })
}

use crate::proto::t;

pub fn init_descriptors() {
    register_message::<t::RpcRequestHandShake>();
}

 

在main.rs里面这样使用:

mod proto;

use proto::factory::*;
use proto::rpc::*;

fn main() {
    let desc = get_descriptor("t111.RpcRequestHandShake".to_string()).unwrap();
    println!("{}", desc.full_name());
    let msg = desc.new_instance();
    println!("msg: {:?}", msg);
}

 

cargo run:

就可以看到

t111.RpcRequestHandShake

msg:

 

这时候就可以拿到MessageDesriptor, 通过这个对象可以new一个instance

 

还未整理代码, 到时候上传到crates.io上面去

 

 

关键字: MessageFactory, Protobuf, Rust