Skip to content

Commit

Permalink
chore: pull _check_idle (single step) from juju#1104
Browse files Browse the repository at this point in the history
  • Loading branch information
dimaqq committed Nov 28, 2024
1 parent 039a9cb commit 2e19281
Showing 1 changed file with 151 additions and 3 deletions.
154 changes: 151 additions & 3 deletions juju/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@

if TYPE_CHECKING:
from .application import Application
from .client._definitions import FullStatus
from .client._definitions import FullStatus, UnitStatus
from .constraints import StorageConstraintDict
from .machine import Machine
from .relation import Relation
Expand Down Expand Up @@ -2663,14 +2663,21 @@ async def get_action_status(self, uuid_or_prefix=None, name=None):
results[tag.untag("action-", a.action.tag)] = a.status
return results

async def get_status(self, filters=None, utc=False) -> FullStatus:
async def get_status(self, filters=None, utc: bool = False) -> FullStatus:
"""Return the status of the model.
:param str filters: Optional list of applications, units, or machines
to include, which can use wildcards ('*').
:param bool utc: Display time as UTC in RFC3339 format
:param bool utc: Deprecated, display time as UTC in RFC3339 format
"""
if utc:
warnings.warn(
"Model.get_status() utc= parameter is deprecated",
DeprecationWarning,
stacklevel=2,
)

client_facade = client.ClientFacade.from_connection(self.connection())
return await client_facade.FullStatus(patterns=filters)

Expand Down Expand Up @@ -3212,6 +3219,147 @@ def _raise_for_status(entities: dict[str, list[str]], status: Any):
last_log_time = datetime.now()
await jasyncio.sleep(check_freq)

async def _check_idle(
self,
*,
apps: list[str],
raise_on_error: bool,
raise_on_blocked: bool,
status: str | None,
wait_for_at_least_units: int | None,
wait_for_exact_units: int | None,
timeout: float | None,
idle_period: float,
_wait_for_units: int,
idle_times: dict[str, datetime],
units_ready: set[str],
last_log_time: list[datetime | None],
start_time: datetime,
) -> bool:
now = datetime.now()
expected_idle_since = now - timedelta(seconds=idle_period)
full_status = await self.get_status()
# import pdb; pdb.set_trace()

# FIXME check this precedence
for app_name in apps:
if not full_status.applications.get(app_name):
logging.info("Waiting for app %r", app_name)
return False

# Order of errors:
#
# Machine error (any unit of any app from apps)
# Agent error (-"-)
# Workload error (-"-)
# App error (any app from apps)
#
# Workload blocked (any unit of any app from apps)
# App blocked (any app from apps)
units: dict[str, UnitStatus] = {}

for app_name in apps:
# assert full_status.applications[app_name]
app = full_status.applications[app_name]
assert app
for unit_name, unit in app.units.items():
assert unit
units[unit_name] = unit

for unit_name, unit in units.items():
if unit.machine:
machine = full_status.machines[unit.machine]
assert machine
assert machine.instance_status
if machine.instance_status.status == "error" and raise_on_error:
raise JujuMachineError(
f"{unit_name!r} machine {unit.machine!r} has errored: {machine.instance_status.info!r}"
)

for unit_name, unit in units.items():
assert unit.agent_status
if unit.agent_status.status == "error" and raise_on_error:
raise JujuAgentError(
f"{unit_name!r} agent has errored: {unit.agent_status.info!r}"
)

for unit_name, unit in units.items():
assert unit.workload_status
if unit.workload_status.status == "error" and raise_on_error:
raise JujuUnitError(
f"{unit_name!r} workload has errored: {unit.workload_status.info!r}"
)

for app_name in apps:
app = full_status.applications[app_name]
assert app
assert app.status
if app.status.status == "error" and raise_on_error:
raise JujuAppError(f"{app_name!r} has errored: {app.status.info!r}")

for unit_name, unit in units.items():
assert unit.workload_status
if unit.workload_status.status == "blocked" and raise_on_blocked:
raise JujuUnitError(
f"{unit_name!r} workload is blocked: {unit.workload_status.info!r}"
)

for app_name in apps:
app = full_status.applications[app_name]
assert app
assert app.status
if app.status.status == "blocked" and raise_on_blocked:
raise JujuAppError(f"{app_name!r} is blocked: {app.status.info!r}")

for unit_name, unit in units.items():
assert unit.agent_status
idle_times.setdefault(unit_name, now)
if unit.agent_status.status != "idle":
idle_times[unit_name] = now

for app_name in apps:
ready_units = []
app = full_status.applications[app_name]
assert app
for unit in app.units.values():
assert unit
assert unit.agent_status
assert unit.workload_status

if unit.agent_status.status != "idle":
continue
if status and unit.workload_status.status != status:
continue

ready_units.append(unit)

if wait_for_exact_units is None and len(ready_units) < _wait_for_units:
logging.info(
"Waiting for app %r units %s/%s",
app_name,
len(ready_units),
_wait_for_units,
)
return False

if (
wait_for_exact_units is not None
and len(ready_units) != wait_for_exact_units
):
logging.info(
"Waiting for app %r units %s/%s",
app_name,
len(ready_units),
_wait_for_units,
)
return False

if busy := [n for n, t in idle_times.items() if expected_idle_since < t]:
logging.info("Waiting for %s to be idle enough", busy)
return False

return True


def _create_consume_args(offer, macaroon, controller_info):
"""Convert a typed object that has been normalised to a overridden typed
Expand Down

0 comments on commit 2e19281

Please sign in to comment.