Update ProxyNode and UserProxyNodeOccupancyRecord models

apps/proxy/migrations/0022_proxynode_cost_price_relaynode_cost_price_and_more.py

@@ -1,8 +1,8 @@
-# Generated by Django 4.2.6 on 2023-12-20 07:30
+# Generated by Django 4.2.6 on 2023-12-20 08:17
 
+import django.db.models.deletion
 from django.conf import settings
 from django.db import migrations, models
-import django.db.models.deletion
 
 
 class Migration(migrations.Migration):
@@ -82,14 +82,15 @@ class Migration(migrations.Migration):
                     "start_time",
                     models.DateTimeField(auto_now_add=True, verbose_name="开始占用时间"),
                 ),
-                (
-                    "end_time",
-                    models.DateTimeField(db_index=True, verbose_name="结束占用时间"),
-                ),
+                ("end_time", models.DateTimeField(verbose_name="结束占用时间")),
                 (
                     "traffic_used",
                     models.BigIntegerField(default=0, verbose_name="已用流量"),
                 ),
+                (
+                    "out_of_traffic",
+                    models.BooleanField(default=False, verbose_name="是否超出流量"),
+                ),
                 (
                     "occupancy_config_snapshot",
                     models.JSONField(default=dict, verbose_name="快照"),
@@ -114,7 +115,11 @@ class Migration(migrations.Migration):
             options={
                 "verbose_name": "用户代理节点占用记录",
                 "verbose_name_plural": "用户代理节点占用记录",
-                "index_together": {("proxy_node", "end_time")},
+                "index_together": {
+                    ("out_of_traffic", "user", "end_time"),
+                    ("out_of_traffic", "end_time"),
+                    ("out_of_traffic", "proxy_node", "end_time"),
+                },
             },
         ),
     ]

apps/proxy/models.py

@@ -259,15 +259,12 @@ def get_user_active_nodes(cls, user):
         base_query = cls.get_active_nodes()
         query = base_query.filter(level__gte=user.level)
         # 2. filter out nodes that has been occupied by other users
-        occupied_node_ids = UserProxyNodeOccupancyRecord.objects.filter(
-            end_time__gte=utils.get_current_datetime(),
-        ).values("proxy_node_id")
+        occupied_node_ids = UserProxyNodeOccupancyRecord.get_occupied_node_ids()
         query = query.exclude(id__in=occupied_node_ids)
         # 3. add nodes that has been occupied by this user
-        user_occupied_node_ids = UserProxyNodeOccupancyRecord.objects.filter(
-            end_time__gte=utils.get_current_datetime(),
-            user=user,
-        ).values("proxy_node_id")
+        user_occupied_node_ids = (
+            UserProxyNodeOccupancyRecord.get_user_occupied_node_ids(user)
+        )
         query = query | base_query.filter(id__in=user_occupied_node_ids)
         return query
 
@@ -286,29 +283,29 @@ def get_node_users(self):
         if not self.enable:
             return []
         # 2. node occupied by users, return users
-        records = UserProxyNodeOccupancyRecord.objects.filter(
-            proxy_node=self,
-            end_time__gte=utils.get_current_datetime(),
-        )
-        if records.exists():
-            user_ids = records.values("user_id")
+        occupancies_query = UserProxyNodeOccupancyRecord.get_node_occupancies(self)
+        if occupancies_query.count() > 0:
+            user_ids = occupancies_query.values("user_id")
             return User.objects.filter(id__in=user_ids)
         # 3. shared node filter user that level >= node.level
         return User.objects.filter(level__gte=self.level)
 
     def get_proxy_configs(self):
-        data = {}
         if self.node_type == self.NODE_TYPE_SS:
-            data = self.ss_config.to_node_config(self)
+            proxy_cfg = self.ss_config
         elif self.node_type == self.NODE_TYPE_TROJAN:
-            data = self.trojan_config.to_node_config()
+            proxy_cfg = self.trojan_config
+        else:
+            raise Exception("not support node type")
+
+        configs = proxy_cfg.to_node_config(self)
         if not self.enable:
-            data["users"] = []
+            configs["users"] = []
         else:
-            data["users"] = [
-                user.to_user_config(self, user) for user in User.get_active_users()
+            configs["users"] = [
+                proxy_cfg.to_user_config(self, user) for user in self.get_node_users()
             ]
-        return data
+        return configs
 
     def get_ehco_server_config(self):
         if self.enable_ehco_tunnel:
@@ -862,24 +859,24 @@ def get_by_proxy_node(cls, node: ProxyNode):
 
 
 class UserProxyNodeOccupancyRecord(BaseModel):
-    user = models.ForeignKey(
-        User, on_delete=models.CASCADE, verbose_name="用户", db_index=True
-    )
+    user = models.ForeignKey(User, on_delete=models.CASCADE, verbose_name="用户")
     proxy_node = models.ForeignKey(
-        ProxyNode, on_delete=models.CASCADE, verbose_name="代理节点", db_index=True
+        ProxyNode, on_delete=models.CASCADE, verbose_name="代理节点"
     )
     start_time = models.DateTimeField(auto_now_add=True, verbose_name="开始占用时间")
-    end_time = models.DateTimeField(
-        null=False, blank=False, verbose_name="结束占用时间", db_index=True
-    )
+    end_time = models.DateTimeField(null=False, blank=False, verbose_name="结束占用时间")
     traffic_used = models.BigIntegerField(default=0, verbose_name="已用流量")
+    out_of_traffic = models.BooleanField(default=False, verbose_name="是否超出流量")
     occupancy_config_snapshot = models.JSONField(verbose_name="快照", default=dict)
 
     class Meta:
         verbose_name = "用户代理节点占用记录"
         verbose_name_plural = "用户代理节点占用记录"
-        index_together = ["user", "end_time"]
-        index_together = ["proxy_node", "end_time"]
+        index_together = (
+            ["out_of_traffic", "end_time"],
+            ["out_of_traffic", "user", "end_time"],
+            ["out_of_traffic", "proxy_node", "end_time"],
+        )
 
     @classmethod
     @transaction.atomic
@@ -906,3 +903,44 @@ def create_by_occupancy_config(
             traffic_used=occupancy_config.occupancy_traffic,
             occupancy_config_snapshot=occupancy_config.to_dict(),
         )
+
+    @classmethod
+    def get_node_occupancy_user_ids(cls, node: ProxyNode):
+        return cls.objects.filter(
+            out_of_traffic=False,
+            proxy_node=node,
+            end_time__gte=utils.get_current_datetime(),
+        ).values("user_id")
+
+    @classmethod
+    def get_occupied_node_ids(cls):
+        occupied_node_ids = cls.objects.filter(
+            out_of_traffic=False,
+            end_time__gte=utils.get_current_datetime(),
+        ).values("proxy_node_id")
+        return occupied_node_ids
+
+    @classmethod
+    def get_user_occupied_node_ids(cls, user: User):
+        user_occupied_node_ids = cls.objects.filter(
+            out_of_traffic=False,
+            user=user,
+            end_time__gte=utils.get_current_datetime(),
+        ).values("proxy_node_id")
+        return user_occupied_node_ids
+
+    @classmethod
+    def get_node_occupies(cls, node: ProxyNode):
+        return UserProxyNodeOccupancyRecord.objects.filter(
+            out_of_traffic=False,
+            proxy_node=node,
+            end_time__gte=utils.get_current_datetime(),
+        )
+
+    @classmethod
+    def check_and_incr_traffic(cls, user_id, proxy_node_id, traffic):
+        r = cls.objects.get(user__id=user_id, proxy_node__id=proxy_node_id)
+        r.traffic_used += traffic
+        if r.traffic_used > r.occupancy_config_snapshot["occupancy_traffic"]:
+            r.out_of_traffic = True
+        r.save()

apps/sspanel/tasks.py

@@ -4,7 +4,7 @@
 from django.core.mail import send_mail
 
 from apps import celery_app
-from apps.proxy.models import ProxyNode, UserTrafficLog
+from apps.proxy.models import ProxyNode, UserProxyNodeOccupancyRecord, UserTrafficLog
 from apps.sspanel import models as m
 from apps.utils import get_current_datetime
 
@@ -28,6 +28,9 @@ def sync_user_traffic_task(node_id, data):
     node: ProxyNode = ProxyNode.get_or_none(node_id)
     if not node:
         return
+    node_occurred_user_ids = UserProxyNodeOccupancyRecord.get_node_occupancy_user_ids(
+        node
+    )
     node_total_traffic = 0
     log_time = get_current_datetime()
     user_model_list = []
@@ -66,6 +69,12 @@ def sync_user_traffic_task(node_id, data):
         # 节点流量增量
         node_total_traffic += u + d
 
+        # 记录用户占用节点流量
+        if user_id in node_occurred_user_ids:
+            UserProxyNodeOccupancyRecord.check_and_incr_traffic(
+                user_id=user_id, node_id=node_id, traffic=d + u
+            )
+
     if not traffic_data:
         # NOTE add blank log to show node is online
         trafficlog_model_list.append(UserTrafficLog(proxy_node=node))