feat(volo-build): Split grpc generated files into several files (#540)

* Cleanup: better variable names in thrift_backend.rs * Split grpc generated files into several * Cleanup: move common functions from grpc_backend.rs and thrift_backend.rs to util.rs * Cleanup: cargo fmt * Cleanup: remove unused imports * Cleanup: imports order * Cleanup: imports order v2 * Cleanup: better variable names in thrift_backend.rs * Split grpc generated files into several * Cleanup: move common functions from grpc_backend.rs and thrift_backend.rs to util.rs * Cleanup: cargo fmt * Cleanup: remove unused imports * Cleanup: imports order * Cleanup: imports order v2 * Use streaming.proto for tests * Remove obsolete echo.proto
cloudwego · Dec 3, 2024 · 03e22de · 03e22de
1 parent 0d39960
commit 03e22de
Show file tree

Hide file tree

Showing 7 changed files with 233 additions and 68 deletions.
diff --git a/tests/code-generation-workspace-split-grpc/.gitignore b/tests/code-generation-workspace-split-grpc/.gitignore
@@ -0,0 +1,2 @@
+rpc_streaming
+Cargo.lock
diff --git a/tests/code-generation-workspace-split-grpc/Cargo.toml b/tests/code-generation-workspace-split-grpc/Cargo.toml
@@ -0,0 +1,38 @@
+[[bin]]
+bench = false
+name = "gen"
+test = false
+
+[dependencies.pilota-build]
+version = "*"
+
+[dependencies.volo-build]
+workspace = true
+
+[package]
+edition = "2021"
+name = "code-generation-workspace-split-grpc"
+publish = false
+version = "0.0.0"
+
+[workspace]
+members = []
+
+[workspace.dependencies]
+anyhow = "1"
+async-trait = "0.1"
+lazy_static = "1"
+serde = "1"
+volo-grpc = "*"
+
+[workspace.dependencies.pilota]
+version = "*"
+
+[workspace.dependencies.volo]
+path = "../../volo"
+
+[workspace.dependencies.volo-build]
+path = "../../volo-build"
+
+[workspace.dependencies.volo-thrift]
+path = "../../volo-thrift"
diff --git a/tests/code-generation-workspace-split-grpc/src/bin/gen.rs b/tests/code-generation-workspace-split-grpc/src/bin/gen.rs
@@ -0,0 +1,7 @@
+use volo_build::plugin::SerdePlugin;
+
+fn main() {
+    volo_build::workspace::Builder::protobuf()
+        .plugin(SerdePlugin)
+        .gen()
+}
diff --git a/tests/code-generation-workspace-split-grpc/volo.workspace.yml b/tests/code-generation-workspace-split-grpc/volo.workspace.yml
@@ -0,0 +1,21 @@
+common_crate_name: "common" # common_crate_name = "common" by default
+touch_all: true
+dedups: [] # remove the repeated structure generation in one entry
+special_namings: []
+split_generated_files: true
+# repos: # exsit if non-local
+# { repo }: # repo = extract the name from url by default
+# url: { git } # url = repo git url
+# ref: { ref } # ref = "HEAD" by default
+# lock: { commit hash } # lock is the last commit hash
+services:
+  - idl:
+      source: local
+      path: ../../examples/proto/streaming.proto
+      includes:
+      - ../../examples/proto
+    codegen_option:
+      touch: []
+      keep_unknown_fields: false # A->B->C, B could use this to transfer the unknown fields which is needed by A and C.
+      config:
+        crate_name: rpc_streaming
diff --git a/volo-build/src/grpc_backend.rs b/volo-build/src/grpc_backend.rs
@@ -8,6 +8,8 @@ use pilota_build::{
 };
 use volo::FastStr;
 
+use crate::util::{get_base_dir, write_file, write_item};
+
 pub struct MkGrpcBackend;
 
 impl pilota_build::MakeBackend for MkGrpcBackend {
@@ -254,6 +256,16 @@ impl CodegenBackend for VoloGrpcBackend {
         let resp_enum_name_send = format!("{}ResponseSend", service_name);
         let req_enum_name_recv = format!("{}RequestRecv", service_name);
         let resp_enum_name_recv = format!("{}ResponseRecv", service_name);
+
+        let path = self.cx().item_path(def_id);
+        let path = path.as_ref();
+        let buf = get_base_dir(self.cx().mode.as_ref(), self.cx().names.get(&def_id), path);
+        let base_dir = buf.as_path();
+
+        if self.cx().split {
+            std::fs::create_dir_all(base_dir).expect("Failed to create base directory");
+        }
+
         let paths = s
             .methods
             .iter()
@@ -426,8 +438,9 @@ impl CodegenBackend for VoloGrpcBackend {
             }}"
         );
 
-        stream.push_str(&format! {
-            r#"pub enum {req_enum_name_send} {{
+        let req_enum_send_impl = format! {
+            r#"
+            pub enum {req_enum_name_send} {{
                 {req_enum_send_variants}
             }}
 
@@ -437,8 +450,11 @@ impl CodegenBackend for VoloGrpcBackend {
                         {req_send_into_body}
                     }}
                 }}
-            }}
+            }}"#
+        };
 
