Merge pull request #952 from cderici/optimize-teardown

#952

#### Description

This solves a couple of issues about the connection teardown. The main observation is that we have a rogue connection that we open to the controller during model connection to get some model info but never close. That connection has its own pinger/receiver tasks and an internal websocket object (and that object has its own managing tasks etc), and because this connection is totally forgotten, these tasks etc are not handled properly. For a long time I thought the task errors we were intermittently receiving were about the tasks belonging to the main connection. This fixes it by just using the `ConnectedController` context manager to properly disconnect after used, which wraps up the tasks gracefully.

Another change is that this reduces the websocket close timeout to 1 second (default is 10). With the first change, we essentially won't need this, but in any case, we might wanna experiment with this timeout to make sure we don't get false closes.

#### QA Steps

I used @jameinel 's annotate script to reproduce the issues and investigate this, so I think it would make a good use for QAing this as well:

https://github.com/jameinel/test-annotate/blob/main/annotate.py

Author: jujubot

juju/client/connection.py

@@ -441,11 +441,15 @@ async def close(self, to_reconnect=False):
         self.monitor.close_called.set()
 
         # Cancel all the tasks (that we started):
+        tasks_need_to_be_gathered = []
         if self._pinger_task:
+            tasks_need_to_be_gathered.append(self._pinger_task)
             self._pinger_task.cancel()
         if self._receiver_task:
+            tasks_need_to_be_gathered.append(self._receiver_task)
             self._receiver_task.cancel()
         if self._debug_log_task:
+            tasks_need_to_be_gathered.append(self._debug_log_task)
             self._debug_log_task.cancel()
 
         if self._ws and not self._ws.closed:
@@ -454,9 +458,6 @@ async def close(self, to_reconnect=False):
         if not to_reconnect:
             try:
                 log.debug('Gathering all tasks for connection close')
-
-                # Avoid gathering the current task
-                tasks_need_to_be_gathered = [task for task in jasyncio.all_tasks() if task != jasyncio.current_task()]
                 await jasyncio.gather(*tasks_need_to_be_gathered)
             except jasyncio.CancelledError:
                 pass
