feat(blocklist): blocklist api (#190)

Author: jianxiaoguo

charts/controller/templates/controller-clusterrole.yaml

@@ -51,6 +51,12 @@ rules:
 - apiGroups: ["extensions", "apps"]
   resources: ["deployments/scale", "replicasets/scale"]
   verbs: ["get", "update"]
+- apiGroups: ["apps"]
+  resources: ["statefulsets"]
+  verbs: ["get", "list", "patch"]
+- apiGroups: ["extensions", "apps"]
+  resources: ["daemonsets"]
+  verbs: ["get", "list", "patch"]
 - apiGroups: ["extensions", "autoscaling"]
   resources: ["horizontalpodautoscalers"]
   verbs: ["get", "list", "create", "update", "delete"]

rootfs/api/backend.py

@@ -21,7 +21,7 @@ class DryccOIDC(OpenIdConnectAuth):
     USERINFO_URL = settings.SOCIAL_AUTH_DRYCC_USERINFO_URL
     JWKS_URI = settings.SOCIAL_AUTH_DRYCC_JWKS_URI
     OIDC_ENDPOINT = settings.SOCIAL_AUTH_DRYCC_OIDC_ENDPOINT
-    DEFAULT_SCOPE = ['openid']
+    DEFAULT_SCOPE = ['openid', 'profile', 'email']
     EXTRA_DATA = [
         ('id', 'id'),
         ('access_token', 'access_token'),

rootfs/api/migrations/0025_alter_app_id_alter_gateway_name.py

@@ -0,0 +1,25 @@
+# Generated by Django 4.2.22 on 2025-08-22 02:48
+
+import api.models.app
+import api.utils
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('api', '0024_limitplan_container_volume_mounts_and_more'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='app',
+            name='id',
+            field=models.SlugField(max_length=63, null=True, unique=True, validators=[api.models.app.validate_app_id]),
+        ),
+        migrations.AlterField(
+            model_name='gateway',
+            name='name',
+            field=models.CharField(db_index=True, max_length=63, validators=[api.utils.validate_label]),
+        ),
+    ]

rootfs/api/migrations/0026_app_suspended_state.py

@@ -0,0 +1,18 @@
+# Generated by Django 4.2.22 on 2025-08-22 02:56
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('api', '0025_alter_app_id_alter_gateway_name'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='app',
+            name='suspended_state',
+            field=models.JSONField(blank=True, default=dict),
+        ),
+    ]

rootfs/api/models/app.py

@@ -114,6 +114,7 @@ class App(UuidAuditedModel):
                           validators=[validate_app_id])
     structure = models.JSONField(
         default=dict, blank=True, validators=[validate_app_structure])
+    suspended_state = models.JSONField(default=dict, blank=True)
 
     class Meta:
         verbose_name = 'Application'

rootfs/api/models/blocklist.py

@@ -1,9 +1,16 @@
+import logging
+import functools
 from django.db import models
 from django.contrib.auth import get_user_model
+from scheduler import KubeHTTPException
+from api.exceptions import ServiceUnavailable
+from api.utils import apply_tasks
 from .app import App
 from .base import UuidAuditedModel
 
+
 User = get_user_model()
+logger = logging.getLogger(__name__)
 
 
 class Blocklist(UuidAuditedModel):
@@ -41,3 +48,227 @@ def get_blocklist(cls, app: App):
     class Meta:
         ordering = ['-created']
         unique_together = (("id", "type"),)
