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