From 45999001be29c4be88fbb2e489baf74e3fd50ef0 Mon Sep 17 00:00:00 2001
From: painor <pi.oussama@gmail.com>
Date: Sun, 28 Oct 2018 10:55:58 +0100
Subject: [PATCH] Added retry_delay parameter for auto-reconnection (#1031)

---
 telethon/client/telegrambaseclient.py | 6 ++++++
 telethon/network/mtprotosender.py     | 6 +++++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/telethon/client/telegrambaseclient.py b/telethon/client/telegrambaseclient.py
index 22a5ac29..33ebfa6f 100644
--- a/telethon/client/telegrambaseclient.py
+++ b/telethon/client/telegrambaseclient.py
@@ -89,6 +89,9 @@ class TelegramBaseClient(abc.ABC):
             retries, but this is not recommended, since the program can
             get stuck in an infinite loop.
 
+        retry_delay (`int` | `float`, optional):
+            The delay in seconds to sleep between automatic reconnections.
+
         auto_reconnect (`bool`, optional):
             Whether reconnection should be retried `connection_retries`
             times automatically if Telegram disconnects us or not.
@@ -150,6 +153,7 @@ class TelegramBaseClient(abc.ABC):
                  timeout=10,
                  request_retries=5,
                  connection_retries=5,
+                 retry_delay=1,
                  auto_reconnect=True,
                  sequential_updates=False,
                  flood_sleep_threshold=60,
@@ -210,6 +214,7 @@ class TelegramBaseClient(abc.ABC):
 
         self._request_retries = request_retries or sys.maxsize
         self._connection_retries = connection_retries or sys.maxsize
+        self._retry_delay = retry_delay or 0
         self._proxy = proxy
         self._timeout = timeout
         self._auto_reconnect = auto_reconnect
@@ -237,6 +242,7 @@ class TelegramBaseClient(abc.ABC):
         self._sender = MTProtoSender(
             self.session.auth_key, self._loop,
             retries=self._connection_retries,
+            delay=self._retry_delay,
             auto_reconnect=self._auto_reconnect,
             connect_timeout=self._timeout,
             update_callback=self._handle_update,
diff --git a/telethon/network/mtprotosender.py b/telethon/network/mtprotosender.py
index 8c4eeb4b..b626abe6 100644
--- a/telethon/network/mtprotosender.py
+++ b/telethon/network/mtprotosender.py
@@ -60,11 +60,12 @@ class MTProtoSender:
     key exists yet.
     """
     def __init__(self, auth_key, loop, *,
-                 retries=5, auto_reconnect=True, connect_timeout=None,
+                 retries=5, delay=1, auto_reconnect=True, connect_timeout=None,
                  update_callback=None, auto_reconnect_callback=None):
         self._connection = None
         self._loop = loop
         self._retries = retries
+        self._delay = delay
         self._auto_reconnect = auto_reconnect
         self._connect_timeout = connect_timeout
         self._update_callback = update_callback
@@ -215,6 +216,7 @@ class MTProtoSender:
             except (ConnectionError, asyncio.TimeoutError) as e:
                 __log__.warning('Attempt {} at connecting failed: {}: {}'
                                 .format(retry, type(e).__name__, e))
+                await asyncio.sleep(self._delay)
             else:
                 break
         else:
@@ -234,6 +236,7 @@ class MTProtoSender:
                 except (SecurityError, AssertionError) as e:
                     __log__.warning('Attempt {} at new auth_key failed: {}'
                                     .format(retry, e))
+                    await asyncio.sleep(self._delay)
             else:
                 e = ConnectionError('auth_key generation failed {} times'
                                     .format(self._retries))
@@ -314,6 +317,7 @@ class MTProtoSender:
                 await self._connect()
             except ConnectionError:
                 __log__.info('Failed reconnection retry %d/%d', retry, retries)
+                await asyncio.sleep(self._delay)
             else:
                 self._send_queue.extend(self._pending_state.values())
                 self._pending_state.clear()