feat(ps): add pod logs support

Author: duanhongyi

charts/controller/templates/controller-clusterrole.yaml

@@ -82,5 +82,8 @@ rules:
 - apiGroups: ["gateway.networking.k8s.io"]
   resources: ["gateways", "httproutes", "grocroutes", "tcproutes", "udproutes", "tlsroutes"]
   verbs: ["get", "patch", "list", "create", "update", "delete"]
+- apiGroups: ["batch"]
+  resources: ["jobs"]
+  verbs: ["create"]
 {{- end -}}
 

charts/controller/templates/controller-webhook-register.yaml

@@ -6,20 +6,27 @@
 apiVersion: admissionregistration.k8s.io/v1
 kind: MutatingWebhookConfiguration
 metadata:
-  name: drycc-controller-webhook
+  name: {{ $altName2 }}
 webhooks:
-- name: "{{ $altName2 }}"
+- name: {{ $altName2 }}
   sideEffects: None
   admissionReviewVersions: ["v1"]
   clientConfig:
     caBundle: {{ b64enc $ca.Cert }}
     service:
       name: drycc-controller-webhook
       namespace: "{{ .Release.Namespace }}"
-      path: "{{ printf "/v2/webhooks/scale/%s/" $token }}"
+      path: "{{ printf "/v2/webhooks/%s/" $token }}"
       port: 8443
   failurePolicy: Fail
+  objectSelector:
+    matchLabels:
+      heritage: drycc
   rules:
+  - operations: ["UPDATE"]
+    apiGroups: ["batch"]
+    apiVersions: ["*"]
+    resources: ["jobs/status"]
   - operations: ["UPDATE"]
     apiGroups: ["apps"]
     apiVersions: ["*"]

rootfs/api/admissions.py

@@ -0,0 +1,83 @@
+from django.db.models import F, Func, Value, JSONField
+from rest_framework.request import Request
+
+from api import models
+
+
+class BaseHandler(object):
+
+    def detect(self, request: Request) -> bool:
+        raise NotImplementedError()
+
+    def handle(self, request: Request) -> bool:
+        raise NotImplementedError()
+
+
+class JobsStatusHandler(BaseHandler):
+
+    def detect(self, request: Request) -> bool:
+        group = request.get("resource", {}).get("group", None)
+        resource = "/".join([
+            request.get("resource", {}).get("resource", None),
+            request.get("subResource", ""),
+        ])
+        if (group, resource) == ("batch", "jobs/status"):
+            return True
+        return False
+
+    def handle(self, request: Request) -> bool:
+        app_id = request["object"]["metadata"]["namespace"]
+        app = models.app.App.objects.filter(id=app_id).first()
+        container_type = request["object"]["metadata"].get("labels", {}).get("type", "")
+        if app and container_type:
+            status = request["object"]["status"]
+            replicas = request["object"]["spec"].get("replicas", 0)
+            if "active" in status:
+                replicas += 1
+            elif "succeeded" in status or "failed" in status:
+                replicas -= 1
+            replicas = 0 if replicas < 0 else replicas
+            if app.structure.get(container_type, 0) != replicas:
+                models.app.App.objects.filter(id=app.id).update(
+                    structure=Func(
+                        F("structure"),
+                        Value([container_type]),
+                        Value(replicas, JSONField()),
+                        function="jsonb_set",
+                    )
+                )
+        return True
+
+
+class DeploymentsScaleHandler(BaseHandler):
+
+    def detect(self, request: Request) -> bool:
+        group = request.get("resource", {}).get("group", None)
+        resource = "/".join([
+            request.get("resource", {}).get("resource", None),
+            request.get("subResource", ""),
+        ])
+        if (group, resource) == ("apps", "deployments/scale"):
+            return True
+        return False
+
+    def handle(self, request: Request) -> bool:
+        app_id = request["object"]["metadata"]["namespace"]
+        app = models.app.App.objects.filter(id=app_id).first()
+        container_type = None
+        for item in request["object"]["status"]["selector"].split(","):
+            key, value = item.split("=")
+            if key == "type":
+                container_type = value
+        if app and container_type:
+            replicas = request["object"]["spec"].get("replicas", 0)
+            if app.structure.get(container_type, 0) != replicas:
+                models.app.App.objects.filter(id=app.id).update(
+                    structure=Func(
+                        F("structure"),
+                        Value([container_type]),
+                        Value(replicas, JSONField()),
+                        function="jsonb_set",
+                    )
+                )
+        return True

rootfs/api/consumers.py

