From 135fb7c6f9c08add35cb03032fbb2a2f55a31d55 Mon Sep 17 00:00:00 2001
From: wangruidong <940853815@qq.com>
Date: Mon, 19 Feb 2024 14:47:32 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E7=BB=88=E6=96=AD=E6=89=B9=E9=87=8F?=
 =?UTF-8?q?=E5=BF=AB=E6=8D=B7=E5=91=BD=E4=BB=A4=E6=89=A7=E8=A1=8C=E7=9A=84?=
 =?UTF-8?q?=E4=BB=BB=E5=8A=A1?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/ops/ansible/callback.py | 14 ++++++++++++++
 apps/ops/api/job.py          | 13 ++++++++++++-
 apps/ops/models/job.py       |  9 +++++++++
 apps/ops/serializers/job.py  |  7 +++++++
 4 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/apps/ops/ansible/callback.py b/apps/ops/ansible/callback.py
index d5fb3a35e..d05158725 100644
--- a/apps/ops/ansible/callback.py
+++ b/apps/ops/ansible/callback.py
@@ -1,3 +1,4 @@
+import os
 from collections import defaultdict
 from functools import reduce
 
@@ -29,6 +30,8 @@ class DefaultCallback:
         )
         self.status = 'running'
         self.finished = False
+        self.local_pid = 0
+        self.private_data_dir = None
 
     @property
     def host_results(self):
@@ -45,6 +48,9 @@ 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', '')
@@ -152,3 +158,11 @@ class DefaultCallback:
     def status_handler(self, data, **kwargs):
         status = data.get('status', '')
         self.status = self.STATUS_MAPPER.get(status, 'unknown')
+
+        rc = kwargs.get('runner_config', None)
+        self.private_data_dir = rc.private_data_dir if rc else '/tmp/'
+
+    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))
diff --git a/apps/ops/api/job.py b/apps/ops/api/job.py
index eb908d04c..4c60956fd 100644
--- a/apps/ops/api/job.py
+++ b/apps/ops/api/job.py
@@ -16,7 +16,7 @@ from common.const.http import POST
 from common.permissions import IsValidUser
 from ops.const import Types
 from ops.models import Job, JobExecution
-from ops.serializers.job import JobSerializer, JobExecutionSerializer, FileSerializer
+from ops.serializers.job import JobSerializer, JobExecutionSerializer, FileSerializer, JobTaskStopSerializer
 
 __all__ = [
     'JobViewSet', 'JobExecutionViewSet', 'JobRunVariableHelpAPIView',
@@ -187,6 +187,17 @@ class JobExecutionViewSet(OrgBulkModelViewSet):
         queryset = queryset.filter(creator=self.request.user)
         return queryset
 
+    @action(methods=[POST], detail=False, serializer_class=JobTaskStopSerializer, permission_classes=[IsValidUser, ],
+            url_path='stop')
+    def stop(self, request, *args, **kwargs):
+        serializer = self.get_serializer(data=request.data)
+        if not serializer.is_valid():
+            return Response({'error': serializer.errors}, status=400)
+        task_id = serializer.validated_data['task_id']
+        instance = get_object_or_404(JobExecution, task_id=task_id, creator=request.user)
+        instance.stop()
+        return Response({'task_id': task_id}, status=200)
+
 
 class JobAssetDetail(APIView):
     rbac_perms = {
diff --git a/apps/ops/models/job.py b/apps/ops/models/job.py
index 04c7fa519..f7831251b 100644
--- a/apps/ops/models/job.py
+++ b/apps/ops/models/job.py
@@ -554,6 +554,15 @@ class JobExecution(JMSOrgBaseModel):
         finally:
             ssh_tunnel.local_gateway_clean(runner)
 
+    def stop(self):
+        with open(os.path.join(self.private_dir, 'local.pid')) as f:
+            try:
+                pid = f.read()
+                os.kill(int(pid), 9)
+            except Exception as e:
+                print(e)
+        self.set_error('Job stop by "kill -9 {}"'.format(pid))
+
     class Meta:
         verbose_name = _("Job Execution")
         ordering = ['-date_created']
diff --git a/apps/ops/serializers/job.py b/apps/ops/serializers/job.py
index 75729f988..ce4faee55 100644
--- a/apps/ops/serializers/job.py
+++ b/apps/ops/serializers/job.py
@@ -57,6 +57,13 @@ class FileSerializer(serializers.Serializer):
         ref_name = "JobFileSerializer"
 
 
+class JobTaskStopSerializer(serializers.Serializer):
+    task_id = serializers.CharField(max_length=128)
+
+    class Meta:
+        ref_name = "JobTaskStopSerializer"
+
+
 class JobExecutionSerializer(BulkOrgResourceModelSerializer):
     creator = ReadableHiddenField(default=serializers.CurrentUserDefault())
     job_type = serializers.ReadOnlyField(label=_("Job type"))