CDH: fix ttrpc cdh memory occupation

Fixes confidential-containers#688 This commit will copy only the pointer of CDH to implement different APIs. In old version, each CDH instance will serve for one API, thus occupies multiple times of memory. Signed-off-by: Xynnn007 <[email protected]>
Xynnn007 · Sep 27, 2024 · 387e3cb · 387e3cb
1 parent 49476b5
commit 387e3cb
Show file tree

Hide file tree

Showing 11 changed files with 60 additions and 93 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -65,7 +65,7 @@ tokio = "1.39"
 toml = "0.8.19"
 tonic = "0.9"
 tonic-build = "0.9"
-ttrpc = "0.8.0"
+ttrpc = "0.8.2"
 ttrpc-codegen = "0.4.2"
 url = "2.5.2"
 uuid = "1"

diff --git a/api-server-rest/src/ttrpc_proto/attestation_agent_ttrpc.rs b/api-server-rest/src/ttrpc_proto/attestation_agent_ttrpc.rs
@@ -1,4 +1,4 @@
-// This file is generated by ttrpc-compiler 0.6.2. Do not edit
+// This file is generated by ttrpc-compiler 0.6.3. Do not edit
 // @generated
 
 #![cfg_attr(rustfmt, rustfmt_skip)]
@@ -43,7 +43,7 @@ impl AttestationAgentServiceClient {
 }
 
 struct GetEvidenceMethod {
-    service: Arc<Box<dyn AttestationAgentService + Send + Sync>>,
+    service: Arc<dyn AttestationAgentService + Send + Sync>,
 }
 
 #[async_trait]
@@ -54,7 +54,7 @@ impl ::ttrpc::r#async::MethodHandler for GetEvidenceMethod {
 }
 
 struct GetTokenMethod {
-    service: Arc<Box<dyn AttestationAgentService + Send + Sync>>,
+    service: Arc<dyn AttestationAgentService + Send + Sync>,
 }
 
 #[async_trait]
@@ -74,7 +74,7 @@ pub trait AttestationAgentService: Sync {
     }
 }
 
-pub fn create_attestation_agent_service(service: Arc<Box<dyn AttestationAgentService + Send + Sync>>) -> HashMap<String, ::ttrpc::r#async::Service> {
+pub fn create_attestation_agent_service(service: Arc<dyn AttestationAgentService + Send + Sync>) -> HashMap<String, ::ttrpc::r#async::Service> {
     let mut ret = HashMap::new();
     let mut methods = HashMap::new();
     let streams = HashMap::new();

diff --git a/api-server-rest/src/ttrpc_proto/confidential_data_hub_ttrpc.rs b/api-server-rest/src/ttrpc_proto/confidential_data_hub_ttrpc.rs
@@ -1,4 +1,4 @@
-// This file is generated by ttrpc-compiler 0.6.2. Do not edit
+// This file is generated by ttrpc-compiler 0.6.3. Do not edit
 // @generated
 
 #![cfg_attr(rustfmt, rustfmt_skip)]
@@ -38,7 +38,7 @@ impl GetResourceServiceClient {
 }
 
 struct GetResourceMethod {
-    service: Arc<Box<dyn GetResourceService + Send + Sync>>,
+    service: Arc<dyn GetResourceService + Send + Sync>,
 }
 
 #[async_trait]
@@ -55,7 +55,7 @@ pub trait GetResourceService: Sync {
     }
 }
 
-pub fn create_get_resource_service(service: Arc<Box<dyn GetResourceService + Send + Sync>>) -> HashMap<String, ::ttrpc::r#async::Service> {
+pub fn create_get_resource_service(service: Arc<dyn GetResourceService + Send + Sync>) -> HashMap<String, ::ttrpc::r#async::Service> {
     let mut ret = HashMap::new();
     let mut methods = HashMap::new();
     let streams = HashMap::new();

diff --git a/attestation-agent/kbs_protocol/src/token_provider/aa/attestation_agent_ttrpc.rs b/attestation-agent/kbs_protocol/src/token_provider/aa/attestation_agent_ttrpc.rs
@@ -1,4 +1,4 @@
-// This file is generated by ttrpc-compiler 0.6.2. Do not edit
+// This file is generated by ttrpc-compiler 0.6.3. Do not edit
 // @generated
 
 #![cfg_attr(rustfmt, rustfmt_skip)]
@@ -63,7 +63,7 @@ impl AttestationAgentServiceClient {
 }
 
 struct GetEvidenceMethod {
-    service: Arc<Box<dyn AttestationAgentService + Send + Sync>>,
+    service: Arc<dyn AttestationAgentService + Send + Sync>,
 }
 
 #[async_trait]
@@ -74,7 +74,7 @@ impl ::ttrpc::r#async::MethodHandler for GetEvidenceMethod {
 }
 
 struct GetTokenMethod {
-    service: Arc<Box<dyn AttestationAgentService + Send + Sync>>,
+    service: Arc<dyn AttestationAgentService + Send + Sync>,
 }
 
 #[async_trait]
@@ -85,7 +85,7 @@ impl ::ttrpc::r#async::MethodHandler for GetTokenMethod {
 }
 
 struct ExtendRuntimeMeasurementMethod {
-    service: Arc<Box<dyn AttestationAgentService + Send + Sync>>,
+    service: Arc<dyn AttestationAgentService + Send + Sync>,
 }
 
 #[async_trait]
@@ -96,7 +96,7 @@ impl ::ttrpc::r#async::MethodHandler for ExtendRuntimeMeasurementMethod {
 }
 
 struct CheckInitDataMethod {
-    service: Arc<Box<dyn AttestationAgentService + Send + Sync>>,
+    service: Arc<dyn AttestationAgentService + Send + Sync>,
 }
 
 #[async_trait]
@@ -107,7 +107,7 @@ impl ::ttrpc::r#async::MethodHandler for CheckInitDataMethod {
 }
 
 struct UpdateConfigurationMethod {
-    service: Arc<Box<dyn AttestationAgentService + Send + Sync>>,
+    service: Arc<dyn AttestationAgentService + Send + Sync>,
 }
 
 #[async_trait]
@@ -118,7 +118,7 @@ impl ::ttrpc::r#async::MethodHandler for UpdateConfigurationMethod {
 }
 
 struct GetTeeTypeMethod {
-    service: Arc<Box<dyn AttestationAgentService + Send + Sync>>,
+    service: Arc<dyn AttestationAgentService + Send + Sync>,
 }
 
 #[async_trait]
@@ -150,7 +150,7 @@ pub trait AttestationAgentService: Sync {
     }
 }
 
-pub fn create_attestation_agent_service(service: Arc<Box<dyn AttestationAgentService + Send + Sync>>) -> HashMap<String, ::ttrpc::r#async::Service> {
+pub fn create_attestation_agent_service(service: Arc<dyn AttestationAgentService + Send + Sync>) -> HashMap<String, ::ttrpc::r#async::Service> {
     let mut ret = HashMap::new();
     let mut methods = HashMap::new();
     let streams = HashMap::new();

diff --git a/confidential-data-hub/hub/src/bin/protos/api_ttrpc.rs b/confidential-data-hub/hub/src/bin/protos/api_ttrpc.rs
@@ -1,4 +1,4 @@
-// This file is generated by ttrpc-compiler 0.6.2. Do not edit
+// This file is generated by ttrpc-compiler 0.6.3. Do not edit
 // @generated
 
 #![cfg_attr(rustfmt, rustfmt_skip)]
@@ -38,7 +38,7 @@ impl SealedSecretServiceClient {
 }
 
 struct UnsealSecretMethod {
-    service: Arc<Box<dyn SealedSecretService + Send + Sync>>,
+    service: Arc<dyn SealedSecretService + Send + Sync>,
 }
 
 #[async_trait]
@@ -55,7 +55,7 @@ pub trait SealedSecretService: Sync {
     }
 }
 
-pub fn create_sealed_secret_service(service: Arc<Box<dyn SealedSecretService + Send + Sync>>) -> HashMap<String, ::ttrpc::r#async::Service> {
+pub fn create_sealed_secret_service(service: Arc<dyn SealedSecretService + Send + Sync>) -> HashMap<String, ::ttrpc::r#async::Service> {
     let mut ret = HashMap::new();
     let mut methods = HashMap::new();
     let streams = HashMap::new();
@@ -86,7 +86,7 @@ impl GetResourceServiceClient {
 }
 
 struct GetResourceMethod {
-    service: Arc<Box<dyn GetResourceService + Send + Sync>>,
+    service: Arc<dyn GetResourceService + Send + Sync>,
 }
 
 #[async_trait]
@@ -103,7 +103,7 @@ pub trait GetResourceService: Sync {
     }
 }
 
-pub fn create_get_resource_service(service: Arc<Box<dyn GetResourceService + Send + Sync>>) -> HashMap<String, ::ttrpc::r#async::Service> {
+pub fn create_get_resource_service(service: Arc<dyn GetResourceService + Send + Sync>) -> HashMap<String, ::ttrpc::r#async::Service> {
     let mut ret = HashMap::new();
     let mut methods = HashMap::new();
     let streams = HashMap::new();
@@ -134,7 +134,7 @@ impl SecureMountServiceClient {
 }
 
 struct SecureMountMethod {
-    service: Arc<Box<dyn SecureMountService + Send + Sync>>,
+    service: Arc<dyn SecureMountService + Send + Sync>,
 }
 
 #[async_trait]
@@ -151,7 +151,7 @@ pub trait SecureMountService: Sync {
     }
 }
 