+        let req_enum_recv_impl = format! {
+            r#"
             pub enum {req_enum_name_recv} {{
                 {req_enum_recv_variants}
             }}
@@ -450,8 +466,11 @@ impl CodegenBackend for VoloGrpcBackend {
                         _ => ::std::result::Result::Err(::volo_grpc::Status::new(::volo_grpc::Code::Unimplemented, "Method not found.")),
                     }}
                 }}
-            }}
+            }}"#
+        };
 
+        let resp_enum_send_impl = format! {
+            r#"
             pub enum {resp_enum_name_send} {{
                 {resp_enum_send_variants}
             }}
@@ -462,8 +481,11 @@ impl CodegenBackend for VoloGrpcBackend {
                         {resp_send_into_body}
                     }}
                 }}
-            }}
+            }}"#
+        };
 
+        let resp_enum_recv_impl = format! {
+            r#"
             pub enum {resp_enum_name_recv} {{
                 {resp_enum_recv_variants}
             }}
@@ -478,8 +500,11 @@ impl CodegenBackend for VoloGrpcBackend {
                         _ => ::std::result::Result::Err(::volo_grpc::Status::new(::volo_grpc::Code::Unimplemented, "Method not found.")),
                     }}
                 }}
-            }}
+            }}"#
+        };
 
+        let client_impl = format! {
+            r#"
             pub struct {client_builder_name} {{}}
             impl {client_builder_name} {{
                 pub fn new(
@@ -522,8 +547,11 @@ impl CodegenBackend for VoloGrpcBackend {
 
             impl<S: ::volo::client::OneShotService<::volo_grpc::context::ClientContext,::volo_grpc::Request<{req_enum_name_send}>, Response=::volo_grpc::Response<{resp_enum_name_recv}>, Error = ::volo_grpc::Status> + Send + Sync + 'static> {oneshot_client_name}<S> {{
                 {oneshot_client_methods}
-            }}
+            }}"#
+        };
 
+        let server_impl = format! {
+            r#"
             pub struct {server_name}<S> {{
                 inner: ::std::sync::Arc<S>,
             }}
@@ -570,7 +598,69 @@ impl CodegenBackend for VoloGrpcBackend {
             impl<S: {service_name}> ::volo_grpc::server::NamedService for {server_name}<S> {{
                 const NAME: &'static str = "{name}";
             }}"#
-        });
+        };
+
+        if self.cx().split {
+            let mut mod_rs_stream = String::new();
+            write_item(
+                &mut mod_rs_stream,
+                base_dir,
+                format!("enum_{}.rs", req_enum_name_send),
+                req_enum_send_impl,
+            );
+            write_item(
+                &mut mod_rs_stream,
+                base_dir,
+                format!("enum_{}.rs", req_enum_name_recv),
+                req_enum_recv_impl,
+            );
+            write_item(
+                &mut mod_rs_stream,
+                base_dir,
+                format!("enum_{}.rs", resp_enum_name_send),
+                resp_enum_send_impl,
+            );
+            write_item(
+                &mut mod_rs_stream,
+                base_dir,
+                format!("enum_{}.rs", resp_enum_name_recv),
+                resp_enum_recv_impl,
+            );
+
+            write_item(
+                &mut mod_rs_stream,
+                base_dir,
+                format!("client_{}.rs", client_name),
+                client_impl,
+            );
+            write_item(
+                &mut mod_rs_stream,
+                base_dir,
+                format!("server_{}.rs", server_name),
+                server_impl,
+            );
+
+            let mod_rs_file_path = base_dir.join("mod.rs");
+            write_file(&mod_rs_file_path, mod_rs_stream);
+            stream.push_str(
+                format!(
+                    "include!(\"{}/mod.rs\");",
+                    base_dir.file_name().unwrap().to_str().unwrap()
+                )
+                .as_str(),
+            );
+        } else {
+            stream.push_str(&format! {
+                r#"
+            {req_enum_send_impl}
+            {req_enum_recv_impl}
+            {resp_enum_send_impl}
+            {resp_enum_recv_impl}
+
+            {client_impl}
+            {server_impl}
+            "#});
+        }
     }
 
     fn codegen_service_method(&self, _service_def_id: DefId, method: &rir::Method) -> String {

diff --git a/volo-build/src/thrift_backend.rs b/volo-build/src/thrift_backend.rs
@@ -1,17 +1,18 @@
-use std::{io::Write, path::Path};
+use std::path::Path;
 
 use itertools::Itertools;
 use pilota_build::{
     codegen::thrift::DecodeHelper,
     db::RirDatabase,
-    middle::context::Mode,
     rir::{self, Method},
     tags::RustWrapperArc,
     CodegenBackend, Context, DefId, IdentName, Symbol, ThriftBackend,
 };
 use quote::format_ident;
 use volo::FastStr;
 
+use crate::util::{get_base_dir, write_file, write_item};
+
 #[derive(Clone)]
 pub struct VoloThriftBackend {
     inner: ThriftBackend,
@@ -320,25 +321,25 @@ impl VoloThriftBackend {
             "#
             };
 
-            Self::write_item(
+            write_item(
                 stream,
                 base_dir,
                 format!("enum_{}.rs", &req_recv_name),
                 req_recv_stream,
             );
-            Self::write_item(
+            write_item(
                 stream,
                 base_dir,
                 format!("enum_{}.rs", &res_recv_name),
                 res_recv_stream,
             );
-            Self::write_item(
+            write_item(
                 stream,
                 base_dir,
                 format!("enum_{}.rs", &req_send_name),
                 req_send_stream,
             );
-            Self::write_item(
+            write_item(
                 stream,
                 base_dir,
                 format!("enum_{}.rs", &res_send_name),
@@ -375,20 +376,6 @@ impl VoloThriftBackend {
         }
     }
 
-    fn write_item(stream: &mut String, base_dir: &Path, name: String, impl_str: String) {
-        let req_recv_buf = base_dir.join(&name);
-        let req_recv_file = req_recv_buf.as_path();
-        Self::write_file(req_recv_file, impl_str);
-        stream.push_str(format!("include!(\"{}\");", &name).as_str());
-    }
-
-    fn write_file(path: &Path, stream: String) {
-        let mut req_file_writer = std::io::BufWriter::new(std::fs::File::create(path).unwrap());
-        req_file_writer.write_all(stream.as_bytes()).unwrap();
-        req_file_writer.flush().unwrap();
-        pilota_build::fmt::fmt_file(path);
-    }
-
     fn method_ty_path(&self, service_name: &Symbol, method: &Method, suffix: &str) -> FastStr {
         match method.source {
             rir::MethodSource::Extend(def_id) => {
@@ -469,43 +456,8 @@ impl pilota_build::CodegenBackend for VoloThriftBackend {
 
         let path = self.cx().item_path(def_id);
         let path = path.as_ref();
-
-        // Locate directory based on the full item path
-        let base_dir = match self.cx().mode.as_ref() {
-            // In a workspace mode, the base directory is next to the `.rs` file for the service
-            Mode::Workspace(info) => {
-                let mut dir = info.dir.clone();
-                if path.is_empty() {
-                    dir
-                } else {
-                    dir.push(path[0].0.as_str());
-                    if path.len() > 1 {
-                        dir.push("src");
-                        for segment in path.iter().skip(1) {
-                            dir.push(Path::new(segment.0.as_str()));
-                        }
-                    }
-                    dir
-                }
-            }
-            // In single file mode, the files directory is the root
-            // The base directory path is the root + the item path
-            Mode::SingleFile { file_path } => {
-                let mut dir = file_path.clone();
-                dir.pop();
-                for segment in path {
-                    dir.push(Path::new(segment.0.as_str()));
-                }
-                dir
-            }
-        };
-
-        let base_dir = if let Some(suffix) = self.cx().names.get(&def_id) {
-            format!("{}_{suffix}", base_dir.display())
-        } else {
-            base_dir.display().to_string()
-        };
-        let base_dir = Path::new(&base_dir);
+        let buf = get_base_dir(self.cx().mode.as_ref(), self.cx().names.get(&def_id), path);
+        let base_dir = buf.as_path();
 
         if self.cx().split {
             std::fs::create_dir_all(base_dir).expect("Failed to create base directory");
@@ -747,13 +699,13 @@ impl pilota_build::CodegenBackend for VoloThriftBackend {
         };
 
         if self.cx().split {
-            Self::write_item(
+            write_item(
                 &mut mod_rs_stream,
                 base_dir,
                 format!("service_{}Server.rs", service_name),
                 server_string,
             );
-            Self::write_item(
+            write_item(
                 &mut mod_rs_stream,
                 base_dir,
                 format!("service_{}Client.rs", service_name),
@@ -772,7 +724,7 @@ impl pilota_build::CodegenBackend for VoloThriftBackend {
 
         if self.cx().split {
             let mod_rs_file_path = base_dir.join("mod.rs");
-            Self::write_file(&mod_rs_file_path, mod_rs_stream);
+            write_file(&mod_rs_file_path, mod_rs_stream);
             stream.push_str(
                 format!(
                     "include!(\"{}/mod.rs\");",