diff --git a/app/main.py b/app/main.py index 2f6d743..ec55e1a 100644 --- a/app/main.py +++ b/app/main.py @@ -1,5 +1,6 @@ import json import os +import signal import socket import sqlite3 import subprocess @@ -15,6 +16,7 @@ TOOLS_DIR = Path("/tools") WORKER_ID = os.getenv("APPFACTORY_WORKER_ID", socket.gethostname()) POLL_SECONDS = int(os.getenv("APPFACTORY_WORKER_POLL_SECONDS", "3")) HEARTBEAT_SECONDS = int(os.getenv("APPFACTORY_WORKER_HEARTBEAT_SECONDS", "10")) +CANCEL_CHECK_SECONDS = int(os.getenv("APPFACTORY_CANCEL_CHECK_SECONDS", "2")) app = FastAPI(title="AppFactory Worker") @@ -42,6 +44,20 @@ def append_job_log(job_id: int, stream: str, message: str): con.close() +def get_job_status(job_id: int) -> str | None: + con = get_connection() + row = con.execute( + "SELECT status FROM jobs WHERE id = ?", + (job_id,), + ).fetchone() + con.close() + + if not row: + return None + + return row["status"] + + def update_worker_status(current_job_id=None): con = get_connection() @@ -77,6 +93,7 @@ def update_worker_status(current_job_id=None): { "poll_seconds": POLL_SECONDS, "heartbeat_seconds": HEARTBEAT_SECONDS, + "cancel_check_seconds": CANCEL_CHECK_SECONDS, }, ensure_ascii=False, ), @@ -104,7 +121,7 @@ def recover_stale_jobs(): """ SELECT id FROM jobs - WHERE status = 'running' + WHERE status IN ('running', 'cancelled_requested') AND started_at < datetime('now', '-15 minutes') """ ).fetchall() @@ -197,7 +214,7 @@ def get_next_job(): """ SELECT id FROM jobs - WHERE status = 'running' + WHERE status IN ('running', 'cancelled_requested') AND target_type = ? AND target_id = ? AND id != ? @@ -257,6 +274,22 @@ def command_for_job(job: dict): raise ValueError(f"Unsupported job type: {job_type}") +def terminate_process(process: subprocess.Popen, job_id: int): + append_job_log(job_id, "system", "Cancel requested. Terminating process...") + + try: + process.terminate() + process.wait(timeout=10) + append_job_log(job_id, "system", "Process terminated.") + except subprocess.TimeoutExpired: + append_job_log(job_id, "system", "Process did not terminate in time. Killing process...") + process.kill() + process.wait(timeout=10) + append_job_log(job_id, "system", "Process killed.") + except Exception as exc: + append_job_log(job_id, "stderr", f"Error while terminating process: {exc}") + + def execute_job(job: dict): job_id = int(job["id"]) target = job["target_id"] @@ -277,8 +310,38 @@ def execute_job(job: dict): stderr=subprocess.PIPE, text=True, env=env, + start_new_session=True, ) + while process.poll() is None: + status = get_job_status(job_id) + + if status == "cancelled_requested": + try: + os.killpg(os.getpgid(process.pid), signal.SIGTERM) + except Exception: + terminate_process(process, job_id) + else: + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + try: + os.killpg(os.getpgid(process.pid), signal.SIGKILL) + except Exception: + process.kill() + process.wait(timeout=10) + + update_job_status( + job_id, + "cancelled", + error_text="Job cancelled by user request", + result={"cancelled": True}, + ) + append_job_log(job_id, "system", "Job cancelled.") + return + + time.sleep(CANCEL_CHECK_SECONDS) + stdout, stderr = process.communicate() returncode = process.returncode @@ -288,6 +351,16 @@ def execute_job(job: dict): if stderr: append_job_log(job_id, "stderr", stderr) + if get_job_status(job_id) == "cancelled_requested": + update_job_status( + job_id, + "cancelled", + error_text="Job cancelled by user request", + result={"returncode": returncode, "cancelled": True}, + ) + append_job_log(job_id, "system", "Job cancelled after process finished.") + return + if returncode == 0: update_job_status( job_id, @@ -361,4 +434,5 @@ def health(): "worker_id": WORKER_ID, "poll_seconds": POLL_SECONDS, "heartbeat_seconds": HEARTBEAT_SECONDS, + "cancel_check_seconds": CANCEL_CHECK_SECONDS, }