From 594b81e463303a75cafa4b3af0c80b97b88fcfd7 Mon Sep 17 00:00:00 2001
From: Noam Meltzer <tsnoam@gmail.com>
Date: Sun, 13 Mar 2016 21:44:59 +0200
Subject: [PATCH] start_polling(): new argument - bootstrap_retries

refs #196
---
 telegram/ext/updater.py | 40 ++++++++++++++++++++++++++++++++++------
 1 file changed, 34 insertions(+), 6 deletions(-)

diff --git a/telegram/ext/updater.py b/telegram/ext/updater.py
index 0b1965ce9..f90eba37d 100644
--- a/telegram/ext/updater.py
+++ b/telegram/ext/updater.py
@@ -30,6 +30,7 @@ import subprocess
 from signal import signal, SIGINT, SIGTERM, SIGABRT
 from telegram import Bot, TelegramError, NullHandler
 from telegram.ext import dispatcher, Dispatcher, JobQueue, UpdateQueue
+from telegram.error import Unauthorized, InvalidToken
 from telegram.utils.webhookhandler import (WebhookServer, WebhookHandler)
 
 logging.getLogger(__name__).addHandler(NullHandler())
@@ -94,7 +95,7 @@ class Updater:
         """:type: list[Thread]"""
 
     def start_polling(self, poll_interval=0.0, timeout=10, network_delay=2,
-                      clean=False):
+                      clean=False, bootstrap_retries=0):
         """
         Starts polling updates from Telegram.
 
@@ -106,6 +107,11 @@ class Updater:
             clean (Optional[bool]): Whether to clean any pending updates on
                 Telegram servers before actually starting to poll. Default is
                 False.
+            bootstrap_retries (Optional[int[): Whether the bootstrapping phase
+                of the `Updater` will retry on failures on the Telegram server.
+                < 0 - retry indefinitely
+                  0 - no retries (default)
+                > 0 - retry up to X times
 
         Returns:
             Queue: The update queue that can be filled from the main thread
@@ -120,7 +126,8 @@ class Updater:
                 # Create & start threads
                 self._init_thread(self.dispatcher.start, "dispatcher")
                 self._init_thread(self._start_polling, "updater",
-                                  poll_interval, timeout, network_delay)
+                                  poll_interval, timeout, network_delay,
+                                  bootstrap_retries)
 
                 # Return the update queue so the main thread can insert updates
                 return self.update_queue
@@ -185,18 +192,18 @@ class Updater:
                 # Return the update queue so the main thread can insert updates
                 return self.update_queue
 
-    def _start_polling(self, poll_interval, timeout, network_delay):
+    def _start_polling(self, poll_interval, timeout, network_delay,
+                       bootstrap_retries):
         """
         Thread target of thread 'updater'. Runs in background, pulls
         updates from Telegram and inserts them in the update queue of the
         Dispatcher.
-        """
 
+        """
         cur_interval = poll_interval
         self.logger.debug('Updater thread started')
 
-        # Remove webhook
-        self.bot.setWebhook(webhook_url=None)
+        self._set_webhook(None, bootstrap_retries)
 
         while self.running:
             try:
@@ -228,6 +235,27 @@ class Updater:
 
             sleep(cur_interval)
 
+    def _set_webhook(self, webhook_url, max_retries):
+        retries = 0
+        while 1:
+            try:
+                # Remove webhook
+                self.bot.setWebhook(webhook_url=webhook_url)
+            except (Unauthorized, InvalidToken):
+                raise
+            except TelegramError:
+                msg = 'failed to set webhook; try={0} max_retries={1}'.format(
+                    retries, max_retries)
+                if max_retries < 0 or retries < max_retries:
+                    self.logger.info(msg)
+                    retries += 1
+                else:
+                    self.logger.exception(msg)
+                    raise
+            else:
+                break
+            sleep(1)
+
     @staticmethod
     def _increase_poll_interval(current_interval):
         # increase waiting times on subsequent errors up to 30secs