perf: 优化org获取逻辑 - 采用redis订阅机制实现orgs_mapping数据的维护;删除get_org_by_id等方法;

perf: 优化get_instance接口

apps/orgs/migrations/0010_auto_20210226_1536.py

@@ -1,18 +0,0 @@
-# Generated by Django 3.1 on 2021-02-26 07:36
-
-from django.db import migrations, models
-
-
-class Migration(migrations.Migration):
-
-    dependencies = [
-        ('orgs', '0009_auto_20201023_1628'),
-    ]
-
-    operations = [
-        migrations.AlterField(
-            model_name='organizationmember',
-            name='role',
-            field=models.CharField(choices=[('Admin', 'Organization administrator'), ('Auditor', 'Organization auditor'), ('User', 'User')], db_index=True, default='User', max_length=16, verbose_name='Role'),
-        ),
-    ]

apps/orgs/mixins/models.py

@@ -7,8 +7,7 @@ from django.core.exceptions import ValidationError
 
 from common.utils import get_logger
 from ..utils import (
-    set_current_org, get_current_org, current_org,
+    set_current_org, get_current_org, current_org, filter_org_queryset
-    filter_org_queryset, get_org_by_id, get_org_name_by_id
 )
 from ..models import Organization
 
@@ -59,11 +58,11 @@ class OrgModelMixin(models.Model):
 
     @property
     def org(self):
-        return get_org_by_id(self.org_id)
+        return Organization.get_instance(self.org_id)
 
     @property
     def org_name(self):
-        return get_org_name_by_id(self.org_id)
+        return self.org.name
 
     @property
     def fullname(self, attr=None):

apps/orgs/models.py

@@ -24,63 +24,53 @@ class Organization(models.Model):
     created_by = models.CharField(max_length=32, null=True, blank=True, verbose_name=_('Created by'))
     date_created = models.DateTimeField(auto_now_add=True, null=True, blank=True, verbose_name=_('Date created'))
     comment = models.TextField(default='', blank=True, verbose_name=_('Comment'))
-    members = models.ManyToManyField('users.User', related_name='orgs', through='orgs.OrganizationMember',
+    members = models.ManyToManyField('users.User', related_name='orgs', through='orgs.OrganizationMember', through_fields=('org', 'user'))
-                                     through_fields=('org', 'user'))
 
-    orgs = None
-    CACHE_PREFIX = 'JMS_ORG_{}'
     ROOT_ID = '00000000-0000-0000-0000-000000000000'
     ROOT_NAME = _('GLOBAL')
     DEFAULT_ID = '00000000-0000-0000-0000-000000000001'
     DEFAULT_NAME = 'DEFAULT'
-    _user_admin_orgs = None
+    orgs_mapping = None
 
     class Meta:
         verbose_name = _("Organization")
 
     def __str__(self):
-        return self.name
+        return str(self.name)
-
-    def set_to_cache(self):
-        if self.__class__.orgs is None:
-            self.__class__.orgs = {}
-        self.__class__.orgs[str(self.id)] = self
-
-    def expire_cache(self):
-        self.__class__.orgs.pop(str(self.id), None)
 
     @classmethod
-    def get_instance_from_cache(cls, oid):
+    def get_instance(cls, id_or_name, default=None):
-        if not cls.orgs or not isinstance(cls.orgs, dict):
+        assert default is None or isinstance(default, cls), (
-            return None
+            '`default` must be None or `Organization` instance'
-        return cls.orgs.get(str(oid))
+        )
-
+        org = cls.get_instance_from_memory(id_or_name)
-    @classmethod
+        org = org or default
-    def get_instance(cls, id_or_name, default=False):
-        cached = cls.get_instance_from_cache(id_or_name)
-        if cached:
-            return cached
-
-        if id_or_name is None:
-            return cls.default() if default else None
-        elif id_or_name in [cls.DEFAULT_ID, cls.DEFAULT_NAME, '']:
-            return cls.default()
-        elif id_or_name in [cls.ROOT_ID, cls.ROOT_NAME]:
-            return cls.root()
-
-        try:
-            if is_uuid(id_or_name):
-                org = cls.objects.get(id=id_or_name)
-            else:
-                org = cls.objects.get(name=id_or_name)
-            org.set_to_cache()
-        except cls.DoesNotExist as e:
-            if default:
-                return cls.default()
-            else:
-                raise e
         return org
 
