Use deployment locks in worker

app/main.py

@@ -114,6 +114,105 @@ def heartbeat_loop():
         time.sleep(HEARTBEAT_SECONDS)
 
 
+def cleanup_stale_locks():
+    con = get_connection()
+
+    con.execute(
+        """
+        DELETE FROM deployment_locks
+        WHERE expires_at IS NOT NULL
+          AND expires_at < CURRENT_TIMESTAMP
+        """
+    )
+
+    con.commit()
+    con.close()
+
+
+def acquire_lock(job: dict) -> bool:
+    con = get_connection()
+
+    try:
+        con.execute(
+            """
+            INSERT INTO deployment_locks (
+                target_type,
+                target_id,
+                job_id,
+                worker_id,
+                acquired_at,
+                expires_at
+            )
+            VALUES (
+                ?,
+                ?,
+                ?,
+                ?,
+                CURRENT_TIMESTAMP,
+                datetime('now', '+30 minutes')
+            )
+            """,
+            (
+                job["target_type"],
+                job["target_id"],
+                job["id"],
+                WORKER_ID,
+            ),
+        )
+
+        con.commit()
+        return True
+
+    except sqlite3.IntegrityError:
+        return False
+
+    finally:
+        con.close()
+
+
+def release_lock(job: dict):
+    con = get_connection()
+
+    con.execute(
+        """
+        DELETE FROM deployment_locks
+        WHERE target_type = ?
+          AND target_id = ?
+          AND job_id = ?
+        """,
+        (
+            job["target_type"],
+            job["target_id"],
+            job["id"],
+        ),
+    )
+
+    con.commit()
+    con.close()
+
+
+def extend_lock(job: dict):
+    con = get_connection()
+
+    con.execute(
+        """
+        UPDATE deployment_locks
+        SET expires_at = datetime('now', '+30 minutes')
+        WHERE target_type = ?
+          AND target_id = ?
+          AND job_id = ?
+        """,
+        (
+            job["target_type"],
+            job["target_id"],
+            job["id"],
+        ),
+    )
+
+    con.commit()
+    con.close()
+
+
 def recover_stale_jobs():
     con = get_connection()
 
@@ -148,6 +247,14 @@ def recover_stale_jobs():
             (job_id,),
         )
 
+        con.execute(
+            """
+            DELETE FROM deployment_locks
+            WHERE job_id = ?
+            """,
+            (job_id,),
+        )
+
     con.commit()
     con.close()
 
@@ -210,24 +317,6 @@ def get_next_job():
 
     job_id = row["id"]
 
-    running_same_target = con.execute(
-        """
-        SELECT id
-        FROM jobs
-        WHERE status IN ('running', 'cancelled_requested')
-          AND target_type = ?
-          AND target_id = ?
-          AND id != ?
-        LIMIT 1
-        """,
-        (row["target_type"], row["target_id"], job_id),
-    ).fetchone()
-
-    if running_same_target:
-        con.commit()
-        con.close()
-        return None
-
     con.execute(
         """
         UPDATE jobs
@@ -274,20 +363,28 @@ def command_for_job(job: dict):
     raise ValueError(f"Unsupported job type: {job_type}")
 
 
-def terminate_process(process: subprocess.Popen, job_id: int):
+def terminate_process_group(process: subprocess.Popen, job_id: int):
-    append_job_log(job_id, "system", "Cancel requested. Terminating process...")
+    append_job_log(job_id, "system", "Cancel requested. Terminating process group...")
 
     try:
-        process.terminate()
+        os.killpg(os.getpgid(process.pid), signal.SIGTERM)
         process.wait(timeout=10)
-        append_job_log(job_id, "system", "Process terminated.")
+        append_job_log(job_id, "system", "Process group terminated.")
     except subprocess.TimeoutExpired:
-        append_job_log(job_id, "system", "Process did not terminate in time. Killing process...")
+        append_job_log(job_id, "system", "Process group did not terminate in time. Killing process group...")
-        process.kill()
+        try:
+            os.killpg(os.getpgid(process.pid), signal.SIGKILL)
+        except Exception:
+            process.kill()
         process.wait(timeout=10)
-        append_job_log(job_id, "system", "Process killed.")
+        append_job_log(job_id, "system", "Process group killed.")
     except Exception as exc:
-        append_job_log(job_id, "stderr", f"Error while terminating process: {exc}")
+        append_job_log(job_id, "stderr", f"Error while terminating process group: {exc}")
+        try:
+            process.kill()
+            process.wait(timeout=10)
+        except Exception as kill_exc:
+            append_job_log(job_id, "stderr", f"Error while killing process: {kill_exc}")
 
 
 def execute_job(job: dict):
@@ -314,22 +411,12 @@ def execute_job(job: dict):
         )
 
         while process.poll() is None:
+            extend_lock(job)
+
             status = get_job_status(job_id)
 
             if status == "cancelled_requested":
-                try:
+                terminate_process_group(process, job_id)
-                    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,
@@ -392,12 +479,40 @@ def execute_job(job: dict):
 def worker_loop():
     while True:
         try:
+            cleanup_stale_locks()
             recover_stale_jobs()
 
             job = get_next_job()
 
             if job:
-                execute_job(job)
+                if not acquire_lock(job):
+                    append_job_log(
+                        int(job["id"]),
+                        "system",
+                        "Deployment lock is already held for this target. Re-queueing job.",
+                    )
+
+                    con = get_connection()
+                    con.execute(
+                        """
+                        UPDATE jobs
+                        SET status = 'queued',
+                            started_at = NULL,
+                            worker_id = NULL
+                        WHERE id = ?
+                        """,
+                        (job["id"],),
+                    )
+                    con.commit()
+                    con.close()
+
+                    time.sleep(POLL_SECONDS)
+                    continue
+
+                try:
+                    execute_job(job)
+                finally:
+                    release_lock(job)
             else:
                 time.sleep(POLL_SECONDS)
 
@@ -415,6 +530,7 @@ def start_worker():
     if _worker_started:
         return
 
+    cleanup_stale_locks()
     recover_stale_jobs()
     update_worker_status(current_job_id=None)