From 48dead76aca3fa9afc7735de7af7331f1fc0cb42 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Mon, 11 Sep 2017 10:52:36 +0200
Subject: [PATCH] Retry up to five times before giving up on a request

---
 telethon/telegram_bare_client.py | 19 +++++++++++++------
 telethon/telegram_client.py      |  7 ++++---
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/telethon/telegram_bare_client.py b/telethon/telegram_bare_client.py
index 1bfa0ade..e7e769d2 100644
--- a/telethon/telegram_bare_client.py
+++ b/telethon/telegram_bare_client.py
@@ -282,7 +282,7 @@ class TelegramBareClient:
 
     # region Invoking Telegram requests
 
-    def invoke(self, request, updates=None, call_receive=True):
+    def invoke(self, request, call_receive=True, retries=5):
         """Invokes (sends) a MTProtoRequest and returns (receives) its result.
 
            If 'updates' is not None, all read update object will be put
@@ -298,6 +298,9 @@ class TelegramBareClient:
         if not self._sender:
             raise ValueError('You must be connected to invoke requests!')
 
+        if retries <= 0:
+            raise ValueError('Number of retries reached 0.')
+
         try:
             # Ensure that we start with no previous errors (i.e. resending)
             request.confirm_received.clear()
@@ -314,20 +317,24 @@ class TelegramBareClient:
                 while not request.confirm_received.is_set():
                     self._sender.receive(update_state=self.updates)
 
-            if request.rpc_error:
-                raise request.rpc_error
-            return request.result
-
         except ConnectionResetError:
             self._logger.debug('Server disconnected us. Reconnecting and '
                                'resending request...')
             self.reconnect()
-            return self.invoke(request)
 
         except FloodWaitError:
             self.disconnect()
             raise
 
+        if request.rpc_error:
+            raise request.rpc_error
+        if request.result is None:
+            return self.invoke(
+                request, call_receive=call_receive, retries=(retries - 1)
+            )
+        else:
+            return request.result
+
     # Let people use client(SomeRequest()) instead client.invoke(...)
     __call__ = invoke
 
diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py
index 1bdd6d04..8dce22f5 100644
--- a/telethon/telegram_client.py
+++ b/telethon/telegram_client.py
@@ -218,8 +218,9 @@ class TelegramClient(TelegramBareClient):
 
     # region Telegram requests functions
 
-    def invoke(self, request, *args):
+    def invoke(self, request, *args, **kwargs):
         """Invokes (sends) a MTProtoRequest and returns (receives) its result.
+           An optional 'retries' parameter can be set.
 
            *args will be ignored.
         """
@@ -233,9 +234,9 @@ class TelegramClient(TelegramBareClient):
             # will be the one which should be reading (but is invoking the
             # request) thus not being available to read it "in the background"
             # and it's needed to call receive.
-            # TODO Retry if 'result' is None?
             return super().invoke(
-                request, call_receive=self._recv_thread is None
+                request, call_receive=self._recv_thread is None,
+                retries=kwargs.get('retries', 5)
             )
 
         except (PhoneMigrateError, NetworkMigrateError, UserMigrateError) as e: