n2vc/juju_watcher.py

   1 # Copyright 2020 Canonical Ltd.
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at
   6 #
   7 #     http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 #     Unless required by applicable law or agreed to in writing, software
  10 #     distributed under the License is distributed on an "AS IS" BASIS,
  11 #     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 #     See the License for the specific language governing permissions and
  13 #     limitations under the License.
  14
  15 import asyncio
  16 import time
  17 from juju.client import client
  18 from n2vc.exceptions import EntityInvalidException
  19 from n2vc.n2vc_conn import N2VCConnector
  20 from juju.model import ModelEntity, Model
  21 from juju.client.overrides import Delta
  22 from juju.status import derive_status
  23 from juju.application import Application
  24 from websockets.exceptions import ConnectionClosed
  25 import logging
  26
  27 logger = logging.getLogger("__main__")
  28
  29
  30 def status(application: Application) -> str:
  31     unit_status = []
  32     for unit in application.units:
  33         unit_status.append(unit.workload_status)
  34     return derive_status(unit_status)
  35
  36
  37 def entity_ready(entity: ModelEntity) -> bool:
  38     """
  39     Check if the entity is ready
  40
  41     :param: entity: Model entity. It can be a machine, action, or application.
  42
  43     :returns: boolean saying if the entity is ready or not
  44     """
  45     entity_type = entity.entity_type
  46     if entity_type == "machine":
  47         return entity.agent_status in ["started"]
  48     elif entity_type == "action":
  49         return entity.status in ["completed", "failed", "cancelled"]
  50     elif entity_type == "application":
  51         # Workaround for bug: https://github.com/juju/python-libjuju/issues/441
  52         return entity.status in ["active", "blocked"]
  53     else:
  54         raise EntityInvalidException("Unknown entity type: {}".format(entity_type))
  55
  56
  57 def application_ready(application: Application) -> bool:
  58     """
  59     Check if an application has a leader
  60
  61     :param: application: Application entity.
  62
  63     :returns: boolean saying if the application has a unit that is a leader.
  64     """
  65     ready_status_list = ["active", "blocked"]
  66     application_ready = application.status in ready_status_list
  67     units_ready = all(
  68         unit.workload_status in ready_status_list for unit in application.units
  69     )
  70     return application_ready and units_ready
  71
  72
  73 class JujuModelWatcher:
  74     @staticmethod
  75     async def wait_for_model(
  76         model: Model,
  77         timeout: float = 3600
  78     ):
  79         """
  80         Wait for all entities in model to reach its final state.
  81
  82         :param: model:              Model to observe
  83         :param: timeout:            Timeout for the model applications to be active
  84
  85         :raises: asyncio.TimeoutError when timeout reaches
  86         """
  87
  88         if timeout is None:
  89             timeout = 3600.0
  90
  91         # Coroutine to wait until the entity reaches the final state
  92         async def wait_until_model_ready():
  93             wait_for_entity = asyncio.ensure_future(
  94                 asyncio.wait_for(
  95                     model.block_until(
  96                         lambda: all(
  97                             application_ready(application)
  98                             for application in model.applications.values()
  99                         ),
 100                     ),
 101                     timeout=timeout,
 102                 )
 103             )
 104
 105             tasks = [wait_for_entity]
 106             try:
 107                 await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
 108             finally:
 109                 # Cancel tasks
 110                 for task in tasks:
 111                     task.cancel()
 112
 113         await wait_until_model_ready()
 114         # Check model is still ready after 10 seconds
 115
 116         await asyncio.sleep(10)
 117         await wait_until_model_ready()
 118
 119     @staticmethod
 120     async def wait_for(
 121         model: Model,
 122         entity: ModelEntity,
 123         progress_timeout: float = 3600,
 124         total_timeout: float = 3600,
 125         db_dict: dict = None,
 126         n2vc: N2VCConnector = None,
 127         vca_id: str = None,
 128     ):
 129         """
 130         Wait for entity to reach its final state.
 131
 132         :param: model:              Model to observe
 133         :param: entity:             Entity object
 134         :param: progress_timeout:   Maximum time between two updates in the model
 135         :param: total_timeout:      Timeout for the entity to be active
 136         :param: db_dict:            Dictionary with data of the DB to write the updates
 137         :param: n2vc:               N2VC Connector objector
 138         :param: vca_id:             VCA ID
 139
 140         :raises: asyncio.TimeoutError when timeout reaches
 141         """
 142
 143         if progress_timeout is None:
 144             progress_timeout = 3600.0
 145         if total_timeout is None:
 146             total_timeout = 3600.0
 147
 148         entity_type = entity.entity_type
 149         if entity_type not in ["application", "action", "machine"]:
 150             raise EntityInvalidException("Unknown entity type: {}".format(entity_type))
 151
 152         # Coroutine to wait until the entity reaches the final state
 153         wait_for_entity = asyncio.ensure_future(
 154             asyncio.wait_for(
 155                 model.block_until(lambda: entity_ready(entity)),
 156                 timeout=total_timeout,
 157             )
 158         )
 159
 160         # Coroutine to watch the model for changes (and write them to DB)
 161         watcher = asyncio.ensure_future(
 162             JujuModelWatcher.model_watcher(
 163                 model,
 164                 entity_id=entity.entity_id,
 165                 entity_type=entity_type,
 166                 timeout=progress_timeout,
 167                 db_dict=db_dict,
 168                 n2vc=n2vc,
 169                 vca_id=vca_id,
 170             )
 171         )
 172
 173         tasks = [wait_for_entity, watcher]
 174         try:
 175             # Execute tasks, and stop when the first is finished
 176             # The watcher task won't never finish (unless it timeouts)
 177             await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
 178         finally:
 179             # Cancel tasks
 180             for task in tasks:
 181                 task.cancel()
 182
 183     @staticmethod
 184     async def model_watcher(
 185         model: Model,
 186         entity_id: str,
 187         entity_type: str,
 188         timeout: float,
 189         db_dict: dict = None,
 190         n2vc: N2VCConnector = None,
 191         vca_id: str = None,
 192     ):
 193         """
 194         Observes the changes related to an specific entity in a model
 195
 196         :param: model:          Model to observe
 197         :param: entity_id:      ID of the entity to be observed
 198         :param: entity_type:    Entity Type (p.e. "application", "machine, and "action")
 199         :param: timeout:        Maximum time between two updates in the model
 200         :param: db_dict:        Dictionary with data of the DB to write the updates
 201         :param: n2vc:           N2VC Connector objector
 202         :param: vca_id:         VCA ID
 203
 204         :raises: asyncio.TimeoutError when timeout reaches
 205         """
 206
 207         allwatcher = client.AllWatcherFacade.from_connection(model.connection())
 208
 209         # Genenerate array with entity types to listen
 210         entity_types = (
 211             [entity_type, "unit"]
 212             if entity_type == "application"  # TODO: Add "action" too
 213             else [entity_type]
 214         )
 215
 216         # Get time when it should timeout
 217         timeout_end = time.time() + timeout
 218
 219         try:
 220             while True:
 221                 change = await allwatcher.Next()
 222                 for delta in change.deltas:
 223                     write = False
 224                     delta_entity = None
 225
 226                     # Get delta EntityType
 227                     delta_entity = delta.entity
 228
 229                     if delta_entity in entity_types:
 230                         # Get entity id
 231                         if entity_type == "application":
 232                             id = (
 233                                 delta.data["application"]
 234                                 if delta_entity == "unit"
 235                                 else delta.data["name"]
 236                             )
 237                         else:
 238                             id = delta.data["id"]
 239
 240                         # Write if the entity id match
 241                         write = True if id == entity_id else False
 242
 243                         # Update timeout
 244                         timeout_end = time.time() + timeout
 245                         (
 246                             status,
 247                             status_message,
 248                             vca_status,
 249                         ) = JujuModelWatcher.get_status(delta)
 250
 251                         if write and n2vc is not None and db_dict:
 252                             # Write status to DB
 253                             status = n2vc.osm_status(delta_entity, status)
 254                             await n2vc.write_app_status_to_db(
 255                                 db_dict=db_dict,
 256                                 status=status,
 257                                 detailed_status=status_message,
 258                                 vca_status=vca_status,
 259                                 entity_type=delta_entity,
 260                                 vca_id=vca_id,
 261                             )
 262                 # Check if timeout
 263                 if time.time() > timeout_end:
 264                     raise asyncio.TimeoutError()
 265         except ConnectionClosed:
 266             pass
 267             # This is expected to happen when the
 268             # entity reaches its final state, because
 269             # the model connection is closed afterwards
 270
 271     @staticmethod
 272     def get_status(delta: Delta) -> (str, str, str):
 273         """
 274         Get status from delta
 275
 276         :param: delta:          Delta generated by the allwatcher
 277         :param: entity_type:    Entity Type (p.e. "application", "machine, and "action")
 278
 279         :return (status, message, vca_status)
 280         """
 281         if delta.entity == "machine":
 282             return (
 283                 delta.data["agent-status"]["current"],
 284                 delta.data["instance-status"]["message"],
 285                 delta.data["instance-status"]["current"],
 286             )
 287         elif delta.entity == "action":
 288             return (
 289                 delta.data["status"],
 290                 delta.data["status"],
 291                 delta.data["status"],
 292             )
 293         elif delta.entity == "application":
 294             return (
 295                 delta.data["status"]["current"],
 296                 delta.data["status"]["message"],
 297                 delta.data["status"]["current"],
 298             )
 299         elif delta.entity == "unit":
 300             return (
 301                 delta.data["workload-status"]["current"],
 302                 delta.data["workload-status"]["message"],
 303                 delta.data["workload-status"]["current"],
 304             )