[DPE-4952] Add add/remove broker capability to rebalance action (#122)

Author: Batalex

actions.yaml

@@ -50,3 +50,8 @@ rebalance:
       description: Only generate the partition rebalance proposals and estimated result, without executing
       type: boolean
       default: true
+    brokerid:
+      description: Broker ID newly added to the cluster or to be removed. The broker ID is the unit number, e.g. kafka/0 is broker 0.
+      type: integer
+      minimum: 0
+  required: [mode]

src/events/balancer.py

@@ -20,6 +20,8 @@
     BALANCER_WEBSERVER_USER,
     CONTAINER,
     GROUP,
+    MODE_ADD,
+    MODE_REMOVE,
     USER,
     Status,
 )
@@ -200,16 +202,25 @@ def rebalance(self, event: ActionEvent) -> None:
                 not self.balancer_manager.cruise_control.ready,
                 "CruiseControl balancer service has not yet collected enough data to provide a partition reallocation proposal",
             ),
+            (
+                event.params.get("brokerid", None) is None
+                and event.params["mode"] in (MODE_ADD, MODE_REMOVE),
+                "'add' and 'remove' rebalance action require passing the 'brokerid' parameter",
+            ),
+            (
+                event.params["mode"] in (MODE_ADD, MODE_REMOVE)
+                and event.params.get("brokerid")
+                not in [broker.unit_id for broker in self.charm.state.brokers],
+                "invalid brokerid",
+            ),
         ]
 
         for check, msg in failure_conditions:
             if check:
                 event.fail(msg)
                 return
 
-        response, user_task_id = self.balancer_manager.rebalance(
-            mode=event.params["mode"], dryrun=event.params["dryrun"]
-        )
+        response, user_task_id = self.balancer_manager.rebalance(**event.params)
         logger.debug(f"rebalance - {vars(response)=}")
 
         if response.status_code != 200:

src/literals.py

@@ -173,6 +173,11 @@ def __eq__(self, value: object, /) -> bool:
 ]
 
 
+MODE_FULL = "full"
+MODE_ADD = "add"
+MODE_REMOVE = "remove"
+
+
 @dataclass
 class StatusLevel:
     """Status object helper."""

src/managers/balancer.py

@@ -7,12 +7,12 @@
 import json
 import logging
 import time
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any
 
 import requests
 
 from core.models import JSON
-from literals import BALANCER, BALANCER_TOPICS, STORAGE
+from literals import BALANCER, BALANCER_TOPICS, MODE_FULL, STORAGE
 
 if TYPE_CHECKING:
     from charm import KafkaCharm
@@ -22,8 +22,6 @@
 
 logger = logging.getLogger(__name__)
 
-RebalanceMode = Literal["full", "add", "remove"]
-
 
 class CruiseControlClient:
     """Client wrapper for CruiseControl."""
