Use sunbeam for determining number of machines need controllers

anvil-python/anvil/commands/juju.py

@@ -13,18 +13,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
 import logging
 from os import environ
 import os.path
+import random
 import subprocess
 
 from rich.status import Status
-
 from sunbeam.commands.juju import JujuStepHelper
 from sunbeam.jobs.common import BaseStep, Result, ResultType
+from sunbeam.jobs.juju import JujuHelper, run_sync
 
-from anvil.utils import machines_missing_juju_controllers
+from anvil.jobs.juju import CONTROLLER
 
 LOG = logging.getLogger(__name__)
 MAX_JUJU_CONTROLLERS = 3
@@ -65,29 +65,43 @@ def run(self, status: Status | None) -> Result:
         return Result(ResultType.COMPLETED)
 
 
-class ScaleUpJujuStep(BaseStep, JujuStepHelper):
+class ScaleJujuStep(BaseStep, JujuStepHelper):
     """Enable Juju HA."""
 
     def __init__(
         self,
-        controller: str,
-        joining: bool,
-        n: int = MAX_JUJU_CONTROLLERS,
-        extra_args: list[str] | None = None,
+        jhelper: JujuHelper,
+        model: str,
     ):
         super().__init__("Juju HA", "Enable Juju High Availability")
-        self.controller = controller
-        self.joining = joining
-        self.n = n
-        self.extra_args = extra_args or []
+
+        self.jhelper = jhelper
+        self.model = model
+
+        self.controller_machines = self.get_controller(CONTROLLER)[
+            "controller-machines"
+        ].keys()
+        self.machines = run_sync(self.jhelper.get_machines(self.model)).keys()
 
     def run(self, status: Status | None = None) -> Result:
+        """Run the step to completion."""
+
+        available_machines = list(self.machines ^ self.controller_machines)
+        n_machines_to_join = min(
+            len(available_machines),
+            MAX_JUJU_CONTROLLERS - len(self.controller_machines),
+        )
+
         cmd = [
             self._get_juju_binary(),
             "enable-ha",
             "-n",
-            str(self.n),
-            *self.extra_args,
+            str(len(self.controller_machines) + n_machines_to_join),
+            "--to",
+            ",".join(
+                str(s)
+                for s in random.sample(available_machines, n_machines_to_join)
+            ),
         ]
         LOG.debug(f'Running command {" ".join(cmd)}')
         process = subprocess.run(
@@ -119,37 +133,22 @@ def run(self, status: Status | None = None) -> Result:
 
     def is_skip(self, status: Status | None = None) -> Result:
         """Determines if the step should be skipped or not."""
-        machines_res = subprocess.run(
-            ["juju", "machines", "--format", "json"], capture_output=True
-        )
-        machines = json.loads(machines_res.stdout)["machines"]
-        n_machines = len(machines)
-        machines_to_join = machines_missing_juju_controllers()
-        n_machines_no_controller = len(machines_to_join)
-        n_controller_machines = n_machines - n_machines_no_controller
-        if (
-            self.joining
-            and n_controller_machines < MAX_JUJU_CONTROLLERS
-            and n_machines == 3
-        ):
-            self.extra_args.extend(("--to", ",".join(machines_to_join)))
+
+        available_machines = self.machines ^ self.controller_machines
+
+        if len(self.controller_machines) == MAX_JUJU_CONTROLLERS:
             LOG.debug(
-                f"Will enable Juju controller on machines {machines_to_join}"
+                "Number of machines with controllers must not be greater than "
+                f"{MAX_JUJU_CONTROLLERS}, skipping scaling Juju controllers"
             )
-            return Result(ResultType.COMPLETED)
-        elif (
-            not self.joining
-            and n_controller_machines < MAX_JUJU_CONTROLLERS
-            and n_machines >= MAX_JUJU_CONTROLLERS
-        ):
-            # a controller machine has been removed, need to pick a new one
-            machines_to_join = machines_to_join[
-                : (MAX_JUJU_CONTROLLERS - n_controller_machines)
-            ]
-            self.extra_args.extend(("--to", ",".join(machines_to_join)))
-            return Result(ResultType.COMPLETED)
-        LOG.debug(
-            "Number of machines with controllers must not be greater than "
-            f"{MAX_JUJU_CONTROLLERS}, skipping scaling Juju controllers"
-        )
-        return Result(ResultType.SKIPPED)
+            return Result(ResultType.SKIPPED)
+        if len(available_machines) == 0:
+            LOG.debug(
+                "No available machines, skipping scaling Juju controllers"
+            )
+            return Result(ResultType.SKIPPED)
+        if len(self.machines) < 3:
+            LOG.debug("Number of machines must be at least 3")
+            return Result(ResultType.SKIPPED)
+
+        return Result(ResultType.COMPLETED)

anvil-python/anvil/provider/local/commands.py

@@ -73,7 +73,7 @@
     RemoveHAProxyUnitStep,
     haproxy_install_steps,
 )
-from anvil.commands.juju import ScaleUpJujuStep, JujuAddSSHKeyStep
+from anvil.commands.juju import JujuAddSSHKeyStep, ScaleJujuStep
 from anvil.commands.maas_agent import (
     RemoveMAASAgentUnitStep,
     maas_agent_install_steps,
@@ -498,7 +498,7 @@ def join(
                 name,
             )
         )
-    plan2.append(ScaleUpJujuStep(controller, True))
+    plan2.append(ScaleJujuStep(jhelper, deployment.infrastructure_model))
     run_plan(plan2, console)
 
     click.echo(f"Node joined cluster with roles: {pretty_roles}")
@@ -590,7 +590,7 @@ def remove(ctx: click.Context, name: str) -> None:
         # Cannot remove user as the same user name cannot be reused,
         # so commenting the RemoveJujuUserStep
         # RemoveJujuUserStep(name),
-        ScaleUpJujuStep(CONTROLLER, False),
+        ScaleJujuStep(jhelper, deployment.infrastructure_model),
         ClusterRemoveNodeStep(client, name),
     ]
     run_plan(plan, console)

anvil-python/anvil/utils.py

@@ -13,16 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
 import logging
-import subprocess
 import sys
 
 import click
 from sunbeam.plugins.interface.v1.base import PluginError
 
-from anvil.jobs.juju import CONTROLLER
-
 LOG = logging.getLogger(__name__)
 LOCAL_ACCESS = "local"
 REMOTE_ACCESS = "remote"
@@ -47,20 +43,3 @@ def __call__(self, *args, **kwargs):  # type: ignore[no-untyped-def]
             LOG.warn(message)
             LOG.error("Error: %s", e)
             sys.exit(1)
-
-
-def machines_missing_juju_controllers() -> list[str]:
-    result = subprocess.run(
-        ["juju", "show-controller", CONTROLLER, "--format", "json"],
-        capture_output=True,
-    )
-    controllers = json.loads(result.stdout)
-    controller_machines = set(
-        controllers[CONTROLLER]["controller-machines"].keys()
-    )
-
-    machines_res = subprocess.run(
-        ["juju", "machines", "--format", "json"], capture_output=True
-    )
-    machines = set(json.loads(machines_res.stdout)["machines"].keys())
-    return list(machines.difference(controller_machines))