完善了删除过期额度的方案，避免第二天继承上一天没使用完的额度

lss233 · Sep 8, 2023 · e6bfabc · e6bfabc
1 parent 145282e
commit e6bfabc
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/manager/ratelimit.py b/manager/ratelimit.py
@@ -62,7 +62,7 @@ def get_draw_usage(self, _type: str, _id: str) -> Document:
 
         # 初始化
         if usage is None:
-            usage = {'type': _type, 'id': _id, 'count': 0, 'time': current_time}
+            usage = {'type': _type, 'id': _id, 'count': 0, 'time': current_time, 'day': current_day}
             self.draw_usage_db.insert(usage)
 
         return usage
@@ -73,15 +73,18 @@ def get_usage(self, _type: str, _id: str) -> Document:
         q = Query()
         usage = self.usage_db.get(q.fragment({"type": _type, "id": _id}))
         current_time = time.localtime(time.time()).tm_hour
+        current_day = time.localtime(time.time()).tm_mday
 
         # 删除过期的记录
-        if usage is not None and usage['time'] != current_time:
+        time_diff_dondition = (usage is not None and usage['time'] != current_time)
+        day_diff_condition = (usage is not None and usage['time'] == current_time and usage['day'] != current_day)
+        if time_diff_dondition or day_diff_condition:
             self.usage_db.remove(doc_ids=[usage.doc_id])
             usage = None
-
+            
         # 初始化
         if usage is None:
-            usage = {'type': _type, 'id': _id, 'count': 0, 'time': current_time}
+            usage = {'type': _type, 'id': _id, 'count': 0, 'time': current_time, 'day': current_day}
             self.usage_db.insert(usage)
 
         return usage