From 4351881f2e25374c2a3cfc6cf0c42d1fd23672cd Mon Sep 17 00:00:00 2001
From: John Smith <xfgryujk@126.com>
Date: Fri, 15 Sep 2023 22:14:25 +0800
Subject: [PATCH] =?UTF-8?q?=E5=8F=AF=E9=85=8D=E7=BD=AE=E9=87=8D=E8=BF=9E?=
 =?UTF-8?q?=E9=97=B4=E9=9A=94=E6=97=B6=E9=97=B4?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 blivedm/clients/ws_base.py | 15 ++++++++++++++-
 blivedm/utils.py           | 15 +++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/blivedm/clients/ws_base.py b/blivedm/clients/ws_base.py
index 5795afe..bae3464 100644
--- a/blivedm/clients/ws_base.py
+++ b/blivedm/clients/ws_base.py
@@ -76,6 +76,9 @@ class AuthError(Exception):
     """认证失败"""
 
 
+DEFAULT_RECONNECT_POLICY = utils.make_constant_retry_policy(1)
+
+
 class WebSocketClientBase:
     """
     基于WebSocket的客户端
@@ -102,6 +105,8 @@ class WebSocketClientBase:
         self._need_init_room = True
         self._handler: Optional[handlers.HandlerInterface] = None
         """消息处理器"""
+        self._get_reconnect_interval: Callable[[int], float] = DEFAULT_RECONNECT_POLICY
+        """重连间隔时间增长策略"""
 
         # 在调用init_room后初始化的字段
         self._room_id: Optional[int] = None
@@ -139,6 +144,14 @@ class WebSocketClientBase:
         """
         self._handler = handler
 
+    def set_reconnect_policy(self, get_reconnect_interval: Callable[[int], float]):
+        """
+        设置重连间隔时间增长策略
+
+        :param get_reconnect_interval: 一个可调用对象，输入重试次数，返回间隔时间
+        """
+        self._get_reconnect_interval = get_reconnect_interval
+
     def start(self):
         """
         启动本客户端
@@ -280,7 +293,7 @@ class WebSocketClientBase:
             # 准备重连
             retry_count += 1
             logger.warning('room=%d is reconnecting, retry_count=%d', self.room_id, retry_count)
-            await asyncio.sleep(1)
+            await asyncio.sleep(self._get_reconnect_interval(retry_count))
 
     async def _on_before_ws_connect(self, retry_count):
         """
diff --git a/blivedm/utils.py b/blivedm/utils.py
index 088aedb..3ee78d1 100644
--- a/blivedm/utils.py
+++ b/blivedm/utils.py
@@ -2,3 +2,18 @@
 USER_AGENT = (
     'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/102.0.0.0 Safari/537.36'
 )
+
+
+def make_constant_retry_policy(interval: float):
+    def get_interval(_retry_count: int):
+        return interval
+    return get_interval
+
+
+def make_linear_retry_policy(start_interval: float, interval_step: float, max_interval: float):
+    def get_interval(retry_count: int):
+        return min(
+            start_interval + (retry_count - 1) * interval_step,
+            max_interval
+        )
+    return get_interval