@@ -60,10 +58,14 @@ def post(self, endpoint: str, dryrun: bool = False, **kwargs) -> requests.Respon
             dryrun: flag to decide whether to return only proposals (True), or execute (False)
             **kwargs: any REST API query parameters provided by that endpoint
         """
+        payload = {"dryrun": str(dryrun)}
+        if brokerid := kwargs.get("brokerid", None) is not None:
+            payload |= {"brokerid": brokerid}
+
         r = requests.post(
             url=f"{self.address}/{endpoint}",
             auth=(self.username, self.password),
-            params=kwargs | {"dryrun": str(dryrun)} | self.default_params,
+            params=kwargs | payload | self.default_params,
         )
         logger.debug(f"POST {endpoint} - {vars(r)}")
 
@@ -178,15 +180,18 @@ def create_internal_topics(self) -> None:
                 )
                 logger.info(f"Created topic {topic}")
 
-    def rebalance(self, mode: str, dryrun: bool = True) -> tuple[requests.Response, str]:
+    def rebalance(
+        self, mode: str, dryrun: bool = True, brokerid: int | None = None
+    ) -> tuple[requests.Response, str]:
         """Triggers a full Kafka cluster partition rebalance.
 
         Returns:
             Tuple of requests.Response and string of the CruiseControl User-Task-ID for the rebalance
         """
+        mode = f"{mode}_broker" if mode != MODE_FULL else mode
         rebalance_request = self.cruise_control.post(
-            endpoint=mode, dryrun=dryrun
-        )  # FIXME: will need updating to include add/remove_broker when it works
+            endpoint=mode, dryrun=dryrun, brokerid=brokerid
+        )
 
         return (rebalance_request, rebalance_request.headers.get("User-Task-ID", ""))
 

tests/integration/test_balancer.py

@@ -13,6 +13,7 @@
 
 from .helpers import (
     APP_NAME,
+    KAFKA_CONTAINER,
     ZK_NAME,
     balancer_exporter_is_up,
     balancer_is_ready,
@@ -38,27 +39,23 @@ async def balancer_app(ops_test: OpsTest, request):
 class TestBalancer:
     @pytest.mark.abort_on_fail
     async def test_build_and_deploy(self, ops_test: OpsTest, kafka_charm, balancer_app):
-        await ops_test.model.add_machine(series="jammy")
-        machine_ids = await ops_test.model.get_machines()
 
         await asyncio.gather(
             ops_test.model.deploy(
                 kafka_charm,
                 application_name=APP_NAME,
                 num_units=1,
-                series="jammy",
-                to=machine_ids[0],
                 config={"roles": "broker,balancer" if balancer_app == APP_NAME else "broker"},
+                resources={"kafka-image": KAFKA_CONTAINER},
             ),
             ops_test.model.deploy(
-                ZK_NAME, channel="edge", application_name=ZK_NAME, num_units=1, series="jammy"
+                ZK_NAME, channel="3/edge", application_name=ZK_NAME, num_units=3, series="jammy"
             ),
             ops_test.model.deploy(
                 "kafka-test-app",
                 application_name=PRODUCER_APP,
                 channel="edge",
                 num_units=1,
-                series="jammy",
                 config={
                     "topic_name": "HOT-TOPIC",
                     "num_messages": 100000,
@@ -74,12 +71,15 @@ async def test_build_and_deploy(self, ops_test: OpsTest, kafka_charm, balancer_a
                 kafka_charm,
                 application_name=balancer_app,
                 num_units=1,
-                series="jammy",
                 config={"roles": balancer_app},
+                resources={"kafka-image": KAFKA_CONTAINER},
             )
 
         await ops_test.model.wait_for_idle(
-            apps=list({APP_NAME, ZK_NAME, balancer_app}), idle_period=30, timeout=3600
+            apps=list({APP_NAME, ZK_NAME, balancer_app}),
+            idle_period=30,
+            timeout=3600,
+            raise_on_error=False,
         )
         assert ops_test.model.applications[APP_NAME].status == "blocked"
         assert ops_test.model.applications[ZK_NAME].status == "active"
@@ -237,6 +237,132 @@ async def test_remove_unit_full_rebalance(self, ops_test: OpsTest, balancer_app)
             # verify that post-rebalance, surviving units increased replica counts
             assert int(value) < int(post_rebalance_replica_counts.get(key, 0))
 
+    @pytest.mark.abort_on_fail
+    async def test_add_unit_targeted_rebalance(self, ops_test: OpsTest, balancer_app):
+        await ops_test.model.applications[APP_NAME].add_units(
+            count=1  # up to 4, new unit won't have any partitions
+        )
+        await ops_test.model.block_until(
+            lambda: len(ops_test.model.applications[APP_NAME].units) == 4
+        )
+        await ops_test.model.wait_for_idle(
+            apps=list({APP_NAME, ZK_NAME, PRODUCER_APP, balancer_app}),
+            status="active",
+            timeout=1800,
+            idle_period=30,
+        )
+        async with ops_test.fast_forward(fast_interval="20s"):
+            await asyncio.sleep(120)  # ensure update-status adds broker-capacities if missed
+
+        assert balancer_is_ready(ops_test=ops_test, app_name=balancer_app)
+
+        await asyncio.sleep(30)  # let the API breathe after so many requests
+
+        # verify CC can find the new broker_id 3, with no replica partitions allocated
+        broker_replica_count = get_replica_count_by_broker_id(ops_test, balancer_app)
+        new_broker_id = max(map(int, broker_replica_count.keys()))
+        pre_rebalance_replica_counts = {
+            key: value for key, value in broker_replica_count.items() if key != str(new_broker_id)
+        }
+        new_broker_replica_count = int(broker_replica_count.get(str(new_broker_id), 0))
+
+        assert not new_broker_replica_count
+
+        for unit in ops_test.model.applications[balancer_app].units:
+            if await unit.is_leader_from_status():
+                leader_unit = unit
+
+        rebalance_action_dry_run = await leader_unit.run_action(
+            "rebalance", mode="add", brokerid=new_broker_id, dryrun=True, timeout=600, block=True
+        )
+        response = await rebalance_action_dry_run.wait()
+        assert response.results
+
+        rebalance_action = await leader_unit.run_action(
+            "rebalance",
+            mode="add",
+            brokerid=new_broker_id,
+            dryrun=False,
+            timeout=600,
+            block=True,
+        )
+        response = await rebalance_action.wait()
+        assert response.results
+
+        post_rebalance_replica_counts = get_replica_count_by_broker_id(ops_test, balancer_app)
+
+        # Partition only were moved from existing brokers to the new one
+        for existing_broker, previous_replica_count in pre_rebalance_replica_counts.items():
+            assert previous_replica_count >= post_rebalance_replica_counts.get(
+                str(existing_broker)
+            )
+
+        # New broker has partition(s)
+        assert int(
+            get_replica_count_by_broker_id(ops_test, balancer_app).get(str(new_broker_id), 0)
+        )  # replicas were successfully moved
+
+        # Total sum of partition conserved
+        assert sum(pre_rebalance_replica_counts.values()) == sum(
+            post_rebalance_replica_counts.values()
+        )
+
+    @pytest.mark.abort_on_fail
+    async def test_balancer_prepare_unit_removal(self, ops_test: OpsTest, balancer_app):
+        broker_replica_count = get_replica_count_by_broker_id(ops_test, balancer_app)
+        new_broker_id = max(map(int, broker_replica_count.keys()))
+
+        # storing the current replica counts of 0, 1, 2 - they will persist
+        pre_rebalance_replica_counts = {
+            key: value
+            for key, value in get_replica_count_by_broker_id(ops_test, balancer_app).items()
+            if key != str(new_broker_id)
+        }
+
+        for unit in ops_test.model.applications[balancer_app].units:
+            if await unit.is_leader_from_status():
+                leader_unit = unit
+
+        rebalance_action_dry_run = await leader_unit.run_action(
+            "rebalance",
+            mode="remove",
+            brokerid=new_broker_id,
+            dryrun=True,
+            timeout=600,
+            block=True,
+        )
+        response = await rebalance_action_dry_run.wait()
+        assert response.results
+
+        rebalance_action = await leader_unit.run_action(
+            "rebalance",
+            mode="remove",
+            brokerid=[new_broker_id],
+            dryrun=False,
+            timeout=600,
+            block=True,
+        )
+        response = await rebalance_action.wait()
+        assert response.results
+
+        post_rebalance_replica_counts = get_replica_count_by_broker_id(ops_test, balancer_app)
+
+        # Partition only were moved from the removed broker to the other ones
+        for existing_broker, previous_replica_count in pre_rebalance_replica_counts.items():
+            assert previous_replica_count <= post_rebalance_replica_counts.get(
+                str(existing_broker)
+            )
+
+        # Replicas were successfully moved
+        assert not int(
+            get_replica_count_by_broker_id(ops_test, balancer_app).get(str(new_broker_id), 0)
+        )
+
+        # Total sum of partition conserved
+        assert sum(pre_rebalance_replica_counts.values()) == sum(
+            post_rebalance_replica_counts.values()
+        )
+
     @pytest.mark.abort_on_fail
     async def test_tls(self, ops_test: OpsTest, balancer_app):
         # deploy and integrate tls

tests/integration/test_provider.py

@@ -243,7 +243,8 @@ async def test_connection_updated_on_tls_enabled(ops_test: OpsTest, app_charm: P
 
     # deploying tls
     tls_config = {"ca-common-name": "kafka"}
-    await ops_test.model.deploy(TLS_NAME, channel="edge", config=tls_config)
+    # FIXME (certs): Unpin the revision once the charm is fixed
+    await ops_test.model.deploy(TLS_NAME, channel="edge", config=tls_config, revision=163)
     await ops_test.model.wait_for_idle(
         apps=[TLS_NAME], idle_period=30, timeout=1800, status="active"
     )

tests/integration/test_tls.py

@@ -41,7 +41,8 @@ async def test_deploy_tls(ops_test: OpsTest, app_charm):
     tls_config = {"ca-common-name": "kafka"}
 
     await asyncio.gather(
-        ops_test.model.deploy(TLS_NAME, channel="edge", config=tls_config),
+        # FIXME (certs): Unpin the revision once the charm is fixed
+        ops_test.model.deploy(TLS_NAME, channel="edge", config=tls_config, revision=163),
         ops_test.model.deploy(ZK_NAME, channel="3/edge", num_units=3),
         ops_test.model.deploy(
             kafka_charm,

tests/unit/test_balancer.py

@@ -12,6 +12,7 @@
 
 from charm import KafkaCharm
 from literals import BALANCER_TOPICS, CHARM_KEY, CONTAINER, SUBSTRATE
+from managers.balancer import CruiseControlClient
 
 logger = logging.getLogger(__name__)
 
@@ -54,7 +55,7 @@ def harness():
     return harness
 
 
-def test_client_get_args(client):
+def test_client_get_args(client: CruiseControlClient):
     with patch("managers.balancer.requests.get") as patched_get:
         client.get("silmaril")
 
@@ -67,7 +68,7 @@ def test_client_get_args(client):
         assert kwargs["auth"] == ("Beren", "Luthien")
 
 
-def test_client_post_args(client):
+def test_client_post_args(client: CruiseControlClient):
     with patch("managers.balancer.requests.post") as patched_post:
         client.post("silmaril")
 
@@ -82,25 +83,25 @@ def test_client_post_args(client):
         assert kwargs["auth"] == ("Beren", "Luthien")
 
 
-def test_client_get_task_status(client, user_tasks):
+def test_client_get_task_status(client: CruiseControlClient, user_tasks: dict):
     with patch("managers.balancer.requests.get", return_value=MockResponse(user_tasks)):
         assert (
             client.get_task_status(user_task_id="e4256bcb-93f7-4290-ab11-804a665bf011")
             == "Completed"
         )
 
 
-def test_client_monitoring(client, state):
+def test_client_monitoring(client: CruiseControlClient, state: dict):
     with patch("managers.balancer.requests.get", return_value=MockResponse(state)):
         assert client.monitoring
 
 
-def test_client_executing(client, state):
+def test_client_executing(client: CruiseControlClient, state: dict):
     with patch("managers.balancer.requests.get", return_value=MockResponse(state)):
         assert not client.executing
 
 
-def test_client_ready(client, state):
+def test_client_ready(client: CruiseControlClient, state: dict):
     with patch("managers.balancer.requests.get", return_value=MockResponse(state)):
         assert client.ready
 
@@ -111,7 +112,7 @@ def test_client_ready(client, state):
         assert not client.ready
 
 
-def test_balancer_manager_create_internal_topics(harness):
+def test_balancer_manager_create_internal_topics(harness: Harness[KafkaCharm]):
     with (
         patch("core.models.PeerCluster.broker_uris", new_callable=PropertyMock, return_value=""),
         patch(
@@ -146,8 +147,14 @@ def test_balancer_manager_create_internal_topics(harness):
 @pytest.mark.parametrize("executing", [True, False])
 @pytest.mark.parametrize("ready", [True, False])
 @pytest.mark.parametrize("status", [200, 404])
-def test_balancer_manager_rebalance(
-    harness, proposal, leader, monitoring, executing, ready, status
+def test_balancer_manager_rebalance_full(
+    harness: Harness[KafkaCharm],
+    proposal: dict,
+    leader: bool,
+    monitoring: bool,
+    executing: bool,
+    ready: bool,
+    status: int,
 ):
     mock_event = MagicMock()
     mock_event.params = {"mode": "full", "dryrun": True}
@@ -189,7 +196,89 @@ def test_balancer_manager_rebalance(
             assert mock_event._mock_children.get("set_results")  # event.set_results was called
 
 
-def test_balancer_manager_clean_results(harness, proposal):
+@pytest.mark.parametrize("mode", ["add", "remove"])
+@pytest.mark.parametrize("brokerid", [None, 0])
+def test_rebalance_add_remove_broker_id_length(
+    harness: Harness[KafkaCharm], proposal: dict, mode: str, brokerid: int | None
+):
+    mock_event = MagicMock()
+    payload = {"mode": mode, "dryrun": True}
+    payload = payload | {"brokerid": brokerid} if brokerid is not None else payload
+    mock_event.params = payload
+
+    with (
+        harness.hooks_disabled(),
+        patch(
+            "managers.balancer.CruiseControlClient.monitoring",
+            new_callable=PropertyMock,
+            return_value=True,
+        ),
+        patch(
+            "managers.balancer.CruiseControlClient.executing",
+            new_callable=PropertyMock,
+            return_value=not True,
+        ),
+        patch(
+            "managers.balancer.CruiseControlClient.ready",
+            new_callable=PropertyMock,
+            return_value=True,
+        ),
+        patch(
+            "managers.balancer.BalancerManager.rebalance",
+            new_callable=None,
+            return_value=(MockResponse(content=proposal, status_code=200), "foo"),
+        ),
+        patch(
+            "managers.balancer.BalancerManager.wait_for_task",
+            new_callable=None,
+        ) as patched_wait_for_task,
+    ):
+        harness.set_leader(True)
+
+        # When
+        harness.charm.balancer.rebalance(mock_event)
+
+        # Then
+        if brokerid is None:
+            assert mock_event._mock_children.get("fail")  # event.fail was called
+        else:
+            assert patched_wait_for_task.call_count
+            assert mock_event._mock_children.get("set_results")  # event.set_results was called
+
+
+def test_rebalance_broker_id_not_found(harness: Harness[KafkaCharm]):
+    mock_event = MagicMock()
+    payload = {"mode": "add", "dryrun": True, "brokerid": 999}
+    mock_event.params = payload
+
+    with (
+        harness.hooks_disabled(),
+        patch(
+            "managers.balancer.CruiseControlClient.monitoring",
+            new_callable=PropertyMock,
+            return_value=True,
+        ),
+        patch(
+            "managers.balancer.CruiseControlClient.executing",
+            new_callable=PropertyMock,
+            return_value=not True,
+        ),
+        patch(
+            "managers.balancer.CruiseControlClient.ready",
+            new_callable=PropertyMock,
+            return_value=True,
+        ),
+    ):
+        harness.set_leader(True)
+
+        # When
+        harness.charm.balancer.rebalance(mock_event)
+
+        # Then
+        assert mock_event._mock_children.get("fail")  # event.fail was called
+
+
+def test_balancer_manager_clean_results(harness: Harness[KafkaCharm], proposal: dict):
     cleaned_results = harness.charm.balancer.balancer_manager.clean_results(value=proposal)
 
     def _check_cleaned_results(value) -> bool: