Add watchdog functionalities

cmake/UtilTargets.cmake

@@ -73,6 +73,9 @@ if(TEST_MODE)
   add_custom_target(
     run-sdk-tests COMMAND ${TEACLAVE_COMMON_ENVS}
                           ${MT_SCRIPT_DIR}/test.sh sdk)
+  add_custom_target(
+    run-cancel-test COMMAND ${TEACLAVE_COMMON_ENVS}
+                          ${MT_SCRIPT_DIR}/test.sh cancel)
 else()
   add_custom_target(
     run-tests

cmake/scripts/test.sh

@@ -97,7 +97,7 @@ run_functional_tests() {
   pushd ${TEACLAVE_SERVICE_INSTALL_DIR}
   ./teaclave_authentication_service &
   ./teaclave_storage_service &
-  sleep 3    # wait for authentication and storage service
+  sleep 10    # wait for authentication and storage service
   ./teaclave_management_service &
   ./teaclave_scheduler_service &
   sleep 3    # wait for management service and scheduler_service
@@ -165,7 +165,7 @@ run_sdk_tests() {
   pushd ${TEACLAVE_SERVICE_INSTALL_DIR}
   ./teaclave_authentication_service &
   ./teaclave_storage_service &
-  sleep 3    # wait for authentication and storage service
+  sleep 10    # wait for authentication and storage service
   ./teaclave_management_service &
   ./teaclave_scheduler_service &
   sleep 3    # wait for management service and scheduler_service
@@ -258,6 +258,70 @@ run_examples() {
   cleanup
 }
 
+run_cancel_test() {
+  trap cleanup INT TERM ERR
+
+  echo_title "cancel"
+  mkdir -p /tmp/fusion_data
+  pushd ${TEACLAVE_CLI_INSTALL_DIR}
+  ./teaclave_cli verify \
+                 --enclave-info ../examples/enclave_info.toml \
+                 --public-keys $(find ../examples -name "*.public.pem") \
+                 --signatures $(find ../examples -name "*.sign.sha256")
+  popd
+
+  echo "initiating Teaclave with 2 executors..."
+
+  pushd ${TEACLAVE_SERVICE_INSTALL_DIR}
+  ./teaclave_authentication_service &
+  ./teaclave_storage_service &
+  sleep 3    # wait for authentication and storage service
+  ./teaclave_management_service &
+  ./teaclave_scheduler_service &
+  sleep 3    # wait for management service and scheduler_service
+  ./teaclave_access_control_service &
+  ./teaclave_frontend_service &
+  sleep 3    # wait for other services
+
+  start_storage_server
+
+  # Run of execution services separately
+  ./teaclave_execution_service & exe_pid1=$!
+  ./teaclave_execution_service & exe_pid2=$!
+  sleep 10    # wait for execution services
+  popd
+
+  echo "executor 1 pid: $exe_pid1"
+  echo "executor 2 pid: $exe_pid2"
+
+  pushd ${TEACLAVE_PROJECT_ROOT}/examples/python
+  export PYTHONPATH=${TEACLAVE_PROJECT_ROOT}/sdk/python
+  python3 mesapy_deadloop_cancel.py
+  popd
+
+  sleep 3
+
+  live_pids=0
+  if ps -p $exe_pid1 > /dev/null
+  then
+    live_pids=$((live_pids+1))
+  fi
+
+  if ps -p $exe_pid2 > /dev/null
+  then
+    live_pids=$((live_pids+1))
+  fi
+
+  if [ $live_pids -eq 1 ]
+  then
+    echo "only one executor is killed, test passed"
+  else
+    echo "Some unexpected happens, test failed"
+  fi
+
+  cleanup
+}
+
 case "$1" in
     "unit")
         run_unit_tests
@@ -274,11 +338,15 @@ case "$1" in
     "example")
         run_examples
         ;;
+    "cancel")
+        run_cancel_test
+        ;;
     *)
         run_unit_tests
         run_integration_tests
         run_functional_tests
         run_sdk_tests
         run_examples
+        run_cancel_test
         ;;
 esac

examples/python/builtin_gbdt_train.py

@@ -103,6 +103,13 @@ def gbdt(self):
         print("[+] invoking task")
         client.invoke_task(task_id)
 
+        # If you feel like you don't want to wait for the task to finish
+        # or something unexpected happens,
+        # you may uncomment the following lines to cancel the task
+        # time.sleep(3)
+        # print("[+] canceling task")
+        # client.cancel_task(task_id)
+
         print("[+] getting result")
         result = client.get_task_result(task_id)
         print("[+] done")

examples/python/mesapy_deadloop_cancel.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import sys
+import time
+
+from teaclave import FunctionInput, FunctionOutput, OwnerList, DataMap
+from utils import USER_ID, USER_PASSWORD, connect_authentication_service, connect_frontend_service, PlatformAdmin
+
+
+class MesaPyEchoExample:
+    def __init__(self, user_id, user_password):
+        self.user_id = user_id
+        self.user_password = user_password
+
+    def deadloop(self, payload_file="mesapy_deadloop_payload.py"):
+        with connect_authentication_service() as client:
+            print(f"[+] {self.user_id} login")
+            token = client.user_login(self.user_id, self.user_password)
+
+        client = connect_frontend_service()
+        metadata = {"id": self.user_id, "token": token}
+        client.metadata = metadata
+
+        print("[+] registering function")
+
+        with open(payload_file, "rb") as f:
+            payload = f.read()
+        function_id = client.register_function(
+            name="mesapy-deadloop",
+            description="A deadloop function to test task cancellation",
+            executor_type="python",
+            payload=list(payload),
+            arguments=[])
+
+        print("[+] creating task")
+        task_id = client.create_task(function_id=function_id,
+                                     function_arguments={},
+                                     executor="mesapy")
+
+        print("[+] invoking task")
+        client.invoke_task(task_id)
+
+        print("[+] canceling task")
+        time.sleep(5)
+        client.cancel_task(task_id)
+
+        print("[+] getting result")
+
+        try:
+            result = client.get_task_result(task_id)
+        except Exception as e:
+            print(f"[+] result: {str(e)}")
+            result = str(e)
+
+        return result
+
+
+def main():
+    example = MesaPyEchoExample(USER_ID, USER_PASSWORD)
+    rv = example.deadloop()
+
+    print("[+] function return: ", rv)
+
+
+if __name__ == '__main__':
+    main()

examples/python/mesapy_deadloop_payload.py

@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+
+def entrypoint(argv):
+    while True:
+        pass

sdk/python/teaclave.py

@@ -371,6 +371,13 @@ def __init__(self, metadata: Metadata, task_id: str):
         self.task_id = task_id
 
 
+class CancelTaskRequest(Request):
+    def __init__(self, metadata: Metadata, task_id: str):
+        self.request = "cancel_task"
+        self.metadata = metadata
+        self.task_id = task_id
+
+
 class GetTaskRequest(Request):
     def __init__(self, metadata: Metadata, task_id: str):
         self.request = "get_task"
@@ -557,6 +564,17 @@ def invoke_task(self, task_id: str):
         else:
             raise TeaclaveException("Failed to invoke task")
 
+    def cancel_task(self, task_id: str):
+        self.check_metadata()
+        self.check_channel()
+        request = CancelTaskRequest(self.metadata, task_id)
+        _write_message(self.channel, request)
+        response = _read_message(self.channel)
+        if response["result"] == "ok":
+            pass
+        else:
+            raise TeaclaveException("Failed to cancel task")
+
     def get_task_result(self, task_id: str):
         self.check_metadata()
         self.check_channel()
@@ -568,6 +586,15 @@ def get_task_result(self, task_id: str):
             if response["result"] != "ok":
                 raise TeaclaveException("Failed to get task result")
             time.sleep(1)
+            if response["content"]["status"] == 20:
+                print(response["content"]["result"]["result"])
+                raise TeaclaveException(
+                    "Task Canceled, Error: " +
+                    response["content"]["result"]["result"]["Err"]["reason"])
+            if response["content"]["status"] == 99:
+                raise TeaclaveException(
+                    "Task Failed, Error: " +
+                    response["content"]["result"]["result"]["Err"]["reason"])
             if response["content"]["status"] == 10:
                 break
 

services/execution/enclave/src/lib.rs

@@ -77,18 +77,17 @@ fn start_service(config: &RuntimeConfig) -> Result<()> {
 
     let mut service =
         service::TeaclaveExecutionService::new(scheduler_service_endpoint, fusion_base)?;
-    let _ = service.start();
-
-    Ok(())
+    service.start()
 }
 
 #[handle_ecall]
 fn handle_start_service(input: &StartServiceInput) -> TeeServiceResult<StartServiceOutput> {
     match start_service(&input.config) {
         Ok(_) => Ok(StartServiceOutput),
+        // terminate the enclave for executor
         Err(e) => {
-            log::error!("Failed to start the service: {}", e);
-            Err(TeeServiceError::ServiceError)
+            log::error!("Service shutdown, reason: {}", e);
+            Err(TeeServiceError::EnclaveForceTermination)
         }
     }
 }