From c5fd92e9be5654d5335c474ace1218dfa9c6d199 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=98=B3=E5=85=89=E5=B0=91=E5=B9=B4?= <849317537@qq.com>
Date: Thu, 31 Oct 2024 14:49:58 +0000
Subject: [PATCH] =?UTF-8?q?=E8=B6=85=E6=97=B6=E5=A4=84=E7=90=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 test/service/handler.py | 10 ++++++----
 test/service/server.py  |  9 ++++-----
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/test/service/handler.py b/test/service/handler.py
index 00245d6..9f6253d 100644
--- a/test/service/handler.py
+++ b/test/service/handler.py
@@ -5,12 +5,13 @@ import traceback
 
 TARGET_LANGUAGE = set(["python", "cpp", "go", "rust", "node", "javascript"])
 PROJECT_PATH = "/home/rocker/project"
+BUILD_TIMEOUT_S = 3
 
 def go_handler():
     source_path = f"{PROJECT_PATH}/go_project/main.go"
     subprocess.run(["cp", "/tmp/code", source_path], capture_output=True)
     main_path = f"{PROJECT_PATH}/go_project/main"
-    ret = subprocess.run(["go", "build", "-o", main_path], cwd=f"{PROJECT_PATH}/go_project", capture_output=True)
+    ret = subprocess.run(["/usr/bin/timeout", f"{BUILD_TIMEOUT_S}s", "go", "build", "-o", main_path], cwd=f"{PROJECT_PATH}/go_project", capture_output=True)
     std_err = ret.stderr.decode("utf-8")
     if len(std_err) > 0:
         return std_err, ""
@@ -20,7 +21,7 @@ def cpp_handler():
     source_path = f"{PROJECT_PATH}/cpp_project/main.cpp"
     subprocess.run(["cp", "/tmp/code", source_path], capture_output=True)
     main_path = f"{PROJECT_PATH}/cpp_project/main"
-    ret = subprocess.run(["g++", source_path, "-o", main_path], capture_output=True)
+    ret = subprocess.run(["/usr/bin/timeout", f"{BUILD_TIMEOUT_S}s", "g++", source_path, "-o", main_path], capture_output=True)
     std_err = ret.stderr.decode("utf-8")
     if len(std_err) > 0:
         return std_err, ""
@@ -34,7 +35,7 @@ def node_handler():
 def rust_handler():
     source_path = f"{PROJECT_PATH}/rust_project/src/main.rs"
     subprocess.run(["cp", "/tmp/code", source_path], capture_output=True)
-    ret = subprocess.run(["cargo", "build"], cwd=f"{PROJECT_PATH}/rust_project", capture_output=True)
+    ret = subprocess.run(["/usr/bin/timeout", f"{BUILD_TIMEOUT_S}s", "cargo", "build"], cwd=f"{PROJECT_PATH}/rust_project", capture_output=True)
     std_err = ret.stderr.decode("utf-8")
     if len(std_err) > 0 and ret.returncode != 0:
         return std_err, ""
@@ -50,12 +51,13 @@ def python_handler():
 
 def main():
     language = sys.argv[1]
+    timeout_s = sys.argv[2]
     if language not in TARGET_LANGUAGE:
         print(f"不支持的编程语言: {language}")
         return
     
     # 编译和执行
-    cmd = ["/usr/bin/time", "-f", "'%M %U %S %e'"]
+    cmd = ["/usr/bin/time", "-f", "'%M %U %S %e'", "/usr/bin/timeout", f"{timeout_s}s"]
     err = ""
     if language == "python":
         err, main_path = python_handler()
diff --git a/test/service/server.py b/test/service/server.py
index ae2bd9e..b84be7c 100644
--- a/test/service/server.py
+++ b/test/service/server.py
@@ -8,6 +8,7 @@ import traceback
 
 app = Flask(__name__)
 USER_NAME = "ubuntu"
+TIMEOUT_S = 10
 CGROUP_PATH = "/sys/fs/cgroup/rocker"
 ROCKER_PATH = f"/home/{USER_NAME}/rocker/target/debug/rocker"
 HANDLER_FILE = "handler.py"
@@ -40,7 +41,7 @@ def main():
         _ = subprocess.run(["sudo", ROCKER_PATH, 
                         "--id", _id,
                         "--image", "base_ubuntu_image", 
-                        "--run", f"python3 -u /tmp/handler.py {language}",
+                        "--run", f"python3 -u /tmp/handler.py {language} {TIMEOUT_S}",
                         "--env", f"/home/{USER_NAME}/rocker/quant.env", 
                         "--volume", f"{user_path}:/tmp",
                         "--log"])
@@ -50,7 +51,7 @@ def main():
     
     # 20s 超时
     procs_path = f"{CGROUP_PATH}/{_id}/cgroup.procs"
-    for _ in range(2000):
+    for _ in range(int(TIMEOUT_S * 2 / 0.01)):  # *2 留出冗余
         time.sleep(0.01)
         try:
             with open(procs_path) as f:
@@ -59,13 +60,11 @@ def main():
                     break
         except Exception as e:
             break
-    else:
-        resp["msg"] = "Timeout"
     resp["rt"] = time.time() - starat
     
     # 读取文件
+    log_path = f"/home/rocker/containers/{_id}/upper/logs/log"
     try:
-        log_path = f"/home/rocker/containers/{_id}/upper/logs/log"
         with open(log_path) as f:
             log_str = f"{f.read()}\n"
             resp["msg"] += log_str