fix: improve strategy

Author: exiaohu

volcenginesdkarkruntime/_utils/_model_breaker.py

@@ -1,56 +1,78 @@
 import asyncio
-import math
 import time
+import uuid
+
+
+class _QuerySet(object):
+    def __init__(self):
+        self._items = list()
+        self._index = dict()
+
+    def add(self, item: int) -> None:
+        if item in self._index:
+            return
+
+        self._items.append(item)
+        self._index[item] = len(self._items) - 1
+
+    def remove(self, item: int) -> None:
+        if item not in self._index:
+            return
+
+        index = self._index[item]
+        self._items[index] = self._items[-1]
+        self._index[self._items[-1]] = index
+        self._items.pop()
+        del self._index[item]
+
+    def query(self, item: int) -> int:
+        return self._index[item]
 
 
 class ModelBreaker(object):
     def __init__(self):
         # 初始化 allow_time 为当前时间
         self._allow_time = time.perf_counter()
+        self._waiters = _QuerySet()
 
-        self._waiters = 0
-        self._permits = 0
-        self._base = 0
-
-    def _allow(self) -> bool:
-        # 检查当前时间是否在 allow_time 之后
-        return time.perf_counter() > self._allow_time
+    def _allow(self, id: int) -> bool:
+        cur = time.perf_counter()
+        # 如果当前时间小于等于 allow_time，不允许通过
+        if cur <= self._allow_time:
+            return 0
+        # 如果当前时间与 allow_time 的差值大于 10，允许通过
+        if cur - self._allow_time > 10:
+            return True
+        # 如果当前时间与 allow_time 的差值小于等于 10，慢启动通过
+        return self._waiters.query(id) < 2 ** (cur - self._allow_time)
 
     def _get_allowed_duration(self) -> float:
         # 计算当前时间与 allow_time 之间的持续时间
         allow_duration = self._allow_time - time.perf_counter()
 
-        # 如果持续时间为负，返回零
-        if allow_duration < 0:
-            return 0
-        return allow_duration
+        # 至少有 1 秒的等待时间
+        return max(allow_duration, 1)
 
     def _acquire(self) -> int:
-        self._waiters += 1
-        return self._waiters
-
-    def _release(self) -> None:
-        self._waiters -= 1
-        self._permits += 1
+        id = uuid.uuid4().int
+        self._waiters.add(id)
+        return id
 
-    def _jitter(self, i: int) -> float:
-        if i <= self._base:
-            return 0
-        return math.log2(i - self._base)
+    def _release(self, id: int) -> None:
+        self._waiters.remove(id)
 
     def reset(self, duration: float) -> None:
         # 将 allow_time 重置为当前时间加上指定的持续时间
         self._allow_time = time.perf_counter() + duration
-        self._base = self._permits
 
     def wait(self) -> None:
-        i = self._acquire()
-        while not self._allow():
-            time.sleep(self._get_allowed_duration() + self._jitter(i))
-        self._release()
+        id = self._acquire()
+        while not self._allow(id):
+            time.sleep(self._get_allowed_duration())
+        self._release(id)
 
     async def asyncwait(self) -> None:
-        i = self._acquire()
-        while not self._allow():
-            await asyncio.sleep(self._get_allowed_duration() + self._jitter(i))
-        self._release()
+        id = self._acquire()
+        while not self._allow(id):
+            await asyncio.sleep(self._get_allowed_duration())
+        self._release(id)