From 639828f5508dff9e3bd744bf0c5ea43a04a6e8f0 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Wed, 18 Apr 2018 15:17:46 +0200
Subject: [PATCH] Retry on internal server errors, up to MAX_RETRIES times Also
 add support for custom retry count

---
 pyrogram/session/session.py | 30 ++++++++++++++++--------------
 1 file changed, 16 insertions(+), 14 deletions(-)

diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
index 5be2eaec..8ae70ccc 100644
--- a/pyrogram/session/session.py
+++ b/pyrogram/session/session.py
@@ -32,7 +32,7 @@ from pyrogram import __copyright__, __license__, __version__
 from pyrogram.api import functions, types, core
 from pyrogram.api.all import layer
 from pyrogram.api.core import Message, Object, MsgContainer, Long, FutureSalt, Int
-from pyrogram.api.errors import Error
+from pyrogram.api.errors import Error, InternalServerError
 from pyrogram.connection import Connection
 from pyrogram.crypto import AES, KDF
 from .internals import MsgId, MsgFactory, DataCenter
@@ -399,17 +399,19 @@ class Session:
             else:
                 return result
 
-    def send(self, data: Object):
-        for i in range(self.MAX_RETRIES):
-            self.is_connected.wait(self.WAIT_TIMEOUT)
+    def send(self, data: Object, retries: int = MAX_RETRIES):
+        self.is_connected.wait(self.WAIT_TIMEOUT)
 
-            try:
-                return self._send(data)
-            except (OSError, TimeoutError):
-                (log.warning if i > 2 else log.info)(
-                    "{}: {} Retrying {}".format(i, datetime.now(), type(data))
-                )
-                time.sleep(1)
-                continue
-        else:
-            return None
+        try:
+            return self._send(data)
+        except (OSError, TimeoutError, InternalServerError) as e:
+            if retries == 0:
+                raise e from None
+
+            (log.warning if retries < 3 else log.info)(
+                "{}: {} Retrying {}".format(
+                    Session.MAX_RETRIES - retries,
+                    datetime.now(), type(data)))
+
+            time.sleep(0.5)
+            self.send(data, retries - 1)