You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@teaclave.apache.org by ms...@apache.org on 2021/08/11 17:37:09 UTC

[incubator-teaclave] branch master updated: Use struture builder to construct complicated structures (#542)

This is an automated email from the ASF dual-hosted git repository.

mssun pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-teaclave.git


The following commit(s) were added to refs/heads/master by this push:
     new 6baa3ad  Use struture builder to construct complicated structures (#542)
6baa3ad is described below

commit 6baa3add34d73e8664a7164b81e0c00594122eaa
Author: He Sun <su...@gmail.com>
AuthorDate: Thu Aug 12 01:37:03 2021 +0800

    Use struture builder to construct complicated structures (#542)
---
 sdk/rust/src/lib.rs                                | 17 ++--
 services/execution/enclave/src/service.rs          | 15 ++--
 services/management/enclave/src/service.rs         | 35 ++++----
 services/proto/src/teaclave_frontend_service.rs    | 98 +++++++++++-----------
 services/proto/src/teaclave_management_service.rs  |  2 +
 .../enclave/src/end_to_end/builtin_echo.rs         |  5 +-
 .../enclave/src/end_to_end/builtin_gbdt_train.rs   |  5 +-
 .../enclave/src/end_to_end/mesapy_data_fusion.rs   | 10 ++-
 .../enclave/src/end_to_end/mesapy_echo.rs          |  5 +-
 tests/functional/enclave/src/execution_service.rs  |  5 +-
 tests/functional/enclave/src/management_service.rs | 10 ++-
 tests/functional/enclave/src/scheduler_service.rs  | 10 ++-
 tests/integration/enclave/src/teaclave_worker.rs   |  7 +-
 types/src/function.rs                              | 76 +++++++++--------
 types/src/staged_function.rs                       | 69 +++++++--------
 types/src/staged_task.rs                           | 84 +++++++++----------
 16 files changed, 246 insertions(+), 207 deletions(-)

diff --git a/sdk/rust/src/lib.rs b/sdk/rust/src/lib.rs
index 95c768c..80fad99 100644
--- a/sdk/rust/src/lib.rs
+++ b/sdk/rust/src/lib.rs
@@ -36,8 +36,8 @@ pub use teaclave_proto::teaclave_frontend_service::{
     ApproveTaskRequest, ApproveTaskResponse, AssignDataRequest, AssignDataResponse,
     CreateTaskRequest, CreateTaskResponse, GetFunctionRequest, GetFunctionResponse, GetTaskRequest,
     GetTaskResponse, InvokeTaskRequest, InvokeTaskResponse, RegisterFunctionRequest,
-    RegisterFunctionResponse, RegisterInputFileRequest, RegisterInputFileResponse,
-    RegisterOutputFileRequest, RegisterOutputFileResponse,
+    RegisterFunctionRequestBuilder, RegisterFunctionResponse, RegisterInputFileRequest,
+    RegisterInputFileResponse, RegisterOutputFileRequest, RegisterOutputFileResponse,
 };
 pub use teaclave_types::{
     EnclaveInfo, Executor, FileCrypto, FunctionInput, FunctionOutput, TaskResult,
@@ -202,22 +202,25 @@ impl FrontendClient {
         outputs: Option<Vec<FunctionOutput>>,
     ) -> Result<String> {
         let executor_type = executor_type.try_into()?;
-        let mut request = RegisterFunctionRequest::new()
+        let mut builder = RegisterFunctionRequestBuilder::new()
             .name(name)
             .description(description)
             .executor_type(executor_type);
+
         if let Some(payload) = payload {
-            request = request.payload(payload.into());
+            builder = builder.payload(payload.into());
         }
         if let Some(arguments) = arguments {
-            request = request.arguments(arguments);
+            builder = builder.arguments(arguments);
         }
         if let Some(inputs) = inputs {
-            request = request.inputs(inputs);
+            builder = builder.inputs(inputs);
         }
         if let Some(outputs) = outputs {
-            request = request.outputs(outputs);
+            builder = builder.outputs(outputs);
         }
+
+        let request = builder.build();
         let response = self.register_function_with_request(request)?;
 
         Ok(response.function_id.to_string())
diff --git a/services/execution/enclave/src/service.rs b/services/execution/enclave/src/service.rs
index 116e9af..ac348ea 100644
--- a/services/execution/enclave/src/service.rs
+++ b/services/execution/enclave/src/service.rs
@@ -157,7 +157,7 @@ fn prepare_task(task: &StagedTask, file_mgr: &TaskFileManager) -> Result<StagedF
     let input_files = file_mgr.prepare_staged_inputs()?;
     let output_files = file_mgr.prepare_staged_outputs()?;
 
-    let staged_function = StagedFunction::new()
+    let staged_function = StagedFunctionBuilder::new()
         .executor_type(task.executor_type)
         .executor(task.executor)
         .name(&task.function_name)
@@ -165,7 +165,8 @@ fn prepare_task(task: &StagedTask, file_mgr: &TaskFileManager) -> Result<StagedF
         .payload(task.function_payload.clone())
         .input_files(input_files)
         .output_files(output_files)
-        .runtime_name("default");
+        .runtime_name("default")
+        .build();
     Ok(staged_function)
 }
 
@@ -186,11 +187,12 @@ pub mod tests {
         let task_id = Uuid::new_v4();
         let function_arguments =
             FunctionArguments::from_json(json!({"message": "Hello, Teaclave!"})).unwrap();
-        let staged_task = StagedTask::new()
+        let staged_task = StagedTaskBuilder::new()
             .task_id(task_id)
             .executor(Executor::Builtin)
             .function_name("builtin-echo")
-            .function_arguments(function_arguments);
+            .function_arguments(function_arguments)
+            .build();
 
         let file_mgr = TaskFileManager::new(
             WORKER_BASE_DIR,
@@ -243,13 +245,14 @@ pub mod tests {
         let input_data = hashmap!("training_data" => training_input_data);
         let output_data = hashmap!("trained_model" => model_output_data);
 
-        let staged_task = StagedTask::new()
+        let staged_task = StagedTaskBuilder::new()
             .task_id(task_id)
             .executor(Executor::Builtin)
             .function_name("builtin-gbdt-train")
             .function_arguments(function_arguments)
             .input_data(input_data)
-            .output_data(output_data);
+            .output_data(output_data)
+            .build();
 
         let file_mgr = TaskFileManager::new(
             WORKER_BASE_DIR,
diff --git a/services/management/enclave/src/service.rs b/services/management/enclave/src/service.rs
index 550037e..e51f693 100644
--- a/services/management/enclave/src/service.rs
+++ b/services/management/enclave/src/service.rs
@@ -255,9 +255,10 @@ impl TeaclaveManagement for TeaclaveManagementService {
     ) -> TeaclaveServiceResponseResult<RegisterFunctionResponse> {
         let user_id = self.get_request_user_id(request.metadata())?;
 
-        let function = Function::from(request.message)
+        let function = FunctionBuilder::from(request.message)
             .id(Uuid::new_v4())
-            .owner(user_id);
+            .owner(user_id)
+            .build();
 
         self.write_to_db(&function)
             .map_err(|_| TeaclaveManagementServiceError::StorageError)?;
@@ -607,7 +608,7 @@ impl TeaclaveManagementService {
         let function_input2 = FunctionInput::new("input2", "input_desc");
         let function_output2 = FunctionOutput::new("output2", "output_desc");
 
-        let function = Function::new()
+        let function = FunctionBuilder::new()
             .id(Uuid::parse_str("00000000-0000-0000-0000-000000000001").unwrap())
             .name("mock-func-1")
             .description("mock-desc")
@@ -616,12 +617,13 @@ impl TeaclaveManagementService {
             .arguments(vec!["arg1".to_string(), "arg2".to_string()])
             .inputs(vec![function_input, function_input2])
             .outputs(vec![function_output, function_output2])
-            .owner("teaclave".to_string());
+            .owner("teaclave".to_string())
+            .build();
 
         self.write_to_db(&function)?;
 
         let function_output = FunctionOutput::new("output", "output_desc");
-        let function = Function::new()
+        let function = FunctionBuilder::new()
             .id(Uuid::parse_str("00000000-0000-0000-0000-000000000002").unwrap())
             .name("mock-func-2")
             .description("mock-desc")
@@ -629,7 +631,8 @@ impl TeaclaveManagementService {
             .public(true)
             .arguments(vec!["arg1".to_string()])
             .outputs(vec![function_output])
-            .owner("teaclave".to_string());
+            .owner("teaclave".to_string())
+            .build();
 
         self.write_to_db(&function)?;
         Ok(())
@@ -670,7 +673,7 @@ pub mod tests {
     pub fn handle_function() {
         let function_input = FunctionInput::new("input", "input_desc");
         let function_output = FunctionOutput::new("output", "output_desc");
-        let function = Function::new()
+        let function = FunctionBuilder::new()
             .id(Uuid::new_v4())
             .name("mock_function")
             .description("mock function")
@@ -679,7 +682,8 @@ pub mod tests {
             .inputs(vec![function_input])
             .outputs(vec![function_output])
             .public(true)
-            .owner("mock_user");
+            .owner("mock_user")
+            .build();
         assert!(Function::match_prefix(&function.key_string()));
         let value = function.to_vec().unwrap();
         let deserialized_function = Function::from_slice(&value).unwrap();
@@ -687,14 +691,15 @@ pub mod tests {
     }
 
     pub fn handle_task() {
-        let function = Function::new()
+        let function = FunctionBuilder::new()
             .id(Uuid::new_v4())
             .name("mock_function")
             .description("mock function")
             .payload(b"python script".to_vec())
             .arguments(vec!["arg".to_string()])
             .public(true)
-            .owner("mock_user");
+            .owner("mock_user")
+            .build();
         let function_arguments = FunctionArguments::from_json(json!({"arg": "data"})).unwrap();
 
         let task = Task::<Create>::new(
@@ -714,26 +719,28 @@ pub mod tests {
     }
 
     pub fn handle_staged_task() {
-        let function = Function::new()
+        let function = FunctionBuilder::new()
             .id(Uuid::new_v4())
             .name("mock_function")
             .description("mock function")
             .payload(b"python script".to_vec())
             .public(true)
-            .owner("mock_user");
+            .owner("mock_user")
+            .build();
 
         let url = Url::parse("s3://bucket_id/path?token=mock_token").unwrap();
         let cmac = FileAuthTag::mock();
         let input_data = FunctionInputFile::new(url.clone(), cmac, FileCrypto::default());
         let output_data = FunctionOutputFile::new(url, FileCrypto::default());
 
-        let staged_task = StagedTask::new()
+        let staged_task = StagedTaskBuilder::new()
             .task_id(Uuid::new_v4())
             .executor(Executor::MesaPy)
             .function_payload(function.payload)
             .function_arguments(hashmap!("arg" => "data"))
             .input_data(hashmap!("input" => input_data))
-            .output_data(hashmap!("output" => output_data));
+            .output_data(hashmap!("output" => output_data))
+            .build();
 
         let value = staged_task.to_vec().unwrap();
         let deserialized_data = StagedTask::from_slice(&value).unwrap();
diff --git a/services/proto/src/teaclave_frontend_service.rs b/services/proto/src/teaclave_frontend_service.rs
index 711f2d4..c68d7a1 100644
--- a/services/proto/src/teaclave_frontend_service.rs
+++ b/services/proto/src/teaclave_frontend_service.rs
@@ -26,12 +26,11 @@ use std::collections::HashMap;
 use std::prelude::v1::*;
 use teaclave_rpc::into_request;
 use teaclave_types::{
-    Executor, ExecutorType, ExternalID, FileAuthTag, FileCrypto, Function, FunctionArguments,
-    FunctionInput, FunctionOutput, OwnerList, TaskFileOwners, TaskResult, TaskStatus, UserID,
-    UserList,
+    Executor, ExecutorType, ExternalID, FileAuthTag, FileCrypto, FunctionArguments,
+    FunctionBuilder, FunctionInput, FunctionOutput, OwnerList, TaskFileOwners, TaskResult,
+    TaskStatus, UserID, UserList,
 };
 use url::Url;
-use uuid::Uuid;
 
 pub use proto::TeaclaveFrontend;
 pub use proto::TeaclaveFrontendClient;
@@ -276,78 +275,81 @@ pub struct RegisterFunctionRequest {
     pub outputs: Vec<FunctionOutput>,
 }
 
-impl RegisterFunctionRequest {
+pub struct RegisterFunctionRequestBuilder {
+    request: RegisterFunctionRequest,
+}
+
+impl RegisterFunctionRequestBuilder {
     pub fn new() -> Self {
-        Self {
+        let request = RegisterFunctionRequest {
             executor_type: ExecutorType::Builtin,
             public: true,
             ..Default::default()
-        }
+        };
+
+        Self { request }
     }
 
-    pub fn name(self, name: impl ToString) -> Self {
-        Self {
-            name: name.to_string(),
-            ..self
-        }
+    pub fn name(mut self, name: impl ToString) -> Self {
+        self.request.name = name.to_string();
+        self
     }
 
-    pub fn description(self, description: impl ToString) -> Self {
-        Self {
-            description: description.to_string(),
-            ..self
-        }
+    pub fn description(mut self, description: impl ToString) -> Self {
+        self.request.description = description.to_string();
+        self
     }
 
-    pub fn executor_type(self, executor_type: ExecutorType) -> Self {
-        Self {
-            executor_type,
-            ..self
-        }
+    pub fn executor_type(mut self, executor_type: ExecutorType) -> Self {
+        self.request.executor_type = executor_type;
+        self
     }
 
-    pub fn payload(self, payload: Vec<u8>) -> Self {
-        Self { payload, ..self }
+    pub fn payload(mut self, payload: Vec<u8>) -> Self {
+        self.request.payload = payload;
+        self
     }
 
-    pub fn public(self, public: bool) -> Self {
-        Self { public, ..self }
+    pub fn public(mut self, public: bool) -> Self {
+        self.request.public = public;
+        self
     }
 
-    pub fn arguments<T: IntoIterator>(self, args: T) -> Self
+    pub fn arguments<T: IntoIterator>(mut self, args: T) -> Self
     where
         <T as IntoIterator>::Item: ToString,
     {
-        Self {
-            arguments: args.into_iter().map(|x| x.to_string()).collect(),
-            ..self
-        }
+        self.request.arguments = args.into_iter().map(|x| x.to_string()).collect();
+        self
+    }
+
+    pub fn inputs(mut self, inputs: Vec<FunctionInput>) -> Self {
+        self.request.inputs = inputs;
+        self
     }
 
-    pub fn inputs(self, inputs: Vec<FunctionInput>) -> Self {
-        Self { inputs, ..self }
+    pub fn outputs(mut self, outputs: Vec<FunctionOutput>) -> Self {
+        self.request.outputs = outputs;
+        self
     }
 
-    pub fn outputs(self, outputs: Vec<FunctionOutput>) -> Self {
-        Self { outputs, ..self }
+    pub fn build(self) -> RegisterFunctionRequest {
+        self.request
     }
 }
 
 // We explicitly construct Function here in case of missing any field
-impl From<RegisterFunctionRequest> for Function {
+impl From<RegisterFunctionRequest> for FunctionBuilder {
     fn from(request: RegisterFunctionRequest) -> Self {
-        Function {
-            id: Uuid::default(),
-            owner: UserID::default(),
-            name: request.name,
-            description: request.description,
-            public: request.public,
-            executor_type: request.executor_type,
-            payload: request.payload,
-            arguments: request.arguments,
-            inputs: request.inputs,
-            outputs: request.outputs,
-        }
+        FunctionBuilder::new()
+            .name(request.name)
+            .description(request.description)
+            .public(request.public)
+            .executor_type(request.executor_type)
+            .payload(request.payload)
+            .arguments(request.arguments)
+            .inputs(request.inputs)
+            .outputs(request.outputs)
     }
 }
 
diff --git a/services/proto/src/teaclave_management_service.rs b/services/proto/src/teaclave_management_service.rs
index 73787b4..9dd8fb1 100644
--- a/services/proto/src/teaclave_management_service.rs
+++ b/services/proto/src/teaclave_management_service.rs
@@ -43,6 +43,8 @@ pub type GetInputFileResponse = crate::teaclave_frontend_service::GetInputFileRe
 pub type GetOutputFileRequest = crate::teaclave_frontend_service::GetOutputFileRequest;
 pub type GetOutputFileResponse = crate::teaclave_frontend_service::GetOutputFileResponse;
 pub type RegisterFunctionRequest = crate::teaclave_frontend_service::RegisterFunctionRequest;
+pub type RegisterFunctionRequestBuilder =
+    crate::teaclave_frontend_service::RegisterFunctionRequestBuilder;
 pub type RegisterFunctionResponse = crate::teaclave_frontend_service::RegisterFunctionResponse;
 pub type GetFunctionRequest = crate::teaclave_frontend_service::GetFunctionRequest;
 pub type GetFunctionResponse = crate::teaclave_frontend_service::GetFunctionResponse;
diff --git a/tests/functional/enclave/src/end_to_end/builtin_echo.rs b/tests/functional/enclave/src/end_to_end/builtin_echo.rs
index 67242bc..15a1f68 100644
--- a/tests/functional/enclave/src/end_to_end/builtin_echo.rs
+++ b/tests/functional/enclave/src/end_to_end/builtin_echo.rs
@@ -28,10 +28,11 @@ pub fn test_echo_task_success() {
         create_frontend_client(shared_enclave_info(), FRONTEND_SERVICE_ADDR, cred).unwrap();
 
     // Register Function
-    let request = RegisterFunctionRequest::new()
+    let request = RegisterFunctionRequestBuilder::new()
         .name("builtin-echo")
         .description("Native Echo Function")
-        .arguments(vec!["message"]);
+        .arguments(vec!["message"])
+        .build();
 
     let response = client.register_function(request).unwrap();
 
diff --git a/tests/functional/enclave/src/end_to_end/builtin_gbdt_train.rs b/tests/functional/enclave/src/end_to_end/builtin_gbdt_train.rs
index b69e5f7..675c61d 100644
--- a/tests/functional/enclave/src/end_to_end/builtin_gbdt_train.rs
+++ b/tests/functional/enclave/src/end_to_end/builtin_gbdt_train.rs
@@ -61,12 +61,13 @@ fn register_gbdt_function(client: &mut TeaclaveFrontendClient) -> ExternalID {
     ];
 
     // Register Function
-    let request = RegisterFunctionRequest::new()
+    let request = RegisterFunctionRequestBuilder::new()
         .name("builtin-gbdt-train")
         .description("Native Gbdt Training Function")
         .arguments(fn_args)
         .inputs(vec![fn_input])
-        .outputs(vec![fn_output]);
+        .outputs(vec![fn_output])
+        .build();
 
     let response = client.register_function(request).unwrap();
     log::debug!("Register function: {:?}", response);
diff --git a/tests/functional/enclave/src/end_to_end/mesapy_data_fusion.rs b/tests/functional/enclave/src/end_to_end/mesapy_data_fusion.rs
index dbbc228..d8c901c 100644
--- a/tests/functional/enclave/src/end_to_end/mesapy_data_fusion.rs
+++ b/tests/functional/enclave/src/end_to_end/mesapy_data_fusion.rs
@@ -55,14 +55,15 @@ def entrypoint(argv):
     let input1 = FunctionInput::new("InPartyA", "Input from party A");
     let input2 = FunctionInput::new("InPartyB", "Input from party B");
     let fusion_output = FunctionOutput::new("OutFusionData", "Output fusion data");
-    let request = RegisterFunctionRequest::new()
+    let request = RegisterFunctionRequestBuilder::new()
         .name("mesapy_data_fusion_demo")
         .description("Mesapy Data Fusion Function")
         .payload(script.into())
         .executor_type(ExecutorType::Python)
         .public(true)
         .inputs(vec![input1, input2])
-        .outputs(vec![fusion_output]);
+        .outputs(vec![fusion_output])
+        .build();
     let response = client.register_function(request).unwrap();
     log::debug!("Resgister function: {:?}", response);
     response.function_id
@@ -230,14 +231,15 @@ def entrypoint(argv):
 "#;
 
     let input_spec = FunctionInput::new("InputData", "Lines of Data");
-    let request = RegisterFunctionRequest::new()
+    let request = RegisterFunctionRequestBuilder::new()
         .name("wlc")
         .description("Mesapy Word Line Count Function")
         .arguments(vec!["query"])
         .payload(script.into())
         .executor_type(ExecutorType::Python)
         .public(true)
-        .inputs(vec![input_spec]);
+        .inputs(vec![input_spec])
+        .build();
     let response = client.register_function(request).unwrap();
     log::debug!("Resgister function: {:?}", response);
     response.function_id
diff --git a/tests/functional/enclave/src/end_to_end/mesapy_echo.rs b/tests/functional/enclave/src/end_to_end/mesapy_echo.rs
index 53058a5..42d1fb6 100644
--- a/tests/functional/enclave/src/end_to_end/mesapy_echo.rs
+++ b/tests/functional/enclave/src/end_to_end/mesapy_echo.rs
@@ -34,13 +34,14 @@ def entrypoint(argv):
     return argv[1]
 ";
     // Register Function
-    let request = RegisterFunctionRequest::new()
+    let request = RegisterFunctionRequestBuilder::new()
         .name("mesapy_echo_demo")
         .description("Mesapy Echo Function")
         .payload(script.into())
         .executor_type(ExecutorType::Python)
         .public(true)
-        .arguments(vec!["message"]);
+        .arguments(vec!["message"])
+        .build();
 
     let response = client.register_function(request).unwrap();
 
diff --git a/tests/functional/enclave/src/execution_service.rs b/tests/functional/enclave/src/execution_service.rs
index fba4880..443ca9e 100644
--- a/tests/functional/enclave/src/execution_service.rs
+++ b/tests/functional/enclave/src/execution_service.rs
@@ -34,14 +34,15 @@ fn test_execute_function() {
 
     let function_id = Uuid::new_v4();
 
-    let staged_task = StagedTask::new()
+    let staged_task = StagedTaskBuilder::new()
         .task_id(task_id)
         .function_id(function_id)
         .function_name("builtin-echo")
         .executor(Executor::Builtin)
         .function_arguments(hashmap!(
             "message" => "Hello, Teaclave Tests!"
-        ));
+        ))
+        .build();
 
     let mut storage_client = get_storage_client();
     let enqueue_request = EnqueueRequest::new(
diff --git a/tests/functional/enclave/src/management_service.rs b/tests/functional/enclave/src/management_service.rs
index cc95f7f..0839fa3 100644
--- a/tests/functional/enclave/src/management_service.rs
+++ b/tests/functional/enclave/src/management_service.rs
@@ -127,14 +127,15 @@ fn test_get_input_file() {
 fn test_register_function() {
     let function_input = FunctionInput::new("input", "input_desc");
     let function_output = FunctionOutput::new("output", "output_desc");
-    let request = RegisterFunctionRequest::new()
+    let request = RegisterFunctionRequestBuilder::new()
         .name("mock_function")
         .executor_type(ExecutorType::Python)
         .payload(b"def entrypoint:\n\treturn".to_vec())
         .public(true)
         .arguments(vec!["arg"])
         .inputs(vec![function_input])
-        .outputs(vec![function_output]);
+        .outputs(vec![function_output])
+        .build();
 
     let mut client = authorized_client("mock_user");
     let response = client.register_function(request);
@@ -146,14 +147,15 @@ fn test_register_function() {
 fn test_get_function() {
     let function_input = FunctionInput::new("input", "input_desc");
     let function_output = FunctionOutput::new("output", "output_desc");
-    let request = RegisterFunctionRequest::new()
+    let request = RegisterFunctionRequestBuilder::new()
         .name("mock_function")
         .executor_type(ExecutorType::Python)
         .payload(b"def entrypoint:\n\treturn".to_vec())
         .public(false)
         .arguments(vec!["arg"])
         .inputs(vec![function_input])
-        .outputs(vec![function_output]);
+        .outputs(vec![function_output])
+        .build();
 
     let mut client = authorized_client("mock_user");
     let response = client.register_function(request).unwrap();
diff --git a/tests/functional/enclave/src/scheduler_service.rs b/tests/functional/enclave/src/scheduler_service.rs
index 3308f11..a96a08f 100644
--- a/tests/functional/enclave/src/scheduler_service.rs
+++ b/tests/functional/enclave/src/scheduler_service.rs
@@ -27,11 +27,12 @@ use uuid::Uuid;
 #[test_case]
 fn test_pull_task() {
     let function_id = Uuid::new_v4();
-    let staged_task = StagedTask::new()
+    let staged_task = StagedTaskBuilder::new()
         .task_id(Uuid::new_v4())
         .function_name("builtin-echo")
         .function_id(function_id)
-        .executor(Executor::Builtin);
+        .executor(Executor::Builtin)
+        .build();
 
     let mut storage_client = get_storage_client();
     let enqueue_request = EnqueueRequest::new(
@@ -53,11 +54,12 @@ fn test_update_task_status_result() {
     let task_id = Uuid::new_v4();
     let function_id = Uuid::new_v4();
 
-    let staged_task = StagedTask::new()
+    let staged_task = StagedTaskBuilder::new()
         .task_id(task_id)
         .function_name("builtin-echo")
         .function_id(function_id)
-        .executor(Executor::Builtin);
+        .executor(Executor::Builtin)
+        .build();
 
     let mut storage_client = get_storage_client();
     let enqueue_request = EnqueueRequest::new(
diff --git a/tests/integration/enclave/src/teaclave_worker.rs b/tests/integration/enclave/src/teaclave_worker.rs
index 35b0547..c376624 100644
--- a/tests/integration/enclave/src/teaclave_worker.rs
+++ b/tests/integration/enclave/src/teaclave_worker.rs
@@ -21,7 +21,7 @@ use serde_json::json;
 use teaclave_crypto::TeaclaveFile128Key;
 use teaclave_types::{
     hashmap, read_all_bytes, Executor, ExecutorType, FileAuthTag, FunctionArguments,
-    StagedFileInfo, StagedFiles, StagedFunction,
+    StagedFileInfo, StagedFiles, StagedFunctionBuilder,
 };
 use teaclave_worker::Worker;
 
@@ -56,14 +56,15 @@ fn test_start_worker() {
     let output_files = StagedFiles::new(hashmap!(
         "trained_model" => output_info.clone()));
 
-    let staged_function = StagedFunction::new()
+    let staged_function = StagedFunctionBuilder::new()
         .executor_type(ExecutorType::Builtin)
         .executor(Executor::Builtin)
         .name("builtin-gbdt-train")
         .arguments(arguments)
         .input_files(input_files)
         .output_files(output_files)
-        .runtime_name("default");
+        .runtime_name("default")
+        .build();
 
     let worker = Worker::default();
 
diff --git a/types/src/function.rs b/types/src/function.rs
index b8d8fe2..abb727d 100644
--- a/types/src/function.rs
+++ b/types/src/function.rs
@@ -66,61 +66,69 @@ pub struct Function {
     pub owner: UserID,
 }
 
-impl Function {
+pub struct FunctionBuilder {
+    function: Function,
+}
+
+impl FunctionBuilder {
     pub fn new() -> Self {
-        Self::default()
+        Self {
+            function: Function::default(),
+        }
     }
 
-    pub fn id(self, id: Uuid) -> Self {
-        Self { id, ..self }
+    pub fn id(mut self, id: Uuid) -> Self {
+        self.function.id = id;
+        self
     }
 
-    pub fn executor_type(self, executor_type: ExecutorType) -> Self {
-        Self {
-            executor_type,
-            ..self
-        }
+    pub fn executor_type(mut self, executor_type: ExecutorType) -> Self {
+        self.function.executor_type = executor_type;
+        self
     }
 
-    pub fn name(self, name: impl ToString) -> Self {
-        Self {
-            name: name.to_string(),
-            ..self
-        }
+    pub fn name(mut self, name: impl ToString) -> Self {
+        self.function.name = name.to_string();
+        self
     }
 
-    pub fn description(self, description: impl ToString) -> Self {
-        Self {
-            description: description.to_string(),
-            ..self
-        }
+    pub fn description(mut self, description: impl ToString) -> Self {
+        self.function.description = description.to_string();
+        self
     }
 
-    pub fn payload(self, payload: Vec<u8>) -> Self {
-        Self { payload, ..self }
+    pub fn payload(mut self, payload: Vec<u8>) -> Self {
+        self.function.payload = payload;
+        self
     }
 
-    pub fn public(self, public: bool) -> Self {
-        Self { public, ..self }
+    pub fn public(mut self, public: bool) -> Self {
+        self.function.public = public;
+        self
     }
 
-    pub fn arguments(self, arguments: Vec<String>) -> Self {
-        Self { arguments, ..self }
+    pub fn arguments(mut self, arguments: Vec<String>) -> Self {
+        self.function.arguments = arguments;
+        self
     }
 
-    pub fn inputs(self, inputs: Vec<FunctionInput>) -> Self {
-        Self { inputs, ..self }
+    pub fn inputs(mut self, inputs: Vec<FunctionInput>) -> Self {
+        self.function.inputs = inputs;
+        self
     }
 
-    pub fn outputs(self, outputs: Vec<FunctionOutput>) -> Self {
-        Self { outputs, ..self }
+    pub fn outputs(mut self, outputs: Vec<FunctionOutput>) -> Self {
+        self.function.outputs = outputs;
+        self
     }
 
-    pub fn owner(self, owner: impl Into<UserID>) -> Self {
-        Self {
-            owner: owner.into(),
-            ..self
-        }
+    pub fn owner(mut self, owner: impl Into<UserID>) -> Self {
+        self.function.owner = owner.into();
+        self
+    }
+
+    pub fn build(self) -> Function {
+        self.function
     }
 }
 
diff --git a/types/src/staged_function.rs b/types/src/staged_function.rs
index a8133bc..3c706e0 100644
--- a/types/src/staged_function.rs
+++ b/types/src/staged_function.rs
@@ -116,55 +116,58 @@ pub struct StagedFunction {
     pub runtime_name: String,
 }
 
-impl StagedFunction {
-    pub fn new() -> Self {
-        Self::default()
-    }
+pub struct StagedFunctionBuilder {
+    function: StagedFunction,
+}
 
-    pub fn name(self, name: impl ToString) -> Self {
+impl StagedFunctionBuilder {
+    pub fn new() -> Self {
         Self {
-            name: name.to_string(),
-            ..self
+            function: StagedFunction::default(),
         }
     }
 
-    pub fn executor(self, executor: Executor) -> Self {
-        Self { executor, ..self }
+    pub fn name(mut self, name: impl ToString) -> Self {
+        self.function.name = name.to_string();
+        self
     }
 
-    pub fn payload(self, payload: Vec<u8>) -> Self {
-        Self { payload, ..self }
+    pub fn executor(mut self, executor: Executor) -> Self {
+        self.function.executor = executor;
+        self
     }
 
-    pub fn arguments(self, arguments: FunctionArguments) -> Self {
-        Self { arguments, ..self }
+    pub fn payload(mut self, payload: Vec<u8>) -> Self {
+        self.function.payload = payload;
+        self
     }
 
-    pub fn input_files(self, input_files: StagedFiles) -> Self {
-        Self {
-            input_files,
-            ..self
-        }
+    pub fn arguments(mut self, arguments: FunctionArguments) -> Self {
+        self.function.arguments = arguments;
+        self
     }
 
-    pub fn output_files(self, output_files: StagedFiles) -> Self {
-        Self {
-            output_files,
-            ..self
-        }
+    pub fn input_files(mut self, input_files: StagedFiles) -> Self {
+        self.function.input_files = input_files;
+        self
     }
 
-    pub fn runtime_name(self, runtime_name: impl ToString) -> Self {
-        Self {
-            runtime_name: runtime_name.to_string(),
-            ..self
-        }
+    pub fn output_files(mut self, output_files: StagedFiles) -> Self {
+        self.function.output_files = output_files;
+        self
     }
 
-    pub fn executor_type(self, executor_type: ExecutorType) -> Self {
-        Self {
-            executor_type,
-            ..self
-        }
+    pub fn runtime_name(mut self, runtime_name: impl ToString) -> Self {
+        self.function.runtime_name = runtime_name.to_string();
+        self
+    }
+
+    pub fn executor_type(mut self, executor_type: ExecutorType) -> Self {
+        self.function.executor_type = executor_type;
+        self
+    }
+
+    pub fn build(self) -> StagedFunction {
+        self.function
     }
 }
diff --git a/types/src/staged_task.rs b/types/src/staged_task.rs
index 6cb930c..595386b 100644
--- a/types/src/staged_task.rs
+++ b/types/src/staged_task.rs
@@ -199,68 +199,68 @@ impl Storable for StagedTask {
 }
 
 impl StagedTask {
-    pub fn new() -> Self {
-        Self::default()
+    pub fn get_queue_key() -> &'static str {
+        QUEUE_KEY
     }
+}
 
-    pub fn task_id(self, task_id: Uuid) -> Self {
-        Self { task_id, ..self }
-    }
+pub struct StagedTaskBuilder {
+    task: StagedTask,
+}
 
-    pub fn function_id(self, function_id: Uuid) -> Self {
+impl StagedTaskBuilder {
+    pub fn new() -> Self {
         Self {
-            function_id,
-            ..self
+            task: StagedTask::default(),
         }
     }
 
-    pub fn executor(self, executor: Executor) -> Self {
-        Self { executor, ..self }
+    pub fn task_id(mut self, task_id: Uuid) -> Self {
+        self.task.task_id = task_id;
+        self
     }
 
-    pub fn function_name(self, name: impl ToString) -> Self {
-        Self {
-            function_name: name.to_string(),
-            ..self
-        }
+    pub fn function_id(mut self, function_id: Uuid) -> Self {
+        self.task.function_id = function_id;
+        self
     }
 
-    pub fn function_arguments(self, function_arguments: impl Into<FunctionArguments>) -> Self {
-        Self {
-            function_arguments: function_arguments.into(),
-            ..self
-        }
+    pub fn executor(mut self, executor: Executor) -> Self {
+        self.task.executor = executor;
+        self
     }
 
-    pub fn function_payload(self, function_payload: Vec<u8>) -> Self {
-        Self {
-            function_payload,
-            ..self
-        }
+    pub fn function_name(mut self, name: impl ToString) -> Self {
+        self.task.function_name = name.to_string();
+        self
     }
 
-    pub fn input_data(self, input_data: impl Into<FunctionInputFiles>) -> Self {
-        Self {
-            input_data: input_data.into(),
-            ..self
-        }
+    pub fn function_arguments(mut self, function_arguments: impl Into<FunctionArguments>) -> Self {
+        self.task.function_arguments = function_arguments.into();
+        self
     }
 
-    pub fn output_data(self, output_data: impl Into<FunctionOutputFiles>) -> Self {
-        Self {
-            output_data: output_data.into(),
-            ..self
-        }
+    pub fn function_payload(mut self, function_payload: Vec<u8>) -> Self {
+        self.task.function_payload = function_payload;
+        self
     }
 
-    pub fn executor_type(self, executor_type: ExecutorType) -> Self {
-        Self {
-            executor_type,
-            ..self
-        }
+    pub fn input_data(mut self, input_data: impl Into<FunctionInputFiles>) -> Self {
+        self.task.input_data = input_data.into();
+        self
     }
 
-    pub fn get_queue_key() -> &'static str {
-        QUEUE_KEY
+    pub fn output_data(mut self, output_data: impl Into<FunctionOutputFiles>) -> Self {
+        self.task.output_data = output_data.into();
+        self
+    }
+
+    pub fn executor_type(mut self, executor_type: ExecutorType) -> Self {
+        self.task.executor_type = executor_type;
+        self
+    }
+
+    pub fn build(self) -> StagedTask {
+        self.task
     }
 }

---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@teaclave.apache.org
For additional commands, e-mail: commits-help@teaclave.apache.org