+
+    def related_resource_deployments(self, app: App):
+        "get resource deployments"
+        try:
+            deployments = self.scheduler.deployment.get(app.id, labels={}).json()['items']  # noqa
+            if not deployments:
+                deployments = []
+            data = []
+            for d in deployments:
+                item = {
+                    'name': d['metadata']['name'],
+                    'replicas': d['spec'].get("replicas", 0),
+                }
+                data.append(item)
+            data.sort(key=lambda x: x['name'])
+            return data
+        except KubeHTTPException:
+            pass
+        except Exception as e:
+            err = f'(list resource deployments): {e}'
+            logger.info(err)
+            raise ServiceUnavailable(err) from e
+
+    def related_resource_statefulsets(self, app: App):
+        "get resource statefulsets"
+        try:
+            statefulsets = self.scheduler.statefulset.get(app.id, labels={}).json()['items']  # noqa
+            if not statefulsets:
+                statefulsets = []
+            data = []
+            for s in statefulsets:
+                item = {
+                    'name': s['metadata']['name'],
+                    'replicas': s['spec'].get("replicas", 0),
+                }
+                data.append(item)
+            data.sort(key=lambda x: x['name'])
+            return data
+        except KubeHTTPException:
+            pass
+        except Exception as e:
+            err = f'(list resource statefulsets): {e}'
+            logger.info(err)
+            raise ServiceUnavailable(err) from e
+
+    def related_resource_daemonsets(self, app: App):
+        "get resource daemonsets"
+        try:
+            daemonsets = self.scheduler.daemonset.get(app.id, labels={}).json()['items']  # noqa
+            if not daemonsets:
+                daemonsets = []
+            data = []
+            for d in daemonsets:
+                item = {
+                    'name': d['metadata']['name'],
+                    'affinity': d['spec']['template']['spec'].get('affinity', {}),
+                }
+                data.append(item)
+            data.sort(key=lambda x: x['name'])
+            return data
+        except KubeHTTPException:
+            pass
+        except Exception as e:
+            err = f'(list resource daemonsets): {e}'
+            logger.info(err)
+            raise ServiceUnavailable(err) from e
+
+    def suspended_state(self, app: App):
+        """
+        Get deployments/statefulsets/daemonsets with labels app.kubernetes.io/managed-by=Helm
+        Store in the suspended_state field of the app model
+        Format:
+        {"deployments": [{"name":"sample1", "replicas": 1}],
+        "statefulsets": [{"name":"sample1", "replicas": 1}],
+        "daemonsets": [{"name":"sample1", "affinity": {}}]}
+        """
+        suspended_state = {}
+        suspended_state['deployments'] = self.related_resource_deployments(app)
+        suspended_state['statefulsets'] = self.related_resource_statefulsets(app)
+        suspended_state['daemonsets'] = self.related_resource_daemonsets(app)
+        return suspended_state
+
+    def scale_resource_deployments(self, app: App, deployment_name: str, replicas=0):
+        """scale deployments"""
+        try:
+            deployment = self.scheduler.deployment.get(app.id, deployment_name).json()
+            self.scheduler.scales.update(app.id, deployment_name, replicas, deployment)
+        except KubeHTTPException:
+            pass
+        except Exception as e:
+            err = f'(scale resource deployments): {e}'
+            logger.info(err)
+            raise ServiceUnavailable(err) from e
+
+    def scale_resource_statefulsets(self, app: App, statefulset_name: str, replicas=0):
+        """scale statefulsets"""
+        try:
+            self.scheduler.statefulset.get(app.id, statefulset_name).json()
+            manifest = {
+                'spec': {
+                    'persistentVolumeClaimRetentionPolicy': {
+                        'whenScaled': 'Retain'
+                    },
+                    'replicas': replicas
+                }
+            }
+            self.scheduler.statefulset.patch(app.id, statefulset_name, manifest)
+        except KubeHTTPException:
+            pass
+        except Exception as e:
+            err = f'(scale resource statefulsets): {e}'
+            logger.info(err)
+            raise ServiceUnavailable(err) from e
+
+    def scale_resource_daemonsets(self, app: App, daemonset_name: str, manifest: dict):
+        """set affinity"""
+        try:
+            self.scheduler.daemonset.get(app.id, daemonset_name).json()
+            self.scheduler.daemonset.patch(app.id, daemonset_name, manifest)
+        except KubeHTTPException:
+            pass
+        except Exception as e:
+            err = f'(scale resource daemonsets): {e}'
+            logger.info(err)
+            raise ServiceUnavailable(err) from e
+
+    def scale_resources(self, app: App, suspended_state: dict, scale_type="block"):
+        "scale resources tasks"
+        tasks = []
+        deployments = suspended_state.get('deployments', [])
+        if not deployments:
+            deployments = []
+        for d in deployments:
+            if scale_type == "unblock":
+                replicas = d['replicas']
+            else:
+                replicas = 0
+            tasks.append((
+                functools.partial(
+                    self.scale_resource_deployments,
+                    app=app,
+                    deployment_name=d['name'],
+                    replicas=replicas
+                ),
+                lambda future, name=d["name"]: app.log(
+                    f'{scale_type} scale deployment {name} callback: {future.result()}',
+                )
+            ))
+
+        statefulsets = suspended_state.get('statefulsets', [])
+        if not statefulsets:
+            statefulsets = []
+        for s in statefulsets:
+            if scale_type == "unblock":
+                replicas = s['replicas']
+            else:
+                replicas = 0
+            tasks.append((
+                functools.partial(
+                    self.scale_resource_statefulsets,
+                    app=app,
+                    statefulset_name=s['name'],
+                    replicas=replicas
+                ),
+                lambda future, name=s["name"]: app.log(
+                    f'{scale_type} scale statefulset {name} callback: {future.result()}',
+                )
+            ))
+
+        daemonsets = suspended_state.get('daemonsets', [])
+        if not daemonsets:
+            daemonsets = []
+        for d in daemonsets:
+            if scale_type == "unblock":
+                manifest = {
+                    "spec": {
+                        "template": {
+                            "spec": {
+                                "affinity": d['affinity']
+                            }
+                        }
+                    }
+                }
+            else:
+                manifest = {
+                    "spec": {
+                        "template": {
+                            "spec": {
+                                "affinity": {
+                                    "nodeAffinity": {
+                                        "requiredDuringSchedulingIgnoredDuringExecution": {
+                                            "nodeSelectorTerms": [{
+                                                "matchExpressions": [{
+                                                    "key": "kubernetes.io/hostname",
+                                                    "operator": "In",
+                                                    "values": [
+                                                        "nohostname"
+                                                    ]
+                                                }]
+                                            }]
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            tasks.append((
+                functools.partial(
+                    self.scale_resource_daemonsets,
+                    app=app,
+                    daemonset_name=d['name'],
+                    manifest=manifest
+                ),
+                lambda future, name=d["name"]: app.log(
+                    f'{scale_type} scale daemonset {name} callback: {future.result()}',
+                )
+            ))
+        try:
+            apply_tasks(tasks)
+        except Exception as e:
+            err = f'({scale_type} scale resources): {e}'
+            logger.info(err)
+            raise ServiceUnavailable(err) from e

rootfs/api/tasks.py

@@ -158,3 +158,19 @@ def downstream_model_owner(app, old_owner, new_owner):
         raise e
     else:
         signals.request_finished.send(sender=task_id)
+
+
+@shared_task(
+    autoretry_for=(ServiceUnavailable, ),
+    retry_kwargs={'max_retries': None}
+)
+def scale_resources(blocklist, app, suspended_state, scale_type):
+    task_id = uuid.uuid4().hex
+    signals.request_started.send(sender=task_id)
+    try:
+        blocklist.scale_resources(app, suspended_state, scale_type)
+    except Exception as e:
+        signals.got_request_exception.send(sender=task_id)
+        raise e
+    else:
+        signals.request_finished.send(sender=task_id)

rootfs/api/views.py

@@ -35,7 +35,7 @@
 
 from api import monitor, models, permissions, serializers, viewsets, authentication, __version__
 from api.tasks import scale_app, restart_app, mount_app, downstream_model_owner, \
-    delete_pod
+    delete_pod, scale_resources
 from api.exceptions import AlreadyExists, ServiceUnavailable, DryccException
 
 from django.views.decorators.cache import never_cache
@@ -183,24 +183,35 @@ class WorkflowManagerViewset(GenericViewSet):
 
     def block(self, request,  **kwargs):
         try:
-            blocklist, _ = models.blocklist.Blocklist.objects.get_or_create(
+            blocklist, created = models.blocklist.Blocklist.objects.get_or_create(
                 id=kwargs['id'],
                 type=models.blocklist.Blocklist.get_type(kwargs["type"]),
                 defaults={"remark": request.data.get("remark")}
             )
             for app in blocklist.related_apps:
-                scale_app.delay(app, app.owner, {key: 0 for key in app.structure.keys()})
+                if created:
+                    app.suspended_state = blocklist.suspended_state(app)
+                    app.save()
+                # scale to 0
+                scale_resources.delay(blocklist, app, app.suspended_state, "block")
             return HttpResponse(status=201)
         except ValueError as e:
             logger.info(e)
             raise DryccException("Unsupported block type: %s" % kwargs["type"])
 
     def unblock(self, request,  **kwargs):
         try:
-            models.blocklist.Blocklist.objects.filter(
+            blocklists = models.blocklist.Blocklist.objects.filter(
                 id=kwargs['id'],
                 type=models.blocklist.Blocklist.get_type(kwargs["type"])
-            ).delete()
+            )
+            for blocklist in blocklists:
+                for app in blocklist.related_apps:
+                    # scale to up
+                    scale_resources.delay(blocklist, app, app.suspended_state, "unblock")
+                    app.suspended_state = {}
+                    app.save()
+                blocklist.delete()
             return HttpResponse(status=204)
         except ValueError as e:
             logger.info(e)

rootfs/dev_requirements.txt

@@ -8,3 +8,6 @@ flake8==7.2.0
 
 # mock out python-requests, mostly k8s
 requests-mock==1.12.1
+
+# tracebacks
+tblib==3.1.0