-pub fn create_secure_mount_service(service: Arc<Box<dyn SecureMountService + Send + Sync>>) -> HashMap<String, ::ttrpc::r#async::Service> {
+pub fn create_secure_mount_service(service: Arc<dyn SecureMountService + Send + Sync>) -> HashMap<String, ::ttrpc::r#async::Service> {
     let mut ret = HashMap::new();
     let mut methods = HashMap::new();
     let streams = HashMap::new();
@@ -182,7 +182,7 @@ impl ImagePullServiceClient {
 }
 
 struct PullImageMethod {
-    service: Arc<Box<dyn ImagePullService + Send + Sync>>,
+    service: Arc<dyn ImagePullService + Send + Sync>,
 }
 
 #[async_trait]
@@ -199,7 +199,7 @@ pub trait ImagePullService: Sync {
     }
 }
 
-pub fn create_image_pull_service(service: Arc<Box<dyn ImagePullService + Send + Sync>>) -> HashMap<String, ::ttrpc::r#async::Service> {
+pub fn create_image_pull_service(service: Arc<dyn ImagePullService + Send + Sync>) -> HashMap<String, ::ttrpc::r#async::Service> {
     let mut ret = HashMap::new();
     let mut methods = HashMap::new();
     let streams = HashMap::new();

diff --git a/confidential-data-hub/hub/src/bin/protos/keyprovider_ttrpc.rs b/confidential-data-hub/hub/src/bin/protos/keyprovider_ttrpc.rs
@@ -1,4 +1,4 @@
-// This file is generated by ttrpc-compiler 0.6.2. Do not edit
+// This file is generated by ttrpc-compiler 0.6.3. Do not edit
 // @generated
 
 #![cfg_attr(rustfmt, rustfmt_skip)]
@@ -43,7 +43,7 @@ impl KeyProviderServiceClient {
 }
 
 struct WrapKeyMethod {
-    service: Arc<Box<dyn KeyProviderService + Send + Sync>>,
+    service: Arc<dyn KeyProviderService + Send + Sync>,
 }
 
 #[async_trait]
@@ -54,7 +54,7 @@ impl ::ttrpc::r#async::MethodHandler for WrapKeyMethod {
 }
 
 struct UnWrapKeyMethod {
-    service: Arc<Box<dyn KeyProviderService + Send + Sync>>,
+    service: Arc<dyn KeyProviderService + Send + Sync>,
 }
 
 #[async_trait]
@@ -74,7 +74,7 @@ pub trait KeyProviderService: Sync {
     }
 }
 
-pub fn create_key_provider_service(service: Arc<Box<dyn KeyProviderService + Send + Sync>>) -> HashMap<String, ::ttrpc::r#async::Service> {
+pub fn create_key_provider_service(service: Arc<dyn KeyProviderService + Send + Sync>) -> HashMap<String, ::ttrpc::r#async::Service> {
     let mut ret = HashMap::new();
     let mut methods = HashMap::new();
     let streams = HashMap::new();

diff --git a/confidential-data-hub/hub/src/bin/ttrpc-cdh.rs b/confidential-data-hub/hub/src/bin/ttrpc-cdh.rs
@@ -41,14 +41,6 @@ struct Cli {
     config: Option<String>,
 }
 
-macro_rules! ttrpc_service {
-    ($func: expr, $conf: expr) => {{
-        let server = Server::new($conf).await?;
-        let server = Arc::new(Box::new(server) as _);
-        $func(server)
-    }};
-}
-
 #[tokio::main]
 async fn main() -> Result<()> {
     env_logger::init_from_env(env_logger::Env::new().default_filter_or("info"));
@@ -64,20 +56,17 @@ async fn main() -> Result<()> {
     create_socket_parent_directory(unix_socket_path).await?;
     clean_previous_sock_file(unix_socket_path).await?;
 
-    let sealed_secret_service = ttrpc_service!(create_sealed_secret_service, &config);
-    let get_resource_service = ttrpc_service!(create_get_resource_service, &config);
-    let key_provider_service = ttrpc_service!(create_key_provider_service, &config);
-    let secure_mount_service = ttrpc_service!(create_secure_mount_service, &config);
-    let image_pull_service = ttrpc_service!(create_image_pull_service, &config);
+    let server = Server::new(&config).await.context("create CDH instance")?;
+    let server = Arc::new(server);
 
     let mut server = TtrpcServer::new()
         .bind(&config.socket)
         .context("cannot bind cdh ttrpc service")?
-        .register_service(sealed_secret_service)
-        .register_service(get_resource_service)
-        .register_service(secure_mount_service)
-        .register_service(key_provider_service)
-        .register_service(image_pull_service);
+        .register_service(create_sealed_secret_service(server.clone() as _))
+        .register_service(create_get_resource_service(server.clone() as _))
+        .register_service(create_key_provider_service(server.clone() as _))
+        .register_service(create_secure_mount_service(server.clone() as _))
+        .register_service(create_image_pull_service(server.clone() as _));
 
     info!(
         "[ttRPC] Confidential Data Hub starts to listen to request: {}",