ttrpc-rust: Support the new version of codegen

The ttrpc provides a "prost" feature to support the new version of
codegen. An "example2" has been added to demonstrate how to use the
codegen.

Signed-off-by: Xuewei Niu <niuxuewei.nxw@antgroup.com>

Author: justxuewei

Cargo.toml

@@ -11,7 +11,8 @@ homepage = "https://github.com/containerd/ttrpc-rust"
 description = "A Rust version of ttrpc."
 
 [dependencies]
-protobuf = { version = "3.1.0" }
+prost = { version = "0.11", optional = true }
+protobuf = "3.1.0"
 libc = { version = "0.2.59", features = [ "extra_traits" ] }
 nix = "0.23.0"
 log = "0.4"
@@ -27,11 +28,13 @@ tokio-vsock = { version = "0.3.1", optional = true }
 
 [build-dependencies]
 protobuf-codegen = "3.1.0"
+prost-build = "0.11"
 
 [features]
 default = ["sync"]
-async = ["async-trait", "tokio", "futures", "tokio-vsock"]
+async = ["dep:async-trait", "dep:tokio", "dep:futures", "dep:tokio-vsock"]
 sync = []
+prost = ["dep:prost"]
 
 [package.metadata.docs.rs]
 all-features = true

README.md

@@ -136,3 +136,19 @@ cargo update
 cargo install --force --path .
 ```
 3. Build your project.
+
+# ttrpc-rust with the Prost
+
+The new version of the ttrpc-rust is built with the Prost crate, a modern
+protobuf compiler written by Rust. There are certain different behaviors from
+the Rust-protobuf version:
+
+1. The protoc should be installed.
+2. Enabling "prost" feature for the ttrpc-rust.
+3. The Rust files are named based on their package name, rather than the proto
+   filename, e.g. `ttrpc = { version = "1.0", features = ["prost"] }`.
+4. Some variable names are different, e.g. for "cpu", "CPU" is generated by the
+   Rust-protobuf, and "Cpu" is genereated by the Prost.
+
+The "example" is an example with the Rust-protobuf version, and the "example2"
+is an example with the Prost version.

build.rs

@@ -7,6 +7,11 @@ fn main() {
     let path: PathBuf = [out_dir.clone(), "mod.rs".to_string()].iter().collect();
     fs::write(path, "pub mod ttrpc;").unwrap();
 
+    generate_ttrpc(&out_dir);
+}
+
+#[cfg(not(feature = "prost"))]
+fn generate_ttrpc(out_dir: &str) {
     let customize = protobuf_codegen::Customize::default()
         .gen_mod_rs(false)
         .generate_accessors(true);
@@ -20,3 +25,13 @@ fn main() {
         .run()
         .expect("Codegen failed.");
 }
+
+#[cfg(feature = "prost")]
+fn generate_ttrpc(out_dir: &str) {
+    prost_build::Config::new()
+        .out_dir(out_dir)
+        .compile_well_known_types()
+        .protoc_arg("--experimental_allow_proto3_optional")
+        .compile_protos(&["src/ttrpc.proto"], &["src"])
+        .expect("Codegen failed")
+}

codegen/Cargo.toml

@@ -2,6 +2,14 @@
 name = "ttrpc-codegen"
 version = "1.0.0"
 edition = "2021"
+authors = ["The Ant Group Kata Team <kata@list.alibaba-inc.com>"]
+license = "Apache-2.0"
+keywords = ["codegen", "ttrpc", "protobuf"]
+description = "Rust codegen for ttrpc using prost crate"
+categories = ["network-programming", "development-tools::build-utils"]
+repository = "https://github.com/containerd/ttrpc-rust/tree/master/codegen"
+homepage = "https://github.com/containerd/ttrpc-rust/tree/master/codegen"
+readme = "README.md"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 

codegen/README.md

@@ -0,0 +1,34 @@
+# Ttrpc-rust Codegen
+
+## Getting started
+
+Please ensure that the protoc has been installed on your local environment. Then
+write the following code into "build.rs".
+
+```rust
+let mut protos = vec![
+    "protocols/protos/health.proto",
+    "protocols/protos/agent.proto",
+    "protocols/protos/oci.proto",
+];
+
+let includes = vec!["protocols/protos"];
+
+let codegen = CodegenBuilder::new()
+    .set_out_dir(&"protocols/sync")
+    .set_protos(&protos)
+    .set_includes(&includes)
+    .set_serde(true)
+    .set_async_mode(AsyncMode::None)
+    .set_generate_service(true)
+    .build()
+    .unwrap();
+codegen.generate().unwrap();
+```
+
+Add ttrpc-codegen to "build-dependencies" section in "Cargo.toml".
+
+```toml
+[build-dependencies]
+ttrpc-codegen = "1.0"
+```

codegen/src/codegen.rs

@@ -1,20 +1,26 @@
 use anyhow::{anyhow, Context, Result};
+use prost::Message;
 use prost_build::Config;
+use prost_types::FileDescriptorSet;
 use std::{
     fs::{self, File},
-    io::{Read, Write},
-    path::Path,
+    io::{BufReader, Read, Write},
+    path::{Path, PathBuf},
 };
 
 use crate::svcgen::{AsyncMode, TtrpcServiceGenerator};
 
+const FILE_DESCRIPTOR_SET: &str = "fd_set.bin";
+
 pub struct Codegen<'a, P: AsRef<Path>> {
     out_dir: &'a P,
     protos: &'a [P],
     includes: &'a [P],
     /// Whether to enable serde
     serde: bool,
     async_mode: AsyncMode,
+    /// Whether to generate service
+    generate_service: bool,
 }
 
 impl<'a, P> Codegen<'a, P>
@@ -24,27 +30,30 @@ where
     pub fn generate(&self) -> Result<()> {
         self.compile_protos().context("Compile protos")?;
         self.write_header().context("Write header")?;
+        self.clean_up().context("Clean up")?;
 
         Ok(())
     }
 
+    // TODO: Do not write header to the files that already has the header
+    // TODO: Write header to the files generated by the codegen
     fn write_header(&self) -> Result<()> {
-        let dir = fs::read_dir(self.out_dir.as_ref()).context("Read out_dir")?;
-
-        for child in dir {
-            let entry = child.context("Unable to get entry")?;
-            let path = entry.path();
-            let ext = match path.extension() {
-                Some(ext) => match ext.to_str() {
-                    Some(ext) => ext,
-                    None => "",
-                },
-                None => "",
+        // Read fd_set.bin
+        let f = File::open(PathBuf::from(self.out_dir.as_ref()).join(FILE_DESCRIPTOR_SET))
+            .context("Open fd_set.bin")?;
+        let mut reader = BufReader::new(f);
+        let mut buffer = Vec::new();
+        reader.read_to_end(&mut buffer).context("Read fd_set.bin")?;
+
+        let fd_set = FileDescriptorSet::decode(&buffer as &[u8]).context("Decode fd_set")?;
+
+        for fd in fd_set.file.iter() {
+            let rs_path =
+                PathBuf::from(self.out_dir.as_ref()).join(&format!("{}.rs", fd.package()));
+            let mut f = match File::open(&rs_path) {
+                Ok(f) => f,
+                _ => continue,
             };
-            if entry.file_name() == "mod.rs" || ext != "rs" {
-                continue;
-            }
-
             let header = format!(
                 r#"// This file is generated by ttrpc-codegen {}. Do not edit
 // @generated
@@ -55,12 +64,11 @@ where
 
             let mut buf = Vec::<u8>::new();
             buf.write(header.as_bytes()).context("Write header")?;
-            let mut f = File::open(&path).context(format!("Open rust file {:?}", path))?;
             f.read_to_end(&mut buf)
-                .context(format!("Read from rust file {:?}", path))?;
-            let mut f = File::create(&path).context(format!("Open rust file {:?}", path))?;
+                .context(format!("Read from rust file {:?}", rs_path))?;
+            let mut f = File::create(&rs_path).context(format!("Open rust file {:?}", rs_path))?;
             f.write_all(buf.as_slice())
-                .context(format!("Write to rust file {:?}", path))?;
+                .context(format!("Write to rust file {:?}", rs_path))?;
         }
 
         Ok(())
@@ -72,39 +80,47 @@ where
         config.service_generator(Box::new(TtrpcServiceGenerator::new(self.async_mode)));
         config.protoc_arg("--experimental_allow_proto3_optional");
         config.compile_well_known_types();
-        config.include_file("_include.rs");
-        if self.serde {
-            config.message_attribute(".", "#[derive(::serde::Serialize, ::serde::Deserialize)]");
+        config.file_descriptor_set_path(
+            PathBuf::from(self.out_dir.as_ref()).join(FILE_DESCRIPTOR_SET),
+        );
+        if self.generate_service {
+            config.include_file("_include.rs");
+            if self.serde {
+                config
+                    .message_attribute(".", "#[derive(::serde::Serialize, ::serde::Deserialize)]");
+            }
         }
         config
             .compile_protos(self.protos, self.includes)
             .context("Compile protos by prost")?;
         Ok(())
     }
+
+    fn clean_up(&self) -> Result<()> {
+        fs::remove_file(PathBuf::from(self.out_dir.as_ref()).join(FILE_DESCRIPTOR_SET))
+            .context("Remove fd_set.bin")?;
+        Ok(())
+    }
 }
 
 #[derive(Default)]
-pub struct CodegenBuilder<'a, P: AsRef<Path>> {
+pub struct CodegenBuilder<'a, P: AsRef<Path> + Default> {
     out_dir: Option<&'a P>,
     protos: Option<&'a [P]>,
     includes: Option<&'a [P]>,
     /// Whether to enable serde
     serde: Option<bool>,
     async_mode: Option<AsyncMode>,
+    /// Whether to generate service
+    generate_service: Option<bool>,
 }
 
 impl<'a, P> CodegenBuilder<'a, P>
 where
-    P: AsRef<Path>,
+    P: AsRef<Path> + Default,
 {
     pub fn new() -> Self {
-        Self {
-            out_dir: None,
-            protos: None,
-            includes: None,
-            serde: None,
-            async_mode: None,
-        }
+        Default::default()
     }
 
     pub fn set_out_dir(mut self, out_dir: &'a P) -> Self {
@@ -132,6 +148,11 @@ where
         self
     }
 
+    pub fn set_generate_service(mut self, generate_service: bool) -> Self {
+        self.generate_service = Some(generate_service);
+        self
+    }
+
     pub fn build(&self) -> Result<Codegen<'a, P>> {
         let out_dir = match self.out_dir {
             Some(out_dir) => out_dir,
@@ -158,12 +179,18 @@ where
             None => AsyncMode::None,
         };
 
+        let generate_service = match self.generate_service {
+            Some(gen) => gen,
+            None => false,
+        };
+
         Ok(Codegen {
             out_dir,
             protos,
             includes,
             serde,
             async_mode,
+            generate_service,
         })
     }
 }