Merge pull request #9512 from jumpserver/pr@dev@perf_asset_task

perf: 优化 applet account 释放

apps/authentication/api/connection_token.py

@@ -14,16 +14,16 @@ from rest_framework.request import Request
 from rest_framework.response import Response
 from rest_framework.serializers import ValidationError
 
+from assets.const import CloudTypes
 from common.api import JMSModelViewSet
 from common.exceptions import JMSException
-from common.utils import random_string
+from common.utils import random_string, get_logger
 from common.utils.django import get_request_os
 from common.utils.http import is_true
 from orgs.mixins.api import RootOrgViewMixin
 from perms.models import ActionChoices
 from terminal.connect_methods import NativeClient, ConnectMethodUtil
 from terminal.models import EndpointRule
-from assets.const import CloudTypes
 from ..models import ConnectionToken
 from ..serializers import (
     ConnectionTokenSerializer, ConnectionTokenSecretSerializer,
@@ -31,6 +31,7 @@ from ..serializers import (
 )
 
 __all__ = ['ConnectionTokenViewSet', 'SuperConnectionTokenViewSet']
+logger = get_logger(__name__)
 
 
 class RDPFileClientProtocolURLMixin:
@@ -366,5 +367,11 @@ class SuperConnectionTokenViewSet(ConnectionTokenViewSet):
     @action(methods=['DELETE', 'POST'], detail=False, url_path='applet-account/release')
     def release_applet_account(self, *args, **kwargs):
         account_id = self.request.data.get('id')
-        msg = ConnectionToken.release_applet_account(account_id)
-        return Response({'msg': msg})
+        released = ConnectionToken.release_applet_account(account_id)
+
+        if released:
+            logger.debug('Release applet account success: {}'.format(account_id))
+            return Response({'msg': 'released'})
+        else:
+            logger.error('Release applet account error: {}'.format(account_id))
+            return Response({'error': 'not found or expired'}, status=400)

apps/authentication/models/connection_token.py

@@ -11,7 +11,8 @@ from rest_framework.exceptions import PermissionDenied
 
 from assets.const import Protocol
 from common.db.fields import EncryptCharField
-from common.utils import lazyproperty, pretty_string, bulk_get, reverse
+from common.exceptions import JMSException
+from common.utils import lazyproperty, pretty_string, bulk_get
 from common.utils.timezone import as_current_tz
 from orgs.mixins.models import JMSOrgBaseModel
 from terminal.models import Applet
@@ -172,7 +173,7 @@ class ConnectionToken(JMSOrgBaseModel):
 
         host_account = applet.select_host_account()
         if not host_account:
-            return None
+            raise JMSException({'error': 'No host account available'})
 
         host, account, lock_key, ttl = bulk_get(host_account, ('host', 'account', 'lock_key', 'ttl'))
         gateway = host.gateway.select_gateway() if host.domain else None
@@ -196,8 +197,7 @@ class ConnectionToken(JMSOrgBaseModel):
         if lock_key:
             cache.delete(lock_key)
             cache.delete(token_account_relate_key)
-            return 'released'
-        return 'not found or expired'
+            return True
 
     @lazyproperty
     def account_object(self):

apps/orgs/models.py

@@ -134,12 +134,12 @@ class Organization(OrgRoleMixin, JMSBaseModel):
         return self.id
 
     @classmethod
-    def get_or_create_builtin(cls, name, **kwargs):
-        _id = kwargs.get('id')
-        org = cls.get_instance(cls.DEFAULT_ID)
+    def get_or_create_builtin(cls, **kwargs):
+        _id = kwargs['id']
+        org = cls.get_instance(_id)
         if org:
             return org
-        org, created = cls.objects.get_or_create(name=name, defaults=kwargs)
+        org, created = cls.objects.get_or_create(id=_id, defaults=kwargs)
         if created:
             org.builtin = True
             org.save()

apps/terminal/models/applet/applet.py

@@ -11,7 +11,9 @@ from django.utils.translation import gettext_lazy as _
 from rest_framework.serializers import ValidationError
 
 from common.db.models import JMSBaseModel
-from common.utils import lazyproperty
+from common.utils import lazyproperty, get_logger
+
+logger = get_logger(__name__)
 
 __all__ = ['Applet', 'AppletPublication']
 
@@ -109,19 +111,27 @@ class Applet(JMSBaseModel):
         if not hosts:
             return None
 
+        key_tmpl = 'applet_host_accounts_{}_{}'
         host = random.choice(hosts)
-        using_keys = cache.keys('host_accounts_{}_*'.format(host.id)) or []
-        accounts_used = cache.get_many(using_keys)
-        accounts = host.accounts.all().exclude(username__in=accounts_used)
+        using_keys = cache.keys(key_tmpl.format(host.id, '*')) or []
+        accounts_username_used = list(cache.get_many(using_keys).values())
+        logger.debug('Applet host account using: {}: {}'.format(host.name, accounts_username_used))
+        accounts = host.accounts.all() \
+            .filter(is_active=True, privileged=False) \
+            .exclude(username__in=accounts_username_used)
+
+        msg = 'Applet host remain accounts: {}: {}'.format(host.name, len(accounts))
+        if len(accounts) == 0:
+            logger.error(msg)
+        else:
+            logger.debug(msg)
 
-        if not accounts:
-            accounts = host.accounts.all()
         if not accounts:
             return None
 
         account = random.choice(accounts)
         ttl = 60 * 60 * 24
-        lock_key = 'applet_host_accounts_{}_{}'.format(host.id, account.username)
+        lock_key = key_tmpl.format(host.id, account.username)
         cache.set(lock_key, account.username, ttl)
 
         return {
@@ -131,11 +141,6 @@ class Applet(JMSBaseModel):
             'ttl': ttl
         }
 
-    @staticmethod
-    def release_host_and_account(host_id, username):
-        key = 'applet_host_accounts_{}_{}'.format(host_id, username)
-        cache.delete(key)
-
 
 class AppletPublication(JMSBaseModel):
     applet = models.ForeignKey('Applet', on_delete=models.CASCADE, related_name='publications',