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     entity_type = entity.entity_type
  39     if entity_type == "machine":
  40         return entity.agent_status in ["started"]
  41     elif entity_type == "action":
  42         return entity.status in ["completed", "failed", "cancelled"]
  43     elif entity_type == "application":
  44         # Workaround for bug: https://github.com/juju/python-libjuju/issues/441
  45         return status(entity) in ["active", "blocked"]
  46     else:
  47         raise EntityInvalidException("Unknown entity type: {}".format(entity_type))
  48
  49
  50 class JujuModelWatcher:
  51     @staticmethod
  52     async def wait_for_model(model: Model, timeout: float = 3600):
  53         """
  54         Wait for all entities in model to reach its final state.
  55
  56         :param: model:              Model to observe
  57         :param: timeout:            Timeout for the model applications to be active
  58
  59         :raises: asyncio.TimeoutError when timeout reaches
  60         """
  61
  62         if timeout is None:
  63             timeout = 3600.0
  64
  65         # Coroutine to wait until the entity reaches the final state
  66         wait_for_entity = asyncio.ensure_future(
  67             asyncio.wait_for(
  68                 model.block_until(
  69                     lambda: all(
  70                         entity_ready(entity) for entity in model.applications.values()
  71                     )
  72                 ),
  73                 timeout=timeout,
  74             )
  75         )
  76
  77         tasks = [wait_for_entity]
  78         try:
  79             await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
  80         finally:
  81             # Cancel tasks
  82             for task in tasks:
  83                 task.cancel()
  84
  85     @staticmethod
  86     async def wait_for(
  87         model: Model,
  88         entity: ModelEntity,
  89         progress_timeout: float = 3600,
  90         total_timeout: float = 3600,
  91         db_dict: dict = None,
  92         n2vc: N2VCConnector = None,
  93     ):
  94         """
  95         Wait for entity to reach its final state.
  96
  97         :param: model:              Model to observe
  98         :param: entity:             Entity object
  99         :param: progress_timeout:   Maximum time between two updates in the model
 100         :param: total_timeout:      Timeout for the entity to be active
 101         :param: db_dict:            Dictionary with data of the DB to write the updates
 102         :param: n2vc:               N2VC Connector objector
 103
 104         :raises: asyncio.TimeoutError when timeout reaches
 105         """
 106
 107         if progress_timeout is None:
 108             progress_timeout = 3600.0
 109         if total_timeout is None:
 110             total_timeout = 3600.0
 111
 112         entity_type = entity.entity_type
 113         if entity_type not in ["application", "action", "machine"]:
 114             raise EntityInvalidException("Unknown entity type: {}".format(entity_type))
 115
 116         # Coroutine to wait until the entity reaches the final state
 117         wait_for_entity = asyncio.ensure_future(
 118             asyncio.wait_for(
 119                 model.block_until(lambda: entity_ready(entity)), timeout=total_timeout,
 120             )
 121         )
 122
 123         # Coroutine to watch the model for changes (and write them to DB)
 124         watcher = asyncio.ensure_future(
 125             JujuModelWatcher.model_watcher(
 126                 model,
 127                 entity_id=entity.entity_id,
 128                 entity_type=entity_type,
 129                 timeout=progress_timeout,
 130                 db_dict=db_dict,
 131                 n2vc=n2vc,
 132             )
 133         )
 134
 135         tasks = [wait_for_entity, watcher]
 136         try:
 137             # Execute tasks, and stop when the first is finished
 138             # The watcher task won't never finish (unless it timeouts)
 139             await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
 140         finally:
 141             # Cancel tasks
 142             for task in tasks:
 143                 task.cancel()
 144
 145     @staticmethod
 146     async def model_watcher(
 147         model: Model,
 148         entity_id: str,
 149         entity_type: str,
 150         timeout: float,
 151         db_dict: dict = None,
 152         n2vc: N2VCConnector = None,
 153     ):
 154         """
 155         Observes the changes related to an specific entity in a model
 156
 157         :param: model:          Model to observe
 158         :param: entity_id:      ID of the entity to be observed
 159         :param: entity_type:    Entity Type (p.e. "application", "machine, and "action")
 160         :param: timeout:        Maximum time between two updates in the model
 161         :param: db_dict:        Dictionary with data of the DB to write the updates
 162         :param: n2vc:           N2VC Connector objector
 163
 164         :raises: asyncio.TimeoutError when timeout reaches
 165         """
 166
 167         allwatcher = client.AllWatcherFacade.from_connection(model.connection())
 168
 169         # Genenerate array with entity types to listen
 170         entity_types = (
 171             [entity_type, "unit"]
 172             if entity_type == "application"  # TODO: Add "action" too
 173             else [entity_type]
 174         )
 175
 176         # Get time when it should timeout
 177         timeout_end = time.time() + timeout
 178
 179         try:
 180             while True:
 181                 change = await allwatcher.Next()
 182                 for delta in change.deltas:
 183                     write = False
 184                     delta_entity = None
 185
 186                     # Get delta EntityType
 187                     delta_entity = delta.entity
 188
 189                     if delta_entity in entity_types:
 190                         # Get entity id
 191                         if entity_type == "application":
 192                             id = (
 193                                 delta.data["application"]
 194                                 if delta_entity == "unit"
 195                                 else delta.data["name"]
 196                             )
 197                         else:
 198                             id = delta.data["id"]
 199
 200                         # Write if the entity id match
 201                         write = True if id == entity_id else False
 202
 203                         # Update timeout
 204                         timeout_end = time.time() + timeout
 205                         (
 206                             status,
 207                             status_message,
 208                             vca_status,
 209                         ) = JujuModelWatcher.get_status(delta)
 210
 211                         if write and n2vc is not None and db_dict:
 212                             # Write status to DB
 213                             status = n2vc.osm_status(delta_entity, status)
 214                             await n2vc.write_app_status_to_db(
 215                                 db_dict=db_dict,
 216                                 status=status,
 217                                 detailed_status=status_message,
 218                                 vca_status=vca_status,
 219                                 entity_type=delta_entity,
 220                             )
 221                 # Check if timeout
 222                 if time.time() > timeout_end:
 223                     raise asyncio.TimeoutError()
 224         except ConnectionClosed:
 225             pass
 226             # This is expected to happen when the
 227             # entity reaches its final state, because
 228             # the model connection is closed afterwards
 229
 230     @staticmethod
 231     def get_status(delta: Delta) -> (str, str, str):
 232         """
 233         Get status from delta
 234
 235         :param: delta:          Delta generated by the allwatcher
 236         :param: entity_type:    Entity Type (p.e. "application", "machine, and "action")
 237
 238         :return (status, message, vca_status)
 239         """
 240         if delta.entity == "machine":
 241             return (
 242                 delta.data["agent-status"]["current"],
 243                 delta.data["instance-status"]["message"],
 244                 delta.data["instance-status"]["current"],
 245             )
 246         elif delta.entity == "action":
 247             return (
 248                 delta.data["status"],
 249                 delta.data["status"],
 250                 delta.data["status"],
 251             )
 252         elif delta.entity == "application":
 253             return (
 254                 delta.data["status"]["current"],
 255                 delta.data["status"]["message"],
 256                 delta.data["status"]["current"],
 257             )
 258         elif delta.entity == "unit":
 259             return (
 260                 delta.data["workload-status"]["current"],
 261                 delta.data["workload-status"]["message"],
 262                 delta.data["workload-status"]["current"],
 263             )