Add Juju Public Key

[osm/N2VC.git] / modules / libjuju / juju / client / connection.py

diff --git a/modules/libjuju/juju/client/connection.py b/modules/libjuju/juju/client/connection.py
index bdd1c3f..f2150b7 100644 (file)
--- a/modules/libjuju/juju/client/connection.py
+++ b/modules/libjuju/juju/client/connection.py
@@ -104,25 +104,34 @@ class Connection:
             bakery_client=None,
             loop=None,
             max_frame_size=None,
+            retries=3,
+            retry_backoff=10,
     ):
         """Connect to the websocket.
 
         If uuid is None, the connection will be to the controller. Otherwise it
         will be to the model.
-        :param str endpoint The hostname:port of the controller to connect to.
-        :param str uuid The model UUID to connect to (None for a
+
+        :param str endpoint: The hostname:port of the controller to connect to.
+        :param str uuid: The model UUID to connect to (None for a
             controller-only connection).
-        :param str username The username for controller-local users (or None
+        :param str username: The username for controller-local users (or None
             to use macaroon-based login.)
-        :param str password The password for controller-local users.
-        :param str cacert The CA certificate of the controller (PEM formatted).
-        :param httpbakery.Client bakery_client The macaroon bakery client to
+        :param str password: The password for controller-local users.
+        :param str cacert: The CA certificate of the controller
+            (PEM formatted).
+        :param httpbakery.Client bakery_client: The macaroon bakery client to
             to use when performing macaroon-based login. Macaroon tokens
             acquired when logging will be saved to bakery_client.cookies.
             If this is None, a default bakery_client will be used.
-        :param loop asyncio.BaseEventLoop The event loop to use for async
+        :param asyncio.BaseEventLoop loop: The event loop to use for async
             operations.
-        :param max_frame_size The maximum websocket frame size to allow.
+        :param int max_frame_size: The maximum websocket frame size to allow.
+        :param int retries: When connecting or reconnecting, and all endpoints
+            fail, how many times to retry the connection before giving up.
+        :param int retry_backoff: Number of seconds to increase the wait
+            between connection retry attempts (a backoff of 10 with 3 retries
+            would wait 10s, 20s, and 30s).
         """
         self = cls()
         if endpoint is None:
@@ -157,6 +166,9 @@ class Connection:
         self._pinger_task = _Task(self._pinger, self.loop)
         self._receiver_task = _Task(self._receiver, self.loop)
 
+        self._retries = retries
+        self._retry_backoff = retry_backoff
+
         self.facades = {}
         self.messages = IdQueue(loop=self.loop)
         self.monitor = Monitor(connection=self)
@@ -444,16 +456,30 @@ class Connection:
         tasks = [self.loop.create_task(_try_endpoint(endpoint, cacert,
                                                      0.1 * i))
                  for i, (endpoint, cacert) in enumerate(endpoints)]
-        for task in asyncio.as_completed(tasks, loop=self.loop):
-            try:
-                result = await task
-                break
-            except ConnectionError:
-                continue  # ignore; try another endpoint
-        else:
-            raise errors.JujuConnectionError(
-                'Unable to connect to any endpoint: {}'.format(', '.join([
-                    endpoint for endpoint, cacert in endpoints])))
+        for attempt in range(self._retries + 1):
+            for task in asyncio.as_completed(tasks, loop=self.loop):
+                try:
+                    result = await task
+                    break
+                except ConnectionError:
+                    continue  # ignore; try another endpoint
+            else:
+                _endpoints_str = ', '.join([endpoint
+                                            for endpoint, cacert in endpoints])
+                if attempt < self._retries:
+                    log.debug('Retrying connection to endpoints: {}; '
+                              'attempt {} of {}'.format(_endpoints_str,
+                                                        attempt + 1,
+                                                        self._retries + 1))
+                    await asyncio.sleep((attempt + 1) * self._retry_backoff)
+                    continue
+                else:
+                    raise errors.JujuConnectionError(
+                        'Unable to connect to any endpoint: '
+                        '{}'.format(_endpoints_str))
+            # only executed if inner loop's else did not continue
+            # (i.e., inner loop did break due to successful connection)
+            break
         for task in tasks:
             task.cancel()
         self.ws = result[0]
@@ -461,7 +487,7 @@ class Connection:
         self.endpoint = result[2]
         self.cacert = result[3]
         self._receiver_task.start()
-        log.info("Driver connected to juju %s", self.addr)
+        log.debug("Driver connected to juju %s", self.addr)
         self.monitor.close_called.clear()
 
     async def _connect_with_login(self, endpoints):
@@ -519,8 +545,8 @@ class Connection:
 
     async def login(self):
         params = {}
+        params['auth-tag'] = self.usertag
         if self.password:
-            params['auth-tag'] = self.usertag
             params['credentials'] = self.password
         else:
             macaroons = _macaroons_for_domain(self.bakery_client.cookies,
@@ -560,7 +586,7 @@ class _Task:
     def start(self):
         async def run():
             try:
-                return await(self.task())
+                return await self.task()
             finally:
                 self.stopped.set()
         self.stopped.clear()