+    @classmethod
+    def get_instance_from_memory(cls, id_or_name):
+        if not isinstance(cls.orgs_mapping, dict):
+            cls.orgs_mapping = cls.construct_orgs_mapping()
+        return cls.orgs_mapping.get(str(id_or_name))
+
+    @classmethod
+    def construct_orgs_mapping(cls):
+        orgs_mapping = {}
+        for org in cls.objects.all():
+            orgs_mapping[str(org.id)] = org
+            orgs_mapping[str(org.name)] = org
+        root_org = cls.root()
+        orgs_mapping.update({
+            root_org.id: root_org,
+            'GLOBAL': root_org,
+            '全局组织': root_org
+        })
+        return orgs_mapping
+
+    @classmethod
+    def expire_orgs_mapping(cls):
+        cls.orgs_mapping = None
+
     def get_org_members_by_role(self, role):
         from users.models import User
         if not self.is_root():
@@ -184,7 +174,7 @@ class Organization(models.Model):
 
     @classmethod
     def default(cls):
-        defaults = dict(name=cls.DEFAULT_NAME, id=cls.DEFAULT_ID)
+        defaults = dict(id=cls.DEFAULT_ID, name=cls.DEFAULT_NAME)
         obj, created = cls.objects.get_or_create(defaults=defaults, id=cls.DEFAULT_ID)
         return obj
 
@@ -411,7 +401,7 @@ class OrganizationMember(models.Model):
     id = models.UUIDField(default=uuid.uuid4, primary_key=True)
     org = models.ForeignKey(Organization, related_name='m2m_org_members', on_delete=models.CASCADE, verbose_name=_('Organization'))
     user = models.ForeignKey('users.User', related_name='m2m_org_members', on_delete=models.CASCADE, verbose_name=_('User'))
-    role = models.CharField(db_index=True, max_length=16, choices=ROLE.choices, default=ROLE.USER, verbose_name=_("Role"))
+    role = models.CharField(max_length=16, choices=ROLE.choices, default=ROLE.USER, verbose_name=_("Role"))
     date_created = models.DateTimeField(auto_now_add=True, verbose_name=_("Date created"))
     date_updated = models.DateTimeField(auto_now=True, verbose_name=_("Date updated"))
     created_by = models.CharField(max_length=128, null=True, verbose_name=_('Created by'))

apps/orgs/signals_handler/common.py

@@ -1,11 +1,13 @@
 # -*- coding: utf-8 -*-
 #
+import threading
 from collections import defaultdict
 from functools import partial
 
-from django.db.models.signals import m2m_changed
-from django.db.models.signals import post_save
 from django.dispatch import receiver
+from django.utils.functional import LazyObject
+from django.db.models.signals import m2m_changed
+from django.db.models.signals import post_save, post_delete
 
 from orgs.utils import tmp_to_org
 from orgs.models import Organization, OrganizationMember
@@ -13,10 +15,51 @@ from orgs.hands import set_current_org, Node, get_current_org
 from perms.models import (AssetPermission, ApplicationPermission)
 from users.models import UserGroup, User
 from common.const.signals import PRE_REMOVE, POST_REMOVE
+from common.signals import django_ready
+from common.utils import get_logger
+from common.utils.connection import RedisPubSub
+
+
+logger = get_logger(__file__)
+
+
+def get_orgs_mapping_for_memory_pub_sub():
+    return RedisPubSub('fm.orgs_mapping')
+
+
+class OrgsMappingForMemoryPubSub(LazyObject):
+    def _setup(self):
+        self._wrapped = get_orgs_mapping_for_memory_pub_sub()
+
+
+orgs_mapping_for_memory_pub_sub = OrgsMappingForMemoryPubSub()
+
+
+def expire_orgs_mapping_for_memory():
+    orgs_mapping_for_memory_pub_sub.publish('expire_orgs_mapping')
+
+
+@receiver(django_ready)
+def subscribe_orgs_mapping_expire(sender, **kwargs):
+    logger.debug("Start subscribe for expire orgs mapping from memory")
+
+    def keep_subscribe():
+        subscribe = orgs_mapping_for_memory_pub_sub.subscribe()
+        for message in subscribe.listen():
+            if message['type'] != 'message':
+                continue
+            Organization.expire_orgs_mapping()
+            logger.debug('Expire orgs mapping')
+
+    t = threading.Thread(target=keep_subscribe)
+    t.daemon = True
+    t.start()
 
 
 @receiver(post_save, sender=Organization)
 def on_org_create_or_update(sender, instance=None, created=False, **kwargs):
+    # 必须放到最开始, 因为下面调用Node.save方法时会获取当前组织的org_id(即instance.org_id), 如果不过期会找不到
+    expire_orgs_mapping_for_memory()
     if instance:
         old_org = get_current_org()
         set_current_org(instance)
@@ -26,8 +69,10 @@ def on_org_create_or_update(sender, instance=None, created=False, **kwargs):
             node_root.save()
         set_current_org(old_org)
 
-    if instance and not created:
+
-        instance.expire_cache()
+@receiver(post_delete, sender=Organization)
+def on_org_delete(sender, **kwargs):
+    expire_orgs_mapping_for_memory()
 
 
 def _remove_users(model, users, org):

