Skip to content

Commit

Permalink
Refactor occupancy traffic handling and add out of usage check
Browse files Browse the repository at this point in the history
  • Loading branch information
@Ehco1996
Ehco1996 committed Dec 22, 2023
1 parent 595a6b4 commit 4af759b
Show file tree
Hide file tree
Showing 3 changed files with 184 additions and 54 deletions.
26 changes: 18 additions & 8 deletions apps/proxy/admin.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand Down Expand Up @@ -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)
Expand Down
108 changes: 108 additions & 0 deletions apps/proxy/migrations/0023_alter_userproxynodeoccupancy_index_together_and_more.py
View file
Original file line number Diff line number Diff line change
@@ -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",
),
]
104 changes: 58 additions & 46 deletions apps/proxy/models.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -875,85 +876,96 @@ 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
if occupancy_config.occupancy_user_limit <= 0:
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()
)

0 comments on commit 4af759b

Please sign in to comment.