From 4af759bcbc8e26055c42ebb34ccffc73611a0a26 Mon Sep 17 00:00:00 2001 From: Ehco1996 Date: Fri, 22 Dec 2023 09:37:19 +0800 Subject: [PATCH] Refactor occupancy traffic handling and add out of usage check --- apps/proxy/admin.py | 26 +++-- ...xynodeoccupancy_index_together_and_more.py | 108 ++++++++++++++++++ apps/proxy/models.py | 104 +++++++++-------- 3 files changed, 184 insertions(+), 54 deletions(-) create mode 100644 apps/proxy/migrations/0023_alter_userproxynodeoccupancy_index_together_and_more.py diff --git a/apps/proxy/admin.py b/apps/proxy/admin.py index 154bead621..0b8eafd8d0 100644 --- a/apps/proxy/admin.py +++ b/apps/proxy/admin.py @@ -34,12 +34,14 @@ class OccupancyConfigInline(admin.StackedInline): def get_formset(self, request, obj=None, **kwargs): if obj: - traffic = traffic_format(obj.occupancy_config.occupancy_traffic) - help_texts = { - "occupancy_traffic": f"={traffic}", - } - print("obj", obj, kwargs) - kwargs.update({"help_texts": help_texts}) + try: + traffic = traffic_format(obj.occupancy_config.occupancy_traffic) + help_texts = { + "occupancy_traffic": f"={traffic}", + } + kwargs.update({"help_texts": help_texts}) + except models.OccupancyConfig.DoesNotExist: + pass return super().get_formset(request, obj, **kwargs) @@ -218,14 +220,22 @@ class UserProxyNodeOccupancyAdmin(admin.ModelAdmin): "user", "start_time", "end_time", - "traffic_used", - "out_of_traffic", + "traffic_info", + "out_of_usage", ] search_fields = ["user__username"] list_filter = ["proxy_node", "user"] list_per_page = 10 show_full_result_count = False + @admin.display(description="已用/总流量") + def traffic_info(self, instance): + return f"{traffic_format(instance.used_traffic)}/{traffic_format(instance.total_traffic)}" + + @admin.display(description="是否超出") + def out_of_usage(self, instance): + return instance.out_of_usage() + # Register your models here. admin.site.register(models.ProxyNode, ProxyNodeAdmin) diff --git a/apps/proxy/migrations/0023_alter_userproxynodeoccupancy_index_together_and_more.py b/apps/proxy/migrations/0023_alter_userproxynodeoccupancy_index_together_and_more.py new file mode 100644 index 0000000000..70c5403d2b --- /dev/null +++ b/apps/proxy/migrations/0023_alter_userproxynodeoccupancy_index_together_and_more.py @@ -0,0 +1,108 @@ +# Generated by Django 4.2.6 on 2023-12-22 01:05 + +import django.db.models.deletion +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("proxy", "0022_proxynode_cost_price_relaynode_cost_price_and_more"), + ] + + operations = [ + migrations.AlterIndexTogether( + name="userproxynodeoccupancy", + index_together=set(), + ), + migrations.AddField( + model_name="userproxynodeoccupancy", + name="total_traffic", + field=models.BigIntegerField(default=1073741824, verbose_name="总流量(单位字节)"), + ), + migrations.AddField( + model_name="userproxynodeoccupancy", + name="used_traffic", + field=models.BigIntegerField(default=0, verbose_name="已用流量(单位字节)"), + ), + migrations.AlterField( + model_name="occupancyconfig", + name="occupancy_price", + field=models.DecimalField( + decimal_places=2, max_digits=10, verbose_name="价格" + ), + ), + migrations.AlterField( + model_name="occupancyconfig", + name="occupancy_traffic", + field=models.BigIntegerField(default=0, verbose_name="流量(单位字节)"), + ), + migrations.AlterField( + model_name="occupancyconfig", + name="occupancy_user_limit", + field=models.PositiveIntegerField(default=0, verbose_name="用户数"), + ), + migrations.AlterField( + model_name="occupancyconfig", + name="proxy_node", + field=models.OneToOneField( + on_delete=django.db.models.deletion.CASCADE, + related_name="occupancy_config", + to="proxy.proxynode", + verbose_name="代理节点", + ), + ), + migrations.AlterField( + model_name="proxynode", + name="cost_price", + field=models.DecimalField( + decimal_places=2, default=0, max_digits=10, verbose_name="成本" + ), + ), + migrations.AlterField( + model_name="proxynode", + name="total_traffic", + field=models.BigIntegerField(default=1073741824, verbose_name="总流量(单位字节)"), + ), + migrations.AlterField( + model_name="proxynode", + name="used_traffic", + field=models.BigIntegerField(default=0, verbose_name="已用流量(单位字节)"), + ), + migrations.AlterField( + model_name="relaynode", + name="cost_price", + field=models.DecimalField( + decimal_places=2, default=0, max_digits=10, verbose_name="成本" + ), + ), + migrations.AddIndex( + model_name="userproxynodeoccupancy", + index=models.Index( + fields=["end_time"], name="proxy_userp_end_tim_ba27c4_idx" + ), + ), + migrations.AddIndex( + model_name="userproxynodeoccupancy", + index=models.Index( + fields=["user", "end_time"], name="proxy_userp_user_id_a195e2_idx" + ), + ), + migrations.AddIndex( + model_name="userproxynodeoccupancy", + index=models.Index( + fields=["proxy_node", "end_time"], name="proxy_userp_proxy_n_ffb646_idx" + ), + ), + migrations.RemoveField( + model_name="userproxynodeoccupancy", + name="occupancy_config_snapshot", + ), + migrations.RemoveField( + model_name="userproxynodeoccupancy", + name="out_of_traffic", + ), + migrations.RemoveField( + model_name="userproxynodeoccupancy", + name="traffic_used", + ), + ] diff --git a/apps/proxy/models.py b/apps/proxy/models.py index bf6ce91160..ada8f76558 100644 --- a/apps/proxy/models.py +++ b/apps/proxy/models.py @@ -10,6 +10,7 @@ import pendulum from django.conf import settings from django.db import models, transaction +from django.db.models import F from apps import constants as c from apps import utils @@ -875,25 +876,30 @@ class UserProxyNodeOccupancy(BaseModel): ) start_time = models.DateTimeField(auto_now_add=True, verbose_name="开始占用时间") 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) + used_traffic = models.BigIntegerField("已用流量(单位字节)", default=0) + total_traffic = models.BigIntegerField("总流量(单位字节)", default=settings.GB) class Meta: verbose_name = "占用记录" verbose_name_plural = "占用记录" - index_together = ( - ["out_of_traffic", "end_time"], - ["out_of_traffic", "user", "end_time"], - ["out_of_traffic", "proxy_node", "end_time"], - ) + indexes = [ + models.Index(fields=["end_time"]), + models.Index(fields=["user", "end_time"]), + models.Index(fields=["proxy_node", "end_time"]), + ] def __str__(self) -> str: return f"用户占用配置:{self.id}" + @classmethod + def _valid_occupancy_query(cls): + return cls.objects.filter(end_time__gt=utils.get_current_datetime()).filter( + used_traffic__lt=F("total_traffic") + ) + @classmethod @transaction.atomic - def create_by_occupancy_config( + def create_occupancy( cls, user: User, proxy_node: ProxyNode, occupancy_config: OccupancyConfig ): # check user limit first @@ -901,59 +907,65 @@ def create_by_occupancy_config( raise Exception("not allow to create occupancy record with user limit 0") if occupancy_config.occupancy_user_limit > 0: if ( - cls.objects.filter( - proxy_node=proxy_node, - end_time__gte=utils.get_current_datetime(), - ).count() + cls.get_node_occupancies(proxy_node).count() >= occupancy_config.occupancy_user_limit ): raise Exception("occupancy user limit exceed") - return cls.objects.create( - user=user, - proxy_node=proxy_node, - start_time=utils.get_current_datetime(), - end_time=utils.get_current_datetime().add(days=30), - traffic_used=occupancy_config.occupancy_traffic, - occupancy_config_snapshot=occupancy_config.to_snapshot(), - ) - @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") + # check user balance + if user.balance < occupancy_config.occupancy_price: + raise Exception("user balance not enough") + + # check if user already occupied this node + o = cls.objects.filter(user=user, proxy_node=proxy_node).first() + if o: + if o.out_of_usage(): + # reset traffic and time when out of usage + o.end_time = utils.get_current_datetime().add(days=30) + o.start_time = utils.get_current_datetime() + o.used_traffic = 0 + o.total_traffic = occupancy_config.occupancy_traffic + o.save() + else: + # incr traffic and time + o.end_time = o.end_time.add(days=30) + o.total_traffic += occupancy_config.occupancy_traffic + o.save() + else: + return cls.objects.create( + user=user, + proxy_node=proxy_node, + start_time=utils.get_current_datetime(), + end_time=utils.get_current_datetime().add(days=30), + total_traffic=occupancy_config.occupancy_traffic, + ) @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") + occupied_node_ids = cls._valid_occupancy_query().values("proxy_node_id") return occupied_node_ids + @classmethod + def get_node_occupancy_user_ids(cls, node: ProxyNode): + return cls._valid_occupancy_query().filter(proxy_node=node).values("user_id") + + @classmethod @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 + return cls._valid_occupancy_query().filter(user=user).values("proxy_node_id") @classmethod def get_node_occupancies(cls, node: ProxyNode): - return UserProxyNodeOccupancy.objects.filter( - out_of_traffic=False, - proxy_node=node, - end_time__gte=utils.get_current_datetime(), - ) + return cls._valid_occupancy_query().filter(proxy_node=node) @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.used_traffic += traffic r.save() + + def out_of_usage(self): + return ( + self.used_traffic >= self.total_traffic + or self.end_time < utils.get_current_datetime() + )