apps/orgs/utils.py

@@ -30,7 +30,7 @@ def get_org_from_request(request):
         oid = Organization.DEFAULT_ID
     elif oid.lower() == "root":
         oid = Organization.ROOT_ID
-    org = Organization.get_instance(oid, True)
+    org = Organization.get_instance(oid, default=Organization.default())
     return org
 
 
@@ -54,9 +54,7 @@ def _find(attr):
 
 def get_current_org():
     org_id = get_current_org_id()
-    if org_id is None:
+    org = Organization.get_instance(org_id, default=Organization.root())
-        return Organization.root()
-    org = Organization.get_instance(org_id)
     return org
 
 
@@ -65,43 +63,6 @@ def get_current_org_id():
     return org_id
 
 
-def construct_org_mapper():
-    orgs = Organization.objects.all()
-    org_mapper = {str(org.id): org for org in orgs}
-    org_mapper.update({
-        Organization.ROOT_ID: Organization.root(),
-    })
-    return org_mapper
-
-
-def set_org_mapper(org_mapper):
-    setattr(thread_local, 'org_mapper', org_mapper)
-
-
-def get_org_mapper():
-    org_mapper = _find('org_mapper')
-    if org_mapper is None:
-        org_mapper = construct_org_mapper()
-        set_org_mapper(org_mapper)
-    return org_mapper
-
-
-def get_org_by_id(org_id):
-    org_id = str(org_id)
-    org_mapper = get_org_mapper()
-    org = org_mapper.get(org_id)
-    return org
-
-
-def get_org_name_by_id(org_id):
-    org = get_org_by_id(org_id)
-    if org:
-        org_name = org.name
-    else:
-        org_name = 'Not Found'
-    return org_name
-
-
 def get_current_org_id_for_serializer():
     org_id = get_current_org_id()
     if org_id == Organization.DEFAULT_ID:

apps/terminal/const.py

@@ -41,7 +41,7 @@ class TerminalTypeChoices(TextChoices):
     koko = 'koko', 'KoKo'
     guacamole = 'guacamole', 'Guacamole'
     omnidb = 'omnidb', 'OmniDB'
-    xrdp = 'xrdp', 'xrdp'
+    xrdp = 'xrdp', 'Xrdp'
 
     @classmethod
     def types(cls):

apps/terminal/migrations/0032_auto_20210302_1853.py

@@ -0,0 +1,18 @@
+# Generated by Django 3.1 on 2021-03-02 10:53
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('terminal', '0031_auto_20210113_1356'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='terminal',
+            name='type',
+            field=models.CharField(choices=[('koko', 'KoKo'), ('guacamole', 'Guacamole'), ('omnidb', 'OmniDB'), ('xrdp', 'Xrdp')], default='koko', max_length=64, verbose_name='type'),
+        ),
+    ]

apps/tickets/api/assignee.py

@@ -5,7 +5,7 @@ from rest_framework import viewsets
 from common.permissions import IsValidUser
 from common.exceptions import JMSException
 from users.models import User
-from orgs.utils import get_org_by_id
+from orgs.models import Organization
 from .. import serializers
 
 
@@ -17,7 +17,7 @@ class AssigneeViewSet(viewsets.ReadOnlyModelViewSet):
 
     def get_org(self):
         org_id = self.request.query_params.get('org_id')
-        org = get_org_by_id(org_id)
+        org = Organization.get_instance(org_id)
         if not org:
             error = ('The organization `{}` does not exist'.format(org_id))
             raise JMSException(error)

apps/tickets/serializers/ticket/ticket.py

@@ -3,8 +3,8 @@
 from django.utils.translation import ugettext_lazy as _
 from rest_framework import serializers
 from common.drf.serializers import MethodSerializer
-from orgs.utils import get_org_by_id
 from orgs.mixins.serializers import OrgResourceModelSerializerMixin
+from orgs.models import Organization
 from users.models import User
 from tickets.models import Ticket
 from .meta import type_serializer_classes_mapping
@@ -104,7 +104,7 @@ class TicketApplySerializer(TicketSerializer):
 
     @staticmethod
     def validate_org_id(org_id):
-        org = get_org_by_id(org_id)
+        org = Organization.get_instance(org_id)
         if not org:
             error = _('The organization `{}` does not exist'.format(org_id))
             raise serializers.ValidationError(error)
@@ -113,7 +113,7 @@ class TicketApplySerializer(TicketSerializer):
     def validate_assignees(self, assignees):
         org_id = self.initial_data.get('org_id')
         self.validate_org_id(org_id)
-        org = get_org_by_id(org_id)
+        org = Organization.get_instance(org_id)
         admins = User.get_super_and_org_admins(org)
         valid_assignees = list(set(assignees) & set(admins))
         if not valid_assignees: