From d63b3c561cbb07fd52c21f5e62fb9bc5663d082c Mon Sep 17 00:00:00 2001
From: LI Qing <geding.lq@antgroup.com>
Date: Fri, 9 Dec 2022 10:54:10 +0800
Subject: [PATCH] Improve the usability of occlum start/exec/stop cmd

---
 src/exec/src/bin/occlum_exec_client.rs |  4 ++--
 tools/occlum                           | 30 ++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/src/exec/src/bin/occlum_exec_client.rs b/src/exec/src/bin/occlum_exec_client.rs
index f131acd4..3a3fef5b 100644
--- a/src/exec/src/bin/occlum_exec_client.rs
+++ b/src/exec/src/bin/occlum_exec_client.rs
@@ -292,14 +292,14 @@ fn main() -> Result<(), i32> {
 
     if let Some(ref _matches) = matches.subcommand_matches("start") {
         if let Err(s) = start_server(&client, DEFAULT_SERVER_FILE) {
-            println!("start_server failed {}", s);
+            println!("start server failed {}", s);
             return Err(-1);
         }
         println!("server is running.");
     } else if let Some(ref matches) = matches.subcommand_matches("stop") {
         let stop_time = matches.value_of("time").unwrap().parse::<u32>().unwrap();
         stop_server(&client, stop_time);
-        println!("server stopped.");
+        println!("server is stopping.");
     } else if let Some(ref matches) = matches.subcommand_matches("exec") {
         let mut cmd_args: Vec<&str> = match matches
             .values_of("args")
diff --git a/tools/occlum b/tools/occlum
index d8117577..ae902a71 100755
--- a/tools/occlum
+++ b/tools/occlum
@@ -120,6 +120,20 @@ check_has_run() {
     fi
 }
 
+check_has_start() {
+    if pgrep --full "build/bin/occlum_exec_server -d $instance_dir" > /dev/null ; then
+        echo "Error: the server of current Occlum instance has been started."
+        exit 1
+    fi
+}
+
+check_has_not_start() {
+    if ! pgrep --full "build/bin/occlum_exec_server -d $instance_dir" > /dev/null ; then
+        echo "Error: the server of current Occlum instance has not been started."
+        exit 1
+    fi
+}
+
 check_aesm_service() {
     # Ignore AESM service status for simulation mode
     if [ "$(cat $instance_dir/.sgx_mode 2>/dev/null)" != "HW" ]; then
@@ -375,6 +389,7 @@ cmd_run() {
 
 cmd_start() {
     check_has_built
+    check_has_start
     check_aesm_service
 
     loop=true
@@ -409,6 +424,7 @@ cmd_start() {
 
 cmd_exec() {
     check_has_built
+    check_has_not_start
 
     SGX_MODE=$(cat $instance_dir/.sgx_mode)
     if [[ -n $SGX_MODE && "$SGX_MODE" != "HW" ]]; then
@@ -426,6 +442,7 @@ cmd_exec() {
 
 cmd_stop() {
     check_has_built
+    check_has_not_start
 
     SGX_MODE=$(cat $instance_dir/.sgx_mode)
     if [[ -n $SGX_MODE && "$SGX_MODE" != "HW" ]]; then
@@ -438,6 +455,19 @@ cmd_stop() {
 
     RUST_BACKTRACE=1 "$instance_dir/build/bin/occlum_exec_client" stop -t 0
 
+    ## Ensure the server is stopped, default timeout is 30s
+    for i in $(seq 0 30); do
+        if pgrep --full "build/bin/occlum_exec_server -d $instance_dir" > /dev/null ; then
+            sleep 1
+        else
+            echo "server stopped."
+            echo "built" > "$status_file"
+            exit 0
+        fi
+    done
+
+    echo "Error: timeout to stop the server, let's kill it."
+    pkill -SIGKILL --full "build/bin/occlum_exec_server -d $instance_dir"
     echo "built" > "$status_file"
 }