perf: 远程应用调度优先调度的上个主机，使用上个账号，并支持同名账号

apps/authentication/models/connection_token.py

@@ -172,7 +172,7 @@ class ConnectionToken(JMSOrgBaseModel):
         if not applet:
             return None
 
-        host_account = applet.select_host_account()
+        host_account = applet.select_host_account(self.user)
         if not host_account:
             raise JMSException({'error': 'No host account available'})
 

apps/terminal/migrations/0050_auto_20220606_1745.py

@@ -4,7 +4,6 @@ from django.db import migrations, models
 
 
 class Migration(migrations.Migration):
-
     dependencies = [
         ('terminal', '0049_endpoint_redis_port'),
     ]
@@ -13,10 +12,10 @@ class Migration(migrations.Migration):
         migrations.AlterField(
             model_name='terminal',
             name='type',
-            field=models.CharField(choices=[
-                ('koko', 'KoKo'), ('guacamole', 'Guacamole'), ('omnidb', 'OmniDB'),
-                ('xrdp', 'Xrdp'), ('lion', 'Lion'), ('core', 'Core'), ('celery', 'Celery'),
-                ('magnus', 'Magnus'), ('razor', 'Razor'), ('tinker', 'Tinker'),
-            ], default='koko', max_length=64, verbose_name='type'),
+            field=models.CharField(
+                choices=[('koko', 'KoKo'), ('guacamole', 'Guacamole'), ('omnidb', 'OmniDB'), ('xrdp', 'Xrdp'),
+                         ('lion', 'Lion'), ('core', 'Core'), ('celery', 'Celery'), ('magnus', 'Magnus'),
+                         ('razor', 'Razor'), ('tinker', 'Tinker'), ('video_worker', 'Video Worker')], default='koko',
+                max_length=64, verbose_name='type'),
         ),
     ]

apps/terminal/migrations/0061_applet_can_concurrent.py

@@ -0,0 +1,18 @@
+# Generated by Django 3.2.17 on 2023-05-09 11:02
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('terminal', '0060_sessionsharing_action_permission'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='applet',
+            name='can_concurrent',
+            field=models.BooleanField(default=True, verbose_name='Can concurrent'),
+        ),
+    ]

apps/terminal/models/applet/applet.py

@@ -32,6 +32,7 @@ class Applet(JMSBaseModel):
     is_active = models.BooleanField(default=True, verbose_name=_('Is active'))
     builtin = models.BooleanField(default=False, verbose_name=_('Builtin'))
     protocols = models.JSONField(default=list, verbose_name=_('Protocol'))
+    can_concurrent = models.BooleanField(default=True, verbose_name=_('Can concurrent'))
     tags = models.JSONField(default=list, verbose_name=_('Tags'))
     comment = models.TextField(default='', blank=True, verbose_name=_('Comment'))
     hosts = models.ManyToManyField(
@@ -134,37 +135,68 @@ class Applet(JMSBaseModel):
         shutil.copytree(path, pkg_path)
         return instance, serializer
 
-    def select_host_account(self):
-        # 选择激活的发布机
+    def select_host(self, user):
         hosts = [
             host for host in self.hosts.filter(is_active=True)
             if host.load != 'offline'
         ]
-
         if not hosts:
             return None
 
-        key_tmpl = 'applet_host_accounts_{}_{}'
-        host = random.choice(hosts)
-        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)
+        prefer_key = 'applet_host_prefer_{}'.format(user.id)
+        prefer_host_id = cache.get(prefer_key, None)
+        pref_host = [host for host in hosts if host.id == prefer_host_id]
+        if pref_host:
+            host = pref_host[0]
+        else:
+            host = random.choice(hosts)
+            cache.set(prefer_key, host.id, timeout=None)
+        return host
 
-        msg = 'Applet host remain accounts: {}: {}'.format(host.name, len(accounts))
+    @staticmethod
+    def random_select_prefer_account(user, host, accounts):
+        msg = 'Applet host remain public accounts: {}: {}'.format(host.name, len(accounts))
         if len(accounts) == 0:
             logger.error(msg)
-        else:
-            logger.debug(msg)
-
-        if not accounts:
             return None
+        prefer_host_account_key = 'applet_host_prefer_account_{}_{}'.format(user.id, host.id)
+        prefer_account_id = cache.get(prefer_host_account_key, None)
+        prefer_account = accounts.filter(id=prefer_account_id).first()
+        if prefer_account:
+            account = prefer_account
+        else:
+            account = random.choice(accounts)
+            cache.set(prefer_host_account_key, account.id, timeout=None)
+        return account
 
-        account = random.choice(accounts)
+    def select_host_account(self, user):
+        # 选择激活的发布机
+        host = self.select_host(user)
+        if not host:
+            return None
+        can_concurrent = self.can_concurrent and self.type == 'general'
+
+        accounts = host.accounts.all().filter(is_active=True, privileged=False)
+        private_account = accounts.filter(username='js_{}'.format(user.username)).first()
+        accounts_using_key_tmpl = 'applet_host_accounts_{}_{}'
+
+        if private_account and can_concurrent:
+            account = private_account
+        else:
+            using_keys = cache.keys(accounts_using_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))
+
+            # 优先使用 private account
+            if private_account and private_account.username not in accounts_username_used:
+                account = private_account
+            else:
+                accounts = accounts.exclude(username__in=accounts_username_used)
+                account = self.random_select_prefer_account(user, host, accounts)
+                if not account:
+                    return
         ttl = 60 * 60 * 24
-        lock_key = key_tmpl.format(host.id, account.username)
+        lock_key = accounts_using_key_tmpl.format(host.id, account.username)
         cache.set(lock_key, account.username, ttl)
 
         return {

apps/terminal/models/applet/host.py

@@ -84,9 +84,13 @@ class AppletHost(Host):
         return random_string(16, special_char=True)
 
     def generate_accounts(self):
-        amount = int(os.getenv('TERMINAL_ACCOUNTS_AMOUNT', 100))
-        now_count = self.accounts.filter(privileged=False).count()
-        need = amount - now_count
+        self.generate_public_accounts()
+        self.generate_private_accounts()
+
+    def generate_public_accounts(self):
+        public_amount = int(os.getenv('TERMINAL_ACCOUNTS_AMOUNT', 100))
+        now_count = self.accounts.filter(privileged=False, username__startswith='jms').count()
+        need = public_amount - now_count
 
         accounts = []
         account_model = self.accounts.model
@@ -99,7 +103,30 @@ class AppletHost(Host):
                 org_id=self.LOCKING_ORG, is_active=False,
             )
             accounts.append(account)
-        bulk_create_with_history(accounts, account_model, batch_size=20)
+        bulk_create_with_history(accounts, account_model, batch_size=20, ignore_conflicts=True)
+
+    def generate_private_accounts_by_usernames(self, usernames):
+        accounts = []
+        account_model = self.accounts.model
+        for username in usernames:
+            password = self.random_password()
+            username = 'js_' + username
+            account = account_model(
+                username=username, secret=password, name=username,
+                asset_id=self.id, secret_type='password', version=1,
+                org_id=self.LOCKING_ORG, is_active=False,
+            )
+            accounts.append(account)
+        bulk_create_with_history(accounts, account_model, batch_size=20, ignore_conflicts=True)
+
+    def generate_private_accounts(self):
+        from users.models import User
+        usernames = User.objects \
+            .filter(is_active=True, is_service_account=False) \
+            .values_list('username', flat=True)
+        account_usernames = self.accounts.all().values_list('username', flat=True)
+        not_exist_users = set(usernames) - set(account_usernames)
+        self.generate_private_accounts_by_usernames(not_exist_users)
 
 
 class AppletHostDeployment(JMSBaseModel):

apps/terminal/signal_handlers/applet.py

@@ -2,11 +2,14 @@ from django.db.models.signals import post_save, post_delete
 from django.dispatch import receiver
 from django.utils.functional import LazyObject
 
+from accounts.models import Account
 from common.signals import django_ready
 from common.utils import get_logger
 from common.utils.connection import RedisPubSub
 from orgs.utils import tmp_to_builtin_org
+from users.models import User
 from ..models import Applet, AppletHost
+from ..tasks import applet_host_generate_accounts
 from ..utils import DBPortManager
 
 db_port_manager: DBPortManager
@@ -19,12 +22,30 @@ def on_applet_host_create(sender, instance, created=False, **kwargs):
         return
     applets = Applet.objects.all()
     instance.applets.set(applets)
-    with tmp_to_builtin_org(system=1):
-        instance.generate_accounts()
 
+    applet_host_generate_accounts.delay(instance.id)
     applet_host_change_pub_sub.publish(True)
 
 
+@receiver(post_save, sender=User)
+def on_user_create_create_account(sender, instance, created=False, **kwargs):
+    if not created:
+        return
+
+    with tmp_to_builtin_org(system=1):
+        applet_hosts = AppletHost.objects.all()
+        for host in applet_hosts:
+            host.generate_private_accounts_by_usernames([instance.username])
+
+
+@receiver(post_delete, sender=User)
+def on_user_delete_remove_account(sender, instance, **kwargs):
+    with tmp_to_builtin_org(system=1):
+        applet_hosts = AppletHost.objects.all().values_list('id', flat=True)
+        accounts = Account.objects.filter(asset_id__in=applet_hosts, username=instance.username)
+        accounts.delete()
+
+
 @receiver(post_delete, sender=AppletHost)
 def on_applet_host_delete(sender, instance, **kwargs):
     applet_host_change_pub_sub.publish(True)

apps/terminal/tasks.py

@@ -16,7 +16,7 @@ from ops.celery.decorator import (
 from orgs.utils import tmp_to_builtin_org
 from .backends import server_replay_storage
 from .models import (
-    Status, Session, Task, AppletHostDeployment
+    Status, Session, Task, AppletHostDeployment, AppletHost
 )
 from .utils import find_session_replay_local
 
@@ -82,7 +82,7 @@ def upload_session_replay_to_external_storage(session_id):
 
 @shared_task(
     verbose_name=_('Run applet host deployment'),
-    activity_callback=lambda self, did, *args, **kwargs: ([did], )
+    activity_callback=lambda self, did, *args, **kwargs: ([did],)
 )
 def run_applet_host_deployment(did):
     with tmp_to_builtin_org(system=1):
@@ -98,3 +98,16 @@ def run_applet_host_deployment_install_applet(did, applet_id):
     with tmp_to_builtin_org(system=1):
         deployment = AppletHostDeployment.objects.get(id=did)
         deployment.install_applet(applet_id)
+
+
+@shared_task(
+    verbose_name=_('Generate applet host accounts'),
+    activity_callback=lambda self, host_id, *args, **kwargs: ([host_id],)
+)
+def applet_host_generate_accounts(host_id):
+    applet_host = AppletHost.objects.filter(id=host_id).first()
+    if not applet_host:
+        return
+
+    with tmp_to_builtin_org(system=1):
+        applet_host.generate_accounts()