Coverage for n2vc/juju_watcher.py: 93%
153 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-05-07 06:04 +0000
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.
15import asyncio
16import time
18from juju.client import client
19from n2vc.exceptions import EntityInvalidException
20from n2vc.n2vc_conn import N2VCConnector
21from juju.model import ModelEntity, Model
22from juju.client.overrides import Delta
23from juju.status import derive_status
24from juju.application import Application
25from websockets.exceptions import ConnectionClosed
26import logging
28logger = logging.getLogger("__main__")
31def status(application: Application) -> str:
32 unit_status = []
33 for unit in application.units:
34 unit_status.append(unit.workload_status)
35 return derive_status(unit_status)
38def entity_ready(entity: ModelEntity) -> bool:
39 """
40 Check if the entity is ready
42 :param: entity: Model entity. It can be a machine, action, or application.
44 :returns: boolean saying if the entity is ready or not
45 """
47 entity_type = entity.entity_type
48 if entity_type == "machine":
49 return entity.agent_status in ["started"]
50 elif entity_type == "action":
51 return entity.status in ["completed", "failed", "cancelled"]
52 elif entity_type == "application":
53 # Workaround for bug: https://github.com/juju/python-libjuju/issues/441
54 return entity.status in ["active", "blocked"]
55 elif entity_type == "unit":
56 return entity.agent_status in ["idle"]
57 else:
58 raise EntityInvalidException("Unknown entity type: {}".format(entity_type))
61def application_ready(application: Application) -> bool:
62 """
63 Check if an application has a leader
65 :param: application: Application entity.
67 :returns: boolean saying if the application has a unit that is a leader.
68 """
69 ready_status_list = ["active", "blocked"]
70 application_ready = application.status in ready_status_list
71 units_ready = all(
72 unit.workload_status in ready_status_list for unit in application.units
73 )
74 return application_ready and units_ready
77class JujuModelWatcher:
78 @staticmethod
79 async def wait_for_model(model: Model, timeout: float = 3600):
80 """
81 Wait for all entities in model to reach its final state.
83 :param: model: Model to observe
84 :param: timeout: Timeout for the model applications to be active
86 :raises: asyncio.TimeoutError when timeout reaches
87 """
89 if timeout is None:
90 timeout = 3600.0
92 # Coroutine to wait until the entity reaches the final state
93 async def wait_until_model_ready():
94 wait_for_entity = asyncio.ensure_future(
95 asyncio.wait_for(
96 model.block_until(
97 lambda: all(
98 application_ready(application)
99 for application in model.applications.values()
100 ),
101 ),
102 timeout=timeout,
103 )
104 )
106 tasks = [wait_for_entity]
107 try:
108 await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
109 finally:
110 # Cancel tasks
111 for task in tasks:
112 task.cancel()
114 await wait_until_model_ready()
115 # Check model is still ready after 10 seconds
117 await asyncio.sleep(10)
118 await wait_until_model_ready()
120 @staticmethod
121 async def wait_for(
122 model: Model,
123 entity: ModelEntity,
124 progress_timeout: float = 3600,
125 total_timeout: float = 3600,
126 db_dict: dict = None,
127 n2vc: N2VCConnector = None,
128 vca_id: str = None,
129 ):
130 """
131 Wait for entity to reach its final state.
133 :param: model: Model to observe
134 :param: entity: Entity object
135 :param: progress_timeout: Maximum time between two updates in the model
136 :param: total_timeout: Timeout for the entity to be active
137 :param: db_dict: Dictionary with data of the DB to write the updates
138 :param: n2vc: N2VC Connector objector
139 :param: vca_id: VCA ID
141 :raises: asyncio.TimeoutError when timeout reaches
142 """
144 if progress_timeout is None:
145 progress_timeout = 3600.0
146 if total_timeout is None:
147 total_timeout = 3600.0
149 entity_type = entity.entity_type
150 if entity_type not in ["application", "action", "machine", "unit"]:
151 raise EntityInvalidException("Unknown entity type: {}".format(entity_type))
153 # Coroutine to wait until the entity reaches the final state
154 wait_for_entity = asyncio.ensure_future(
155 asyncio.wait_for(
156 model.block_until(lambda: entity_ready(entity)),
157 timeout=total_timeout,
158 )
159 )
161 # Coroutine to watch the model for changes (and write them to DB)
162 watcher = asyncio.ensure_future(
163 JujuModelWatcher.model_watcher(
164 model,
165 entity_id=entity.entity_id,
166 entity_type=entity_type,
167 timeout=progress_timeout,
168 db_dict=db_dict,
169 n2vc=n2vc,
170 vca_id=vca_id,
171 )
172 )
174 tasks = [wait_for_entity, watcher]
175 try:
176 # Execute tasks, and stop when the first is finished
177 # The watcher task won't never finish (unless it timeouts)
178 await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
179 finally:
180 # Cancel tasks
181 for task in tasks:
182 task.cancel()
184 @staticmethod
185 async def wait_for_units_idle(
186 model: Model, application: Application, timeout: float = 60
187 ):
188 """
189 Waits for the application and all its units to transition back to idle
191 :param: model: Model to observe
192 :param: application: The application to be observed
193 :param: timeout: Maximum time between two updates in the model
195 :raises: asyncio.TimeoutError when timeout reaches
196 """
198 ensure_units_idle = asyncio.ensure_future(
199 asyncio.wait_for(
200 JujuModelWatcher.ensure_units_idle(model, application), timeout
201 )
202 )
203 tasks = [
204 ensure_units_idle,
205 ]
206 (done, pending) = await asyncio.wait(
207 tasks, timeout=timeout, return_when=asyncio.FIRST_COMPLETED
208 )
210 if ensure_units_idle in pending:
211 ensure_units_idle.cancel()
212 raise TimeoutError(
213 "Application's units failed to return to idle after {} seconds".format(
214 timeout
215 )
216 )
217 if ensure_units_idle.result():
218 pass
220 @staticmethod
221 async def ensure_units_idle(model: Model, application: Application):
222 """
223 Waits forever until the application's units to transition back to idle
225 :param: model: Model to observe
226 :param: application: The application to be observed
227 """
229 try:
230 allwatcher = client.AllWatcherFacade.from_connection(model.connection())
231 unit_wanted_state = "executing"
232 final_state_reached = False
234 units = application.units
235 final_state_seen = {unit.entity_id: False for unit in units}
236 agent_state_seen = {unit.entity_id: False for unit in units}
237 workload_state = {unit.entity_id: False for unit in units}
239 try:
240 while not final_state_reached:
241 change = await allwatcher.Next()
243 # Keep checking to see if new units were added during the change
244 for unit in units:
245 if unit.entity_id not in final_state_seen:
246 final_state_seen[unit.entity_id] = False
247 agent_state_seen[unit.entity_id] = False
248 workload_state[unit.entity_id] = False
250 for delta in change.deltas:
251 await asyncio.sleep(0)
252 if delta.entity != units[0].entity_type:
253 continue
255 final_state_reached = True
256 for unit in units:
257 if delta.data["name"] == unit.entity_id:
258 status = delta.data["agent-status"]["current"]
259 workload_state[unit.entity_id] = delta.data[
260 "workload-status"
261 ]["current"]
263 if status == unit_wanted_state:
264 agent_state_seen[unit.entity_id] = True
265 final_state_seen[unit.entity_id] = False
267 if (
268 status == "idle"
269 and agent_state_seen[unit.entity_id]
270 ):
271 final_state_seen[unit.entity_id] = True
273 final_state_reached = (
274 final_state_reached
275 and final_state_seen[unit.entity_id]
276 and workload_state[unit.entity_id]
277 in [
278 "active",
279 "error",
280 ]
281 )
283 except ConnectionClosed:
284 pass
285 # This is expected to happen when the
286 # entity reaches its final state, because
287 # the model connection is closed afterwards
288 except Exception as e:
289 raise e
291 @staticmethod
292 async def model_watcher(
293 model: Model,
294 entity_id: str,
295 entity_type: str,
296 timeout: float,
297 db_dict: dict = None,
298 n2vc: N2VCConnector = None,
299 vca_id: str = None,
300 ):
301 """
302 Observes the changes related to an specific entity in a model
304 :param: model: Model to observe
305 :param: entity_id: ID of the entity to be observed
306 :param: entity_type: Entity Type (p.e. "application", "machine, and "action")
307 :param: timeout: Maximum time between two updates in the model
308 :param: db_dict: Dictionary with data of the DB to write the updates
309 :param: n2vc: N2VC Connector objector
310 :param: vca_id: VCA ID
312 :raises: asyncio.TimeoutError when timeout reaches
313 """
315 try:
316 allwatcher = client.AllWatcherFacade.from_connection(model.connection())
318 # Genenerate array with entity types to listen
319 entity_types = (
320 [entity_type, "unit"]
321 if entity_type == "application" # TODO: Add "action" too
322 else [entity_type]
323 )
325 # Get time when it should timeout
326 timeout_end = time.time() + timeout
328 try:
329 while True:
330 change = await allwatcher.Next()
331 for delta in change.deltas:
332 write = False
333 delta_entity = None
335 # Get delta EntityType
336 delta_entity = delta.entity
338 if delta_entity in entity_types:
339 # Get entity id
340 id = None
341 if entity_type == "application":
342 id = (
343 delta.data["application"]
344 if delta_entity == "unit"
345 else delta.data["name"]
346 )
347 else:
348 if "id" in delta.data:
349 id = delta.data["id"]
350 else:
351 print("No id {}".format(delta.data))
353 # Write if the entity id match
354 write = True if id == entity_id else False
356 # Update timeout
357 timeout_end = time.time() + timeout
358 (
359 status,
360 status_message,
361 vca_status,
362 ) = JujuModelWatcher.get_status(delta)
364 if write and n2vc is not None and db_dict:
365 # Write status to DB
366 status = n2vc.osm_status(delta_entity, status)
367 await n2vc.write_app_status_to_db(
368 db_dict=db_dict,
369 status=status,
370 detailed_status=status_message,
371 vca_status=vca_status,
372 entity_type=delta_entity,
373 vca_id=vca_id,
374 )
375 # Check if timeout
376 if time.time() > timeout_end:
377 raise asyncio.TimeoutError()
378 except ConnectionClosed:
379 pass
380 # This is expected to happen when the
381 # entity reaches its final state, because
382 # the model connection is closed afterwards
383 except Exception as e:
384 raise e
386 @staticmethod
387 def get_status(delta: Delta) -> (str, str, str):
388 """
389 Get status from delta
391 :param: delta: Delta generated by the allwatcher
392 :param: entity_type: Entity Type (p.e. "application", "machine, and "action")
394 :return (status, message, vca_status)
395 """
396 if delta.entity == "machine":
397 return (
398 delta.data["agent-status"]["current"],
399 delta.data["instance-status"]["message"],
400 delta.data["instance-status"]["current"],
401 )
402 elif delta.entity == "action":
403 return (
404 delta.data["status"],
405 delta.data["status"],
406 delta.data["status"],
407 )
408 elif delta.entity == "application":
409 return (
410 delta.data["status"]["current"],
411 delta.data["status"]["message"],
412 delta.data["status"]["current"],
413 )
414 elif delta.entity == "unit":
415 return (
416 delta.data["workload-status"]["current"],
417 delta.data["workload-status"]["message"],
418 delta.data["workload-status"]["current"],
419 )