Merge pull request #5924 from oyvsyo/gateway_connection_retry

Add reconnection to Gateway (form nb2kg)

Author: kevin-bates

notebook/gateway/handlers.py

@@ -4,6 +4,7 @@
 import os
 import logging
 import mimetypes
+import random
 
 from ..base.handlers import APIHandler, IPythonHandler
 from ..utils import url_path_join
@@ -134,6 +135,7 @@ def __init__(self, **kwargs):
         self.ws = None
         self.ws_future = Future()
         self.disconnected = False
+        self.retry = 0
 
     @gen.coroutine
     def _connect(self, kernel_id):
@@ -155,6 +157,7 @@ def _connect(self, kernel_id):
     def _connection_done(self, fut):
         if not self.disconnected and fut.exception() is None:  # prevent concurrent.futures._base.CancelledError
             self.ws = fut.result()
+            self.retry = 0
             self.log.debug("Connection is ready: ws: {}".format(self.ws))
         else:
             self.log.warning("Websocket connection has been closed via client disconnect or due to error.  "
@@ -189,8 +192,15 @@ def _read_messages(self, callback):
             else:  # ws cancelled - stop reading
                 break
 
-        if not self.disconnected: # if websocket is not disconnected by client, attept to reconnect to Gateway
-            self.log.info("Attempting to re-establish the connection to Gateway: {}".format(self.kernel_id))
+        # NOTE(esevan): if websocket is not disconnected by client, try to reconnect.
+        if not self.disconnected and self.retry < GatewayClient.instance().gateway_retry_max:
+            jitter = random.randint(10, 100) * 0.01
+            retry_interval = min(GatewayClient.instance().gateway_retry_interval * (2 ** self.retry),
+                                 GatewayClient.instance().gateway_retry_interval_max) + jitter
+            self.retry += 1
+            self.log.info("Attempting to re-establish the connection to Gateway in %s secs (%s/%s): %s",
+                          retry_interval, self.retry, GatewayClient.instance().gateway_retry_max, self.kernel_id)
+            yield gen.sleep(retry_interval)
             self._connect(self.kernel_id)
             loop = IOLoop.current()
             loop.add_future(self.ws_future, lambda future: self._read_messages(callback))

notebook/gateway/managers.py

@@ -15,14 +15,14 @@
 from jupyter_client.kernelspec import KernelSpecManager
 from ..utils import url_path_join
 
-from traitlets import Instance, Unicode, Float, Bool, default, validate, TraitError
+from traitlets import Instance, Unicode, Int, Float, Bool, default, validate, TraitError
 from traitlets.config import SingletonConfigurable
 
 
 class GatewayClient(SingletonConfigurable):
     """This class manages the configuration.  It's its own singleton class so that we
        can share these values across all objects.  It also contains some helper methods
-        to build request arguments out of the various config options.
+       to build request arguments out of the various config options.
 
     """
 
@@ -220,6 +220,38 @@ def __init__(self, **kwargs):
     def _env_whitelist_default(self):
         return os.environ.get(self.env_whitelist_env, self.env_whitelist_default_value)
 
+    gateway_retry_interval_default_value = 1.0
+    gateway_retry_interval_env = 'JUPYTER_GATEWAY_RETRY_INTERVAL'
+    gateway_retry_interval = Float(default_value=gateway_retry_interval_default_value, config=True,
+        help="""The time allowed for HTTP reconnection with the Gateway server for the first time.
+                Next will be JUPYTER_GATEWAY_RETRY_INTERVAL multiplied by two in factor of numbers of retries
+                but less than JUPYTER_GATEWAY_RETRY_INTERVAL_MAX.
+                (JUPYTER_GATEWAY_RETRY_INTERVAL env var)""")
+
+    @default('gateway_retry_interval')
+    def gateway_retry_interval_default(self):
+        return float(os.environ.get('JUPYTER_GATEWAY_RETRY_INTERVAL', self.gateway_retry_interval_default_value))
+
+    gateway_retry_interval_max_default_value = 30.0
+    gateway_retry_interval_max_env = 'JUPYTER_GATEWAY_RETRY_INTERVAL_MAX'
+    gateway_retry_interval_max = Float(default_value=gateway_retry_interval_max_default_value, config=True,
+        help="""The maximum time allowed for HTTP reconnection retry with the Gateway server.
+                (JUPYTER_GATEWAY_RETRY_INTERVAL_MAX env var)""")
+
+    @default('gateway_retry_interval_max')
+    def gateway_retry_interval_max_default(self):
+        return float(os.environ.get('JUPYTER_GATEWAY_RETRY_INTERVAL_MAX', self.gateway_retry_interval_max_default_value))
+
+    gateway_retry_max_default_value = 5
+    gateway_retry_max_env = 'JUPYTER_GATEWAY_RETRY_MAX'
+    gateway_retry_max = Int(default_value=gateway_retry_max_default_value, config=True,
+        help="""The maximum retries allowed for HTTP reconnection with the Gateway server.
+                (JUPYTER_GATEWAY_RETRY_MAX env var)""")
+
+    @default('gateway_retry_max')
+    def gateway_retry_max_default(self):
+        return int(os.environ.get('JUPYTER_GATEWAY_RETRY_MAX', self.gateway_retry_max_default_value))
+
     @property
     def gateway_enabled(self):
         return bool(self.url is not None and len(self.url) > 0)
@@ -503,7 +535,6 @@ def shutdown_all(self, now=False):
             self.remove_kernel(kernel_id)
 
 
-
 class GatewayKernelSpecManager(KernelSpecManager):
 
     def __init__(self, **kwargs):