@@ -7,9 +7,9 @@
 import collections
 from django.conf import settings
 
-from asgiref.sync import sync_to_async
+from asgiref.sync import sync_to_async, async_to_sync
 
-from kubernetes.client import Configuration
+from kubernetes.client import Configuration, exceptions
 from kubernetes.client.api import core_v1_api
 from kubernetes.stream import stream
 from kubernetes.stream.ws_client import STDOUT_CHANNEL, STDERR_CHANNEL, ERROR_CHANNEL
@@ -27,12 +27,16 @@
 
 class BaseAppConsumer(AsyncWebsocketConsumer):
 
-    async def has_perm(self):
+    @database_sync_to_async
+    def has_perm(self):
         if self.scope["user"] is None:
             return False, "user not login"
         request = Request(self.scope["user"], "POST")
-        app = await database_sync_to_async(App.objects.get)(id=self.id)
-        return await database_sync_to_async(has_app_permission)(request, app)
+        try:
+            app = App.objects.get(id=self.id)
+            return has_app_permission(request, app)
+        except App.DoesNotExist:
+            return False, "user not exists"
 
     async def connect(self):
         self.id = self.scope["url_route"]["kwargs"]["id"]
@@ -43,6 +47,21 @@ async def connect(self):
             raise DenyConnection(message)
 
 
+class BaseK8sConsumer(BaseAppConsumer):
+
+    @property
+    def kubernetes(self):
+        with open('/var/run/secrets/kubernetes.io/serviceaccount/token') as token_file:
+            token = token_file.read()
+        config = Configuration(host=settings.SCHEDULER_URL)
+        config.api_key = {"authorization": "Bearer " + token}
+        config.verify_ssl = settings.K8S_API_VERIFY_TLS
+        if config.verify_ssl:
+            config.ssl_ca_cert = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
+        Configuration.set_default(config)
+        return core_v1_api.CoreV1Api()
+
+
 class AppLogsConsumer(BaseAppConsumer):
 
     async def receive(self, text_data=None, bytes_data=None):
@@ -68,19 +87,30 @@ async def receive(self, text_data=None, bytes_data=None):
             raise ValueError("text_data cannot be empty!")
 
 
-class AppPodExecConsumer(BaseAppConsumer):
+class AppPodLogsConsumer(BaseK8sConsumer):
 
-    @property
-    def kubernetes(self):
-        with open('/var/run/secrets/kubernetes.io/serviceaccount/token') as token_file:
-            token = token_file.read()
-        config = Configuration(host=settings.SCHEDULER_URL)
-        config.api_key = {"authorization": "Bearer " + token}
-        config.verify_ssl = settings.K8S_API_VERIFY_TLS
-        if config.verify_ssl:
-            config.ssl_ca_cert = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
-        Configuration.set_default(config)
-        return core_v1_api.CoreV1Api()
+    async def connect(self):
+        await super().connect()
+        self.pod_id = self.scope["url_route"]["kwargs"]["pod_id"]
+
+    @sync_to_async
+    def receive(self, text_data=None, bytes_data=None):
+        kwargs = json.loads(text_data)
+        try:
+            stream = self.kubernetes.read_namespaced_pod_log(self.pod_id, self.id, **{
+                "tail_lines": kwargs.get("lines", 100),
+                "follow": kwargs.get("follow", False),
+                "container": kwargs.get("container", ""),
+                "_preload_content": False,
+            }).stream()
+            for line in stream:
+                async_to_sync(self.send)(text_data=line)
+        except exceptions.ApiException as e:
+            async_to_sync(self.send)(text_data=str(e))
+        async_to_sync(self.close)(code=1000)
+
+
+class AppPodExecConsumer(BaseK8sConsumer):
 
     async def connect(self):
         self.stream = None
@@ -121,8 +151,13 @@ async def receive(self, text_data=None, bytes_data=None):
             args = (self.kubernetes.connect_get_namespaced_pod_exec, self.pod_id, self.id)
             kwargs = json.loads(text_data)
             kwargs.update({"stderr": True, "stdout": True, "_preload_content": False})
-            self.stream = stream(*args, **kwargs)
-            asyncio.create_task(self.task())
+            try:
+                self.stream = stream(*args, **kwargs)
+            except exceptions.ApiException as e:
+                await self.send(str(e), STDERR_CHANNEL)
+                await self.close(code=1000)
+            else:
+                asyncio.create_task(self.task())
         elif self.stream is not None:
             data = text_data if text_data else bytes_data
             channel, data = ord(data[0]), data[1:]