@@ -603,6 +604,7 @@ async def _pinger(self):
         '''
         async def _do_ping():
             try:
+                log.debug(f'Pinger {self._pinger_task}: pinging')
                 await pinger_facade.Ping()
             except jasyncio.CancelledError:
                 raise
@@ -642,7 +644,7 @@ async def rpc(self, msg, encoder=None):
         if "version" not in msg:
             msg['version'] = self.facades[msg['type']]
         outgoing = json.dumps(msg, indent=2, cls=encoder)
-        log.debug('connection {} -> {}'.format(id(self), outgoing))
+        log.debug('connection id: {} ---> {}'.format(id(self), outgoing))
         for attempt in range(3):
             if self.monitor.status == Monitor.DISCONNECTED:
                 # closed cleanly; shouldn't try to reconnect
@@ -665,7 +667,7 @@ async def rpc(self, msg, encoder=None):
                     log.error('RPC: Automatic reconnect failed')
                     raise
         result = await self._recv(msg['request-id'])
-        log.debug('connection {} <- {}'.format(id(self), result))
+        log.debug('connection id : {} <--- {}'.format(id(self), result))
 
         if not result:
             return result
@@ -796,6 +798,7 @@ async def reconnect(self):
             if not self.is_debug_log_connection:
                 self._build_facades(res.get('facades', {}))
                 if not self._pinger_task:
+                    log.debug('reconnect: scheduling a pinger task')
                     self._pinger_task = jasyncio.create_task(self._pinger(), name="Task_Pinger")
 
     async def _connect(self, endpoints):
@@ -852,6 +855,7 @@ async def _try_endpoint(endpoint, cacert, delay):
         #  If this is regular connection, and we dont have a
         #  receiver_task yet, then schedule a _receiver_task
         elif not self.is_debug_log_connection and not self._receiver_task:
+            log.debug('_connect: scheduling a receiver task')
             self._receiver_task = jasyncio.create_task(self._receiver(), name="Task_Receiver")
 
         log.debug("Driver connected to juju %s", self.addr)
@@ -907,6 +911,7 @@ async def _connect_with_redirect(self, endpoints):
             login_result = await self._connect_with_login(e.endpoints)
         self._build_facades(login_result.get('facades', {}))
         if not self._pinger_task:
+            log.debug('_connect_with_redirect: scheduling a pinger task')
             self._pinger_task = jasyncio.create_task(self._pinger(), name="Task_Pinger")
 
     # _build_facades takes the facade list that comes from the connection with the controller,

juju/client/connector.py

@@ -73,6 +73,12 @@ async def connect(self, **kwargs):
             assert self._connection
             self._log_connection = await Connection.connect(**kwargs)
         else:
+            # TODO (cderici): we need to investigate how to reuse/share
+            # connections between Model & Controller objects
+            # At least either avoid overwriting self._connection with
+            # different types of connections (model, controller), or
+            # have an indication of which type of entity this is
+            # connected to.
             if self._connection:
                 await self._connection.close()
             self._connection = await Connection.connect(**kwargs)
@@ -84,11 +90,11 @@ async def connect(self, **kwargs):
         if not self._connection.info['server-version'].startswith(SUPPORTED_MAJOR_VERSION):
             raise JujuConnectionError("juju server-version %s not supported" % juju_server_version)
 
-    async def disconnect(self):
+    async def disconnect(self, entity):
         """Shut down the watcher task and close websockets.
         """
         if self._connection:
-            log.debug('Closing model connection')
+            log.debug(f'Connector: closing {entity} connection')
             await self._connection.close()
             self._connection = None
         if self._log_connection:

juju/controller.py

@@ -193,7 +193,7 @@ async def disconnect(self):
         """Shut down the watcher task and close websockets.
 
         """
-        await self._connector.disconnect()
+        await self._connector.disconnect(entity='controller')
 
     async def add_credential(self, name=None, credential=None, cloud=None,
                              owner=None, force=False):

juju/model.py

@@ -771,10 +771,14 @@ async def watch_received_waiter():
                 raise JujuError("AllWatcher task is finished abruptly without an exception.")
             raise self._watcher_task.exception()
 
-        if self.info is None:
-            contr = await self.get_controller()
-            self._info = await contr.get_model_info(model_name, model_uuid)
-            log.debug('Got ModelInfo: %s', vars(self.info))
+        if self._info is None:
+            # TODO (cderici): See if this can be optimized away, or at least
+            # be done lazily (i.e. not everytime after_connect, but whenever
+            # self.info is needed -- which here can be bypassed if model_uuid
+            # is known)
+            async with ConnectedController(self.connection()) as contr:
+                self._info = await contr.get_model_info(model_name, model_uuid)
+                log.debug('Got ModelInfo: %s', vars(self.info))
 
         self.uuid = self.info.uuid
 
@@ -795,8 +799,7 @@ async def disconnect(self):
             self._watch_stopping.clear()
 
         if self.is_connected():
-            log.debug('Closing model connection')
-            await self._connector.disconnect()
+            await self._connector.disconnect(entity='model')
             self._info = None
 
     async def add_local_charm_dir(self, charm_dir, series):

tests/integration/test_model.py

@@ -223,8 +223,6 @@ async def test_deploy_local_charm_channel(event_loop):
     async with base.CleanModel() as model:
         await model.deploy(str(charm_path), channel='stable')
         assert 'charm' in model.applications
-        await model.wait_for_idle(status="active")
-        assert model.units['charm/0'].workload_status == 'active'
 
 
 @base.bootstrapped

tests/integration/test_unit.py

@@ -26,7 +26,8 @@ async def is_leader_elected():
             return any([await u.is_leader_from_status() for u in app.units])
 
         await utils.block_until_with_coroutine(is_leader_elected,
-                                               timeout=480)
+                                               timeout=480,
+                                               wait_period=5)
 
 
 @base.bootstrapped