feat: ansible receptor kill 进程

apps/ops/ansible/callback.py

@@ -48,6 +48,11 @@ class DefaultCallback:
         event = data.get('event', None)
         if not event:
             return
+
+        pid = data.get('pid', None)
+        if pid:
+            self.write_pid(pid)
+
         event_data = data.get('event_data', {})
         host = event_data.get('remote_addr', '')
         task = event_data.get('task', '')
@@ -156,3 +161,8 @@ class DefaultCallback:
         status = data.get('status', '')
         self.status = self.STATUS_MAPPER.get(status, 'unknown')
         self.private_data_dir = data.get("private_data_dir", None)
+
+    def write_pid(self, pid):
+        pid_filepath = os.path.join(self.private_data_dir, 'local.pid')
+        with open(pid_filepath, 'w') as f:
+            f.write(str(pid))

apps/ops/ansible/receptor/receptor_runner.py

@@ -31,6 +31,20 @@ def nodes():
     return receptor_ctl.simple_command("status").get("Advertisements", None)
 
 
+def kill_process(pid):
+    submit_result = receptor_ctl.submit_work(worktype="kill", node="primary", payload=str(pid))
+
+    unit_id = submit_result["unitid"]
+
+    result_socket, result_file = receptor_ctl.get_work_results(unit_id=unit_id, return_sockfile=True,
+                                                               return_socket=True)
+    while not result_socket.close():
+        buf = result_file.read()
+        if not buf:
+            break
+        print(buf.decode('utf8'))
+
+
 def run(**kwargs):
     receptor_runner = AnsibleReceptorRunner(**kwargs)
     return receptor_runner.run()

apps/ops/models/job.py

@@ -524,15 +524,16 @@ class JobExecution(JMSOrgBaseModel):
             ssh_tunnel.local_gateway_clean(runner)
 
     def stop(self):
-        unit_id_path = os.path.join(self.private_dir, "local.unitid")
-        if os.path.exists(unit_id_path):
-            with open(unit_id_path) as f:
+        from ops.signal_handlers import job_execution_stop_pub_sub
+        pid_path = os.path.join(self.private_dir, "local.pid")
+        if os.path.exists(pid_path):
+            with open(pid_path) as f:
                 try:
-                    unit_id = f.read()
-                    receptor_runner.cancel(unit_id)
+                    pid = f.read()
+                    job_execution_stop_pub_sub.publish(int(pid))
                 except Exception as e:
                     print(e)
-            self.set_error('Job stop by "receptor worker cancel, unit id is {}"'.format(unit_id))
+        self.set_error('Job stop by "user cancel"')
 
     class Meta:
         verbose_name = _("Job Execution")

apps/ops/signal_handlers.py

@@ -1,7 +1,6 @@
 import ast
 import time
 
-import psutil
 from celery import signals
 from django.core.cache import cache
 from django.db import transaction
@@ -10,13 +9,13 @@ from django.db.utils import ProgrammingError
 from django.dispatch import receiver
 from django.utils import translation, timezone
 from django.utils.functional import LazyObject
-from psutil import NoSuchProcess
 
 from common.db.utils import close_old_connections, get_logger
 from common.signals import django_ready
 from common.utils.connection import RedisPubSub
 from jumpserver.utils import get_current_request
 from orgs.utils import get_current_org_id, set_current_org
+from .ansible.receptor.receptor_runner import kill_process
 from .celery import app
 from .models import CeleryTaskExecution, CeleryTask, Job
 
@@ -160,24 +159,7 @@ def subscribe_stop_job_execution(sender, **kwargs):
 
     def on_stop(pid):
         logger.info(f"Stop job execution {pid} start")
-        try:
-            current_process = psutil.Process(pid)
-        except NoSuchProcess as e:
-            logger.error(e)
-            return
-
-        children = current_process.children(recursive=True)
-        logger.debug(f"Job execution process children: {children}")
-        for child in children:
-            if child.pid == 1:
-                continue
-            if child.name() != 'ssh':
-                continue
-            try:
-                child.kill()
-                logger.debug(f"Kill job execution process {pid} children process {child.pid} success")
-            except Exception as e:
-                logger.error(e)
+        kill_process(pid)
 
     job_execution_stop_pub_sub.subscribe(on_stop)