From 05833c77f39aa3691854568b8eff2f5bf0819b55 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Mon, 26 Aug 2024 12:58:54 +0300
Subject: [PATCH 01/85] Use syncobj lib directly

---
 poetry.lock    | 12 +++++++++++-
 pyproject.toml |  1 +
 src/cluster.py | 33 +++++++++++++++++----------------
 3 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index 952eeb9cb5..4888ef8369 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1626,6 +1626,16 @@ files = [
 [package.dependencies]
 pytz = "*"
 
+[[package]]
+name = "pysyncobj"
+version = "0.3.12"
+description = "A library for replicating your python class between multiple servers, based on raft protocol"
+optional = false
+python-versions = "*"
+files = [
+    {file = "pysyncobj-0.3.12.tar.gz", hash = "sha256:a4c82f080d3dcbf4a4fd627a6a0c466a60169297d3b7797acff9ff40cc44daf6"},
+]
+
 [[package]]
 name = "pytest"
 version = "8.3.2"
@@ -2380,4 +2390,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools",
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "02a4eff05b211c6940dfb9822fbf06f846034ec9cb8ffa34bce0333144a5ddde"
+content-hash = "1c2fa3ce3f1fbc3d95d1ce9d17001cddabdaa0b07c0d169d6ad8d57d9dc241a8"
diff --git a/pyproject.toml b/pyproject.toml
index d90a8e76c9..4a11b8e7f5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -18,6 +18,7 @@ pydantic = "^1.10.17"
 poetry-core = "^1.9.0"
 pyOpenSSL = "^24.2.1"
 jinja2 = "^3.1.4"
+pysyncobj = "^0.3.12"
 
 [tool.poetry.group.charm-libs.dependencies]
 # data_platform_libs/v0/data_interfaces.py
diff --git a/src/cluster.py b/src/cluster.py
index 0fe668fd3b..05c4e62172 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -15,6 +15,7 @@
 import requests
 from charms.operator_libs_linux.v2 import snap
 from jinja2 import Template
+from pysyncobj.utility import TcpUtility, UtilityException
 from tenacity import (
     AttemptManager,
     RetryError,
@@ -688,27 +689,27 @@ def remove_raft_member(self, member_ip: str) -> None:
                 is not part of the raft cluster.
         """
         # Get the status of the raft cluster.
-        raft_status = subprocess.check_output([
-            "charmed-postgresql.syncobj-admin",
-            "-conn",
-            "127.0.0.1:2222",
-            "-status",
-        ]).decode("UTF-8")
+        syncobj_util = TcpUtility(timeout=3)
+
+        raft_host = "127.0.0.1:2222"
+        try:
+            raft_status = syncobj_util.executeCommand(raft_host, ["status"])
+        except UtilityException:
+            logger.warning("Remove raft member: Cannot connect to raft cluster")
+            raise RemoveRaftMemberFailedError()
 
         # Check whether the member is still part of the raft cluster.
-        if not member_ip or member_ip not in raft_status:
+        if not member_ip or f"partner_node_status_server_{member_ip}:2222" not in raft_status:
             return
 
         # Remove the member from the raft cluster.
-        result = subprocess.check_output([
-            "charmed-postgresql.syncobj-admin",
-            "-conn",
-            "127.0.0.1:2222",
-            "-remove",
-            f"{member_ip}:2222",
-        ]).decode("UTF-8")
-
-        if "SUCCESS" not in result:
+        try:
+            result = syncobj_util.executeCommand(raft_host, ["remove", f"{member_ip}:2222"])
+        except UtilityException:
+            logger.debug("Remove raft member: Remove call failed")
+            raise RemoveRaftMemberFailedError()
+
+        if not result.startswith("SUCCESS"):
             raise RemoveRaftMemberFailedError()
 
     @retry(stop=stop_after_attempt(10), wait=wait_exponential(multiplier=1, min=2, max=10))

From 5753d1f595ecae7435802c89f859c72500b9a447 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Mon, 26 Aug 2024 13:54:33 +0300
Subject: [PATCH 02/85] Bump libs

---
 lib/charms/grafana_agent/v0/cos_agent.py |  8 ++++++--
 lib/charms/tempo_k8s/v1/charm_tracing.py | 17 ++++++++++++++---
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/lib/charms/grafana_agent/v0/cos_agent.py b/lib/charms/grafana_agent/v0/cos_agent.py
index 582b70c079..cc4da25a82 100644
--- a/lib/charms/grafana_agent/v0/cos_agent.py
+++ b/lib/charms/grafana_agent/v0/cos_agent.py
@@ -22,7 +22,7 @@
 Using the `COSAgentProvider` object only requires instantiating it,
 typically in the `__init__` method of your charm (the one which sends telemetry).
 
-The constructor of `COSAgentProvider` has only one required and nine optional parameters:
+The constructor of `COSAgentProvider` has only one required and ten optional parameters:
 
 ```python
     def __init__(
@@ -36,6 +36,7 @@ def __init__(
         log_slots: Optional[List[str]] = None,
         dashboard_dirs: Optional[List[str]] = None,
         refresh_events: Optional[List] = None,
+        tracing_protocols: Optional[List[str]] = None,
         scrape_configs: Optional[Union[List[Dict], Callable]] = None,
     ):
 ```
@@ -65,6 +66,8 @@ def __init__(
 
 - `refresh_events`: List of events on which to refresh relation data.
 
+- `tracing_protocols`: List of requested tracing protocols that the charm requires to send traces.
+
 - `scrape_configs`: List of standard scrape_configs dicts or a callable that returns the list in
     case the configs need to be generated dynamically. The contents of this list will be merged
     with the configs from `metrics_endpoints`.
@@ -108,6 +111,7 @@ def __init__(self, *args):
             log_slots=["my-app:slot"],
             dashboard_dirs=["./src/dashboards_1", "./src/dashboards_2"],
             refresh_events=["update-status", "upgrade-charm"],
+            tracing_protocols=["otlp_http", "otlp_grpc"],
             scrape_configs=[
                 {
                     "job_name": "custom_job",
@@ -249,7 +253,7 @@ class _MetricsEndpointDict(TypedDict):
 
 LIBID = "dc15fa84cef84ce58155fb84f6c6213a"
 LIBAPI = 0
-LIBPATCH = 10
+LIBPATCH = 11
 
 PYDEPS = ["cosl", "pydantic"]
 
diff --git a/lib/charms/tempo_k8s/v1/charm_tracing.py b/lib/charms/tempo_k8s/v1/charm_tracing.py
index fa926539bc..2dbdddd684 100644
--- a/lib/charms/tempo_k8s/v1/charm_tracing.py
+++ b/lib/charms/tempo_k8s/v1/charm_tracing.py
@@ -224,7 +224,6 @@ def _remove_stale_otel_sdk_packages():
 
 _remove_stale_otel_sdk_packages()
 
-
 import functools
 import inspect
 import logging
@@ -271,7 +270,7 @@ def _remove_stale_otel_sdk_packages():
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
 
-LIBPATCH = 14
+LIBPATCH = 15
 
 PYDEPS = ["opentelemetry-exporter-otlp-proto-http==1.21.0"]
 
@@ -281,7 +280,6 @@ def _remove_stale_otel_sdk_packages():
 # set this to 0 if you are debugging/developing this library source
 dev_logger.setLevel(logging.CRITICAL)
 
-
 _CharmType = Type[CharmBase]  # the type CharmBase and any subclass thereof
 _C = TypeVar("_C", bound=_CharmType)
 _T = TypeVar("_T", bound=type)
@@ -333,9 +331,22 @@ def _get_tracer() -> Optional[Tracer]:
     try:
         return tracer.get()
     except LookupError:
+        # fallback: this course-corrects for a user error where charm_tracing symbols are imported
+        # from different paths (typically charms.tempo_k8s... and lib.charms.tempo_k8s...)
         try:
             ctx: Context = copy_context()
             if context_tracer := _get_tracer_from_context(ctx):
+                logger.warning(
+                    "Tracer not found in `tracer` context var. "
+                    "Verify that you're importing all `charm_tracing` symbols from the same module path. \n"
+                    "For example, DO"
+                    ": `from charms.lib...charm_tracing import foo, bar`. \n"
+                    "DONT: \n"
+                    " \t - `from charms.lib...charm_tracing import foo` \n"
+                    " \t - `from lib...charm_tracing import bar` \n"
+                    "For more info: https://python-notes.curiousefficiency.org/en/latest/python"
+                    "_concepts/import_traps.html#the-double-import-trap"
+                )
                 return context_tracer.get()
             else:
                 return None

From e2b082eb92315cebfcb33e2a91ef2e098af140f0 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Mon, 26 Aug 2024 15:03:50 +0300
Subject: [PATCH 03/85] Add raft removal test

---
 tests/unit/test_cluster.py | 165 ++++++++++++++++++++++++++-----------
 1 file changed, 116 insertions(+), 49 deletions(-)

diff --git a/tests/unit/test_cluster.py b/tests/unit/test_cluster.py
index 8662b7b2af..02e1d60e65 100644
--- a/tests/unit/test_cluster.py
+++ b/tests/unit/test_cluster.py
@@ -1,19 +1,25 @@
 # Copyright 2021 Canonical Ltd.
 # See LICENSE file for licensing details.
 
-from unittest import TestCase
 from unittest.mock import MagicMock, Mock, PropertyMock, mock_open, patch, sentinel
 
 import pytest
-import requests as requests
-import tenacity as tenacity
+import requests
 from charms.operator_libs_linux.v2 import snap
 from jinja2 import Template
 from ops.testing import Harness
-from tenacity import RetryError, stop_after_delay, wait_fixed
+from pysyncobj.utility import UtilityException
+from tenacity import (
+    AttemptManager,
+    RetryCallState,
+    RetryError,
+    Retrying,
+    stop_after_delay,
+    wait_fixed,
+)
 
 from charm import PostgresqlOperatorCharm
-from cluster import Patroni
+from cluster import Patroni, RemoveRaftMemberFailedError
 from constants import (
     PATRONI_CONF_PATH,
     PATRONI_LOGS_PATH,
@@ -25,9 +31,6 @@
 PATRONI_SERVICE = "patroni"
 CREATE_CLUSTER_CONF_PATH = "/var/snap/charmed-postgresql/current/etc/postgresql/postgresql.conf"
 
-# used for assert functions
-tc = TestCase()
-
 
 # This method will be used by the mock to replace requests.get
 def mocked_requests_get(*args, **kwargs):
@@ -84,13 +87,13 @@ def patroni(harness, peers_ips):
 
 def test_get_alternative_patroni_url(peers_ips, patroni):
     # Mock tenacity attempt.
-    retry = tenacity.Retrying()
-    retry_state = tenacity.RetryCallState(retry, None, None, None)
-    attempt = tenacity.AttemptManager(retry_state)
+    retry = Retrying()
+    retry_state = RetryCallState(retry, None, None, None)
+    attempt = AttemptManager(retry_state)
 
     # Test the first URL that is returned (it should have the current unit IP).
     url = patroni._get_alternative_patroni_url(attempt)
-    tc.assertEqual(url, f"http://{patroni.unit_ip}:8008")
+    assert url == f"http://{patroni.unit_ip}:8008"
 
     # Test returning the other servers URLs.
     for attempt_number in range(attempt.retry_state.attempt_number + 1, len(peers_ips) + 2):
@@ -106,8 +109,9 @@ def test_get_member_ip(peers_ips, patroni):
     ):
         # Test error on trying to get the member IP.
         _get_alternative_patroni_url.side_effect = "http://server2"
-        with tc.assertRaises(tenacity.RetryError):
+        with pytest.raises(RetryError):
             patroni.get_member_ip(patroni.member_name)
+            assert False
 
         # Test using an alternative Patroni URL.
         _get_alternative_patroni_url.side_effect = [
@@ -116,17 +120,17 @@ def test_get_member_ip(peers_ips, patroni):
             "http://server1",
         ]
         ip = patroni.get_member_ip(patroni.member_name)
-        tc.assertEqual(ip, "1.1.1.1")
+        assert ip == "1.1.1.1"
 
         # Test using the current Patroni URL.
         _get_alternative_patroni_url.side_effect = ["http://server1"]
         ip = patroni.get_member_ip(patroni.member_name)
-        tc.assertEqual(ip, "1.1.1.1")
+        assert ip == "1.1.1.1"
 
         # Test when not having that specific member in the cluster.
         _get_alternative_patroni_url.side_effect = ["http://server1"]
         ip = patroni.get_member_ip("other-member-name")
-        tc.assertIsNone(ip)
+        assert ip is None
 
 
 def test_get_patroni_health(peers_ips, patroni):
@@ -143,12 +147,13 @@ def test_get_patroni_health(peers_ips, patroni):
         _stop_after_delay.assert_called_once_with(60)
         _wait_fixed.assert_called_once_with(7)
 
-        tc.assertEqual(health, {"state": "running"})
+        assert health == {"state": "running"}
 
         # Test when the Patroni API is not reachable.
         _patroni_url.return_value = "http://server2"
-        with tc.assertRaises(tenacity.RetryError):
+        with pytest.raises(RetryError):
             patroni.get_patroni_health()
+            assert False
 
 
 def test_get_postgresql_version(peers_ips, patroni):
@@ -161,7 +166,7 @@ def test_get_postgresql_version(peers_ips, patroni):
         ]
         version = patroni.get_postgresql_version()
 
-        tc.assertEqual(version, "14.0")
+        assert version == "14.0"
         _snap_client.assert_called_once_with()
         _get_installed_snaps.assert_called_once_with()
 
@@ -173,8 +178,9 @@ def test_get_primary(peers_ips, patroni):
     ):
         # Test error on trying to get the member IP.
         _get_alternative_patroni_url.side_effect = "http://server2"
-        with tc.assertRaises(tenacity.RetryError):
+        with pytest.raises(RetryError):
             patroni.get_primary(patroni.member_name)
+            assert False
 
         # Test using an alternative Patroni URL.
         _get_alternative_patroni_url.side_effect = [
@@ -183,17 +189,17 @@ def test_get_primary(peers_ips, patroni):
             "http://server1",
         ]
         primary = patroni.get_primary()
-        tc.assertEqual(primary, "postgresql-0")
+        assert primary == "postgresql-0"
 
         # Test using the current Patroni URL.
         _get_alternative_patroni_url.side_effect = ["http://server1"]
         primary = patroni.get_primary()
-        tc.assertEqual(primary, "postgresql-0")
+        assert primary == "postgresql-0"
 
         # Test requesting the primary in the unit name pattern.
         _get_alternative_patroni_url.side_effect = ["http://server1"]
         primary = patroni.get_primary(unit_name_pattern=True)
-        tc.assertEqual(primary, "postgresql/0")
+        assert primary == "postgresql/0"
 
 
 def test_is_creating_backup(peers_ips, patroni):
@@ -206,13 +212,13 @@ def test_is_creating_backup(peers_ips, patroni):
                 {"name": "postgresql-1", "tags": {"is_creating_backup": True}},
             ]
         }
-        tc.assertTrue(patroni.is_creating_backup)
+        assert patroni.is_creating_backup
 
         # Test when no member is creating a backup.
         response.json.return_value = {
             "members": [{"name": "postgresql-0"}, {"name": "postgresql-1"}]
         }
-        tc.assertFalse(patroni.is_creating_backup)
+        assert not patroni.is_creating_backup
 
 
 def test_is_replication_healthy(peers_ips, patroni):
@@ -223,7 +229,7 @@ def test_is_replication_healthy(peers_ips, patroni):
     ):
         # Test when replication is healthy.
         _get.return_value.status_code = 200
-        tc.assertTrue(patroni.is_replication_healthy)
+        assert patroni.is_replication_healthy
 
         # Test when replication is not healthy.
         _get.side_effect = [
@@ -231,27 +237,27 @@ def test_is_replication_healthy(peers_ips, patroni):
             MagicMock(status_code=200),
             MagicMock(status_code=503),
         ]
-        tc.assertFalse(patroni.is_replication_healthy)
+        assert not patroni.is_replication_healthy
 
 
 def test_is_member_isolated(peers_ips, patroni):
     with (
-        patch("cluster.stop_after_delay", return_value=tenacity.stop_after_delay(0)),
-        patch("cluster.wait_fixed", return_value=tenacity.wait_fixed(0)),
+        patch("cluster.stop_after_delay", return_value=stop_after_delay(0)),
+        patch("cluster.wait_fixed", return_value=wait_fixed(0)),
         patch("requests.get", side_effect=mocked_requests_get) as _get,
         patch("charm.Patroni._patroni_url", new_callable=PropertyMock) as _patroni_url,
     ):
         # Test when it wasn't possible to connect to the Patroni API.
         _patroni_url.return_value = "http://server3"
-        tc.assertFalse(patroni.is_member_isolated)
+        assert not patroni.is_member_isolated
 
         # Test when the member isn't isolated from the cluster.
         _patroni_url.return_value = "http://server1"
-        tc.assertFalse(patroni.is_member_isolated)
+        assert not patroni.is_member_isolated
 
         # Test when the member is isolated from the cluster.
         _patroni_url.return_value = "http://server4"
-        tc.assertTrue(patroni.is_member_isolated)
+        assert patroni.is_member_isolated
 
 
 def test_render_file(peers_ips, patroni):
@@ -275,7 +281,7 @@ def test_render_file(peers_ips, patroni):
             patroni.render_file(filename, "rendered-content", 0o640)
 
         # Check the rendered file is opened with "w+" mode.
-        tc.assertEqual(mock.call_args_list[0][0], (filename, "w+"))
+        assert mock.call_args_list[0][0] == (filename, "w+")
         # Ensure that the correct user is lookup up.
         _pwnam.assert_called_with("snap_daemon")
         # Ensure the file is chmod'd correctly.
@@ -336,7 +342,7 @@ def test_render_patroni_yml_file(peers_ips, patroni):
             patroni.render_patroni_yml_file()
 
         # Check the template is opened read-only in the call to open.
-        tc.assertEqual(mock.call_args_list[0][0], ("templates/patroni.yml.j2", "r"))
+        assert mock.call_args_list[0][0] == ("templates/patroni.yml.j2", "r")
         # Ensure the correct rendered template is sent to _render_file method.
         _render_file.assert_called_once_with(
             "/var/snap/charmed-postgresql/current/etc/patroni/patroni.yaml",
@@ -402,7 +408,7 @@ def test_member_replication_lag(peers_ips, patroni):
 
         # Test when the API call fails.
         _patroni_url.return_value = "http://server2"
-        with patch.object(tenacity.Retrying, "iter", Mock(side_effect=tenacity.RetryError(None))):
+        with patch.object(Retrying, "iter", Mock(side_effect=RetryError(None))):
             lag = patroni.member_replication_lag
             assert lag == "unknown"
 
@@ -443,8 +449,9 @@ def test_update_synchronous_node_count(peers_ips, patroni):
 
         # Test when the request fails.
         response.status_code = 500
-        with tc.assertRaises(RetryError):
+        with pytest.raises(RetryError):
             patroni.update_synchronous_node_count()
+            assert False
 
 
 def test_configure_patroni_on_unit(peers_ips, patroni):
@@ -476,8 +483,8 @@ def test_configure_patroni_on_unit(peers_ips, patroni):
 def test_member_started_true(peers_ips, patroni):
     with (
         patch("cluster.requests.get") as _get,
-        patch("cluster.stop_after_delay", return_value=tenacity.stop_after_delay(0)),
-        patch("cluster.wait_fixed", return_value=tenacity.wait_fixed(0)),
+        patch("cluster.stop_after_delay", return_value=stop_after_delay(0)),
+        patch("cluster.wait_fixed", return_value=wait_fixed(0)),
     ):
         _get.return_value.json.return_value = {"state": "running"}
 
@@ -489,8 +496,8 @@ def test_member_started_true(peers_ips, patroni):
 def test_member_started_false(peers_ips, patroni):
     with (
         patch("cluster.requests.get") as _get,
-        patch("cluster.stop_after_delay", return_value=tenacity.stop_after_delay(0)),
-        patch("cluster.wait_fixed", return_value=tenacity.wait_fixed(0)),
+        patch("cluster.stop_after_delay", return_value=stop_after_delay(0)),
+        patch("cluster.wait_fixed", return_value=wait_fixed(0)),
     ):
         _get.return_value.json.return_value = {"state": "stopped"}
 
@@ -502,8 +509,8 @@ def test_member_started_false(peers_ips, patroni):
 def test_member_started_error(peers_ips, patroni):
     with (
         patch("cluster.requests.get") as _get,
-        patch("cluster.stop_after_delay", return_value=tenacity.stop_after_delay(0)),
-        patch("cluster.wait_fixed", return_value=tenacity.wait_fixed(0)),
+        patch("cluster.stop_after_delay", return_value=stop_after_delay(0)),
+        patch("cluster.wait_fixed", return_value=wait_fixed(0)),
     ):
         _get.side_effect = Exception
 
@@ -515,8 +522,8 @@ def test_member_started_error(peers_ips, patroni):
 def test_member_inactive_true(peers_ips, patroni):
     with (
         patch("cluster.requests.get") as _get,
-        patch("cluster.stop_after_delay", return_value=tenacity.stop_after_delay(0)),
-        patch("cluster.wait_fixed", return_value=tenacity.wait_fixed(0)),
+        patch("cluster.stop_after_delay", return_value=stop_after_delay(0)),
+        patch("cluster.wait_fixed", return_value=wait_fixed(0)),
     ):
         _get.return_value.json.return_value = {"state": "stopped"}
 
@@ -528,8 +535,8 @@ def test_member_inactive_true(peers_ips, patroni):
 def test_member_inactive_false(peers_ips, patroni):
     with (
         patch("cluster.requests.get") as _get,
-        patch("cluster.stop_after_delay", return_value=tenacity.stop_after_delay(0)),
-        patch("cluster.wait_fixed", return_value=tenacity.wait_fixed(0)),
+        patch("cluster.stop_after_delay", return_value=stop_after_delay(0)),
+        patch("cluster.wait_fixed", return_value=wait_fixed(0)),
     ):
         _get.return_value.json.return_value = {"state": "starting"}
 
@@ -541,8 +548,8 @@ def test_member_inactive_false(peers_ips, patroni):
 def test_member_inactive_error(peers_ips, patroni):
     with (
         patch("cluster.requests.get") as _get,
-        patch("cluster.stop_after_delay", return_value=tenacity.stop_after_delay(0)),
-        patch("cluster.wait_fixed", return_value=tenacity.wait_fixed(0)),
+        patch("cluster.stop_after_delay", return_value=stop_after_delay(0)),
+        patch("cluster.wait_fixed", return_value=wait_fixed(0)),
     ):
         _get.side_effect = Exception
 
@@ -600,8 +607,9 @@ def test_get_patroni_restart_condition(patroni):
 
         # Test when there is no restart condition set.
         _open.return_value.__enter__.return_value.read.return_value = ""
-        with tc.assertRaises(RuntimeError):
+        with pytest.raises(RuntimeError):
             patroni.get_patroni_restart_condition()
+            assert False
 
 
 @pytest.mark.parametrize("new_restart_condition", ["on-success", "on-failure"])
@@ -615,3 +623,62 @@ def test_update_patroni_restart_condition(patroni, new_restart_condition):
             f"Restart={new_restart_condition}"
         )
         _run.assert_called_once_with(["/bin/systemctl", "daemon-reload"])
+
+
+def test_remove_raft_member(patroni):
+    with patch("cluster.TcpUtility") as _tcp_utility:
+        # Member already removed
+        _tcp_utility.return_value.executeCommand.return_value = ""
+
+        patroni.remove_raft_member("1.2.3.4")
+
+        _tcp_utility.assert_called_once_with(timeout=3)
+        _tcp_utility.return_value.executeCommand.assert_called_once_with(
+            "127.0.0.1:2222", ["status"]
+        )
+        _tcp_utility.reset_mock()
+
+        # Removing member
+        _tcp_utility.return_value.executeCommand.side_effect = [
+            "partner_node_status_server_1.2.3.4:2222",
+            "SUCCESS",
+        ]
+
+        patroni.remove_raft_member("1.2.3.4")
+
+        _tcp_utility.assert_called_once_with(timeout=3)
+        assert _tcp_utility.return_value.executeCommand.call_count == 2
+        _tcp_utility.return_value.executeCommand.assert_any_call("127.0.0.1:2222", ["status"])
+        _tcp_utility.return_value.executeCommand.assert_any_call(
+            "127.0.0.1:2222", ["remove", "1.2.3.4:2222"]
+        )
+        _tcp_utility.reset_mock()
+
+        # Raises on failed status
+        _tcp_utility.return_value.executeCommand.side_effect = [
+            "partner_node_status_server_1.2.3.4:2222",
+            "FAIL",
+        ]
+
+        with pytest.raises(RemoveRaftMemberFailedError):
+            patroni.remove_raft_member("1.2.3.4")
+            assert False
+
+        # Raises on remove error
+        _tcp_utility.return_value.executeCommand.side_effect = [
+            "partner_node_status_server_1.2.3.4:2222",
+            UtilityException,
+        ]
+
+        with pytest.raises(RemoveRaftMemberFailedError):
+            patroni.remove_raft_member("1.2.3.4")
+            assert False
+
+        # Raises on status error
+        _tcp_utility.return_value.executeCommand.side_effect = [
+            UtilityException,
+        ]
+
+        with pytest.raises(RemoveRaftMemberFailedError):
+            patroni.remove_raft_member("1.2.3.4")
+            assert False

From 538f721bf2c6200c1e863165ea0973fb468efd90 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 3 Sep 2024 21:46:17 +0300
Subject: [PATCH 04/85] Reinitialise RAFT WIP

---
 src/cluster.py             | 63 +++++++++++++++++++++++++++++++++++++-
 tests/unit/test_cluster.py |  6 ++--
 2 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/src/cluster.py b/src/cluster.py
index 6284ee8ec5..4a9aa04107 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -9,7 +9,9 @@
 import os
 import pwd
 import re
+import shutil
 import subprocess
+from pathlib import Path
 from typing import Any, Dict, List, Optional, Set
 
 import requests
@@ -51,6 +53,10 @@
 RUNNING_STATES = ["running", "streaming"]
 
 
+class RaftNotPromotedError(Exception):
+    """Raised when a cluster is not promoted."""
+
+
 class ClusterNotPromotedError(Exception):
     """Raised when a cluster is not promoted."""
 
@@ -558,6 +564,7 @@ def render_patroni_yml_file(
         pitr_target: Optional[str] = None,
         restore_to_latest: bool = False,
         parameters: Optional[dict[str, str]] = None,
+        no_peers: bool = False,
     ) -> None:
         """Render the Patroni configuration file.
 
@@ -572,6 +579,7 @@ def render_patroni_yml_file(
             pitr_target: point-in-time-recovery target for the backup.
             restore_to_latest: restore all the WAL transaction logs from the stanza.
             parameters: PostgreSQL parameters to be added to the postgresql.conf file.
+            no_peers: Don't include peers.
         """
         # Open the template patroni.yml file.
         with open("templates/patroni.yml.j2", "r") as file:
@@ -587,7 +595,7 @@ def render_patroni_yml_file(
             enable_tls=enable_tls,
             member_name=self.member_name,
             partner_addrs=self.charm.async_replication.get_partner_addresses(),
-            peers_ips=self.peers_ips,
+            peers_ips=self.peers_ips if not no_peers else set(),
             pgbackrest_configuration_file=PGBACKREST_CONFIGURATION_FILE,
             scope=self.cluster_name,
             self_ip=self.unit_ip,
@@ -711,6 +719,53 @@ def primary_changed(self, old_primary: str) -> bool:
         primary = self.get_primary()
         return primary != old_primary
 
+    def remove_raft_data(self) -> None:
+        """Stops Patroni and removes the raft journals."""
+        logger.info("Stopping postgresql")
+        self.stop_patroni()
+
+        logger.info("Removing raft data")
+        try:
+            path = Path(f"{PATRONI_CONF_PATH}/raft")
+            if path.exists() and path.is_dir():
+                shutil.rmtree(path)
+        except OSError as e:
+            raise Exception(f"Failed to remove previous cluster information with error: {str(e)}")
+        logger.info("Raft ready to reinitialise")
+
+    def reinitialise_raft_data(self) -> None:
+        """Reinitialise the raft journals and promoting the unit to leader. Should only be run on sync replicas."""
+        # TODO remove magic sleep
+        import time
+
+        time.sleep(30)
+
+        logger.info("Rerendering patroni config without peers")
+        self.render_patroni_yml_file(no_peers=True)
+        logger.info("Starting patroni")
+        self.start_patroni()
+
+        logger.info("Waiting for new raft cluster to initialise")
+        for attempt in Retrying(wait=wait_fixed(5)):
+            with attempt:
+                cluster_status = requests.get(
+                    f"{self._patroni_url}/{PATRONI_CLUSTER_STATUS_ENDPOINT}",
+                    verify=self.verify,
+                    timeout=API_REQUEST_TIMEOUT,
+                    auth=self._patroni_auth,
+                )
+                found_primary = False
+                # logger.error(f"@@@@@@@@@@@@@@@@@@@@@@2{cluster_status.json()}")
+                for member in cluster_status.json()["members"]:
+                    if member["role"] == "leader":
+                        if member["host"] != self.unit_ip:
+                            raise RaftNotPromotedError()
+                if not found_primary:
+                    raise RaftNotPromotedError()
+
+        logger.info("Restarting patroni")
+        self.restart_patroni()
+
     def remove_raft_member(self, member_ip: str) -> None:
         """Remove a member from the raft cluster.
 
@@ -736,6 +791,12 @@ def remove_raft_member(self, member_ip: str) -> None:
         if not member_ip or f"partner_node_status_server_{member_ip}:2222" not in raft_status:
             return
 
+        # If there's no quorum and the leader left raft cluster is stuck
+        # logger.error(f"!!!!!!!!!!!!!!!!!!!!!11{raft_status}")
+        if not raft_status["has_quorum"] and raft_status["leader"].host == member_ip:
+            logger.warning("Remove raft member: Stuck raft cluster detected")
+            self.remove_raft_data()
+            self.reinitialise_raft_data()
         # Remove the member from the raft cluster.
         try:
             result = syncobj_util.executeCommand(raft_host, ["remove", f"{member_ip}:2222"])
diff --git a/tests/unit/test_cluster.py b/tests/unit/test_cluster.py
index 72781aa300..21b84bc67f 100644
--- a/tests/unit/test_cluster.py
+++ b/tests/unit/test_cluster.py
@@ -666,7 +666,7 @@ def test_remove_raft_member(patroni):
 
         # Removing member
         _tcp_utility.return_value.executeCommand.side_effect = [
-            "partner_node_status_server_1.2.3.4:2222",
+            {"partner_node_status_server_1.2.3.4:2222": 0, "has_quorum": True},
             "SUCCESS",
         ]
 
@@ -682,7 +682,7 @@ def test_remove_raft_member(patroni):
 
         # Raises on failed status
         _tcp_utility.return_value.executeCommand.side_effect = [
-            "partner_node_status_server_1.2.3.4:2222",
+            {"partner_node_status_server_1.2.3.4:2222": 0, "has_quorum": True},
             "FAIL",
         ]
 
@@ -692,7 +692,7 @@ def test_remove_raft_member(patroni):
 
         # Raises on remove error
         _tcp_utility.return_value.executeCommand.side_effect = [
-            "partner_node_status_server_1.2.3.4:2222",
+            {"partner_node_status_server_1.2.3.4:2222": 0, "has_quorum": True},
             UtilityException,
         ]
 

From 29aaf1827400fca9daa18340cbaee161cad00509 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 3 Sep 2024 22:08:11 +0300
Subject: [PATCH 05/85] Initial scaling test

---
 tests/integration/ha_tests/test_scaling.py | 65 ++++++++++++++++++++++
 1 file changed, 65 insertions(+)
 create mode 100644 tests/integration/ha_tests/test_scaling.py

diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
new file mode 100644
index 0000000000..3a15c9c5fc
--- /dev/null
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+# Copyright 2023 Canonical Ltd.
+# See LICENSE file for licensing details.
+import logging
+from asyncio import gather
+
+import pytest
+from pytest_operator.plugin import OpsTest
+
+from .. import markers
+from ..helpers import (
+    CHARM_SERIES,
+)
+from .conftest import APPLICATION_NAME
+from .helpers import (
+    app_name,
+    are_writes_increasing,
+    get_primary,
+    start_continuous_writes,
+)
+
+logger = logging.getLogger(__name__)
+
+charm = None
+
+
+@pytest.mark.group(1)
+@markers.juju3
+@pytest.mark.abort_on_fail
+async def test_build_and_deploy(ops_test: OpsTest) -> None:
+    """Build and deploy two PostgreSQL clusters."""
+    # This is a potentially destructive test, so it shouldn't be run against existing clusters
+    charm = await ops_test.build_charm(".")
+    async with ops_test.fast_forward():
+        # Deploy the first cluster with reusable storage
+        await gather(
+            ops_test.model.deploy(
+                charm,
+                num_units=2,
+                series=CHARM_SERIES,
+                config={"profile": "testing"},
+            ),
+            ops_test.model.deploy(
+                APPLICATION_NAME,
+                application_name=APPLICATION_NAME,
+                series=CHARM_SERIES,
+                channel="edge",
+            ),
+        )
+
+        await ops_test.model.wait_for_idle(status="active", timeout=1500)
+
+
+@pytest.mark.group(1)
+@markers.juju3
+async def test_removing_primary(ops_test: OpsTest, continuous_writes) -> None:
+    # Start an application that continuously writes data to the database.
+    app = await app_name(ops_test)
+    await start_continuous_writes(ops_test, app)
+    primary = await get_primary(ops_test, app)
+    await ops_test.model.destroy_unit(primary, force=True, destroy_storage=False, max_wait=1500)
+
+    await ops_test.model.wait_for_idle(status="active", timeout=600)
+
+    await are_writes_increasing(ops_test, primary)

From d050e6eea9f6cdb558f3dd25387cc80200772b50 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 3 Sep 2024 23:21:45 +0300
Subject: [PATCH 06/85] Endless loop

---
 src/cluster.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/cluster.py b/src/cluster.py
index 4a9aa04107..72ac2f4aed 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -760,6 +760,7 @@ def reinitialise_raft_data(self) -> None:
                     if member["role"] == "leader":
                         if member["host"] != self.unit_ip:
                             raise RaftNotPromotedError()
+                        found_primary = True
                 if not found_primary:
                     raise RaftNotPromotedError()
 

From 7626c9901f44dbcf959947389f45d68720dd6a95 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 4 Sep 2024 14:07:24 +0300
Subject: [PATCH 07/85] Try to do fast shutdown

---
 src/cluster.py | 39 +++++++++++++++++++++++----------------
 1 file changed, 23 insertions(+), 16 deletions(-)

diff --git a/src/cluster.py b/src/cluster.py
index 72ac2f4aed..6bb571b49a 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -53,6 +53,10 @@
 RUNNING_STATES = ["running", "streaming"]
 
 
+class RaftPostgresqlStillUpError(Exception):
+    """Raised when a cluster is not promoted."""
+
+
 class RaftNotPromotedError(Exception):
     """Raised when a cluster is not promoted."""
 
@@ -722,6 +726,18 @@ def primary_changed(self, old_primary: str) -> bool:
     def remove_raft_data(self) -> None:
         """Stops Patroni and removes the raft journals."""
         logger.info("Stopping postgresql")
+        subprocess.run([
+            "sudo",
+            "-u",
+            "snap_daemon",
+            "charmed-postgresql.pg-ctl",
+            "stop",
+            "-m",
+            "f",
+            "-D",
+            POSTGRESQL_DATA_PATH,
+        ])
+        logger.info("Stopping patroni")
         self.stop_patroni()
 
         logger.info("Removing raft data")
@@ -735,11 +751,6 @@ def remove_raft_data(self) -> None:
 
     def reinitialise_raft_data(self) -> None:
         """Reinitialise the raft journals and promoting the unit to leader. Should only be run on sync replicas."""
-        # TODO remove magic sleep
-        import time
-
-        time.sleep(30)
-
         logger.info("Rerendering patroni config without peers")
         self.render_patroni_yml_file(no_peers=True)
         logger.info("Starting patroni")
@@ -748,20 +759,17 @@ def reinitialise_raft_data(self) -> None:
         logger.info("Waiting for new raft cluster to initialise")
         for attempt in Retrying(wait=wait_fixed(5)):
             with attempt:
-                cluster_status = requests.get(
-                    f"{self._patroni_url}/{PATRONI_CLUSTER_STATUS_ENDPOINT}",
+                response = requests.get(
+                    f"{self._patroni_url}/health",
                     verify=self.verify,
                     timeout=API_REQUEST_TIMEOUT,
                     auth=self._patroni_auth,
                 )
-                found_primary = False
-                # logger.error(f"@@@@@@@@@@@@@@@@@@@@@@2{cluster_status.json()}")
-                for member in cluster_status.json()["members"]:
-                    if member["role"] == "leader":
-                        if member["host"] != self.unit_ip:
-                            raise RaftNotPromotedError()
-                        found_primary = True
-                if not found_primary:
+                health_status = response.json()
+                if (
+                    health_status["role"] not in ["leader", "master"]
+                    or health_status["state"] != "running"
+                ):
                     raise RaftNotPromotedError()
 
         logger.info("Restarting patroni")
@@ -793,7 +801,6 @@ def remove_raft_member(self, member_ip: str) -> None:
             return
 
         # If there's no quorum and the leader left raft cluster is stuck
-        # logger.error(f"!!!!!!!!!!!!!!!!!!!!!11{raft_status}")
         if not raft_status["has_quorum"] and raft_status["leader"].host == member_ip:
             logger.warning("Remove raft member: Stuck raft cluster detected")
             self.remove_raft_data()

From 774ea9504251b876ee500d44c479c5de80494cc0 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 5 Sep 2024 16:54:56 +0300
Subject: [PATCH 08/85] Correct rerendering of yaml

---
 poetry.lock    | 31 ++++++++++++++++++++++++++++++-
 pyproject.toml |  1 +
 src/charm.py   |  6 +++++-
 src/cluster.py | 28 ++++++++++++++--------------
 4 files changed, 50 insertions(+), 16 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index 649a6de14f..26e64a03aa 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1317,6 +1317,35 @@ files = [
     {file = "protobuf-4.25.4.tar.gz", hash = "sha256:0dc4a62cc4052a036ee2204d26fe4d835c62827c855c8a03f29fe6da146b380d"},
 ]
 
+[[package]]
+name = "psutil"
+version = "6.0.0"
+description = "Cross-platform lib for process and system monitoring in Python."
+optional = false
+python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,>=2.7"
+files = [
+    {file = "psutil-6.0.0-cp27-cp27m-macosx_10_9_x86_64.whl", hash = "sha256:a021da3e881cd935e64a3d0a20983bda0bb4cf80e4f74fa9bfcb1bc5785360c6"},
+    {file = "psutil-6.0.0-cp27-cp27m-manylinux2010_i686.whl", hash = "sha256:1287c2b95f1c0a364d23bc6f2ea2365a8d4d9b726a3be7294296ff7ba97c17f0"},
+    {file = "psutil-6.0.0-cp27-cp27m-manylinux2010_x86_64.whl", hash = "sha256:a9a3dbfb4de4f18174528d87cc352d1f788b7496991cca33c6996f40c9e3c92c"},
+    {file = "psutil-6.0.0-cp27-cp27mu-manylinux2010_i686.whl", hash = "sha256:6ec7588fb3ddaec7344a825afe298db83fe01bfaaab39155fa84cf1c0d6b13c3"},
+    {file = "psutil-6.0.0-cp27-cp27mu-manylinux2010_x86_64.whl", hash = "sha256:1e7c870afcb7d91fdea2b37c24aeb08f98b6d67257a5cb0a8bc3ac68d0f1a68c"},
+    {file = "psutil-6.0.0-cp27-none-win32.whl", hash = "sha256:02b69001f44cc73c1c5279d02b30a817e339ceb258ad75997325e0e6169d8b35"},
+    {file = "psutil-6.0.0-cp27-none-win_amd64.whl", hash = "sha256:21f1fb635deccd510f69f485b87433460a603919b45e2a324ad65b0cc74f8fb1"},
+    {file = "psutil-6.0.0-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:c588a7e9b1173b6e866756dde596fd4cad94f9399daf99ad8c3258b3cb2b47a0"},
+    {file = "psutil-6.0.0-cp36-abi3-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:6ed2440ada7ef7d0d608f20ad89a04ec47d2d3ab7190896cd62ca5fc4fe08bf0"},
+    {file = "psutil-6.0.0-cp36-abi3-manylinux_2_12_x86_64.manylinux2010_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5fd9a97c8e94059b0ef54a7d4baf13b405011176c3b6ff257c247cae0d560ecd"},
+    {file = "psutil-6.0.0-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e2e8d0054fc88153ca0544f5c4d554d42e33df2e009c4ff42284ac9ebdef4132"},
+    {file = "psutil-6.0.0-cp36-cp36m-win32.whl", hash = "sha256:fc8c9510cde0146432bbdb433322861ee8c3efbf8589865c8bf8d21cb30c4d14"},
+    {file = "psutil-6.0.0-cp36-cp36m-win_amd64.whl", hash = "sha256:34859b8d8f423b86e4385ff3665d3f4d94be3cdf48221fbe476e883514fdb71c"},
+    {file = "psutil-6.0.0-cp37-abi3-win32.whl", hash = "sha256:a495580d6bae27291324fe60cea0b5a7c23fa36a7cd35035a16d93bdcf076b9d"},
+    {file = "psutil-6.0.0-cp37-abi3-win_amd64.whl", hash = "sha256:33ea5e1c975250a720b3a6609c490db40dae5d83a4eb315170c4fe0d8b1f34b3"},
+    {file = "psutil-6.0.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:ffe7fc9b6b36beadc8c322f84e1caff51e8703b88eee1da46d1e3a6ae11b4fd0"},
+    {file = "psutil-6.0.0.tar.gz", hash = "sha256:8faae4f310b6d969fa26ca0545338b21f73c6b15db7c4a8d934a5482faa818f2"},
+]
+
+[package.extras]
+test = ["enum34", "ipaddress", "mock", "pywin32", "wmi"]
+
 [[package]]
 name = "psycopg2"
 version = "2.9.9"
@@ -2409,4 +2438,4 @@ type = ["pytest-mypy"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "2b18f22f59c94638fd68e049855db6a027c893c3836bf430346078c9965bec70"
+content-hash = "2e829c2eaa9961b251ed682329f5ee91058e636d721d5cec98185bd7a735b63f"
diff --git a/pyproject.toml b/pyproject.toml
index 863828cec5..6da5792060 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -19,6 +19,7 @@ poetry-core = "^1.9.0"
 pyOpenSSL = "^24.2.1"
 jinja2 = "^3.1.4"
 pysyncobj = "^0.3.12"
+psutil = "^6.0.0"
 
 [tool.poetry.group.charm-libs.dependencies]
 # data_platform_libs/v0/data_interfaces.py
diff --git a/src/charm.py b/src/charm.py
index f0a119d5db..fb7bc91621 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -1619,7 +1619,7 @@ def _can_connect_to_postgresql(self) -> bool:
             return False
         return True
 
-    def update_config(self, is_creating_backup: bool = False) -> bool:
+    def update_config(self, is_creating_backup: bool = False, no_peers: bool = False) -> bool:
         """Updates Patroni config file based on the existence of the TLS files."""
         enable_tls = self.is_tls_enabled
         limit_memory = None
@@ -1645,7 +1645,11 @@ def update_config(self, is_creating_backup: bool = False) -> bool:
                 self.app_peer_data.get("require-change-bucket-after-restore", None)
             ),
             parameters=pg_parameters,
+            no_peers=no_peers,
         )
+        if no_peers:
+            return True
+
         if not self._is_workload_running:
             # If Patroni/PostgreSQL has not started yet and TLS relations was initialised,
             # then mark TLS as enabled. This commonly happens when the charm is deployed
diff --git a/src/cluster.py b/src/cluster.py
index 6bb571b49a..9ee87b2c96 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -14,6 +14,7 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Set
 
+import psutil
 import requests
 from charms.operator_libs_linux.v2 import snap
 from jinja2 import Template
@@ -598,7 +599,9 @@ def render_patroni_yml_file(
             data_path=POSTGRESQL_DATA_PATH,
             enable_tls=enable_tls,
             member_name=self.member_name,
-            partner_addrs=self.charm.async_replication.get_partner_addresses(),
+            partner_addrs=self.charm.async_replication.get_partner_addresses()
+            if not no_peers
+            else [],
             peers_ips=self.peers_ips if not no_peers else set(),
             pgbackrest_configuration_file=PGBACKREST_CONFIGURATION_FILE,
             scope=self.cluster_name,
@@ -725,21 +728,16 @@ def primary_changed(self, old_primary: str) -> bool:
 
     def remove_raft_data(self) -> None:
         """Stops Patroni and removes the raft journals."""
-        logger.info("Stopping postgresql")
-        subprocess.run([
-            "sudo",
-            "-u",
-            "snap_daemon",
-            "charmed-postgresql.pg-ctl",
-            "stop",
-            "-m",
-            "f",
-            "-D",
-            POSTGRESQL_DATA_PATH,
-        ])
         logger.info("Stopping patroni")
         self.stop_patroni()
 
+        logger.info("Wait for postgresql to stop")
+        for attempt in Retrying(wait=wait_fixed(5)):
+            with attempt:
+                for proc in psutil.process_iter(["name"]):
+                    if proc.name() == "postgres":
+                        raise RaftPostgresqlStillUpError()
+
         logger.info("Removing raft data")
         try:
             path = Path(f"{PATRONI_CONF_PATH}/raft")
@@ -752,7 +750,7 @@ def remove_raft_data(self) -> None:
     def reinitialise_raft_data(self) -> None:
         """Reinitialise the raft journals and promoting the unit to leader. Should only be run on sync replicas."""
         logger.info("Rerendering patroni config without peers")
-        self.render_patroni_yml_file(no_peers=True)
+        self.charm.update_config(no_peers=True)
         logger.info("Starting patroni")
         self.start_patroni()
 
@@ -774,6 +772,7 @@ def reinitialise_raft_data(self) -> None:
 
         logger.info("Restarting patroni")
         self.restart_patroni()
+        logger.info("Raft should be unstuck")
 
     def remove_raft_member(self, member_ip: str) -> None:
         """Remove a member from the raft cluster.
@@ -805,6 +804,7 @@ def remove_raft_member(self, member_ip: str) -> None:
             logger.warning("Remove raft member: Stuck raft cluster detected")
             self.remove_raft_data()
             self.reinitialise_raft_data()
+
         # Remove the member from the raft cluster.
         try:
             result = syncobj_util.executeCommand(raft_host, ["remove", f"{member_ip}:2222"])

From 76bbe84ea75464ce4b1685ae9100b8af766e7013 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 5 Sep 2024 16:59:14 +0300
Subject: [PATCH 09/85] Unit tests

---
 tests/unit/test_charm.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/unit/test_charm.py b/tests/unit/test_charm.py
index 761322d95a..6eaab30cec 100644
--- a/tests/unit/test_charm.py
+++ b/tests/unit/test_charm.py
@@ -1288,6 +1288,7 @@ def test_update_config(harness):
             restore_to_latest=False,
             disable_pgbackrest_archiving=False,
             parameters={"test": "test"},
+            no_peers=False,
         )
         _handle_postgresql_restart_need.assert_called_once_with(False)
         assert "tls" not in harness.get_relation_data(rel_id, harness.charm.unit.name)
@@ -1311,6 +1312,7 @@ def test_update_config(harness):
             restore_to_latest=False,
             disable_pgbackrest_archiving=False,
             parameters={"test": "test"},
+            no_peers=False,
         )
         _handle_postgresql_restart_need.assert_called_once()
         assert "tls" not in harness.get_relation_data(

From bdae203f07e2c0b136cbf9038a856bdc1c1d438f Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 5 Sep 2024 17:35:49 +0300
Subject: [PATCH 10/85] Ignore connection error

---
 src/charm.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/charm.py b/src/charm.py
index fb7bc91621..87a18ad94c 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -427,6 +427,11 @@ def _on_peer_relation_departed(self, event: RelationDepartedEvent) -> None:
             )
             event.defer()
             return
+        except ConnectionError:
+            logger.warning(
+                "Early on_peer_relation_departed: unable to connect to the departing unit"
+            )
+            return
 
         # Allow leader to update the cluster members.
         if not self.unit.is_leader():

From fb922e506fd85dc04ef9c8c15865387f12fe3b50 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 5 Sep 2024 18:41:27 +0300
Subject: [PATCH 11/85] Catch exception

---
 src/charm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/charm.py b/src/charm.py
index 87a18ad94c..d028d6dd1c 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -427,7 +427,7 @@ def _on_peer_relation_departed(self, event: RelationDepartedEvent) -> None:
             )
             event.defer()
             return
-        except ConnectionError:
+        except RetryError:
             logger.warning(
                 "Early on_peer_relation_departed: unable to connect to the departing unit"
             )

From e0bbb0887861fe39cf0e98f0384447558cf9885f Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 5 Sep 2024 19:23:43 +0300
Subject: [PATCH 12/85] Sync replica stereo test

---
 tests/integration/ha_tests/test_scaling.py | 31 +++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index 3a15c9c5fc..c77c0823bb 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -10,6 +10,7 @@
 from .. import markers
 from ..helpers import (
     CHARM_SERIES,
+    DATABASE_APP_NAME,
 )
 from .conftest import APPLICATION_NAME
 from .helpers import (
@@ -36,6 +37,7 @@ async def test_build_and_deploy(ops_test: OpsTest) -> None:
         await gather(
             ops_test.model.deploy(
                 charm,
+                application_name=DATABASE_APP_NAME,
                 num_units=2,
                 series=CHARM_SERIES,
                 config={"profile": "testing"},
@@ -53,13 +55,40 @@ async def test_build_and_deploy(ops_test: OpsTest) -> None:
 
 @pytest.mark.group(1)
 @markers.juju3
-async def test_removing_primary(ops_test: OpsTest, continuous_writes) -> None:
+async def test_removing_stereo_primary(ops_test: OpsTest, continuous_writes) -> None:
     # Start an application that continuously writes data to the database.
     app = await app_name(ops_test)
     await start_continuous_writes(ops_test, app)
+    logger.info("Deleting primary")
     primary = await get_primary(ops_test, app)
     await ops_test.model.destroy_unit(primary, force=True, destroy_storage=False, max_wait=1500)
 
     await ops_test.model.wait_for_idle(status="active", timeout=600)
 
     await are_writes_increasing(ops_test, primary)
+
+    logger.info("Scaling back up")
+    await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=1)
+    await ops_test.model.wait_for_idle(status="active", timeout=1500)
+
+
+@pytest.mark.group(1)
+@markers.juju3
+async def test_removing_stereo_sync_replica(ops_test: OpsTest, continuous_writes) -> None:
+    # Start an application that continuously writes data to the database.
+    app = await app_name(ops_test)
+    await start_continuous_writes(ops_test, app)
+    logger.info("Deleting sync replica")
+    primary = await get_primary(ops_test, app)
+    secondary = next(
+        filter(lambda x: x.name != primary, ops_test.model.applications[DATABASE_APP_NAME].units)
+    )
+    await ops_test.model.destroy_unit(secondary, force=True, destroy_storage=False, max_wait=1500)
+
+    await ops_test.model.wait_for_idle(status="active", timeout=600)
+
+    await are_writes_increasing(ops_test, primary)
+
+    logger.info("Scaling back up")
+    await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=1)
+    await ops_test.model.wait_for_idle(status="active", timeout=1500)

From 895a8c4d4bce6fcf59ab38d332a44d92f9a1a175 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 5 Sep 2024 19:56:42 +0300
Subject: [PATCH 13/85] Fix test

---
 tests/integration/ha_tests/test_scaling.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index c77c0823bb..166f96e82f 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-# Copyright 2023 Canonical Ltd.
+# Copyright 2024 Canonical Ltd.
 # See LICENSE file for licensing details.
 import logging
 from asyncio import gather
@@ -74,7 +74,7 @@ async def test_removing_stereo_primary(ops_test: OpsTest, continuous_writes) ->
 
 @pytest.mark.group(1)
 @markers.juju3
-async def test_removing_stereo_sync_replica(ops_test: OpsTest, continuous_writes) -> None:
+async def test_removing_stereo_sync_standby(ops_test: OpsTest, continuous_writes) -> None:
     # Start an application that continuously writes data to the database.
     app = await app_name(ops_test)
     await start_continuous_writes(ops_test, app)
@@ -82,7 +82,7 @@ async def test_removing_stereo_sync_replica(ops_test: OpsTest, continuous_writes
     primary = await get_primary(ops_test, app)
     secondary = next(
         filter(lambda x: x.name != primary, ops_test.model.applications[DATABASE_APP_NAME].units)
-    )
+    ).name
     await ops_test.model.destroy_unit(secondary, force=True, destroy_storage=False, max_wait=1500)
 
     await ops_test.model.wait_for_idle(status="active", timeout=600)

From e739a56681bbfa90d6f5ad05f83cfd30ea319bbe Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 5 Sep 2024 20:33:49 +0300
Subject: [PATCH 14/85] No down unit

---
 tests/integration/ha_tests/test_scaling.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index 166f96e82f..51be38b0d3 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -65,7 +65,7 @@ async def test_removing_stereo_primary(ops_test: OpsTest, continuous_writes) ->
 
     await ops_test.model.wait_for_idle(status="active", timeout=600)
 
-    await are_writes_increasing(ops_test, primary)
+    await are_writes_increasing(ops_test)
 
     logger.info("Scaling back up")
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=1)
@@ -87,7 +87,7 @@ async def test_removing_stereo_sync_standby(ops_test: OpsTest, continuous_writes
 
     await ops_test.model.wait_for_idle(status="active", timeout=600)
 
-    await are_writes_increasing(ops_test, primary)
+    await are_writes_increasing(ops_test)
 
     logger.info("Scaling back up")
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=1)

From ed56f3d3ecdc1efd00484c6c609b95c168edde84 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Mon, 9 Sep 2024 15:39:35 +0300
Subject: [PATCH 15/85] Add check_writes to test

---
 tests/integration/ha_tests/test_scaling.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index 51be38b0d3..9861afc6b4 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -16,6 +16,7 @@
 from .helpers import (
     app_name,
     are_writes_increasing,
+    check_writes,
     get_primary,
     start_continuous_writes,
 )
@@ -71,6 +72,8 @@ async def test_removing_stereo_primary(ops_test: OpsTest, continuous_writes) ->
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=1)
     await ops_test.model.wait_for_idle(status="active", timeout=1500)
 
+    await check_writes(ops_test)
+
 
 @pytest.mark.group(1)
 @markers.juju3
@@ -92,3 +95,5 @@ async def test_removing_stereo_sync_standby(ops_test: OpsTest, continuous_writes
     logger.info("Scaling back up")
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=1)
     await ops_test.model.wait_for_idle(status="active", timeout=1500)
+
+    await check_writes(ops_test)

From 7c8c95c6426ae69c26a66cd648568914a6023fed Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 10 Sep 2024 21:24:41 +0300
Subject: [PATCH 16/85] Unit synchronisation WIP

---
 src/charm.py             | 81 ++++++++++++++++++++++++++++++++++++++--
 src/cluster.py           | 15 +++++++-
 tests/unit/test_charm.py |  2 +
 3 files changed, 92 insertions(+), 6 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index d028d6dd1c..334a758c1e 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -511,20 +511,93 @@ def _on_pgdata_storage_detaching(self, _) -> None:
         if self.primary_endpoint:
             self._update_relation_endpoints()
 
-    def _on_peer_relation_changed(self, event: HookEvent):
-        """Reconfigure cluster members when something changes."""
+    def _collect_raft_cluster_state(self) -> (bool, bool, Optional[Unit]):
+        raft_stuck = False
+        all_units_stuck = True
+        candidate = None
+        for key, data in self._peers.data.items():
+            if key == self.app:
+                continue
+            if "stuck_raft" in data:
+                raft_stuck = True
+            else:
+                all_units_stuck = False
+            if not candidate and "raft_candidate" in data:
+                candidate = key
+        return raft_stuck, all_units_stuck, candidate
+
+    def _stuck_raft_cluster_check(self) -> bool:
+        """Check for stuck raft cluster and reinitialise if safe."""
+        raft_stuck, all_units_stuck, candidate = self._collect_raft_cluster_state()
+
+        if not raft_stuck:
+            return False
+
+        if not all_units_stuck:
+            logger.warning("Stuck raft not yet detected on all units")
+            return True
+
+        if not candidate:
+            logger.warning("Stuck raft has no candidate")
+            return True
+
+        for key, data in self._peers.data.items():
+            if key == self.app:
+                continue
+            data.pop("stuck_raft", None)
+            if key != candidate:
+                data.pop("candidate", None)
+            data["stopping_raft"] = "True"
+        return True
+
+    def _stuck_raft_cluster_rejoin(self) -> bool:
+        """Reconnect cluster to new raft."""
+        # TODO rejoin the cluster
+        return False
+
+    def _raft_reinitialisation(self) -> bool:
+        """Handle raft cluster loss of quorum."""
+        should_exit = False
+        if self.unit.is_leader() and self._stuck_raft_cluster_check():
+            should_exit = True
+
+        if "stopping_raft" in self.unit_peer_data:
+            should_exit = True
+            self._patroni.remove_raft_data()
+            self.unit_peer_data.pop("stopping_raft", None)
+            self.unit_peer_data["stopped_raft"] = "True"
+            if "candidate" in self.unit_peer_data:
+                self._patroni.reinitialise_raft_data()
+                self.unit_peer_data["candidate_raft_up"] = "True"
+
+        if self.unit.is_leader() and self._stuck_raft_cluster_rejoin():
+            should_exit = True
+        return should_exit
+
+    def _peer_relation_changed_checks(self, event: HookEvent) -> bool:
+        """Split of to reduce complexity."""
         # Prevents the cluster to be reconfigured before it's bootstrapped in the leader.
         if "cluster_initialised" not in self._peers.data[self.app]:
             logger.debug("Deferring on_peer_relation_changed: cluster not initialized")
             event.defer()
-            return
+            return False
+
+        # Check whether raft is stuck
+        if hasattr(event, "unit") and event.unit and self._raft_reinitialisation():
+            return False
 
         # If the unit is the leader, it can reconfigure the cluster.
         if self.unit.is_leader() and not self._reconfigure_cluster(event):
             event.defer()
-            return
+            return False
 
         if self._update_member_ip():
+            return False
+        return True
+
+    def _on_peer_relation_changed(self, event: HookEvent):
+        """Reconfigure cluster members when something changes."""
+        if not self._peer_relation_changed_checks(event):
             return
 
         # Don't update this member before it's part of the members list.
diff --git a/src/cluster.py b/src/cluster.py
index 9ee87b2c96..c1a64ca0fd 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -802,8 +802,19 @@ def remove_raft_member(self, member_ip: str) -> None:
         # If there's no quorum and the leader left raft cluster is stuck
         if not raft_status["has_quorum"] and raft_status["leader"].host == member_ip:
             logger.warning("Remove raft member: Stuck raft cluster detected")
-            self.remove_raft_data()
-            self.reinitialise_raft_data()
+            data_flags = {"stuck_raft": "True"}
+            try:
+                health_status = self.get_patroni_health()
+            except Exception:
+                logger.warning("Remove raft member: Unable to get health status")
+                health_status = {}
+            if health_status.get("role") in ("leader", "master", "sync_standby"):
+                data_flags["raft_candidate"] = "True"
+            self.charm.unit_peer_data.update(data_flags)
+            return
+
+            # self.remove_raft_data()
+            # self.reinitialise_raft_data()
 
         # Remove the member from the raft cluster.
         try:
diff --git a/tests/unit/test_charm.py b/tests/unit/test_charm.py
index 6eaab30cec..d71234f401 100644
--- a/tests/unit/test_charm.py
+++ b/tests/unit/test_charm.py
@@ -1488,6 +1488,8 @@ def test_on_peer_relation_changed(harness):
         rel_id = harness.model.get_relation(PEER).id
         # Test an uninitialized cluster.
         mock_event = Mock()
+        mock_event.unit = None
+
         with harness.hooks_disabled():
             harness.update_relation_data(
                 rel_id, harness.charm.app.name, {"cluster_initialised": ""}

From c7b26d2e035f38dbdf913c2f61ec60d3c4736477 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 10 Sep 2024 22:46:39 +0300
Subject: [PATCH 17/85] Tweak synchronisation

---
 src/charm.py   | 23 ++++++++++++++---------
 src/cluster.py |  9 ++++-----
 2 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 334a758c1e..45149f273d 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -518,7 +518,7 @@ def _collect_raft_cluster_state(self) -> (bool, bool, Optional[Unit]):
         for key, data in self._peers.data.items():
             if key == self.app:
                 continue
-            if "stuck_raft" in data:
+            if "raft_stuck" in data:
                 raft_stuck = True
             else:
                 all_units_stuck = False
@@ -544,10 +544,10 @@ def _stuck_raft_cluster_check(self) -> bool:
         for key, data in self._peers.data.items():
             if key == self.app:
                 continue
-            data.pop("stuck_raft", None)
+            data.pop("raft_stuck", None)
             if key != candidate:
-                data.pop("candidate", None)
-            data["stopping_raft"] = "True"
+                data.pop("raft_candidate", None)
+            data["raft_stopping"] = "True"
         return True
 
     def _stuck_raft_cluster_rejoin(self) -> bool:
@@ -561,14 +561,16 @@ def _raft_reinitialisation(self) -> bool:
         if self.unit.is_leader() and self._stuck_raft_cluster_check():
             should_exit = True
 
-        if "stopping_raft" in self.unit_peer_data:
+        if "raft_stopping" in self.unit_peer_data:
             should_exit = True
             self._patroni.remove_raft_data()
-            self.unit_peer_data.pop("stopping_raft", None)
-            self.unit_peer_data["stopped_raft"] = "True"
-            if "candidate" in self.unit_peer_data:
+            self.unit_peer_data.pop("raft_stopping", None)
+            if "candidate_raft" in self.unit_peer_data:
                 self._patroni.reinitialise_raft_data()
-                self.unit_peer_data["candidate_raft_up"] = "True"
+                self.unit_peer_data.pop("raft_candidate", None)
+                self.unit_peer_data["raft_primary"] = "True"
+            else:
+                self.unit_peer_data["raft_stopped"] = "True"
 
         if self.unit.is_leader() and self._stuck_raft_cluster_rejoin():
             should_exit = True
@@ -987,6 +989,9 @@ def _on_leader_elected(self, event: LeaderElectedEvent) -> None:
             if self.get_secret(APP_SCOPE, key) is None:
                 self.set_secret(APP_SCOPE, key, new_password())
 
+        if self._raft_reinitialisation():
+            return
+
         # Update the list of the current PostgreSQL hosts when a new leader is elected.
         # Add this unit to the list of cluster members
         # (the cluster should start with only this member).
diff --git a/src/cluster.py b/src/cluster.py
index c1a64ca0fd..6b3a3b4407 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -802,20 +802,19 @@ def remove_raft_member(self, member_ip: str) -> None:
         # If there's no quorum and the leader left raft cluster is stuck
         if not raft_status["has_quorum"] and raft_status["leader"].host == member_ip:
             logger.warning("Remove raft member: Stuck raft cluster detected")
-            data_flags = {"stuck_raft": "True"}
+            data_flags = {"raft_stuck": "True"}
             try:
                 health_status = self.get_patroni_health()
             except Exception:
                 logger.warning("Remove raft member: Unable to get health status")
                 health_status = {}
-            if health_status.get("role") in ("leader", "master", "sync_standby"):
+            if health_status.get("role") in ("leader", "master") or health_status.get(
+                "sync_standby"
+            ):
                 data_flags["raft_candidate"] = "True"
             self.charm.unit_peer_data.update(data_flags)
             return
 
-            # self.remove_raft_data()
-            # self.reinitialise_raft_data()
-
         # Remove the member from the raft cluster.
         try:
             result = syncobj_util.executeCommand(raft_host, ["remove", f"{member_ip}:2222"])

From a4f0b397226e483417b5451863ff17c885ec039a Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 11 Sep 2024 00:25:14 +0300
Subject: [PATCH 18/85] Add logging

---
 src/charm.py   | 3 ++-
 src/cluster.py | 1 +
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/charm.py b/src/charm.py
index 45149f273d..53f4054cff 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -541,6 +541,7 @@ def _stuck_raft_cluster_check(self) -> bool:
             logger.warning("Stuck raft has no candidate")
             return True
 
+        logger.info("%s selected for new raft leader" % candidate.name)
         for key, data in self._peers.data.items():
             if key == self.app:
                 continue
@@ -565,7 +566,7 @@ def _raft_reinitialisation(self) -> bool:
             should_exit = True
             self._patroni.remove_raft_data()
             self.unit_peer_data.pop("raft_stopping", None)
-            if "candidate_raft" in self.unit_peer_data:
+            if "raft_candidate" in self.unit_peer_data:
                 self._patroni.reinitialise_raft_data()
                 self.unit_peer_data.pop("raft_candidate", None)
                 self.unit_peer_data["raft_primary"] = "True"
diff --git a/src/cluster.py b/src/cluster.py
index 6b3a3b4407..f3cdea6e1a 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -811,6 +811,7 @@ def remove_raft_member(self, member_ip: str) -> None:
             if health_status.get("role") in ("leader", "master") or health_status.get(
                 "sync_standby"
             ):
+                logger.info("%s is raft candidate" % self.charm.unit.name)
                 data_flags["raft_candidate"] = "True"
             self.charm.unit_peer_data.update(data_flags)
             return

From 02d08ff93eb82f91cddf5a90a2543ca1a364e36c Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 11 Sep 2024 01:17:01 +0300
Subject: [PATCH 19/85] Update endpoints after raft nuke

---
 src/charm.py                               | 1 +
 tests/integration/ha_tests/test_scaling.py | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/src/charm.py b/src/charm.py
index 53f4054cff..b01226e9c0 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -575,6 +575,7 @@ def _raft_reinitialisation(self) -> bool:
 
         if self.unit.is_leader() and self._stuck_raft_cluster_rejoin():
             should_exit = True
+        self._update_relation_endpoints()
         return should_exit
 
     def _peer_relation_changed_checks(self, event: HookEvent) -> bool:
diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index 9861afc6b4..9c9d4f964c 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -56,6 +56,7 @@ async def test_build_and_deploy(ops_test: OpsTest) -> None:
 
 @pytest.mark.group(1)
 @markers.juju3
+@pytest.mark.abort_on_fail
 async def test_removing_stereo_primary(ops_test: OpsTest, continuous_writes) -> None:
     # Start an application that continuously writes data to the database.
     app = await app_name(ops_test)
@@ -77,6 +78,7 @@ async def test_removing_stereo_primary(ops_test: OpsTest, continuous_writes) ->
 
 @pytest.mark.group(1)
 @markers.juju3
+@pytest.mark.abort_on_fail
 async def test_removing_stereo_sync_standby(ops_test: OpsTest, continuous_writes) -> None:
     # Start an application that continuously writes data to the database.
     app = await app_name(ops_test)

From b454d069591589d7dfedc4043457a8257ee450fe Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 11 Sep 2024 01:33:32 +0300
Subject: [PATCH 20/85] Fix unit tests

---
 src/charm.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index b01226e9c0..76fc03785b 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -553,7 +553,13 @@ def _stuck_raft_cluster_check(self) -> bool:
 
     def _stuck_raft_cluster_rejoin(self) -> bool:
         """Reconnect cluster to new raft."""
-        # TODO rejoin the cluster
+        for key, data in self._peers.data.items():
+            if key == self.app:
+                continue
+            if "raft_primary" in data:
+                # TODO remove other units' ips so they can rejoin
+                self._update_relation_endpoints()
+                return True
         return False
 
     def _raft_reinitialisation(self) -> bool:
@@ -575,7 +581,6 @@ def _raft_reinitialisation(self) -> bool:
 
         if self.unit.is_leader() and self._stuck_raft_cluster_rejoin():
             should_exit = True
-        self._update_relation_endpoints()
         return should_exit
 
     def _peer_relation_changed_checks(self, event: HookEvent) -> bool:

From a281f7a0448e2a0aa8e798a897d6f4220348ab41 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 11 Sep 2024 02:21:27 +0300
Subject: [PATCH 21/85] Track down unit

---
 tests/integration/ha_tests/helpers.py      | 2 ++
 tests/integration/ha_tests/test_scaling.py | 4 ++--
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/integration/ha_tests/helpers.py b/tests/integration/ha_tests/helpers.py
index ff65047c63..0e9f468729 100644
--- a/tests/integration/ha_tests/helpers.py
+++ b/tests/integration/ha_tests/helpers.py
@@ -100,6 +100,8 @@ async def are_writes_increasing(
         extra_model=extra_model,
     )
     for member, count in writes.items():
+        if member == down_unit:
+            continue
         for attempt in Retrying(stop=stop_after_delay(60 * 3), wait=wait_fixed(3)):
             with attempt:
                 more_writes, _ = await count_writes(
diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index 9c9d4f964c..29a3250a19 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -67,7 +67,7 @@ async def test_removing_stereo_primary(ops_test: OpsTest, continuous_writes) ->
 
     await ops_test.model.wait_for_idle(status="active", timeout=600)
 
-    await are_writes_increasing(ops_test)
+    await are_writes_increasing(ops_test, primary)
 
     logger.info("Scaling back up")
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=1)
@@ -92,7 +92,7 @@ async def test_removing_stereo_sync_standby(ops_test: OpsTest, continuous_writes
 
     await ops_test.model.wait_for_idle(status="active", timeout=600)
 
-    await are_writes_increasing(ops_test)
+    await are_writes_increasing(ops_test, secondary)
 
     logger.info("Scaling back up")
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=1)

From 0f9ec81587a516e58c4ea28aec4a8c91c82c1c75 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 11 Sep 2024 14:44:33 +0300
Subject: [PATCH 22/85] Debug down unit exclusion

---
 tests/integration/ha_tests/helpers.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/integration/ha_tests/helpers.py b/tests/integration/ha_tests/helpers.py
index 0e9f468729..99e7f1c7f7 100644
--- a/tests/integration/ha_tests/helpers.py
+++ b/tests/integration/ha_tests/helpers.py
@@ -100,8 +100,6 @@ async def are_writes_increasing(
         extra_model=extra_model,
     )
     for member, count in writes.items():
-        if member == down_unit:
-            continue
         for attempt in Retrying(stop=stop_after_delay(60 * 3), wait=wait_fixed(3)):
             with attempt:
                 more_writes, _ = await count_writes(
@@ -281,6 +279,7 @@ async def count_writes(
         if model is None:
             continue
         for unit in model.applications[app].units:
+            print(unit.name, down_unit)
             if unit.name != down_unit:
                 members_data = get_patroni_cluster(
                     await (

From c1f95960080d5829aee32e72eaddd01c2d55cd53 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 11 Sep 2024 15:23:31 +0300
Subject: [PATCH 23/85] Reraise

---
 tests/integration/ha_tests/helpers.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/integration/ha_tests/helpers.py b/tests/integration/ha_tests/helpers.py
index 99e7f1c7f7..79d9365194 100644
--- a/tests/integration/ha_tests/helpers.py
+++ b/tests/integration/ha_tests/helpers.py
@@ -100,7 +100,7 @@ async def are_writes_increasing(
         extra_model=extra_model,
     )
     for member, count in writes.items():
-        for attempt in Retrying(stop=stop_after_delay(60 * 3), wait=wait_fixed(3)):
+        for attempt in Retrying(stop=stop_after_delay(60 * 3), wait=wait_fixed(3), reraise=True):
             with attempt:
                 more_writes, _ = await count_writes(
                     ops_test,
@@ -279,7 +279,6 @@ async def count_writes(
         if model is None:
             continue
         for unit in model.applications[app].units:
-            print(unit.name, down_unit)
             if unit.name != down_unit:
                 members_data = get_patroni_cluster(
                     await (

From 47b2b6abc51db921e140c4e001d6495e05458a8b Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 11 Sep 2024 16:20:09 +0300
Subject: [PATCH 24/85] Debug excess key

---
 tests/integration/ha_tests/helpers.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/integration/ha_tests/helpers.py b/tests/integration/ha_tests/helpers.py
index 79d9365194..5172034873 100644
--- a/tests/integration/ha_tests/helpers.py
+++ b/tests/integration/ha_tests/helpers.py
@@ -99,6 +99,7 @@ async def are_writes_increasing(
         use_ip_from_inside=use_ip_from_inside,
         extra_model=extra_model,
     )
+    print(writes)
     for member, count in writes.items():
         for attempt in Retrying(stop=stop_after_delay(60 * 3), wait=wait_fixed(3), reraise=True):
             with attempt:

From ad5b8e85415194efee3619826c3a3cf24b5d31b6 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 11 Sep 2024 16:51:05 +0300
Subject: [PATCH 25/85] Ignore down unit

---
 tests/integration/ha_tests/helpers.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/integration/ha_tests/helpers.py b/tests/integration/ha_tests/helpers.py
index 5172034873..8fbf408c5b 100644
--- a/tests/integration/ha_tests/helpers.py
+++ b/tests/integration/ha_tests/helpers.py
@@ -99,8 +99,10 @@ async def are_writes_increasing(
         use_ip_from_inside=use_ip_from_inside,
         extra_model=extra_model,
     )
-    print(writes)
     for member, count in writes.items():
+        if member.split(".", 1)[-1] == down_unit:
+            continue
+
         for attempt in Retrying(stop=stop_after_delay(60 * 3), wait=wait_fixed(3), reraise=True):
             with attempt:
                 more_writes, _ = await count_writes(

From 39517082ea3a5f05f10c7bd19b6d8bbe6afb05fe Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 11 Sep 2024 17:29:59 +0300
Subject: [PATCH 26/85] Add logging to writes increasing check

---
 tests/integration/ha_tests/helpers.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/integration/ha_tests/helpers.py b/tests/integration/ha_tests/helpers.py
index 8fbf408c5b..b2c7d7a03a 100644
--- a/tests/integration/ha_tests/helpers.py
+++ b/tests/integration/ha_tests/helpers.py
@@ -99,6 +99,7 @@ async def are_writes_increasing(
         use_ip_from_inside=use_ip_from_inside,
         extra_model=extra_model,
     )
+    logger.info(f"Initial writes {writes}")
     for member, count in writes.items():
         if member.split(".", 1)[-1] == down_unit:
             continue
@@ -111,6 +112,7 @@ async def are_writes_increasing(
                     use_ip_from_inside=use_ip_from_inside,
                     extra_model=extra_model,
                 )
+                logger.info(f"Retry writes {more_writes}")
                 assert (
                     more_writes[member] > count
                 ), f"{member}: writes not continuing to DB (current writes: {more_writes[member]} - previous writes: {count})"

From 6afe354c2cd0e9261c701cd11a00e529c4471a11 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 11 Sep 2024 18:17:06 +0300
Subject: [PATCH 27/85] Tweaks writes check

---
 tests/integration/ha_tests/helpers.py | 25 +++++++++++++------------
 1 file changed, 13 insertions(+), 12 deletions(-)

diff --git a/tests/integration/ha_tests/helpers.py b/tests/integration/ha_tests/helpers.py
index b2c7d7a03a..9768467fe2 100644
--- a/tests/integration/ha_tests/helpers.py
+++ b/tests/integration/ha_tests/helpers.py
@@ -100,19 +100,20 @@ async def are_writes_increasing(
         extra_model=extra_model,
     )
     logger.info(f"Initial writes {writes}")
-    for member, count in writes.items():
-        if member.split(".", 1)[-1] == down_unit:
-            continue
 
-        for attempt in Retrying(stop=stop_after_delay(60 * 3), wait=wait_fixed(3), reraise=True):
-            with attempt:
-                more_writes, _ = await count_writes(
-                    ops_test,
-                    down_unit=down_unit,
-                    use_ip_from_inside=use_ip_from_inside,
-                    extra_model=extra_model,
-                )
-                logger.info(f"Retry writes {more_writes}")
+    for attempt in Retrying(stop=stop_after_delay(60 * 3), wait=wait_fixed(3), reraise=True):
+        with attempt:
+            more_writes, _ = await count_writes(
+                ops_test,
+                down_unit=down_unit,
+                use_ip_from_inside=use_ip_from_inside,
+                extra_model=extra_model,
+            )
+            logger.info(f"Retry writes {more_writes}")
+            for member, count in writes.items():
+                print(member, down_unit)
+                if member == down_unit or member.split(".", 1)[-1] == down_unit:
+                    continue
                 assert (
                     more_writes[member] > count
                 ), f"{member}: writes not continuing to DB (current writes: {more_writes[member]} - previous writes: {count})"

From 243f1fb0c0b9ffae639ed5cfc795bdf21fcee3f2 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 11 Sep 2024 20:26:15 +0300
Subject: [PATCH 28/85] Fix down unit skipping

---
 tests/integration/ha_tests/helpers.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/integration/ha_tests/helpers.py b/tests/integration/ha_tests/helpers.py
index 9768467fe2..0cf776c18a 100644
--- a/tests/integration/ha_tests/helpers.py
+++ b/tests/integration/ha_tests/helpers.py
@@ -111,8 +111,7 @@ async def are_writes_increasing(
             )
             logger.info(f"Retry writes {more_writes}")
             for member, count in writes.items():
-                print(member, down_unit)
-                if member == down_unit or member.split(".", 1)[-1] == down_unit:
+                if "/".join(member.split(".", 1)[-1].rsplit("-", 1)) == down_unit:
                     continue
                 assert (
                     more_writes[member] > count

From 15dfeb761c512197a3f29bb0f9abf959b3b26ad0 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 12 Sep 2024 13:42:09 +0300
Subject: [PATCH 29/85] Cleanup and ip cache

---
 src/charm.py | 26 +++++++++++++++++++++++---
 1 file changed, 23 insertions(+), 3 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 76fc03785b..fb6b445c00 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -553,18 +553,38 @@ def _stuck_raft_cluster_check(self) -> bool:
 
     def _stuck_raft_cluster_rejoin(self) -> bool:
         """Reconnect cluster to new raft."""
+        primary = None
+        all_units_down = True
         for key, data in self._peers.data.items():
             if key == self.app:
                 continue
             if "raft_primary" in data:
-                # TODO remove other units' ips so they can rejoin
-                self._update_relation_endpoints()
-                return True
+                primary = key
+            elif "raft_stopped" not in data:
+                all_units_down = False
+        if primary:
+            logger.info("Updating new primary endpoint")
+            self.app_peer_data.pop("members_ips", None)
+            self._add_to_members_ips(self._get_unit_ip(key))
+            self._update_relation_endpoints()
+            if all_units_down:
+                logger.info("Removing stuck raft peer data")
+                for key, data in self._peers.data.items():
+                    if key == self.app:
+                        continue
+                    data.pop("raft_primary", None)
+                    data.pop("raft_stopped", None)
+            return True
         return False
 
     def _raft_reinitialisation(self) -> bool:
         """Handle raft cluster loss of quorum."""
         should_exit = False
+        for key in self.unit_peer_data.keys():
+            if key.startswith("raft_"):
+                should_exit = True
+                break
+
         if self.unit.is_leader() and self._stuck_raft_cluster_check():
             should_exit = True
 

From 1717f27efee50b227bc78e05f6c16c0b6ce62aed Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 12 Sep 2024 14:35:56 +0300
Subject: [PATCH 30/85] Add tests

---
 tests/integration/ha_tests/helpers.py      | 25 ++++++++++
 tests/integration/ha_tests/test_scaling.py | 56 ++++++++++++++++++++++
 2 files changed, 81 insertions(+)

diff --git a/tests/integration/ha_tests/helpers.py b/tests/integration/ha_tests/helpers.py
index 0cf776c18a..a6d23924b6 100644
--- a/tests/integration/ha_tests/helpers.py
+++ b/tests/integration/ha_tests/helpers.py
@@ -627,6 +627,31 @@ async def is_replica(ops_test: OpsTest, unit_name: str, use_ip_from_inside: bool
         return False
 
 
+async def get_cluster_roles(
+    ops_test: OpsTest, unit_name: str, use_ip_from_inside: bool = False
+) -> Dict[str, Union[Optional[str], list[str]]]:
+    """Returns whether the unit a replica in the cluster."""
+    unit_ip = await (
+        get_ip_from_inside_the_unit(ops_test, unit_name)
+        if use_ip_from_inside
+        else get_unit_ip(ops_test, unit_name)
+    )
+
+    members = {"replicas": [], "primary": None, "sync_standbys": []}
+    cluster_info = requests.get(f"http://{unit_ip}:8008/cluster")
+    for member in cluster_info.json()["members"]:
+        role = member["role"]
+        name = "/".join(member["name"].rsplit("-", 1))
+        if role == "leader":
+            members["primary"] = name
+        elif role == "sync_standby":
+            members["sync_standbys"].append(name)
+        else:
+            members["replicas"].append(name)
+
+    return members
+
+
 async def instance_ip(ops_test: OpsTest, instance: str) -> str:
     """Translate juju instance name to IP.
 
diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index 34f18f540c..e9de31f7e7 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -17,6 +17,7 @@
     app_name,
     are_writes_increasing,
     check_writes,
+    get_cluster_roles,
     get_primary,
     start_continuous_writes,
 )
@@ -99,3 +100,58 @@ async def test_removing_stereo_sync_standby(ops_test: OpsTest, continuous_writes
     await ops_test.model.wait_for_idle(status="active", timeout=1500)
 
     await check_writes(ops_test)
+
+
+@pytest.mark.group(1)
+@markers.juju3
+@pytest.mark.abort_on_fail
+async def test_deploy_quatro(ops_test: OpsTest) -> None:
+    await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=2)
+    await ops_test.model.wait_for_idle(status="active", timeout=1500)
+
+
+@pytest.mark.group(1)
+@markers.juju3
+@pytest.mark.abort_on_fail
+async def test_removing_quatro_primary_and_async_replica(
+    ops_test: OpsTest, continuous_writes
+) -> None:
+    # Start an application that continuously writes data to the database.
+    app = await app_name(ops_test)
+    roles = get_cluster_roles(ops_test.model.applications[DATABASE_APP_NAME].units[0])
+    await start_continuous_writes(ops_test, app)
+    logger.info("Deleting primary")
+    await gather(
+        ops_test.model.destroy_unit(
+            roles["primary"], force=True, destroy_storage=False, max_wait=1500
+        ),
+        await ops_test.model.destroy_unit(
+            roles["replicas"][0], force=True, destroy_storage=False, max_wait=1500
+        ),
+    )
+
+    await ops_test.model.wait_for_idle(status="active", timeout=600)
+
+    await are_writes_increasing(ops_test, roles["primary"])
+
+    logger.info("Scaling back up")
+    await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=2)
+    await ops_test.model.wait_for_idle(status="active", timeout=1500)
+
+    await check_writes(ops_test)
+
+
+@pytest.mark.group(1)
+@markers.juju3
+@pytest.mark.abort_on_fail
+async def test_removing_quatro_sync_and_async_replica(
+    ops_test: OpsTest, continuous_writes
+) -> None:
+    pass
+
+
+@pytest.mark.group(1)
+@markers.juju3
+@pytest.mark.abort_on_fail
+async def test_removing_quatro_both_async_replica(ops_test: OpsTest, continuous_writes) -> None:
+    pass

From d08eb1d6fb34eb9c758c988f62ce84b9b6a5b730 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 12 Sep 2024 15:29:02 +0300
Subject: [PATCH 31/85] Add logging on early exit

---
 src/charm.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/charm.py b/src/charm.py
index fb6b445c00..eddf28b535 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -613,6 +613,7 @@ def _peer_relation_changed_checks(self, event: HookEvent) -> bool:
 
         # Check whether raft is stuck
         if hasattr(event, "unit") and event.unit and self._raft_reinitialisation():
+            logger.debug("Early exit on_peer_relation_changed: stuck raft recovery")
             return False
 
         # If the unit is the leader, it can reconfigure the cluster.

From 862660bd9d36f8fe9f210019c25d3823d5bd1b37 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 12 Sep 2024 16:01:55 +0300
Subject: [PATCH 32/85] Protect against None peers

---
 src/charm.py   |  2 +-
 src/cluster.py | 13 +++++++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/src/charm.py b/src/charm.py
index eddf28b535..d4076254db 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -928,7 +928,7 @@ def _update_members_ips(self, ip_to_add: str = None, ip_to_remove: str = None) -
             return
 
         ips = json.loads(self._peers.data[self.app].get("members_ips", "[]"))
-        if ip_to_add and ip_to_add not in ips:
+        if ip_to_add and ip_to_add != "None" and ip_to_add not in ips:
             ips.append(ip_to_add)
         elif ip_to_remove:
             ips.remove(ip_to_remove)
diff --git a/src/cluster.py b/src/cluster.py
index f3cdea6e1a..0fb5693ebc 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -54,6 +54,10 @@
 RUNNING_STATES = ["running", "streaming"]
 
 
+class RaftPostgresqlNotUpError(Exception):
+    """Raised when a cluster is not promoted."""
+
+
 class RaftPostgresqlStillUpError(Exception):
     """Raised when a cluster is not promoted."""
 
@@ -772,6 +776,15 @@ def reinitialise_raft_data(self) -> None:
 
         logger.info("Restarting patroni")
         self.restart_patroni()
+        for attempt in Retrying(wait=wait_fixed(5)):
+            with attempt:
+                found_postgres = False
+                for proc in psutil.process_iter(["name"]):
+                    if proc.name() == "postgres":
+                        found_postgres = True
+                        break
+                if not found_postgres:
+                    raise RaftPostgresqlNotUpError()
         logger.info("Raft should be unstuck")
 
     def remove_raft_member(self, member_ip: str) -> None:

From c3bcac1ea4ed1b899221a4185b70446eef9abe6b Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 12 Sep 2024 17:02:45 +0300
Subject: [PATCH 33/85] Fix quatro test

---
 tests/integration/ha_tests/test_scaling.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index e9de31f7e7..9bb8c0c37c 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -118,7 +118,9 @@ async def test_removing_quatro_primary_and_async_replica(
 ) -> None:
     # Start an application that continuously writes data to the database.
     app = await app_name(ops_test)
-    roles = get_cluster_roles(ops_test.model.applications[DATABASE_APP_NAME].units[0])
+    roles = get_cluster_roles(
+        ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
+    )
     await start_continuous_writes(ops_test, app)
     logger.info("Deleting primary")
     await gather(

From 90e0d83852d02c6a57b34f83aa1fc9825d7dd310 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 12 Sep 2024 18:18:27 +0300
Subject: [PATCH 34/85] Majority removal test

---
 tests/integration/ha_tests/helpers.py      |  4 +-
 tests/integration/ha_tests/test_scaling.py | 62 +++++++++++++---------
 2 files changed, 38 insertions(+), 28 deletions(-)

diff --git a/tests/integration/ha_tests/helpers.py b/tests/integration/ha_tests/helpers.py
index a6d23924b6..6aa47ee2a4 100644
--- a/tests/integration/ha_tests/helpers.py
+++ b/tests/integration/ha_tests/helpers.py
@@ -637,13 +637,13 @@ async def get_cluster_roles(
         else get_unit_ip(ops_test, unit_name)
     )
 
-    members = {"replicas": [], "primary": None, "sync_standbys": []}
+    members = {"replicas": [], "primaries": [], "sync_standbys": []}
     cluster_info = requests.get(f"http://{unit_ip}:8008/cluster")
     for member in cluster_info.json()["members"]:
         role = member["role"]
         name = "/".join(member["name"].rsplit("-", 1))
         if role == "leader":
-            members["primary"] = name
+            members["primaries"].append(name)
         elif role == "sync_standby":
             members["sync_standbys"].append(name)
         else:
diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index 9bb8c0c37c..48888bb943 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -61,6 +61,9 @@ async def test_build_and_deploy(ops_test: OpsTest) -> None:
 async def test_removing_stereo_primary(ops_test: OpsTest, continuous_writes) -> None:
     # Start an application that continuously writes data to the database.
     app = await app_name(ops_test)
+    original_roles = await get_cluster_roles(
+        ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
+    )
     await start_continuous_writes(ops_test, app)
     logger.info("Deleting primary")
     primary = await get_primary(ops_test, app)
@@ -74,6 +77,13 @@ async def test_removing_stereo_primary(ops_test: OpsTest, continuous_writes) ->
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=1)
     await ops_test.model.wait_for_idle(status="active", timeout=1500)
 
+    new_roles = await get_cluster_roles(
+        ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
+    )
+    assert len(new_roles["primaries"]) == 1
+    assert len(new_roles["sync_standbys"]) == 1
+    assert new_roles["primaries"][0] == original_roles["sync_standbys"][0]
+
     await check_writes(ops_test)
 
 
@@ -83,6 +93,9 @@ async def test_removing_stereo_primary(ops_test: OpsTest, continuous_writes) ->
 async def test_removing_stereo_sync_standby(ops_test: OpsTest, continuous_writes) -> None:
     # Start an application that continuously writes data to the database.
     app = await app_name(ops_test)
+    original_roles = await get_cluster_roles(
+        ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
+    )
     await start_continuous_writes(ops_test, app)
     logger.info("Deleting sync replica")
     primary = await get_primary(ops_test, app)
@@ -98,6 +111,12 @@ async def test_removing_stereo_sync_standby(ops_test: OpsTest, continuous_writes
     logger.info("Scaling back up")
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=1)
     await ops_test.model.wait_for_idle(status="active", timeout=1500)
+    new_roles = await get_cluster_roles(
+        ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
+    )
+    assert len(new_roles["primaries"]) == 1
+    assert len(new_roles["sync_standbys"]) == 1
+    assert new_roles["primaries"][0] == original_roles["primaries"][0]
 
     await check_writes(ops_test)
 
@@ -105,55 +124,46 @@ async def test_removing_stereo_sync_standby(ops_test: OpsTest, continuous_writes
 @pytest.mark.group(1)
 @markers.juju3
 @pytest.mark.abort_on_fail
-async def test_deploy_quatro(ops_test: OpsTest) -> None:
-    await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=2)
+async def test_scale_to_five_units(ops_test: OpsTest) -> None:
+    await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=3)
     await ops_test.model.wait_for_idle(status="active", timeout=1500)
 
 
 @pytest.mark.group(1)
 @markers.juju3
 @pytest.mark.abort_on_fail
-async def test_removing_quatro_primary_and_async_replica(
-    ops_test: OpsTest, continuous_writes
-) -> None:
+async def test_removing_raft_majority(ops_test: OpsTest, continuous_writes) -> None:
     # Start an application that continuously writes data to the database.
     app = await app_name(ops_test)
-    roles = get_cluster_roles(
+    original_roles = await get_cluster_roles(
         ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
     )
     await start_continuous_writes(ops_test, app)
     logger.info("Deleting primary")
     await gather(
         ops_test.model.destroy_unit(
-            roles["primary"], force=True, destroy_storage=False, max_wait=1500
+            original_roles["primaries"][0], force=True, destroy_storage=False, max_wait=1500
         ),
-        await ops_test.model.destroy_unit(
-            roles["replicas"][0], force=True, destroy_storage=False, max_wait=1500
+        ops_test.model.destroy_unit(
+            original_roles["replicas"][0], force=True, destroy_storage=False, max_wait=1500
+        ),
+        ops_test.model.destroy_unit(
+            original_roles["sync_standbys"][0], force=True, destroy_storage=False, max_wait=1500
         ),
     )
 
     await ops_test.model.wait_for_idle(status="active", timeout=600)
 
-    await are_writes_increasing(ops_test, roles["primary"])
+    await are_writes_increasing(ops_test, original_roles["primaries"][0])
 
     logger.info("Scaling back up")
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=2)
     await ops_test.model.wait_for_idle(status="active", timeout=1500)
 
     await check_writes(ops_test)
-
-
-@pytest.mark.group(1)
-@markers.juju3
-@pytest.mark.abort_on_fail
-async def test_removing_quatro_sync_and_async_replica(
-    ops_test: OpsTest, continuous_writes
-) -> None:
-    pass
-
-
-@pytest.mark.group(1)
-@markers.juju3
-@pytest.mark.abort_on_fail
-async def test_removing_quatro_both_async_replica(ops_test: OpsTest, continuous_writes) -> None:
-    pass
+    new_roles = await get_cluster_roles(
+        ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
+    )
+    assert len(new_roles["primaries"]) == 1
+    assert len(new_roles["sync_standbys"]) == 2
+    assert new_roles["primaries"][0] == original_roles["sync_standbys"][1]

From 507665d917e2746d1a541b933093625f0bb0a959 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 12 Sep 2024 22:34:16 +0300
Subject: [PATCH 35/85] Log cluster roles

---
 tests/integration/ha_tests/helpers.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/integration/ha_tests/helpers.py b/tests/integration/ha_tests/helpers.py
index 6aa47ee2a4..9763a2f80d 100644
--- a/tests/integration/ha_tests/helpers.py
+++ b/tests/integration/ha_tests/helpers.py
@@ -639,7 +639,9 @@ async def get_cluster_roles(
 
     members = {"replicas": [], "primaries": [], "sync_standbys": []}
     cluster_info = requests.get(f"http://{unit_ip}:8008/cluster")
-    for member in cluster_info.json()["members"]:
+    member_list = cluster_info.json()["members"]
+    logger.info(f"Cluster members are: {member_list}")
+    for member in member_list:
         role = member["role"]
         name = "/".join(member["name"].rsplit("-", 1))
         if role == "leader":

From 0dfa199f0e88ed33ec5f69bcfe176018ed05cace Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Fri, 13 Sep 2024 00:06:34 +0300
Subject: [PATCH 36/85] Test sleeps

---
 tests/integration/ha_tests/test_scaling.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index 48888bb943..7d6fb4eeac 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -2,6 +2,7 @@
 # Copyright 2024 Canonical Ltd.
 # See LICENSE file for licensing details.
 import logging
+import time
 from asyncio import gather
 
 import pytest
@@ -71,12 +72,16 @@ async def test_removing_stereo_primary(ops_test: OpsTest, continuous_writes) ->
 
     await ops_test.model.wait_for_idle(status="active", timeout=600)
 
+    time.sleep(60)
+
     await are_writes_increasing(ops_test, primary)
 
     logger.info("Scaling back up")
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=1)
     await ops_test.model.wait_for_idle(status="active", timeout=1500)
 
+    time.sleep(60)
+
     new_roles = await get_cluster_roles(
         ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
     )
@@ -106,11 +111,16 @@ async def test_removing_stereo_sync_standby(ops_test: OpsTest, continuous_writes
 
     await ops_test.model.wait_for_idle(status="active", timeout=600)
 
+    time.sleep(60)
+
     await are_writes_increasing(ops_test, secondary)
 
     logger.info("Scaling back up")
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=1)
     await ops_test.model.wait_for_idle(status="active", timeout=1500)
+
+    time.sleep(60)
+
     new_roles = await get_cluster_roles(
         ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
     )
@@ -138,6 +148,9 @@ async def test_removing_raft_majority(ops_test: OpsTest, continuous_writes) -> N
     original_roles = await get_cluster_roles(
         ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
     )
+
+    time.sleep(60)
+
     await start_continuous_writes(ops_test, app)
     logger.info("Deleting primary")
     await gather(
@@ -160,6 +173,8 @@ async def test_removing_raft_majority(ops_test: OpsTest, continuous_writes) -> N
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=2)
     await ops_test.model.wait_for_idle(status="active", timeout=1500)
 
+    time.sleep(60)
+
     await check_writes(ops_test)
     new_roles = await get_cluster_roles(
         ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name

From 1b8c2d81e7f3fabfa916f7131178030e42f30230 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Fri, 13 Sep 2024 02:25:39 +0300
Subject: [PATCH 37/85] Try to remove the intialised flag

---
 src/charm.py                               | 13 ++++++++-----
 tests/integration/ha_tests/test_scaling.py | 13 -------------
 2 files changed, 8 insertions(+), 18 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index d4076254db..dcbc645ab2 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -549,6 +549,7 @@ def _stuck_raft_cluster_check(self) -> bool:
             if key != candidate:
                 data.pop("raft_candidate", None)
             data["raft_stopping"] = "True"
+        self.app_peer_data.pop("cluster_initialised", None)
         return True
 
     def _stuck_raft_cluster_rejoin(self) -> bool:
@@ -565,7 +566,9 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
         if primary:
             logger.info("Updating new primary endpoint")
             self.app_peer_data.pop("members_ips", None)
+            logger.error(f"!!!!!!!!!!!!!!!!!1{self._get_unit_ip(key)}")
             self._add_to_members_ips(self._get_unit_ip(key))
+            self.app_peer_data["cluster_initialised"] = "True"
             self._update_relation_endpoints()
             if all_units_down:
                 logger.info("Removing stuck raft peer data")
@@ -605,17 +608,17 @@ def _raft_reinitialisation(self) -> bool:
 
     def _peer_relation_changed_checks(self, event: HookEvent) -> bool:
         """Split of to reduce complexity."""
+        # Check whether raft is stuck
+        if hasattr(event, "unit") and event.unit and self._raft_reinitialisation():
+            logger.debug("Early exit on_peer_relation_changed: stuck raft recovery")
+            return False
+
         # Prevents the cluster to be reconfigured before it's bootstrapped in the leader.
         if "cluster_initialised" not in self._peers.data[self.app]:
             logger.debug("Deferring on_peer_relation_changed: cluster not initialized")
             event.defer()
             return False
 
-        # Check whether raft is stuck
-        if hasattr(event, "unit") and event.unit and self._raft_reinitialisation():
-            logger.debug("Early exit on_peer_relation_changed: stuck raft recovery")
-            return False
-
         # If the unit is the leader, it can reconfigure the cluster.
         if self.unit.is_leader() and not self._reconfigure_cluster(event):
             event.defer()
diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index 7d6fb4eeac..054c3752d0 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -2,7 +2,6 @@
 # Copyright 2024 Canonical Ltd.
 # See LICENSE file for licensing details.
 import logging
-import time
 from asyncio import gather
 
 import pytest
@@ -72,16 +71,12 @@ async def test_removing_stereo_primary(ops_test: OpsTest, continuous_writes) ->
 
     await ops_test.model.wait_for_idle(status="active", timeout=600)
 
-    time.sleep(60)
-
     await are_writes_increasing(ops_test, primary)
 
     logger.info("Scaling back up")
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=1)
     await ops_test.model.wait_for_idle(status="active", timeout=1500)
 
-    time.sleep(60)
-
     new_roles = await get_cluster_roles(
         ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
     )
@@ -111,16 +106,12 @@ async def test_removing_stereo_sync_standby(ops_test: OpsTest, continuous_writes
 
     await ops_test.model.wait_for_idle(status="active", timeout=600)
 
-    time.sleep(60)
-
     await are_writes_increasing(ops_test, secondary)
 
     logger.info("Scaling back up")
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=1)
     await ops_test.model.wait_for_idle(status="active", timeout=1500)
 
-    time.sleep(60)
-
     new_roles = await get_cluster_roles(
         ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
     )
@@ -149,8 +140,6 @@ async def test_removing_raft_majority(ops_test: OpsTest, continuous_writes) -> N
         ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
     )
 
-    time.sleep(60)
-
     await start_continuous_writes(ops_test, app)
     logger.info("Deleting primary")
     await gather(
@@ -173,8 +162,6 @@ async def test_removing_raft_majority(ops_test: OpsTest, continuous_writes) -> N
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=2)
     await ops_test.model.wait_for_idle(status="active", timeout=1500)
 
-    time.sleep(60)
-
     await check_writes(ops_test)
     new_roles = await get_cluster_roles(
         ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name

From ad2d7a98920bbb3f92151e1258d0ed02df6ce8fe Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Fri, 13 Sep 2024 03:22:53 +0300
Subject: [PATCH 38/85] Use private address

---
 src/charm.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index dcbc645ab2..b38d44ea40 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -549,7 +549,6 @@ def _stuck_raft_cluster_check(self) -> bool:
             if key != candidate:
                 data.pop("raft_candidate", None)
             data["raft_stopping"] = "True"
-        self.app_peer_data.pop("cluster_initialised", None)
         return True
 
     def _stuck_raft_cluster_rejoin(self) -> bool:
@@ -566,9 +565,10 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
         if primary:
             logger.info("Updating new primary endpoint")
             self.app_peer_data.pop("members_ips", None)
+            logger.error(f"!!!!!!!!!!!!!!!!!1{key}")
             logger.error(f"!!!!!!!!!!!!!!!!!1{self._get_unit_ip(key)}")
-            self._add_to_members_ips(self._get_unit_ip(key))
-            self.app_peer_data["cluster_initialised"] = "True"
+            logger.error(f"!!!!!!!!!!!!!!!!!1{data['private-address']}")
+            self._add_to_members_ips(data["private-address"])
             self._update_relation_endpoints()
             if all_units_down:
                 logger.info("Removing stuck raft peer data")

From a7bbe67f208dc9f885c4a4f29d6d044a8fbe9192 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Fri, 13 Sep 2024 04:06:53 +0300
Subject: [PATCH 39/85] Log ip

---
 src/charm.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/charm.py b/src/charm.py
index b38d44ea40..86d9dd0499 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -566,6 +566,7 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
             logger.info("Updating new primary endpoint")
             self.app_peer_data.pop("members_ips", None)
             logger.error(f"!!!!!!!!!!!!!!!!!1{key}")
+            logger.error(f"!!!!!!!!!!!!!!!!!1{data}")
             logger.error(f"!!!!!!!!!!!!!!!!!1{self._get_unit_ip(key)}")
             logger.error(f"!!!!!!!!!!!!!!!!!1{data['private-address']}")
             self._add_to_members_ips(data["private-address"])

From 89c368febd7424b63da2ae18da0f401ed3a9e825 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Fri, 13 Sep 2024 04:49:45 +0300
Subject: [PATCH 40/85] Use primary key

---
 src/charm.py | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 86d9dd0499..8e9d7d7363 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -565,11 +565,9 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
         if primary:
             logger.info("Updating new primary endpoint")
             self.app_peer_data.pop("members_ips", None)
-            logger.error(f"!!!!!!!!!!!!!!!!!1{key}")
-            logger.error(f"!!!!!!!!!!!!!!!!!1{data}")
-            logger.error(f"!!!!!!!!!!!!!!!!!1{self._get_unit_ip(key)}")
-            logger.error(f"!!!!!!!!!!!!!!!!!1{data['private-address']}")
-            self._add_to_members_ips(data["private-address"])
+            logger.error(f"!!!!!!!!!!!!!!!!!1{primary}")
+            logger.error(f"!!!!!!!!!!!!!!!!!1{self._get_unit_ip(primary)}")
+            self._add_to_members_ips(self._get_unit_ip(primary))
             self._update_relation_endpoints()
             if all_units_down:
                 logger.info("Removing stuck raft peer data")

From 9efc6caa6cb9b644b19ba2cb5577ba5da2262091 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Fri, 13 Sep 2024 14:39:14 +0300
Subject: [PATCH 41/85] Check for None leader

---
 src/charm.py   | 2 --
 src/cluster.py | 4 +++-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 8e9d7d7363..79981b46af 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -565,8 +565,6 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
         if primary:
             logger.info("Updating new primary endpoint")
             self.app_peer_data.pop("members_ips", None)
-            logger.error(f"!!!!!!!!!!!!!!!!!1{primary}")
-            logger.error(f"!!!!!!!!!!!!!!!!!1{self._get_unit_ip(primary)}")
             self._add_to_members_ips(self._get_unit_ip(primary))
             self._update_relation_endpoints()
             if all_units_down:
diff --git a/src/cluster.py b/src/cluster.py
index 0fb5693ebc..b1ec8d18c9 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -813,7 +813,9 @@ def remove_raft_member(self, member_ip: str) -> None:
             return
 
         # If there's no quorum and the leader left raft cluster is stuck
-        if not raft_status["has_quorum"] and raft_status["leader"].host == member_ip:
+        if not raft_status["has_quorum"] and (
+            not raft_status["leader"] or raft_status["leader"].host == member_ip
+        ):
             logger.warning("Remove raft member: Stuck raft cluster detected")
             data_flags = {"raft_stuck": "True"}
             try:

From ec6f88cbcbe6d81847812290d4f0d8e38800fa59 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Fri, 13 Sep 2024 16:07:16 +0300
Subject: [PATCH 42/85] Try to reuse general health call

---
 src/cluster.py             |   8 +--
 tests/unit/test_cluster.py | 113 +++++++++++++++++++++++++++++++++++--
 2 files changed, 109 insertions(+), 12 deletions(-)

diff --git a/src/cluster.py b/src/cluster.py
index b1ec8d18c9..9ac5fdc3d9 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -761,13 +761,7 @@ def reinitialise_raft_data(self) -> None:
         logger.info("Waiting for new raft cluster to initialise")
         for attempt in Retrying(wait=wait_fixed(5)):
             with attempt:
-                response = requests.get(
-                    f"{self._patroni_url}/health",
-                    verify=self.verify,
-                    timeout=API_REQUEST_TIMEOUT,
-                    auth=self._patroni_auth,
-                )
-                health_status = response.json()
+                health_status = self.get_patroni_health()
                 if (
                     health_status["role"] not in ["leader", "master"]
                     or health_status["state"] != "running"
diff --git a/tests/unit/test_cluster.py b/tests/unit/test_cluster.py
index 21b84bc67f..5326af6fcb 100644
--- a/tests/unit/test_cluster.py
+++ b/tests/unit/test_cluster.py
@@ -1,6 +1,7 @@
 # Copyright 2021 Canonical Ltd.
 # See LICENSE file for licensing details.
 
+from pathlib import Path
 from unittest.mock import MagicMock, Mock, PropertyMock, mock_open, patch, sentinel
 
 import pytest
@@ -136,11 +137,12 @@ def test_get_member_ip(peers_ips, patroni):
 
 
 def test_get_patroni_health(peers_ips, patroni):
-    with patch("cluster.stop_after_delay", new_callable=PropertyMock) as _stop_after_delay, patch(
-        "cluster.wait_fixed", new_callable=PropertyMock
-    ) as _wait_fixed, patch(
-        "charm.Patroni._patroni_url", new_callable=PropertyMock
-    ) as _patroni_url, patch("requests.get", side_effect=mocked_requests_get) as _get:
+    with (
+        patch("cluster.stop_after_delay", new_callable=PropertyMock) as _stop_after_delay,
+        patch("cluster.wait_fixed", new_callable=PropertyMock) as _wait_fixed,
+        patch("charm.Patroni._patroni_url", new_callable=PropertyMock) as _patroni_url,
+        patch("requests.get", side_effect=mocked_requests_get) as _get,
+    ):
         # Test when the Patroni API is reachable.
         _patroni_url.return_value = "http://server1"
         health = patroni.get_patroni_health()
@@ -708,3 +710,104 @@ def test_remove_raft_member(patroni):
         with pytest.raises(RemoveRaftMemberFailedError):
             patroni.remove_raft_member("1.2.3.4")
             assert False
+
+
+def test_remove_raft_member_no_quorum(patroni, harness):
+    with (
+        patch("cluster.TcpUtility") as _tcp_utility,
+        patch("cluster.Patroni.get_patroni_health") as _get_patroni_health,
+        patch(
+            "charm.PostgresqlOperatorCharm.unit_peer_data", new_callable=PropertyMock
+        ) as _unit_peer_data,
+    ):
+        # Async replica
+        _unit_peer_data.return_value = {}
+        _tcp_utility.return_value.executeCommand.return_value = {
+            "partner_node_status_server_1.2.3.4:2222": 0,
+            "has_quorum": False,
+            "leader": None,
+        }
+        _get_patroni_health.return_value = {"role": "replica", "sync_standby": False}
+
+        patroni.remove_raft_member("1.2.3.4")
+        assert harness.charm.unit_peer_data == {"raft_stuck": "True"}
+
+        # No health
+        _unit_peer_data.return_value = {}
+        _tcp_utility.return_value.executeCommand.return_value = {
+            "partner_node_status_server_1.2.3.4:2222": 0,
+            "has_quorum": False,
+            "leader": None,
+        }
+        _get_patroni_health.side_effect = Exception
+
+        patroni.remove_raft_member("1.2.3.4")
+
+        assert harness.charm.unit_peer_data == {"raft_stuck": "True"}
+
+        # Sync replica
+        _unit_peer_data.return_value = {}
+        leader_mock = Mock()
+        leader_mock.host = "1.2.3.4"
+        _tcp_utility.return_value.executeCommand.return_value = {
+            "partner_node_status_server_1.2.3.4:2222": 0,
+            "has_quorum": False,
+            "leader": leader_mock,
+        }
+        _get_patroni_health.side_effect = None
+        _get_patroni_health.return_value = {"role": "replica", "sync_standby": True}
+
+        patroni.remove_raft_member("1.2.3.4")
+
+        assert harness.charm.unit_peer_data == {"raft_candidate": "True", "raft_stuck": "True"}
+
+
+def test_remove_raft_data(patroni):
+    with (
+        patch("cluster.Patroni.stop_patroni") as _stop_patroni,
+        patch("cluster.psutil") as _psutil,
+        patch("cluster.wait_fixed", return_value=wait_fixed(0)),
+        patch("shutil.rmtree") as _rmtree,
+        patch("pathlib.Path.is_dir") as _is_dir,
+        patch("pathlib.Path.exists") as _exists,
+    ):
+        mock_proc_pg = Mock()
+        mock_proc_not_pg = Mock()
+        mock_proc_pg.name.return_value = "postgres"
+        mock_proc_not_pg.name.return_value = "something_else"
+        _psutil.process_iter.side_effect = [[mock_proc_not_pg, mock_proc_pg], [mock_proc_not_pg]]
+
+        patroni.remove_raft_data()
+
+        _stop_patroni.assert_called_once_with()
+        assert _psutil.process_iter.call_count == 2
+        _psutil.process_iter.assert_any_call(["name"])
+        _rmtree.assert_called_once_with(Path(f"{PATRONI_CONF_PATH}/raft"))
+
+
+def test_reinitialise_raft_data(patroni):
+    with (
+        patch("cluster.Patroni.get_patroni_health") as _get_patroni_health,
+        patch("charm.PostgresqlOperatorCharm.update_config") as _update_config,
+        patch("cluster.Patroni.start_patroni") as _start_patroni,
+        patch("cluster.Patroni.restart_patroni") as _restart_patroni,
+        patch("cluster.psutil") as _psutil,
+        patch("cluster.wait_fixed", return_value=wait_fixed(0)),
+    ):
+        mock_proc_pg = Mock()
+        mock_proc_not_pg = Mock()
+        mock_proc_pg.name.return_value = "postgres"
+        mock_proc_not_pg.name.return_value = "something_else"
+        _psutil.process_iter.side_effect = [[mock_proc_not_pg], [mock_proc_not_pg, mock_proc_pg]]
+        _get_patroni_health.side_effect = [
+            {"role": "replica", "state": "streaming"},
+            {"role": "leader", "state": "running"},
+        ]
+
+        patroni.reinitialise_raft_data()
+
+        _update_config.assert_called_once_with(no_peers=True)
+        _start_patroni.assert_called_once_with()
+        _restart_patroni.assert_called_once_with()
+        assert _psutil.process_iter.call_count == 2
+        _psutil.process_iter.assert_any_call(["name"])

From de838b5c8d3483fb51ef0a847b61311cf0ae111f Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Sat, 14 Sep 2024 21:45:12 +0300
Subject: [PATCH 43/85] Force emit event changed

---
 src/cluster.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/cluster.py b/src/cluster.py
index 9ac5fdc3d9..e620f77b25 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -37,6 +37,7 @@
     PATRONI_CONF_PATH,
     PATRONI_LOGS_PATH,
     PATRONI_SERVICE_DEFAULT_PATH,
+    PEER,
     PGBACKREST_CONFIGURATION_FILE,
     POSTGRESQL_CONF_PATH,
     POSTGRESQL_DATA_PATH,
@@ -823,6 +824,7 @@ def remove_raft_member(self, member_ip: str) -> None:
                 logger.info("%s is raft candidate" % self.charm.unit.name)
                 data_flags["raft_candidate"] = "True"
             self.charm.unit_peer_data.update(data_flags)
+            self.charm.on[PEER].relation_changed.emit()
             return
 
         # Remove the member from the raft cluster.

From 1c56adf66e2c87bc8fdcd73d855c6d690e0370e3 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Sat, 14 Sep 2024 22:02:42 +0300
Subject: [PATCH 44/85] Force emit event

---
 src/cluster.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/cluster.py b/src/cluster.py
index e620f77b25..8370687384 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -824,7 +824,13 @@ def remove_raft_member(self, member_ip: str) -> None:
                 logger.info("%s is raft candidate" % self.charm.unit.name)
                 data_flags["raft_candidate"] = "True"
             self.charm.unit_peer_data.update(data_flags)
-            self.charm.on[PEER].relation_changed.emit()
+
+            if self.charm.unit.is_leader():
+                self.charm.on[PEER].relation_changed.emit(
+                    unit=self.charm.unit,
+                    app=self.charm.app,
+                    relation=self.charm.model.get_relation(PEER),
+                )
             return
 
         # Remove the member from the raft cluster.

From 85d0ba32c5d0c382ca23681aeb14c7172b7e6b13 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Mon, 16 Sep 2024 15:36:03 +0300
Subject: [PATCH 45/85] Use app data to sync rejoin

---
 src/charm.py | 61 +++++++++++++++++++++++++++++++---------------------
 1 file changed, 37 insertions(+), 24 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 79981b46af..fb762ac194 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -554,37 +554,41 @@ def _stuck_raft_cluster_check(self) -> bool:
     def _stuck_raft_cluster_rejoin(self) -> bool:
         """Reconnect cluster to new raft."""
         primary = None
+        rejoin = False
         all_units_down = True
+        cleanup = True
+        ret = False
         for key, data in self._peers.data.items():
-            if key == self.app:
+            if key == self.app and "raft_rejoin" in data:
+                rejoin = True
                 continue
             if "raft_primary" in data:
                 primary = key
+                cleanup = False
             elif "raft_stopped" not in data:
                 all_units_down = False
-        if primary:
+            elif "raft_stopped" in data:
+                cleanup = False
+        if self.unit.is_leader() and primary:
             logger.info("Updating new primary endpoint")
             self.app_peer_data.pop("members_ips", None)
             self._add_to_members_ips(self._get_unit_ip(primary))
             self._update_relation_endpoints()
-            if all_units_down:
-                logger.info("Removing stuck raft peer data")
-                for key, data in self._peers.data.items():
-                    if key == self.app:
-                        continue
-                    data.pop("raft_primary", None)
-                    data.pop("raft_stopped", None)
-            return True
-        return False
+            if all_units_down and not rejoin:
+                logger.info("Notify units they can rejoin")
+                self.app_peer_data["raft_rejoin"] = "True"
+                ret = True
+        if rejoin:
+            logger.info("Cleaning up raft unit data")
+            self.unit_peer_data.pop("raft_primary", None)
+            self.unit_peer_data.pop("raft_stopped", None)
+        if self.unit.is_leader() and rejoin and cleanup:
+            logger.info("Cleaning up raft app data")
+            self.app_peer_data.pop("raft_rejoin", None)
+        return ret
 
     def _raft_reinitialisation(self) -> bool:
         """Handle raft cluster loss of quorum."""
-        should_exit = False
-        for key in self.unit_peer_data.keys():
-            if key.startswith("raft_"):
-                should_exit = True
-                break
-
         if self.unit.is_leader() and self._stuck_raft_cluster_check():
             should_exit = True
 
@@ -599,23 +603,32 @@ def _raft_reinitialisation(self) -> bool:
             else:
                 self.unit_peer_data["raft_stopped"] = "True"
 
-        if self.unit.is_leader() and self._stuck_raft_cluster_rejoin():
+        if self._stuck_raft_cluster_rejoin():
             should_exit = True
         return should_exit
 
+    def _has_raft_keys(self):
+        if "raft_rejoin" in self.app_peer_data:
+            return True
+
+        for key in self.unit_peer_data.keys():
+            if key.startswith("raft_"):
+                return True
+        return False
+
     def _peer_relation_changed_checks(self, event: HookEvent) -> bool:
         """Split of to reduce complexity."""
-        # Check whether raft is stuck
-        if hasattr(event, "unit") and event.unit and self._raft_reinitialisation():
-            logger.debug("Early exit on_peer_relation_changed: stuck raft recovery")
-            return False
-
         # Prevents the cluster to be reconfigured before it's bootstrapped in the leader.
         if "cluster_initialised" not in self._peers.data[self.app]:
             logger.debug("Deferring on_peer_relation_changed: cluster not initialized")
             event.defer()
             return False
 
+        # Check whether raft is stuck
+        if self._has_raft_keys() and self._raft_reinitialisation():
+            logger.debug("Early exit on_peer_relation_changed: stuck raft recovery")
+            return False
+
         # If the unit is the leader, it can reconfigure the cluster.
         if self.unit.is_leader() and not self._reconfigure_cluster(event):
             event.defer()
@@ -1017,7 +1030,7 @@ def _on_leader_elected(self, event: LeaderElectedEvent) -> None:
             if self.get_secret(APP_SCOPE, key) is None:
                 self.set_secret(APP_SCOPE, key, new_password())
 
-        if self._raft_reinitialisation():
+        if self._has_raft_keys() and self._raft_reinitialisation():
             return
 
         # Update the list of the current PostgreSQL hosts when a new leader is elected.

From c3171041547c629f92d39a260b1b40798fa838e9 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Mon, 16 Sep 2024 16:04:14 +0300
Subject: [PATCH 46/85] Missing var

---
 src/charm.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/charm.py b/src/charm.py
index fb762ac194..e9ef5bfa3e 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -589,6 +589,7 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
 
     def _raft_reinitialisation(self) -> bool:
         """Handle raft cluster loss of quorum."""
+        should_exit = False
         if self.unit.is_leader() and self._stuck_raft_cluster_check():
             should_exit = True
 

From 887f008050e295ad0ee0b0aed70bce8b36e69c6d Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Mon, 16 Sep 2024 17:48:09 +0300
Subject: [PATCH 47/85] App data sync for stuck cluster detection

---
 src/charm.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index e9ef5bfa3e..104c60cf18 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -542,13 +542,7 @@ def _stuck_raft_cluster_check(self) -> bool:
             return True
 
         logger.info("%s selected for new raft leader" % candidate.name)
-        for key, data in self._peers.data.items():
-            if key == self.app:
-                continue
-            data.pop("raft_stuck", None)
-            if key != candidate:
-                data.pop("raft_candidate", None)
-            data["raft_stopping"] = "True"
+        self.app_peer_data["raft_candidate"] = candidate.name
         return True
 
     def _stuck_raft_cluster_rejoin(self) -> bool:
@@ -593,6 +587,13 @@ def _raft_reinitialisation(self) -> bool:
         if self.unit.is_leader() and self._stuck_raft_cluster_check():
             should_exit = True
 
+        if candidate := self.app_peer_data.get("raft_candidate"):
+            self.unit_peer_data.pop("raft_stuck", None)
+            if self.unit.name != candidate:
+                self.unit_peer_data.pop("raft_candidate", None)
+            self.unit_peer_data["raft_stopping"] = "True"
+            should_exit = True
+
         if "raft_stopping" in self.unit_peer_data:
             should_exit = True
             self._patroni.remove_raft_data()

From 24a3fe03d718010ab7dc8bcb523d34b485ad0f71 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Mon, 16 Sep 2024 18:43:35 +0300
Subject: [PATCH 48/85] Cleanup logic

---
 src/charm.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 104c60cf18..70396f7c96 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -579,6 +579,7 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
         if self.unit.is_leader() and rejoin and cleanup:
             logger.info("Cleaning up raft app data")
             self.app_peer_data.pop("raft_rejoin", None)
+            self.app_peer_data.pop("raft_candidate", None)
         return ret
 
     def _raft_reinitialisation(self) -> bool:
@@ -587,22 +588,25 @@ def _raft_reinitialisation(self) -> bool:
         if self.unit.is_leader() and self._stuck_raft_cluster_check():
             should_exit = True
 
-        if candidate := self.app_peer_data.get("raft_candidate"):
+        if (
+            candidate := self.app_peer_data.get("raft_candidate")
+            and "raft_stuck" not in self.unit_peer_data
+        ):
+            should_exit = True
             self.unit_peer_data.pop("raft_stuck", None)
-            if self.unit.name != candidate:
-                self.unit_peer_data.pop("raft_candidate", None)
+            self.unit_peer_data.pop("raft_candidate", None)
             self.unit_peer_data["raft_stopping"] = "True"
-            should_exit = True
 
         if "raft_stopping" in self.unit_peer_data:
             should_exit = True
             self._patroni.remove_raft_data()
             self.unit_peer_data.pop("raft_stopping", None)
-            if "raft_candidate" in self.unit_peer_data:
+            if candidate == self.unit.name:
+                logger.info("Reinitialising %s as primary" % self.unit.name)
                 self._patroni.reinitialise_raft_data()
-                self.unit_peer_data.pop("raft_candidate", None)
                 self.unit_peer_data["raft_primary"] = "True"
             else:
+                logger.info("Stopping %s" % self.unit.name)
                 self.unit_peer_data["raft_stopped"] = "True"
 
         if self._stuck_raft_cluster_rejoin():

From ac7e21789bf9d252848dd6fa2999292efbf9f4c8 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Mon, 16 Sep 2024 19:36:49 +0300
Subject: [PATCH 49/85] Rename app candidate flag

---
 src/charm.py | 14 +++++---------
 1 file changed, 5 insertions(+), 9 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 70396f7c96..707e2bcf5b 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -511,7 +511,8 @@ def _on_pgdata_storage_detaching(self, _) -> None:
         if self.primary_endpoint:
             self._update_relation_endpoints()
 
-    def _collect_raft_cluster_state(self) -> (bool, bool, Optional[Unit]):
+    def _stuck_raft_cluster_check(self) -> bool:
+        """Check for stuck raft cluster and reinitialise if safe."""
         raft_stuck = False
         all_units_stuck = True
         candidate = None
@@ -524,11 +525,6 @@ def _collect_raft_cluster_state(self) -> (bool, bool, Optional[Unit]):
                 all_units_stuck = False
             if not candidate and "raft_candidate" in data:
                 candidate = key
-        return raft_stuck, all_units_stuck, candidate
-
-    def _stuck_raft_cluster_check(self) -> bool:
-        """Check for stuck raft cluster and reinitialise if safe."""
-        raft_stuck, all_units_stuck, candidate = self._collect_raft_cluster_state()
 
         if not raft_stuck:
             return False
@@ -542,7 +538,7 @@ def _stuck_raft_cluster_check(self) -> bool:
             return True
 
         logger.info("%s selected for new raft leader" % candidate.name)
-        self.app_peer_data["raft_candidate"] = candidate.name
+        self.app_peer_data["raft_selected_candidate"] = candidate.name
         return True
 
     def _stuck_raft_cluster_rejoin(self) -> bool:
@@ -579,7 +575,7 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
         if self.unit.is_leader() and rejoin and cleanup:
             logger.info("Cleaning up raft app data")
             self.app_peer_data.pop("raft_rejoin", None)
-            self.app_peer_data.pop("raft_candidate", None)
+            self.app_peer_data.pop("raft_selected_candidate", None)
         return ret
 
     def _raft_reinitialisation(self) -> bool:
@@ -589,7 +585,7 @@ def _raft_reinitialisation(self) -> bool:
             should_exit = True
 
         if (
-            candidate := self.app_peer_data.get("raft_candidate")
+            candidate := self.app_peer_data.get("raft_selected_candidate")
             and "raft_stuck" not in self.unit_peer_data
         ):
             should_exit = True

From 26a2bc6ebc2f5f658d226b770e7095f3debd7602 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 17 Sep 2024 02:56:40 +0300
Subject: [PATCH 50/85] Wrong condition

---
 src/charm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/charm.py b/src/charm.py
index 707e2bcf5b..30f9858949 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -586,7 +586,7 @@ def _raft_reinitialisation(self) -> bool:
 
         if (
             candidate := self.app_peer_data.get("raft_selected_candidate")
-            and "raft_stuck" not in self.unit_peer_data
+            and "raft_stuck" in self.unit_peer_data
         ):
             should_exit = True
             self.unit_peer_data.pop("raft_stuck", None)

From 63b5d2c475b2b39327808a65788ccf8ee07990be Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 17 Sep 2024 04:44:45 +0300
Subject: [PATCH 51/85] Check all raft_ keys

---
 src/charm.py | 40 +++++++++++++++++++---------------------
 1 file changed, 19 insertions(+), 21 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 30f9858949..83ea4d71b0 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -584,34 +584,32 @@ def _raft_reinitialisation(self) -> bool:
         if self.unit.is_leader() and self._stuck_raft_cluster_check():
             should_exit = True
 
-        if (
-            candidate := self.app_peer_data.get("raft_selected_candidate")
-            and "raft_stuck" in self.unit_peer_data
-        ):
-            should_exit = True
-            self.unit_peer_data.pop("raft_stuck", None)
-            self.unit_peer_data.pop("raft_candidate", None)
-            self.unit_peer_data["raft_stopping"] = "True"
-
-        if "raft_stopping" in self.unit_peer_data:
+        if candidate := self.app_peer_data.get("raft_selected_candidate"):
             should_exit = True
-            self._patroni.remove_raft_data()
-            self.unit_peer_data.pop("raft_stopping", None)
-            if candidate == self.unit.name:
-                logger.info("Reinitialising %s as primary" % self.unit.name)
-                self._patroni.reinitialise_raft_data()
-                self.unit_peer_data["raft_primary"] = "True"
-            else:
-                logger.info("Stopping %s" % self.unit.name)
-                self.unit_peer_data["raft_stopped"] = "True"
+            if "raft_stuck" in self.unit_peer_data:
+                self.unit_peer_data.pop("raft_stuck", None)
+                self.unit_peer_data.pop("raft_candidate", None)
+                self.unit_peer_data["raft_stopping"] = "True"
+
+            if "raft_stopping" in self.unit_peer_data:
+                self._patroni.remove_raft_data()
+                self.unit_peer_data.pop("raft_stopping", None)
+                if candidate == self.unit.name:
+                    logger.info("Reinitialising %s as primary" % self.unit.name)
+                    self._patroni.reinitialise_raft_data()
+                    self.unit_peer_data["raft_primary"] = "True"
+                else:
+                    logger.info("Stopping %s" % self.unit.name)
+                    self.unit_peer_data["raft_stopped"] = "True"
 
         if self._stuck_raft_cluster_rejoin():
             should_exit = True
         return should_exit
 
     def _has_raft_keys(self):
-        if "raft_rejoin" in self.app_peer_data:
-            return True
+        for key in self.app_peer_data.keys():
+            if key.startswith("raft_"):
+                return True
 
         for key in self.unit_peer_data.keys():
             if key.startswith("raft_"):

From 3a4eda7e546130461ea1b2f132903158498c8b67 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 17 Sep 2024 05:49:07 +0300
Subject: [PATCH 52/85] Rework rejoin logic

---
 src/charm.py | 42 ++++++++++++++++++++++++------------------
 1 file changed, 24 insertions(+), 18 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 83ea4d71b0..3b7475d861 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -544,22 +544,16 @@ def _stuck_raft_cluster_check(self) -> bool:
     def _stuck_raft_cluster_rejoin(self) -> bool:
         """Reconnect cluster to new raft."""
         primary = None
-        rejoin = False
         all_units_down = True
-        cleanup = True
-        ret = False
         for key, data in self._peers.data.items():
             if key == self.app and "raft_rejoin" in data:
                 rejoin = True
                 continue
             if "raft_primary" in data:
                 primary = key
-                cleanup = False
             elif "raft_stopped" not in data:
                 all_units_down = False
-            elif "raft_stopped" in data:
-                cleanup = False
-        if self.unit.is_leader() and primary:
+        if primary:
             logger.info("Updating new primary endpoint")
             self.app_peer_data.pop("members_ips", None)
             self._add_to_members_ips(self._get_unit_ip(primary))
@@ -567,16 +561,19 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
             if all_units_down and not rejoin:
                 logger.info("Notify units they can rejoin")
                 self.app_peer_data["raft_rejoin"] = "True"
-                ret = True
-        if rejoin:
-            logger.info("Cleaning up raft unit data")
-            self.unit_peer_data.pop("raft_primary", None)
-            self.unit_peer_data.pop("raft_stopped", None)
-        if self.unit.is_leader() and rejoin and cleanup:
-            logger.info("Cleaning up raft app data")
-            self.app_peer_data.pop("raft_rejoin", None)
-            self.app_peer_data.pop("raft_selected_candidate", None)
-        return ret
+                return True
+        return False
+
+    def _stuck_raft_cluster_cleanup(self) -> None:
+        for key, data in self._peers.data.items():
+            if key == self.app:
+                continue
+            if "raft_primary" in data or "raft_stoppedi" in data:
+                return
+
+        logger.info("Cleaning up raft app data")
+        self.app_peer_data.pop("raft_rejoin", None)
+        self.app_peer_data.pop("raft_selected_candidate", None)
 
     def _raft_reinitialisation(self) -> bool:
         """Handle raft cluster loss of quorum."""
@@ -602,8 +599,17 @@ def _raft_reinitialisation(self) -> bool:
                     logger.info("Stopping %s" % self.unit.name)
                     self.unit_peer_data["raft_stopped"] = "True"
 
-        if self._stuck_raft_cluster_rejoin():
+        if self.unit.is_leader() and self._stuck_raft_cluster_rejoin():
+            should_exit = True
+
+        if "raft_rejoin" in self.app_peer_data:
             should_exit = True
+            logger.info("Cleaning up raft unit data")
+            self.unit_peer_data.pop("raft_primary", None)
+            self.unit_peer_data.pop("raft_stopped", None)
+
+            if self.unit.is_leader() and "raft_rejoin":
+                self._stuck_raft_cluster_cleanup()
         return should_exit
 
     def _has_raft_keys(self):

From bca3e3af04c736e887382ecbcc63bbb573962508 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 17 Sep 2024 13:49:07 +0300
Subject: [PATCH 53/85] Cleanup cleanup

---
 src/charm.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 3b7475d861..dadf38f271 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -546,8 +546,7 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
         primary = None
         all_units_down = True
         for key, data in self._peers.data.items():
-            if key == self.app and "raft_rejoin" in data:
-                rejoin = True
+            if key == self.app:
                 continue
             if "raft_primary" in data:
                 primary = key
@@ -558,7 +557,7 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
             self.app_peer_data.pop("members_ips", None)
             self._add_to_members_ips(self._get_unit_ip(primary))
             self._update_relation_endpoints()
-            if all_units_down and not rejoin:
+            if all_units_down and "raft_rejoin" not in self.app_peer_data:
                 logger.info("Notify units they can rejoin")
                 self.app_peer_data["raft_rejoin"] = "True"
                 return True
@@ -568,7 +567,7 @@ def _stuck_raft_cluster_cleanup(self) -> None:
         for key, data in self._peers.data.items():
             if key == self.app:
                 continue
-            if "raft_primary" in data or "raft_stoppedi" in data:
+            if "raft_primary" in data or "raft_stopped" in data:
                 return
 
         logger.info("Cleaning up raft app data")

From 0e9686a43368578e7df27053091b3eefab0b27e2 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 17 Sep 2024 16:07:07 +0300
Subject: [PATCH 54/85] Delete peers just once

---
 src/charm.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index dadf38f271..63fa2a0ede 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -552,15 +552,16 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
                 primary = key
             elif "raft_stopped" not in data:
                 all_units_down = False
-        if primary:
+        if primary and "raft_rejoin" not in self.app_peer_data:
             logger.info("Updating new primary endpoint")
             self.app_peer_data.pop("members_ips", None)
             self._add_to_members_ips(self._get_unit_ip(primary))
+            self.app_peer_data["raft_reset_primary"] = "True"
             self._update_relation_endpoints()
-            if all_units_down and "raft_rejoin" not in self.app_peer_data:
-                logger.info("Notify units they can rejoin")
-                self.app_peer_data["raft_rejoin"] = "True"
-                return True
+        if all_units_down and "raft_rejoin" not in self.app_peer_data:
+            logger.info("Notify units they can rejoin")
+            self.app_peer_data["raft_rejoin"] = "True"
+            return True
         return False
 
     def _stuck_raft_cluster_cleanup(self) -> None:
@@ -572,6 +573,7 @@ def _stuck_raft_cluster_cleanup(self) -> None:
 
         logger.info("Cleaning up raft app data")
         self.app_peer_data.pop("raft_rejoin", None)
+        self.app_peer_data.pop("raft_reset_primary", None)
         self.app_peer_data.pop("raft_selected_candidate", None)
 
     def _raft_reinitialisation(self) -> bool:

From 3974ea2a54cbc4376d1aabb72d4470a692a17dfe Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 17 Sep 2024 17:11:35 +0300
Subject: [PATCH 55/85] Increase timout

---
 src/charm.py                               |  6 ++---
 tests/integration/ha_tests/test_scaling.py |  2 +-
 tests/unit/test_charm.py                   | 28 ++++++++++++++++++++++
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 63fa2a0ede..9d874e8a0d 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -536,9 +536,9 @@ def _stuck_raft_cluster_check(self) -> bool:
         if not candidate:
             logger.warning("Stuck raft has no candidate")
             return True
-
-        logger.info("%s selected for new raft leader" % candidate.name)
-        self.app_peer_data["raft_selected_candidate"] = candidate.name
+        if "raft_selected_candidate" not in self.app_peer_data:
+            logger.info("%s selected for new raft leader" % candidate.name)
+            self.app_peer_data["raft_selected_candidate"] = candidate.name
         return True
 
     def _stuck_raft_cluster_rejoin(self) -> bool:
diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index 054c3752d0..2bcef61faf 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -154,7 +154,7 @@ async def test_removing_raft_majority(ops_test: OpsTest, continuous_writes) -> N
         ),
     )
 
-    await ops_test.model.wait_for_idle(status="active", timeout=600)
+    await ops_test.model.wait_for_idle(status="active", timeout=1500)
 
     await are_writes_increasing(ops_test, original_roles["primaries"][0])
 
diff --git a/tests/unit/test_charm.py b/tests/unit/test_charm.py
index f1bc570076..4f13b1862e 100644
--- a/tests/unit/test_charm.py
+++ b/tests/unit/test_charm.py
@@ -1847,6 +1847,34 @@ def test_add_cluster_member(harness):
             harness.charm.add_cluster_member("postgresql/0")
 
 
+def test_stuck_raft_cluster_check(harness):
+    # doesn't raise flags if there are no raft flags
+    assert not harness.charm._stuck_raft_cluster_check()
+
+    # Raft is stuck
+    rel_id = harness.model.get_relation(PEER).id
+    with harness.hooks_disabled():
+        harness.set_leader()
+        harness.update_relation_data(rel_id, harness.charm.unit.name, {"raft_stuck": "True"})
+
+    assert harness.charm._stuck_raft_cluster_check()
+    assert "raft_selected_candidate" not in harness.charm.app_peer_data
+
+    # Raft candidate
+    with harness.hooks_disabled():
+        harness.update_relation_data(rel_id, harness.charm.unit.name, {"raft_candidate": "True"})
+    assert harness.charm._stuck_raft_cluster_check()
+    assert harness.charm.app_peer_data["raft_selected_candidate"] == harness.charm.unit.name
+
+    # Don't override existing candidate
+    with harness.hooks_disabled():
+        harness.update_relation_data(
+            rel_id, harness.charm.app.name, {"raft_selected_candidate": "something_else"}
+        )
+    assert harness.charm._stuck_raft_cluster_check()
+    assert harness.charm.app_peer_data["raft_selected_candidate"] != harness.charm.unit.name
+
+
 #
 # Secrets
 #

From 4048015a33093cb4f3a02bad805c5f651e3b997c Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 17 Sep 2024 18:24:46 +0300
Subject: [PATCH 56/85] Force start on rejoin

---
 src/charm.py                               | 3 ++-
 tests/integration/ha_tests/test_scaling.py | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 9d874e8a0d..5baf782618 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -552,7 +552,7 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
                 primary = key
             elif "raft_stopped" not in data:
                 all_units_down = False
-        if primary and "raft_rejoin" not in self.app_peer_data:
+        if primary and "raft_reset_primary" not in self.app_peer_data:
             logger.info("Updating new primary endpoint")
             self.app_peer_data.pop("members_ips", None)
             self._add_to_members_ips(self._get_unit_ip(primary))
@@ -608,6 +608,7 @@ def _raft_reinitialisation(self) -> bool:
             logger.info("Cleaning up raft unit data")
             self.unit_peer_data.pop("raft_primary", None)
             self.unit_peer_data.pop("raft_stopped", None)
+            self._patroni.start_patroni()
 
             if self.unit.is_leader() and "raft_rejoin":
                 self._stuck_raft_cluster_cleanup()
diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index 2bcef61faf..6ef35712fb 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -160,7 +160,7 @@ async def test_removing_raft_majority(ops_test: OpsTest, continuous_writes) -> N
 
     logger.info("Scaling back up")
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=2)
-    await ops_test.model.wait_for_idle(status="active", timeout=1500)
+    await ops_test.model.wait_for_idle(status="active", timeout=600)
 
     await check_writes(ops_test)
     new_roles = await get_cluster_roles(

From 819748fe14b87b58d92741fe50ad12f3fce46a02 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 17 Sep 2024 20:12:45 +0300
Subject: [PATCH 57/85] Supress db reinit

---
 src/charm.py             |  6 ++++--
 tests/unit/test_charm.py | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 5baf782618..0fb7c79963 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -690,7 +690,8 @@ def _on_peer_relation_changed(self, event: HookEvent):
         # Restart the workload if it's stuck on the starting state after a timeline divergence
         # due to a backup that was restored.
         if (
-            not self.is_primary
+            not self._has_raft_keys()
+            and not self.is_primary
             and not self.is_standby_leader
             and (
                 self._patroni.member_replication_lag == "unknown"
@@ -1549,7 +1550,8 @@ def _handle_workload_failures(self) -> bool:
             return False
 
         if (
-            not is_primary
+            not self._has_raft_keys()
+            and not is_primary
             and not is_standby_leader
             and not self._patroni.member_started
             and "postgresql_restarted" in self._peers.data[self.unit]
diff --git a/tests/unit/test_charm.py b/tests/unit/test_charm.py
index 4f13b1862e..ac4130040c 100644
--- a/tests/unit/test_charm.py
+++ b/tests/unit/test_charm.py
@@ -1875,6 +1875,45 @@ def test_stuck_raft_cluster_check(harness):
     assert harness.charm.app_peer_data["raft_selected_candidate"] != harness.charm.unit.name
 
 
+def test_stuck_raft_cluster_cleanup(harness):
+    rel_id = harness.model.get_relation(PEER).id
+
+    # Cleans up app data
+    with harness.hooks_disabled():
+        harness.update_relation_data(
+            rel_id,
+            harness.charm.app.name,
+            {
+                "raft_rejoin": "True",
+                "raft_reset_primary": "True",
+                "raft_selected_candidate": "unit_name",
+            },
+        )
+    harness.charm._stuck_raft_cluster_cleanup()
+
+    assert "raft_rejoin" not in harness.charm.app_peer_data
+    assert "raft_reset_primary" not in harness.charm.app_peer_data
+    assert "raft_selected_candidate" not in harness.charm.app_peer_data
+
+    # Don't clean up if there's unit data flags
+    with harness.hooks_disabled():
+        harness.update_relation_data(rel_id, harness.charm.unit.name, {"raft_primary": "True"})
+        harness.update_relation_data(
+            rel_id,
+            harness.charm.app.name,
+            {
+                "raft_rejoin": "True",
+                "raft_reset_primary": "True",
+                "raft_selected_candidate": "unit_name",
+            },
+        )
+    harness.charm._stuck_raft_cluster_cleanup()
+
+    assert "raft_rejoin" in harness.charm.app_peer_data
+    assert "raft_reset_primary" in harness.charm.app_peer_data
+    assert "raft_selected_candidate" in harness.charm.app_peer_data
+
+
 #
 # Secrets
 #

From 761d28dad34ae8489224892cbb4f2ea8b42c42a6 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 17 Sep 2024 22:19:32 +0300
Subject: [PATCH 58/85] Higher coverage

---
 tests/integration/ha_tests/test_scaling.py |  2 +-
 tests/unit/test_charm.py                   | 60 ++++++++++++++++++++++
 2 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index 6ef35712fb..7da469de50 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -159,7 +159,7 @@ async def test_removing_raft_majority(ops_test: OpsTest, continuous_writes) -> N
     await are_writes_increasing(ops_test, original_roles["primaries"][0])
 
     logger.info("Scaling back up")
-    await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=2)
+    await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=3)
     await ops_test.model.wait_for_idle(status="active", timeout=600)
 
     await check_writes(ops_test)
diff --git a/tests/unit/test_charm.py b/tests/unit/test_charm.py
index ac4130040c..c88a410db6 100644
--- a/tests/unit/test_charm.py
+++ b/tests/unit/test_charm.py
@@ -1914,6 +1914,66 @@ def test_stuck_raft_cluster_cleanup(harness):
     assert "raft_selected_candidate" in harness.charm.app_peer_data
 
 
+def test_stuck_raft_cluster_rejoin(harness):
+    rel_id = harness.model.get_relation(PEER).id
+
+    with (
+        patch(
+            "charm.PostgresqlOperatorCharm._update_relation_endpoints"
+        ) as _update_relation_endpoints,
+        patch("charm.PostgresqlOperatorCharm._add_to_members_ips") as _add_to_members_ips,
+    ):
+        # No data
+        assert not harness.charm._stuck_raft_cluster_rejoin()
+
+        # Raises primary flag
+        with harness.hooks_disabled():
+            harness.update_relation_data(
+                rel_id, harness.charm.unit.name, {"raft_primary": "test_primary"}
+            )
+            harness.update_relation_data(rel_id, harness.charm.app.name, {"members_ips": "values"})
+
+        assert harness.charm._stuck_raft_cluster_rejoin()
+
+        assert "raft_reset_primary" in harness.charm.app_peer_data
+        assert "raft_rejoin" in harness.charm.app_peer_data
+        assert "members_ips" not in harness.charm.app_peer_data
+        _add_to_members_ips.assert_called_once_with("192.0.2.0")
+        _update_relation_endpoints.assert_called_once_with()
+
+
+def test_raft_reinitialisation(harness):
+    rel_id = harness.model.get_relation(PEER).id
+
+    with (
+        patch(
+            "charm.PostgresqlOperatorCharm._stuck_raft_cluster_check"
+        ) as _stuck_raft_cluster_check,
+        patch(
+            "charm.PostgresqlOperatorCharm._stuck_raft_cluster_rejoin"
+        ) as _stuck_raft_cluster_rejoin,
+        patch(
+            "charm.PostgresqlOperatorCharm._stuck_raft_cluster_cleanup"
+        ) as _stuck_raft_cluster_cleanup,
+    ):
+        # No data
+        assert not harness.charm._raft_reinitialisation()
+
+        with harness.hooks_disabled():
+            harness.set_leader()
+            harness.update_relation_data(
+                rel_id, harness.charm.unit.name, {"raft_primary": "True", "raft_stopped": "True"}
+            )
+            harness.update_relation_data(rel_id, harness.charm.app.name, {"raft_rejoin": "True"})
+
+        assert harness.charm._raft_reinitialisation()
+        _stuck_raft_cluster_rejoin.assert_called_once_with()
+        _stuck_raft_cluster_check.assert_called_once_with()
+        _stuck_raft_cluster_cleanup.assert_called_once_with()
+        assert "raft_stopped" not in harness.charm.unit_peer_data
+        assert "raft_primary" not in harness.charm.unit_peer_data
+
+
 #
 # Secrets
 #

From 0163017bade05b381514befc8c0fb450f3d12333 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 17 Sep 2024 22:37:06 +0300
Subject: [PATCH 59/85] Unit tests

---
 tests/unit/test_charm.py | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/tests/unit/test_charm.py b/tests/unit/test_charm.py
index c88a410db6..44594ebbdb 100644
--- a/tests/unit/test_charm.py
+++ b/tests/unit/test_charm.py
@@ -1955,23 +1955,50 @@ def test_raft_reinitialisation(harness):
         patch(
             "charm.PostgresqlOperatorCharm._stuck_raft_cluster_cleanup"
         ) as _stuck_raft_cluster_cleanup,
+        patch("charm.Patroni.remove_raft_data") as _remove_raft_data,
+        patch("charm.Patroni.reinitialise_raft_data") as _reinitialise_raft_data,
+        patch("charm.Patroni.start_patroni") as _start_patroni,
     ):
         # No data
         assert not harness.charm._raft_reinitialisation()
 
+        # Different candidate
         with harness.hooks_disabled():
             harness.set_leader()
             harness.update_relation_data(
-                rel_id, harness.charm.unit.name, {"raft_primary": "True", "raft_stopped": "True"}
+                rel_id, harness.charm.unit.name, {"raft_stuck": "True", "raft_candidate": "True"}
+            )
+            harness.update_relation_data(
+                rel_id,
+                harness.charm.app.name,
+                {"raft_rejoin": "True", "raft_selected_candidate": "test_candidate"},
             )
-            harness.update_relation_data(rel_id, harness.charm.app.name, {"raft_rejoin": "True"})
 
         assert harness.charm._raft_reinitialisation()
         _stuck_raft_cluster_rejoin.assert_called_once_with()
         _stuck_raft_cluster_check.assert_called_once_with()
         _stuck_raft_cluster_cleanup.assert_called_once_with()
+        _remove_raft_data.assert_called_once_with()
+        _start_patroni.assert_called_once_with()
+        assert not _reinitialise_raft_data.called
         assert "raft_stopped" not in harness.charm.unit_peer_data
         assert "raft_primary" not in harness.charm.unit_peer_data
+        _start_patroni.reset_mock()
+
+        # Current candidate
+        with harness.hooks_disabled():
+            harness.set_leader()
+            harness.update_relation_data(
+                rel_id, harness.charm.unit.name, {"raft_stuck": "True", "raft_candidate": "True"}
+            )
+            harness.update_relation_data(
+                rel_id,
+                harness.charm.app.name,
+                {"raft_selected_candidate": "postgresql/0"},
+            )
+
+        assert harness.charm._raft_reinitialisation()
+        _reinitialise_raft_data.assert_called_once_with()
 
 
 #

From db875adb2ee56869cae77cc2934d8ff0f1440a26 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 18 Sep 2024 01:21:09 +0300
Subject: [PATCH 60/85] Try not to reset members cache

---
 src/charm.py             | 2 --
 tests/unit/test_charm.py | 4 ----
 2 files changed, 6 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 0fb7c79963..9b02429374 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -554,8 +554,6 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
                 all_units_down = False
         if primary and "raft_reset_primary" not in self.app_peer_data:
             logger.info("Updating new primary endpoint")
-            self.app_peer_data.pop("members_ips", None)
-            self._add_to_members_ips(self._get_unit_ip(primary))
             self.app_peer_data["raft_reset_primary"] = "True"
             self._update_relation_endpoints()
         if all_units_down and "raft_rejoin" not in self.app_peer_data:
diff --git a/tests/unit/test_charm.py b/tests/unit/test_charm.py
index 44594ebbdb..c6212ba84b 100644
--- a/tests/unit/test_charm.py
+++ b/tests/unit/test_charm.py
@@ -1921,7 +1921,6 @@ def test_stuck_raft_cluster_rejoin(harness):
         patch(
             "charm.PostgresqlOperatorCharm._update_relation_endpoints"
         ) as _update_relation_endpoints,
-        patch("charm.PostgresqlOperatorCharm._add_to_members_ips") as _add_to_members_ips,
     ):
         # No data
         assert not harness.charm._stuck_raft_cluster_rejoin()
@@ -1931,14 +1930,11 @@ def test_stuck_raft_cluster_rejoin(harness):
             harness.update_relation_data(
                 rel_id, harness.charm.unit.name, {"raft_primary": "test_primary"}
             )
-            harness.update_relation_data(rel_id, harness.charm.app.name, {"members_ips": "values"})
 
         assert harness.charm._stuck_raft_cluster_rejoin()
 
         assert "raft_reset_primary" in harness.charm.app_peer_data
         assert "raft_rejoin" in harness.charm.app_peer_data
-        assert "members_ips" not in harness.charm.app_peer_data
-        _add_to_members_ips.assert_called_once_with("192.0.2.0")
         _update_relation_endpoints.assert_called_once_with()
 
 

From 89a1fee3310438fb865c8e877930eff34d1d5301 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 18 Sep 2024 13:04:50 +0300
Subject: [PATCH 61/85] Get selected candidate from peer data

---
 src/charm.py   | 13 +++++++------
 src/cluster.py |  4 ++++
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 9b02429374..6c04e9c556 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -515,7 +515,7 @@ def _stuck_raft_cluster_check(self) -> bool:
         """Check for stuck raft cluster and reinitialise if safe."""
         raft_stuck = False
         all_units_stuck = True
-        candidate = None
+        candidate = self.app_peer_data.get("raft_selected_candidate")
         for key, data in self._peers.data.items():
             if key == self.app:
                 continue
@@ -612,7 +612,8 @@ def _raft_reinitialisation(self) -> bool:
                 self._stuck_raft_cluster_cleanup()
         return should_exit
 
-    def _has_raft_keys(self):
+    def has_raft_keys(self):
+        """Checks for the presence of raft recovery keys in peer data."""
         for key in self.app_peer_data.keys():
             if key.startswith("raft_"):
                 return True
@@ -631,7 +632,7 @@ def _peer_relation_changed_checks(self, event: HookEvent) -> bool:
             return False
 
         # Check whether raft is stuck
-        if self._has_raft_keys() and self._raft_reinitialisation():
+        if self.has_raft_keys() and self._raft_reinitialisation():
             logger.debug("Early exit on_peer_relation_changed: stuck raft recovery")
             return False
 
@@ -688,7 +689,7 @@ def _on_peer_relation_changed(self, event: HookEvent):
         # Restart the workload if it's stuck on the starting state after a timeline divergence
         # due to a backup that was restored.
         if (
-            not self._has_raft_keys()
+            not self.has_raft_keys()
             and not self.is_primary
             and not self.is_standby_leader
             and (
@@ -1037,7 +1038,7 @@ def _on_leader_elected(self, event: LeaderElectedEvent) -> None:
             if self.get_secret(APP_SCOPE, key) is None:
                 self.set_secret(APP_SCOPE, key, new_password())
 
-        if self._has_raft_keys() and self._raft_reinitialisation():
+        if self.has_raft_keys() and self._raft_reinitialisation():
             return
 
         # Update the list of the current PostgreSQL hosts when a new leader is elected.
@@ -1548,7 +1549,7 @@ def _handle_workload_failures(self) -> bool:
             return False
 
         if (
-            not self._has_raft_keys()
+            not self.has_raft_keys()
             and not is_primary
             and not is_standby_leader
             and not self._patroni.member_started
diff --git a/src/cluster.py b/src/cluster.py
index 8370687384..4ee537b846 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -793,6 +793,10 @@ def remove_raft_member(self, member_ip: str) -> None:
             RaftMemberNotFoundError: if the member to be removed
                 is not part of the raft cluster.
         """
+        if self.charm.has_raft_keys():
+            logger.debug("Remove raft member: Raft already in recovery")
+            return
+
         # Get the status of the raft cluster.
         syncobj_util = TcpUtility(password=self.raft_password, timeout=3)
 

From 95c5884514a9197ab7cc7e11bea431c54ab2ad35 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 18 Sep 2024 14:44:42 +0300
Subject: [PATCH 62/85] Reset members

---
 src/charm.py             | 16 ++++++++++------
 tests/unit/test_charm.py |  3 +++
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 6c04e9c556..41ef8a43c2 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -553,7 +553,9 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
             elif "raft_stopped" not in data:
                 all_units_down = False
         if primary and "raft_reset_primary" not in self.app_peer_data:
-            logger.info("Updating new primary endpoint")
+            logger.info("Updating the primary endpoint")
+            self.app_peer_data.pop("members_ips", None)
+            self._add_to_members_ips(self._get_unit_ip(primary))
             self.app_peer_data["raft_reset_primary"] = "True"
             self._update_relation_endpoints()
         if all_units_down and "raft_rejoin" not in self.app_peer_data:
@@ -839,14 +841,16 @@ def add_cluster_member(self, member: str) -> None:
     def _get_unit_ip(self, unit: Unit) -> Optional[str]:
         """Get the IP address of a specific unit."""
         # Check if host is current host.
+        ip = None
         if unit == self.unit:
-            return str(self.model.get_binding(PEER).network.bind_address)
+            ip = self.model.get_binding(PEER).network.bind_address
         # Check if host is a peer.
         elif unit in self._peers.data:
-            return str(self._peers.data[unit].get("private-address"))
+            ip = self._peers.data[unit].get("private-address")
         # Return None if the unit is not a peer neither the current unit.
-        else:
-            return None
+        if ip:
+            return str(ip)
+        return None
 
     @property
     def _hosts(self) -> set:
@@ -949,7 +953,7 @@ def _update_members_ips(self, ip_to_add: str = None, ip_to_remove: str = None) -
             return
 
         ips = json.loads(self._peers.data[self.app].get("members_ips", "[]"))
-        if ip_to_add and ip_to_add != "None" and ip_to_add not in ips:
+        if ip_to_add and ip_to_add not in ips:
             ips.append(ip_to_add)
         elif ip_to_remove:
             ips.remove(ip_to_remove)
diff --git a/tests/unit/test_charm.py b/tests/unit/test_charm.py
index c6212ba84b..7361ea00ef 100644
--- a/tests/unit/test_charm.py
+++ b/tests/unit/test_charm.py
@@ -1921,6 +1921,7 @@ def test_stuck_raft_cluster_rejoin(harness):
         patch(
             "charm.PostgresqlOperatorCharm._update_relation_endpoints"
         ) as _update_relation_endpoints,
+        patch("charm.PostgresqlOperatorCharm._add_to_members_ips") as _add_to_members_ips,
     ):
         # No data
         assert not harness.charm._stuck_raft_cluster_rejoin()
@@ -1935,6 +1936,8 @@ def test_stuck_raft_cluster_rejoin(harness):
 
         assert "raft_reset_primary" in harness.charm.app_peer_data
         assert "raft_rejoin" in harness.charm.app_peer_data
+        assert "members_ips" not in harness.charm.app_peer_data
+        _add_to_members_ips.assert_called_once_with("192.0.2.0")
         _update_relation_endpoints.assert_called_once_with()
 
 

From 09b3fea783ccedec96eb363a08529830d779c421 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 18 Sep 2024 15:32:52 +0300
Subject: [PATCH 63/85] Rework synchronising cluster stop

---
 src/charm.py   | 40 +++++++++++++++++++++++++++-------------
 src/cluster.py |  6 +++---
 2 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 41ef8a43c2..8238b58f2b 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -427,11 +427,6 @@ def _on_peer_relation_departed(self, event: RelationDepartedEvent) -> None:
             )
             event.defer()
             return
-        except RetryError:
-            logger.warning(
-                "Early on_peer_relation_departed: unable to connect to the departing unit"
-            )
-            return
 
         # Allow leader to update the cluster members.
         if not self.unit.is_leader():
@@ -564,6 +559,20 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
             return True
         return False
 
+    def _stuck_raft_cluster_stopped(self) -> bool:
+        """Check that the cluster is stopped."""
+        all_units_down = True
+        for key, data in self._peers.data.items():
+            if key == self.app:
+                continue
+            if "raft_stopped" not in data:
+                all_units_down = False
+        if all_units_down and "raft_followers_stopped" not in self.app_peer_data:
+            logger.info("Cluster is shut down")
+            self.app_peer_data["raft_followers_stopped"] = "True"
+            return True
+        return False
+
     def _stuck_raft_cluster_cleanup(self) -> None:
         for key, data in self._peers.data.items():
             if key == self.app:
@@ -575,6 +584,7 @@ def _stuck_raft_cluster_cleanup(self) -> None:
         self.app_peer_data.pop("raft_rejoin", None)
         self.app_peer_data.pop("raft_reset_primary", None)
         self.app_peer_data.pop("raft_selected_candidate", None)
+        self.app_peer_data.pop("raft_followers_stopped", None)
 
     def _raft_reinitialisation(self) -> bool:
         """Handle raft cluster loss of quorum."""
@@ -592,13 +602,17 @@ def _raft_reinitialisation(self) -> bool:
             if "raft_stopping" in self.unit_peer_data:
                 self._patroni.remove_raft_data()
                 self.unit_peer_data.pop("raft_stopping", None)
-                if candidate == self.unit.name:
-                    logger.info("Reinitialising %s as primary" % self.unit.name)
-                    self._patroni.reinitialise_raft_data()
-                    self.unit_peer_data["raft_primary"] = "True"
-                else:
-                    logger.info("Stopping %s" % self.unit.name)
-                    self.unit_peer_data["raft_stopped"] = "True"
+                logger.info("Stopping %s" % self.unit.name)
+                self.unit_peer_data["raft_stopped"] = "True"
+
+        if self.unit.is_leader() and self._stuck_raft_cluster_stopped():
+            should_exit = True
+
+        if candidate == self.unit.name and "raft_followers_stopped" in self.app_peer_data:
+            should_exit = True
+            logger.info("Reinitialising %s as primary" % self.unit.name)
+            self._patroni.reinitialise_raft_data()
+            self.unit_peer_data["raft_primary"] = "True"
 
         if self.unit.is_leader() and self._stuck_raft_cluster_rejoin():
             should_exit = True
@@ -610,7 +624,7 @@ def _raft_reinitialisation(self) -> bool:
             self.unit_peer_data.pop("raft_stopped", None)
             self._patroni.start_patroni()
 
-            if self.unit.is_leader() and "raft_rejoin":
+            if self.unit.is_leader() and "raft_rejoin" in self.app_peer_data:
                 self._stuck_raft_cluster_cleanup()
         return should_exit
 
diff --git a/src/cluster.py b/src/cluster.py
index 4ee537b846..b2ee36f09c 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -56,15 +56,15 @@
 
 
 class RaftPostgresqlNotUpError(Exception):
-    """Raised when a cluster is not promoted."""
+    """Postgresql not yet started."""
 
 
 class RaftPostgresqlStillUpError(Exception):
-    """Raised when a cluster is not promoted."""
+    """Postgresql not yet down."""
 
 
 class RaftNotPromotedError(Exception):
-    """Raised when a cluster is not promoted."""
+    """Leader not yet set when reinitialising raft."""
 
 
 class ClusterNotPromotedError(Exception):

From 5d253c75cdd5271e0df134611a4f83f564df903b Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 18 Sep 2024 16:55:02 +0300
Subject: [PATCH 64/85] Make sure unit data is clean before cleaning the app
 data

---
 src/charm.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 8238b58f2b..b71762ca66 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -577,8 +577,9 @@ def _stuck_raft_cluster_cleanup(self) -> None:
         for key, data in self._peers.data.items():
             if key == self.app:
                 continue
-            if "raft_primary" in data or "raft_stopped" in data:
-                return
+            for flag in data.keys():
+                if flag.startswith("raft_"):
+                    return
 
         logger.info("Cleaning up raft app data")
         self.app_peer_data.pop("raft_rejoin", None)

From 92dd7041cce89e4b86e668c39376c7cac943a109 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 18 Sep 2024 20:00:41 +0300
Subject: [PATCH 65/85] No double reinitialisation

---
 src/charm.py | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index b71762ca66..20e05415bf 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -559,19 +559,19 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
             return True
         return False
 
-    def _stuck_raft_cluster_stopped(self) -> bool:
+    def _stuck_raft_cluster_stopped(self) -> None:
         """Check that the cluster is stopped."""
-        all_units_down = True
+        if "raft_followers_stopped" in self.app_peer_data:
+            return
+
         for key, data in self._peers.data.items():
             if key == self.app:
                 continue
             if "raft_stopped" not in data:
-                all_units_down = False
-        if all_units_down and "raft_followers_stopped" not in self.app_peer_data:
+                return
             logger.info("Cluster is shut down")
-            self.app_peer_data["raft_followers_stopped"] = "True"
-            return True
-        return False
+
+        self.app_peer_data["raft_followers_stopped"] = "True"
 
     def _stuck_raft_cluster_cleanup(self) -> None:
         for key, data in self._peers.data.items():
@@ -609,7 +609,11 @@ def _raft_reinitialisation(self) -> bool:
         if self.unit.is_leader() and self._stuck_raft_cluster_stopped():
             should_exit = True
 
-        if candidate == self.unit.name and "raft_followers_stopped" in self.app_peer_data:
+        if (
+            candidate == self.unit.name
+            and "raft_primary" not in self.unit_peer_data
+            and "raft_followers_stopped" in self.app_peer_data
+        ):
             should_exit = True
             logger.info("Reinitialising %s as primary" % self.unit.name)
             self._patroni.reinitialise_raft_data()

From 723306dc06e87dde1b60b6900a553eaf6ba8b03b Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 18 Sep 2024 22:42:42 +0300
Subject: [PATCH 66/85] Only fire in relation changed

---
 src/charm.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 20e05415bf..a914554e0e 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -38,6 +38,7 @@
     HookEvent,
     InstallEvent,
     LeaderElectedEvent,
+    RelationChangedEvent,
     RelationDepartedEvent,
     StartEvent,
 )
@@ -569,8 +570,8 @@ def _stuck_raft_cluster_stopped(self) -> None:
                 continue
             if "raft_stopped" not in data:
                 return
-            logger.info("Cluster is shut down")
 
+        logger.info("Cluster is shut down")
         self.app_peer_data["raft_followers_stopped"] = "True"
 
     def _stuck_raft_cluster_cleanup(self) -> None:
@@ -653,7 +654,11 @@ def _peer_relation_changed_checks(self, event: HookEvent) -> bool:
             return False
 
         # Check whether raft is stuck
-        if self.has_raft_keys() and self._raft_reinitialisation():
+        if (
+            isinstance(event, RelationChangedEvent)
+            and self.has_raft_keys()
+            and self._raft_reinitialisation()
+        ):
             logger.debug("Early exit on_peer_relation_changed: stuck raft recovery")
             return False
 

From c86e7c7df7f241e5870c5eb113dd49e39c22efb9 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Wed, 18 Sep 2024 22:49:04 +0300
Subject: [PATCH 67/85] Still exit if secret event and raft flags

---
 src/charm.py | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index a914554e0e..1f747cb32f 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -653,11 +653,9 @@ def _peer_relation_changed_checks(self, event: HookEvent) -> bool:
             event.defer()
             return False
 
-        # Check whether raft is stuck
-        if (
-            isinstance(event, RelationChangedEvent)
-            and self.has_raft_keys()
-            and self._raft_reinitialisation()
+        # Check whether raft is stuck.
+        if self.has_raft_keys() and (
+            isinstance(event, RelationChangedEvent) or self._raft_reinitialisation()
         ):
             logger.debug("Early exit on_peer_relation_changed: stuck raft recovery")
             return False

From 42049d3950d13c6eda768b2369e2fa75e2b29ecf Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 19 Sep 2024 01:30:22 +0300
Subject: [PATCH 68/85] Try to simplifly coordinated stopping

---
 src/charm.py | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 1f747cb32f..24f25d33ac 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -560,7 +560,7 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
             return True
         return False
 
-    def _stuck_raft_cluster_stopped(self) -> None:
+    def _stuck_raft_cluster_stopped_check(self) -> None:
         """Check that the cluster is stopped."""
         if "raft_followers_stopped" in self.app_peer_data:
             return
@@ -596,18 +596,14 @@ def _raft_reinitialisation(self) -> bool:
 
         if candidate := self.app_peer_data.get("raft_selected_candidate"):
             should_exit = True
-            if "raft_stuck" in self.unit_peer_data:
+            if "raft_stopped" not in self.unit_peer_data:
                 self.unit_peer_data.pop("raft_stuck", None)
                 self.unit_peer_data.pop("raft_candidate", None)
-                self.unit_peer_data["raft_stopping"] = "True"
-
-            if "raft_stopping" in self.unit_peer_data:
                 self._patroni.remove_raft_data()
-                self.unit_peer_data.pop("raft_stopping", None)
                 logger.info("Stopping %s" % self.unit.name)
                 self.unit_peer_data["raft_stopped"] = "True"
 
-        if self.unit.is_leader() and self._stuck_raft_cluster_stopped():
+        if self.unit.is_leader() and self._stuck_raft_cluster_stopped_check():
             should_exit = True
 
         if (

From 4fae252a21dd7a25c836a81d6a58ea6136fa9d60 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 19 Sep 2024 02:37:10 +0300
Subject: [PATCH 69/85] Add logging

---
 src/charm.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/charm.py b/src/charm.py
index 24f25d33ac..f0e11b70e5 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -590,6 +590,8 @@ def _stuck_raft_cluster_cleanup(self) -> None:
 
     def _raft_reinitialisation(self) -> bool:
         """Handle raft cluster loss of quorum."""
+        logger.error(f"!!!!!!!!!!!!!!!!!!1{self.unit} {self.unit_peer_data}")
+        logger.error(f"@@@@@@@@@@@@@@@2@{self.unit} {self.app_peer_data}")
         should_exit = False
         if self.unit.is_leader() and self._stuck_raft_cluster_check():
             should_exit = True

From c9b1211dae7aadefb866b1712ceea30f857046a2 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 19 Sep 2024 03:49:43 +0300
Subject: [PATCH 70/85] Try to set status

---
 src/charm.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/charm.py b/src/charm.py
index f0e11b70e5..373a8a017a 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -636,10 +636,12 @@ def has_raft_keys(self):
         """Checks for the presence of raft recovery keys in peer data."""
         for key in self.app_peer_data.keys():
             if key.startswith("raft_"):
+                self.unit.status = MaintenanceStatus("Reinitialising raft...")
                 return True
 
         for key in self.unit_peer_data.keys():
             if key.startswith("raft_"):
+                self.unit.status = MaintenanceStatus("Reinitialising raft...")
                 return True
         return False
 

From cf77dfe4c73dba94eaf7288a3ca787eaaf74ab3f Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 19 Sep 2024 14:42:39 +0300
Subject: [PATCH 71/85] Log event

---
 src/charm.py   | 5 ++---
 src/cluster.py | 1 +
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 373a8a017a..3a933de52b 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -636,12 +636,10 @@ def has_raft_keys(self):
         """Checks for the presence of raft recovery keys in peer data."""
         for key in self.app_peer_data.keys():
             if key.startswith("raft_"):
-                self.unit.status = MaintenanceStatus("Reinitialising raft...")
                 return True
 
         for key in self.unit_peer_data.keys():
             if key.startswith("raft_"):
-                self.unit.status = MaintenanceStatus("Reinitialising raft...")
                 return True
         return False
 
@@ -653,7 +651,8 @@ def _peer_relation_changed_checks(self, event: HookEvent) -> bool:
             event.defer()
             return False
 
-        # Check whether raft is stuck.
+        # Check whether raft is stuck. Secrets events may have stale peer data
+        logger.error(f"@@@@@@@@@@@@@@@2@{self.unit} {event}")
         if self.has_raft_keys() and (
             isinstance(event, RelationChangedEvent) or self._raft_reinitialisation()
         ):
diff --git a/src/cluster.py b/src/cluster.py
index b2ee36f09c..8f9b5fe461 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -829,6 +829,7 @@ def remove_raft_member(self, member_ip: str) -> None:
                 data_flags["raft_candidate"] = "True"
             self.charm.unit_peer_data.update(data_flags)
 
+            # Leader doesn't always trigger when changing it's own peer data.
             if self.charm.unit.is_leader():
                 self.charm.on[PEER].relation_changed.emit(
                     unit=self.charm.unit,

From d59626e32980d5027af80896de401cb24cc6086a Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 19 Sep 2024 16:17:40 +0300
Subject: [PATCH 72/85] Supress transient secrets events

---
 src/charm.py             | 39 +++++++++++++++++----------------------
 tests/unit/test_charm.py | 12 ++++++------
 2 files changed, 23 insertions(+), 28 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 3a933de52b..6c47fbdd02 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -507,7 +507,7 @@ def _on_pgdata_storage_detaching(self, _) -> None:
         if self.primary_endpoint:
             self._update_relation_endpoints()
 
-    def _stuck_raft_cluster_check(self) -> bool:
+    def _stuck_raft_cluster_check(self) -> None:
         """Check for stuck raft cluster and reinitialise if safe."""
         raft_stuck = False
         all_units_stuck = True
@@ -523,19 +523,18 @@ def _stuck_raft_cluster_check(self) -> bool:
                 candidate = key
 
         if not raft_stuck:
-            return False
+            return
 
         if not all_units_stuck:
             logger.warning("Stuck raft not yet detected on all units")
-            return True
+            return
 
         if not candidate:
             logger.warning("Stuck raft has no candidate")
-            return True
+            return
         if "raft_selected_candidate" not in self.app_peer_data:
             logger.info("%s selected for new raft leader" % candidate.name)
             self.app_peer_data["raft_selected_candidate"] = candidate.name
-        return True
 
     def _stuck_raft_cluster_rejoin(self) -> bool:
         """Reconnect cluster to new raft."""
@@ -588,16 +587,14 @@ def _stuck_raft_cluster_cleanup(self) -> None:
         self.app_peer_data.pop("raft_selected_candidate", None)
         self.app_peer_data.pop("raft_followers_stopped", None)
 
-    def _raft_reinitialisation(self) -> bool:
+    def _raft_reinitialisation(self) -> None:
         """Handle raft cluster loss of quorum."""
         logger.error(f"!!!!!!!!!!!!!!!!!!1{self.unit} {self.unit_peer_data}")
         logger.error(f"@@@@@@@@@@@@@@@2@{self.unit} {self.app_peer_data}")
-        should_exit = False
-        if self.unit.is_leader() and self._stuck_raft_cluster_check():
-            should_exit = True
+        if self.unit.is_leader():
+            self._stuck_raft_cluster_check()
 
         if candidate := self.app_peer_data.get("raft_selected_candidate"):
-            should_exit = True
             if "raft_stopped" not in self.unit_peer_data:
                 self.unit_peer_data.pop("raft_stuck", None)
                 self.unit_peer_data.pop("raft_candidate", None)
@@ -605,32 +602,29 @@ def _raft_reinitialisation(self) -> bool:
                 logger.info("Stopping %s" % self.unit.name)
                 self.unit_peer_data["raft_stopped"] = "True"
 
-        if self.unit.is_leader() and self._stuck_raft_cluster_stopped_check():
-            should_exit = True
+        if self.unit.is_leader():
+            self._stuck_raft_cluster_stopped_check()
 
         if (
             candidate == self.unit.name
             and "raft_primary" not in self.unit_peer_data
             and "raft_followers_stopped" in self.app_peer_data
         ):
-            should_exit = True
             logger.info("Reinitialising %s as primary" % self.unit.name)
             self._patroni.reinitialise_raft_data()
             self.unit_peer_data["raft_primary"] = "True"
 
-        if self.unit.is_leader() and self._stuck_raft_cluster_rejoin():
-            should_exit = True
+        if self.unit.is_leader():
+            self._stuck_raft_cluster_rejoin()
 
         if "raft_rejoin" in self.app_peer_data:
-            should_exit = True
             logger.info("Cleaning up raft unit data")
             self.unit_peer_data.pop("raft_primary", None)
             self.unit_peer_data.pop("raft_stopped", None)
             self._patroni.start_patroni()
 
-            if self.unit.is_leader() and "raft_rejoin" in self.app_peer_data:
+            if self.unit.is_leader():
                 self._stuck_raft_cluster_cleanup()
-        return should_exit
 
     def has_raft_keys(self):
         """Checks for the presence of raft recovery keys in peer data."""
@@ -653,9 +647,9 @@ def _peer_relation_changed_checks(self, event: HookEvent) -> bool:
 
         # Check whether raft is stuck. Secrets events may have stale peer data
         logger.error(f"@@@@@@@@@@@@@@@2@{self.unit} {event}")
-        if self.has_raft_keys() and (
-            isinstance(event, RelationChangedEvent) or self._raft_reinitialisation()
-        ):
+        if self.has_raft_keys():
+            if isinstance(event, RelationChangedEvent):
+                self._raft_reinitialisation()
             logger.debug("Early exit on_peer_relation_changed: stuck raft recovery")
             return False
 
@@ -1063,7 +1057,8 @@ def _on_leader_elected(self, event: LeaderElectedEvent) -> None:
             if self.get_secret(APP_SCOPE, key) is None:
                 self.set_secret(APP_SCOPE, key, new_password())
 
-        if self.has_raft_keys() and self._raft_reinitialisation():
+        if self.has_raft_keys():
+            self._raft_reinitialisation()
             return
 
         # Update the list of the current PostgreSQL hosts when a new leader is elected.
diff --git a/tests/unit/test_charm.py b/tests/unit/test_charm.py
index 7361ea00ef..092bd2f94e 100644
--- a/tests/unit/test_charm.py
+++ b/tests/unit/test_charm.py
@@ -1857,13 +1857,13 @@ def test_stuck_raft_cluster_check(harness):
         harness.set_leader()
         harness.update_relation_data(rel_id, harness.charm.unit.name, {"raft_stuck": "True"})
 
-    assert harness.charm._stuck_raft_cluster_check()
+    harness.charm._stuck_raft_cluster_check()
     assert "raft_selected_candidate" not in harness.charm.app_peer_data
 
     # Raft candidate
     with harness.hooks_disabled():
         harness.update_relation_data(rel_id, harness.charm.unit.name, {"raft_candidate": "True"})
-    assert harness.charm._stuck_raft_cluster_check()
+    harness.charm._stuck_raft_cluster_check()
     assert harness.charm.app_peer_data["raft_selected_candidate"] == harness.charm.unit.name
 
     # Don't override existing candidate
@@ -1871,7 +1871,7 @@ def test_stuck_raft_cluster_check(harness):
         harness.update_relation_data(
             rel_id, harness.charm.app.name, {"raft_selected_candidate": "something_else"}
         )
-    assert harness.charm._stuck_raft_cluster_check()
+    harness.charm._stuck_raft_cluster_check()
     assert harness.charm.app_peer_data["raft_selected_candidate"] != harness.charm.unit.name
 
 
@@ -1959,7 +1959,7 @@ def test_raft_reinitialisation(harness):
         patch("charm.Patroni.start_patroni") as _start_patroni,
     ):
         # No data
-        assert not harness.charm._raft_reinitialisation()
+        harness.charm._raft_reinitialisation()
 
         # Different candidate
         with harness.hooks_disabled():
@@ -1973,7 +1973,7 @@ def test_raft_reinitialisation(harness):
                 {"raft_rejoin": "True", "raft_selected_candidate": "test_candidate"},
             )
 
-        assert harness.charm._raft_reinitialisation()
+        harness.charm._raft_reinitialisation()
         _stuck_raft_cluster_rejoin.assert_called_once_with()
         _stuck_raft_cluster_check.assert_called_once_with()
         _stuck_raft_cluster_cleanup.assert_called_once_with()
@@ -1996,7 +1996,7 @@ def test_raft_reinitialisation(harness):
                 {"raft_selected_candidate": "postgresql/0"},
             )
 
-        assert harness.charm._raft_reinitialisation()
+        harness.charm._raft_reinitialisation()
         _reinitialise_raft_data.assert_called_once_with()
 
 

From 45f98e60271e6e6c8977efd182428a79ea9f93ad Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 19 Sep 2024 17:26:59 +0300
Subject: [PATCH 73/85] Multiple down units

---
 tests/integration/ha_tests/helpers.py      | 15 +++++++++++----
 tests/integration/ha_tests/test_scaling.py |  9 ++++++++-
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/tests/integration/ha_tests/helpers.py b/tests/integration/ha_tests/helpers.py
index 9763a2f80d..92f9e41f8d 100644
--- a/tests/integration/ha_tests/helpers.py
+++ b/tests/integration/ha_tests/helpers.py
@@ -7,7 +7,7 @@
 import subprocess
 from pathlib import Path
 from tempfile import mkstemp
-from typing import Dict, Optional, Set, Tuple, Union
+from typing import Dict, List, Optional, Set, Tuple, Union
 
 import psycopg2
 import requests
@@ -90,12 +90,19 @@ async def are_all_db_processes_down(ops_test: OpsTest, process: str, signal: str
 
 
 async def are_writes_increasing(
-    ops_test, down_unit: str = None, use_ip_from_inside: bool = False, extra_model: Model = None
+    ops_test,
+    down_unit: Optional[Union[str, List[str]]] = None,
+    use_ip_from_inside: bool = False,
+    extra_model: Model = None,
 ) -> None:
     """Verify new writes are continuing by counting the number of writes."""
+    if isinstance(down_unit, str):
+        down_units = [down_unit]
+    else:
+        down_units = down_unit
     writes, _ = await count_writes(
         ops_test,
-        down_unit=down_unit,
+        down_unit=down_units[0],
         use_ip_from_inside=use_ip_from_inside,
         extra_model=extra_model,
     )
@@ -111,7 +118,7 @@ async def are_writes_increasing(
             )
             logger.info(f"Retry writes {more_writes}")
             for member, count in writes.items():
-                if "/".join(member.split(".", 1)[-1].rsplit("-", 1)) == down_unit:
+                if "/".join(member.split(".", 1)[-1].rsplit("-", 1)) in down_units:
                     continue
                 assert (
                     more_writes[member] > count
diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index 7da469de50..3fea3ef4bf 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -156,7 +156,14 @@ async def test_removing_raft_majority(ops_test: OpsTest, continuous_writes) -> N
 
     await ops_test.model.wait_for_idle(status="active", timeout=1500)
 
-    await are_writes_increasing(ops_test, original_roles["primaries"][0])
+    await are_writes_increasing(
+        ops_test,
+        [
+            original_roles["primaries"][0],
+            original_roles["replicas"][0],
+            original_roles["sync_standbys"][0],
+        ],
+    )
 
     logger.info("Scaling back up")
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=3)

From c413f7610789c644d2129102c9022e7835548c44 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 19 Sep 2024 19:13:30 +0300
Subject: [PATCH 74/85] Conditional cleanup

---
 src/charm.py             | 36 +++++++++++++++++++-----------------
 tests/unit/test_charm.py | 27 ++++++++++++++++++---------
 2 files changed, 37 insertions(+), 26 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 8438440a19..b1e9e2fbdf 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -593,31 +593,34 @@ def _raft_reinitialisation(self) -> None:
         """Handle raft cluster loss of quorum."""
         logger.error(f"!!!!!!!!!!!!!!!!!!1{self.unit} {self.unit_peer_data}")
         logger.error(f"@@@@@@@@@@@@@@@2@{self.unit} {self.app_peer_data}")
-        if self.unit.is_leader():
-            self._stuck_raft_cluster_check()
+        # Skip to cleanup if rejoining
+        if "raft_rejoin" not in self.app_peer_data:
+            if self.unit.is_leader():
+                self._stuck_raft_cluster_check()
 
-        if candidate := self.app_peer_data.get("raft_selected_candidate"):
-            if "raft_stopped" not in self.unit_peer_data:
+            if (
+                candidate := self.app_peer_data.get("raft_selected_candidate")
+            ) and "raft_stopped" not in self.unit_peer_data:
                 self.unit_peer_data.pop("raft_stuck", None)
                 self.unit_peer_data.pop("raft_candidate", None)
                 self._patroni.remove_raft_data()
                 logger.info("Stopping %s" % self.unit.name)
                 self.unit_peer_data["raft_stopped"] = "True"
 
-        if self.unit.is_leader():
-            self._stuck_raft_cluster_stopped_check()
+            if self.unit.is_leader():
+                self._stuck_raft_cluster_stopped_check()
 
-        if (
-            candidate == self.unit.name
-            and "raft_primary" not in self.unit_peer_data
-            and "raft_followers_stopped" in self.app_peer_data
-        ):
-            logger.info("Reinitialising %s as primary" % self.unit.name)
-            self._patroni.reinitialise_raft_data()
-            self.unit_peer_data["raft_primary"] = "True"
+            if (
+                candidate == self.unit.name
+                and "raft_primary" not in self.unit_peer_data
+                and "raft_followers_stopped" in self.app_peer_data
+            ):
+                logger.info("Reinitialising %s as primary" % self.unit.name)
+                self._patroni.reinitialise_raft_data()
+                self.unit_peer_data["raft_primary"] = "True"
 
-        if self.unit.is_leader():
-            self._stuck_raft_cluster_rejoin()
+            if self.unit.is_leader():
+                self._stuck_raft_cluster_rejoin()
 
         if "raft_rejoin" in self.app_peer_data:
             logger.info("Cleaning up raft unit data")
@@ -648,7 +651,6 @@ def _peer_relation_changed_checks(self, event: HookEvent) -> bool:
             return False
 
         # Check whether raft is stuck. Secrets events may have stale peer data
-        logger.error(f"@@@@@@@@@@@@@@@2@{self.unit} {event}")
         if self.has_raft_keys():
             if isinstance(event, RelationChangedEvent):
                 self._raft_reinitialisation()
diff --git a/tests/unit/test_charm.py b/tests/unit/test_charm.py
index 6eb248c245..6d8525f3ef 100644
--- a/tests/unit/test_charm.py
+++ b/tests/unit/test_charm.py
@@ -1959,7 +1959,6 @@ def test_raft_reinitialisation(harness):
         ) as _stuck_raft_cluster_cleanup,
         patch("charm.Patroni.remove_raft_data") as _remove_raft_data,
         patch("charm.Patroni.reinitialise_raft_data") as _reinitialise_raft_data,
-        patch("charm.Patroni.start_patroni") as _start_patroni,
     ):
         # No data
         harness.charm._raft_reinitialisation()
@@ -1973,35 +1972,45 @@ def test_raft_reinitialisation(harness):
             harness.update_relation_data(
                 rel_id,
                 harness.charm.app.name,
-                {"raft_rejoin": "True", "raft_selected_candidate": "test_candidate"},
+                {"raft_selected_candidate": "test_candidate"},
             )
 
         harness.charm._raft_reinitialisation()
         _stuck_raft_cluster_rejoin.assert_called_once_with()
         _stuck_raft_cluster_check.assert_called_once_with()
-        _stuck_raft_cluster_cleanup.assert_called_once_with()
+        assert not _stuck_raft_cluster_cleanup.called
         _remove_raft_data.assert_called_once_with()
-        _start_patroni.assert_called_once_with()
         assert not _reinitialise_raft_data.called
-        assert "raft_stopped" not in harness.charm.unit_peer_data
-        assert "raft_primary" not in harness.charm.unit_peer_data
-        _start_patroni.reset_mock()
 
         # Current candidate
         with harness.hooks_disabled():
             harness.set_leader()
             harness.update_relation_data(
-                rel_id, harness.charm.unit.name, {"raft_stuck": "True", "raft_candidate": "True"}
+                rel_id,
+                harness.charm.unit.name,
+                {"raft_stuck": "", "raft_candidate": "True", "raft_stopped": "True"},
             )
             harness.update_relation_data(
                 rel_id,
                 harness.charm.app.name,
-                {"raft_selected_candidate": "postgresql/0"},
+                {"raft_selected_candidate": "postgresql/0", "raft_followers_stopped": "True"},
             )
 
         harness.charm._raft_reinitialisation()
         _reinitialise_raft_data.assert_called_once_with()
 
+        # Cleanup
+        with harness.hooks_disabled():
+            harness.set_leader()
+            harness.update_relation_data(
+                rel_id,
+                harness.charm.unit.name,
+                {"raft_stuck": "", "raft_candidate": "True", "raft_stopped": "True"},
+            )
+            harness.update_relation_data(rel_id, harness.charm.app.name, {"raft_rejoin": "True"})
+        harness.charm._raft_reinitialisation()
+        _stuck_raft_cluster_cleanup.assert_called_once_with()
+
 
 #
 # Secrets

From 5a57952e12de3bfc2bc9a8d9e266d06b7f473535 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 19 Sep 2024 20:15:53 +0300
Subject: [PATCH 75/85] Reenable secret hooks check

---
 src/charm.py                               | 6 ++----
 tests/integration/ha_tests/helpers.py      | 2 +-
 tests/integration/ha_tests/test_scaling.py | 2 +-
 3 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index b1e9e2fbdf..cad41b293f 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -38,7 +38,6 @@
     HookEvent,
     InstallEvent,
     LeaderElectedEvent,
-    RelationChangedEvent,
     RelationDepartedEvent,
     StartEvent,
 )
@@ -650,10 +649,9 @@ def _peer_relation_changed_checks(self, event: HookEvent) -> bool:
             event.defer()
             return False
 
-        # Check whether raft is stuck. Secrets events may have stale peer data
+        # Check whether raft is stuck.
         if self.has_raft_keys():
-            if isinstance(event, RelationChangedEvent):
-                self._raft_reinitialisation()
+            self._raft_reinitialisation()
             logger.debug("Early exit on_peer_relation_changed: stuck raft recovery")
             return False
 
diff --git a/tests/integration/ha_tests/helpers.py b/tests/integration/ha_tests/helpers.py
index 92f9e41f8d..ab14e89271 100644
--- a/tests/integration/ha_tests/helpers.py
+++ b/tests/integration/ha_tests/helpers.py
@@ -96,7 +96,7 @@ async def are_writes_increasing(
     extra_model: Model = None,
 ) -> None:
     """Verify new writes are continuing by counting the number of writes."""
-    if isinstance(down_unit, str):
+    if isinstance(down_unit, str) or not down_unit:
         down_units = [down_unit]
     else:
         down_units = down_unit
diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index 3fea3ef4bf..b5d9dced14 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -167,7 +167,7 @@ async def test_removing_raft_majority(ops_test: OpsTest, continuous_writes) -> N
 
     logger.info("Scaling back up")
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=3)
-    await ops_test.model.wait_for_idle(status="active", timeout=600)
+    await ops_test.model.wait_for_idle(status="active", timeout=900)
 
     await check_writes(ops_test)
     new_roles = await get_cluster_roles(

From 2ee87ed58a9ba44a5e2b4b39c26432cf534378e9 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 19 Sep 2024 21:47:36 +0300
Subject: [PATCH 76/85] Sync on reset primary

---
 src/charm.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index cad41b293f..d19eef7de7 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -540,21 +540,23 @@ def _stuck_raft_cluster_check(self) -> None:
     def _stuck_raft_cluster_rejoin(self) -> bool:
         """Reconnect cluster to new raft."""
         primary = None
-        all_units_down = True
         for key, data in self._peers.data.items():
             if key == self.app:
                 continue
             if "raft_primary" in data:
                 primary = key
-            elif "raft_stopped" not in data:
-                all_units_down = False
+                break
         if primary and "raft_reset_primary" not in self.app_peer_data:
             logger.info("Updating the primary endpoint")
             self.app_peer_data.pop("members_ips", None)
             self._add_to_members_ips(self._get_unit_ip(primary))
             self.app_peer_data["raft_reset_primary"] = "True"
             self._update_relation_endpoints()
-        if all_units_down and "raft_rejoin" not in self.app_peer_data:
+        if (
+            "raft_rejoin" not in self.app_peer_data
+            and "raft_followers_stopped" in self.app_peer_data
+            and "raft_reset_primary" in self.app_peer_data
+        ):
             logger.info("Notify units they can rejoin")
             self.app_peer_data["raft_rejoin"] = "True"
             return True
@@ -590,8 +592,6 @@ def _stuck_raft_cluster_cleanup(self) -> None:
 
     def _raft_reinitialisation(self) -> None:
         """Handle raft cluster loss of quorum."""
-        logger.error(f"!!!!!!!!!!!!!!!!!!1{self.unit} {self.unit_peer_data}")
-        logger.error(f"@@@@@@@@@@@@@@@2@{self.unit} {self.app_peer_data}")
         # Skip to cleanup if rejoining
         if "raft_rejoin" not in self.app_peer_data:
             if self.unit.is_leader():

From aa61f3734e58c82376eabdb911b674b53f9fc643 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 19 Sep 2024 21:55:01 +0300
Subject: [PATCH 77/85] Fix tests

---
 src/charm.py             |  4 +---
 tests/unit/test_charm.py | 12 ++++++++++--
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index d19eef7de7..95ce841c24 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -537,7 +537,7 @@ def _stuck_raft_cluster_check(self) -> None:
             logger.info("%s selected for new raft leader" % candidate.name)
             self.app_peer_data["raft_selected_candidate"] = candidate.name
 
-    def _stuck_raft_cluster_rejoin(self) -> bool:
+    def _stuck_raft_cluster_rejoin(self) -> None:
         """Reconnect cluster to new raft."""
         primary = None
         for key, data in self._peers.data.items():
@@ -559,8 +559,6 @@ def _stuck_raft_cluster_rejoin(self) -> bool:
         ):
             logger.info("Notify units they can rejoin")
             self.app_peer_data["raft_rejoin"] = "True"
-            return True
-        return False
 
     def _stuck_raft_cluster_stopped_check(self) -> None:
         """Check that the cluster is stopped."""
diff --git a/tests/unit/test_charm.py b/tests/unit/test_charm.py
index 6d8525f3ef..8bf13e35a4 100644
--- a/tests/unit/test_charm.py
+++ b/tests/unit/test_charm.py
@@ -1927,15 +1927,23 @@ def test_stuck_raft_cluster_rejoin(harness):
         patch("charm.PostgresqlOperatorCharm._add_to_members_ips") as _add_to_members_ips,
     ):
         # No data
-        assert not harness.charm._stuck_raft_cluster_rejoin()
+        harness.charm._stuck_raft_cluster_rejoin()
+
+        assert "raft_reset_primary" not in harness.charm.app_peer_data
+        assert "raft_rejoin" not in harness.charm.app_peer_data
 
         # Raises primary flag
         with harness.hooks_disabled():
             harness.update_relation_data(
                 rel_id, harness.charm.unit.name, {"raft_primary": "test_primary"}
             )
+            harness.update_relation_data(
+                rel_id,
+                harness.charm.app.name,
+                {"raft_followers_stopped": "test_candidate"},
+            )
 
-        assert harness.charm._stuck_raft_cluster_rejoin()
+        harness.charm._stuck_raft_cluster_rejoin()
 
         assert "raft_reset_primary" in harness.charm.app_peer_data
         assert "raft_rejoin" in harness.charm.app_peer_data

From 53966074d3941f4feb1135ed33f60f89e8a04424 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Fri, 20 Sep 2024 13:37:56 +0300
Subject: [PATCH 78/85] Add test to remove both async replicas

---
 tests/integration/ha_tests/test_scaling.py | 55 +++++++++++++++++++++-
 1 file changed, 53 insertions(+), 2 deletions(-)

diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index b5d9dced14..555a26c9ad 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -154,7 +154,7 @@ async def test_removing_raft_majority(ops_test: OpsTest, continuous_writes) -> N
         ),
     )
 
-    await ops_test.model.wait_for_idle(status="active", timeout=1500)
+    await ops_test.model.wait_for_idle(status="active", timeout=900)
 
     await are_writes_increasing(
         ops_test,
@@ -167,7 +167,7 @@ async def test_removing_raft_majority(ops_test: OpsTest, continuous_writes) -> N
 
     logger.info("Scaling back up")
     await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=3)
-    await ops_test.model.wait_for_idle(status="active", timeout=900)
+    await ops_test.model.wait_for_idle(status="active", timeout=1500)
 
     await check_writes(ops_test)
     new_roles = await get_cluster_roles(
@@ -176,3 +176,54 @@ async def test_removing_raft_majority(ops_test: OpsTest, continuous_writes) -> N
     assert len(new_roles["primaries"]) == 1
     assert len(new_roles["sync_standbys"]) == 2
     assert new_roles["primaries"][0] == original_roles["sync_standbys"][1]
+
+
+@pytest.mark.group(1)
+@markers.juju3
+@pytest.mark.abort_on_fail
+async def test_removing_raft_majority_async(ops_test: OpsTest, continuous_writes) -> None:
+    # Start an application that continuously writes data to the database.
+    app = await app_name(ops_test)
+    original_roles = await get_cluster_roles(
+        ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
+    )
+
+    await start_continuous_writes(ops_test, app)
+    logger.info("Deleting primary")
+    await gather(
+        ops_test.model.destroy_unit(
+            original_roles["primaries"][0], force=True, destroy_storage=False, max_wait=1500
+        ),
+        ops_test.model.destroy_unit(
+            original_roles["replicas"][0], force=True, destroy_storage=False, max_wait=1500
+        ),
+        ops_test.model.destroy_unit(
+            original_roles["replicas"][1], force=True, destroy_storage=False, max_wait=1500
+        ),
+    )
+
+    await ops_test.model.wait_for_idle(status="active", timeout=900)
+
+    await are_writes_increasing(
+        ops_test,
+        [
+            original_roles["primaries"][0],
+            original_roles["replicas"][0],
+            original_roles["replicas"][1],
+        ],
+    )
+
+    logger.info("Scaling back up")
+    await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=3)
+    await ops_test.model.wait_for_idle(status="active", timeout=1500)
+
+    await check_writes(ops_test)
+    new_roles = await get_cluster_roles(
+        ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
+    )
+    assert len(new_roles["primaries"]) == 1
+    assert len(new_roles["sync_standbys"]) == 2
+    assert (
+        new_roles["primaries"][0] == original_roles["sync_standbys"][0]
+        or new_roles["primaries"][0] == original_roles["sync_standbys"][1]
+    )

From ca127cbd3ca764d13d4be680e7f95f0cae25557f Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Fri, 20 Sep 2024 18:07:47 +0300
Subject: [PATCH 79/85] Handle missing member in patroni cluster

---
 src/charm.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/charm.py b/src/charm.py
index 95ce841c24..b6650ba541 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -429,6 +429,12 @@ def _on_peer_relation_departed(self, event: RelationDepartedEvent) -> None:
             )
             event.defer()
             return
+        except RetryError:
+            logger.warning(
+                "Early exit on_peer_relation_departed: Cannot get %s member IP"
+                % event.departing_unit.name
+            )
+            return
 
         # Allow leader to update the cluster members.
         if not self.unit.is_leader():

From 88edde1b35fb0ad8fab7449ca3d75d581f2c54e7 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Fri, 20 Sep 2024 22:30:14 +0300
Subject: [PATCH 80/85] increase idle period

---
 tests/integration/ha_tests/test_scaling.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index 555a26c9ad..3d92cd9b2c 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -69,7 +69,7 @@ async def test_removing_stereo_primary(ops_test: OpsTest, continuous_writes) ->
     primary = await get_primary(ops_test, app)
     await ops_test.model.destroy_unit(primary, force=True, destroy_storage=False, max_wait=1500)
 
-    await ops_test.model.wait_for_idle(status="active", timeout=600)
+    await ops_test.model.wait_for_idle(status="active", timeout=600, idle_period=45)
 
     await are_writes_increasing(ops_test, primary)
 
@@ -104,7 +104,7 @@ async def test_removing_stereo_sync_standby(ops_test: OpsTest, continuous_writes
     ).name
     await ops_test.model.destroy_unit(secondary, force=True, destroy_storage=False, max_wait=1500)
 
-    await ops_test.model.wait_for_idle(status="active", timeout=600)
+    await ops_test.model.wait_for_idle(status="active", timeout=600, idle_period=45)
 
     await are_writes_increasing(ops_test, secondary)
 
@@ -154,7 +154,7 @@ async def test_removing_raft_majority(ops_test: OpsTest, continuous_writes) -> N
         ),
     )
 
-    await ops_test.model.wait_for_idle(status="active", timeout=900)
+    await ops_test.model.wait_for_idle(status="active", timeout=900, idle_period=45)
 
     await are_writes_increasing(
         ops_test,
@@ -202,7 +202,7 @@ async def test_removing_raft_majority_async(ops_test: OpsTest, continuous_writes
         ),
     )
 
-    await ops_test.model.wait_for_idle(status="active", timeout=900)
+    await ops_test.model.wait_for_idle(status="active", timeout=900, idle_period=45)
 
     await are_writes_increasing(
         ops_test,

From d11b87e2d4e61fec166006816c89b4b2113393c8 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Sun, 22 Sep 2024 16:02:14 +0300
Subject: [PATCH 81/85] Try to skip hooks that take too long

---
 src/charm.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/charm.py b/src/charm.py
index b6650ba541..5db29cb620 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -418,6 +418,10 @@ def _on_peer_relation_departed(self, event: RelationDepartedEvent) -> None:
             logger.debug("Early exit on_peer_relation_departed: Skipping departing unit")
             return
 
+        if self.has_raft_keys():
+            logger.debug("Early exit on_peer_relation_departed: Raft recovery in progress")
+            return
+
         # Remove the departing member from the raft cluster.
         try:
             departing_member = event.departing_unit.name.replace("/", "-")
@@ -1527,6 +1531,10 @@ def _can_run_on_update_status(self) -> bool:
         if "cluster_initialised" not in self._peers.data[self.app]:
             return False
 
+        if self.has_raft_keys():
+            logger.debug("Early exit on_update_status: Raft recovery in progress")
+            return False
+
         if not self.upgrade.idle:
             logger.debug("Early exit on_update_status: upgrade in progress")
             return False

From 2cb218eea1eef5bea6cddd9b0f00b9ad2de4117f Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Mon, 30 Sep 2024 15:57:23 +0300
Subject: [PATCH 82/85] Two stuck units synchronisation

---
 src/charm.py             | 18 ++++++++++--------
 tests/unit/test_charm.py |  2 ++
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index 5db29cb620..050721a3a6 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -559,7 +559,9 @@ def _stuck_raft_cluster_rejoin(self) -> None:
         if primary and "raft_reset_primary" not in self.app_peer_data:
             logger.info("Updating the primary endpoint")
             self.app_peer_data.pop("members_ips", None)
-            self._add_to_members_ips(self._get_unit_ip(primary))
+            for unit in self._peers.units:
+                self._add_to_members_ips(self._get_unit_ip(unit))
+            self._add_to_members_ips(self._get_unit_ip(self.unit))
             self.app_peer_data["raft_reset_primary"] = "True"
             self._update_relation_endpoints()
         if (
@@ -633,6 +635,7 @@ def _raft_reinitialisation(self) -> None:
             logger.info("Cleaning up raft unit data")
             self.unit_peer_data.pop("raft_primary", None)
             self.unit_peer_data.pop("raft_stopped", None)
+            self.update_config()
             self._patroni.start_patroni()
 
             if self.unit.is_leader():
@@ -670,6 +673,11 @@ def _peer_relation_changed_checks(self, event: HookEvent) -> bool:
 
         if self._update_member_ip():
             return False
+
+        # Don't update this member before it's part of the members list.
+        if self._unit_ip not in self.members_ips:
+            logger.debug("Early exit on_peer_relation_changed: Unit not in the members list")
+            return False
         return True
 
     def _on_peer_relation_changed(self, event: HookEvent):
@@ -677,11 +685,6 @@ def _on_peer_relation_changed(self, event: HookEvent):
         if not self._peer_relation_changed_checks(event):
             return
 
-        # Don't update this member before it's part of the members list.
-        if self._unit_ip not in self.members_ips:
-            logger.debug("Early exit on_peer_relation_changed: Unit not in the members list")
-            return
-
         # Update the list of the cluster members in the replicas to make them know each other.
         try:
             # Update the members of the cluster in the Patroni configuration on this unit.
@@ -716,8 +719,7 @@ def _on_peer_relation_changed(self, event: HookEvent):
         # Restart the workload if it's stuck on the starting state after a timeline divergence
         # due to a backup that was restored.
         if (
-            not self.has_raft_keys()
-            and not self.is_primary
+            not self.is_primary
             and not self.is_standby_leader
             and (
                 self._patroni.member_replication_lag == "unknown"
diff --git a/tests/unit/test_charm.py b/tests/unit/test_charm.py
index d0305dfda8..60be12fbcd 100644
--- a/tests/unit/test_charm.py
+++ b/tests/unit/test_charm.py
@@ -1965,6 +1965,7 @@ def test_raft_reinitialisation(harness):
         ) as _stuck_raft_cluster_cleanup,
         patch("charm.Patroni.remove_raft_data") as _remove_raft_data,
         patch("charm.Patroni.reinitialise_raft_data") as _reinitialise_raft_data,
+        patch("charm.PostgresqlOperatorCharm.update_config") as _update_config,
     ):
         # No data
         harness.charm._raft_reinitialisation()
@@ -2016,6 +2017,7 @@ def test_raft_reinitialisation(harness):
             harness.update_relation_data(rel_id, harness.charm.app.name, {"raft_rejoin": "True"})
         harness.charm._raft_reinitialisation()
         _stuck_raft_cluster_cleanup.assert_called_once_with()
+        _update_config.assert_called_once_with()
 
 
 #

From bd67a39c7a55272d64183cedf99e0cb76f93ceee Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Fri, 18 Oct 2024 13:57:29 +0300
Subject: [PATCH 83/85] Update linting

---
 src/charm.py                          | 18 +++++++-----------
 src/cluster.py                        | 10 ++++++----
 tests/integration/ha_tests/helpers.py |  2 +-
 3 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index d43a53b5d7..bc308ae19b 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -433,8 +433,7 @@ def _on_peer_relation_departed(self, event: RelationDepartedEvent) -> None:
             return
         except RetryError:
             logger.warning(
-                "Early exit on_peer_relation_departed: Cannot get %s member IP"
-                % event.departing_unit.name
+                f"Early exit on_peer_relation_departed: Cannot get {event.departing_unit.name} member IP"
             )
             return
 
@@ -542,7 +541,7 @@ def _stuck_raft_cluster_check(self) -> None:
             logger.warning("Stuck raft has no candidate")
             return
         if "raft_selected_candidate" not in self.app_peer_data:
-            logger.info("%s selected for new raft leader" % candidate.name)
+            logger.info(f"{candidate.name} selected for new raft leader")
             self.app_peer_data["raft_selected_candidate"] = candidate.name
 
     def _stuck_raft_cluster_rejoin(self) -> None:
@@ -588,7 +587,7 @@ def _stuck_raft_cluster_cleanup(self) -> None:
         for key, data in self._peers.data.items():
             if key == self.app:
                 continue
-            for flag in data.keys():
+            for flag in data:
                 if flag.startswith("raft_"):
                     return
 
@@ -611,7 +610,7 @@ def _raft_reinitialisation(self) -> None:
                 self.unit_peer_data.pop("raft_stuck", None)
                 self.unit_peer_data.pop("raft_candidate", None)
                 self._patroni.remove_raft_data()
-                logger.info("Stopping %s" % self.unit.name)
+                logger.info(f"Stopping {self.unit.name}")
                 self.unit_peer_data["raft_stopped"] = "True"
 
             if self.unit.is_leader():
@@ -622,7 +621,7 @@ def _raft_reinitialisation(self) -> None:
                 and "raft_primary" not in self.unit_peer_data
                 and "raft_followers_stopped" in self.app_peer_data
             ):
-                logger.info("Reinitialising %s as primary" % self.unit.name)
+                logger.info(f"Reinitialising {self.unit.name} as primary")
                 self._patroni.reinitialise_raft_data()
                 self.unit_peer_data["raft_primary"] = "True"
 
@@ -641,14 +640,11 @@ def _raft_reinitialisation(self) -> None:
 
     def has_raft_keys(self):
         """Checks for the presence of raft recovery keys in peer data."""
-        for key in self.app_peer_data.keys():
+        for key in self.app_peer_data:
             if key.startswith("raft_"):
                 return True
 
-        for key in self.unit_peer_data.keys():
-            if key.startswith("raft_"):
-                return True
-        return False
+        return any(key.startswith("raft_") for key in self.unit_peer_data)
 
     def _peer_relation_changed_checks(self, event: HookEvent) -> bool:
         """Split of to reduce complexity."""
diff --git a/src/cluster.py b/src/cluster.py
index ed96a94243..88dfab66c8 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -772,7 +772,9 @@ def remove_raft_data(self) -> None:
             if path.exists() and path.is_dir():
                 shutil.rmtree(path)
         except OSError as e:
-            raise Exception(f"Failed to remove previous cluster information with error: {e!s}")
+            raise Exception(
+                f"Failed to remove previous cluster information with error: {e!s}"
+            ) from e
         logger.info("Raft ready to reinitialise")
 
     def reinitialise_raft_data(self) -> None:
@@ -828,7 +830,7 @@ def remove_raft_member(self, member_ip: str) -> None:
             raft_status = syncobj_util.executeCommand(raft_host, ["status"])
         except UtilityException:
             logger.warning("Remove raft member: Cannot connect to raft cluster")
-            raise RemoveRaftMemberFailedError()
+            raise RemoveRaftMemberFailedError() from None
 
         # Check whether the member is still part of the raft cluster.
         if not member_ip or f"partner_node_status_server_{member_ip}:2222" not in raft_status:
@@ -848,7 +850,7 @@ def remove_raft_member(self, member_ip: str) -> None:
             if health_status.get("role") in ("leader", "master") or health_status.get(
                 "sync_standby"
             ):
-                logger.info("%s is raft candidate" % self.charm.unit.name)
+                logger.info(f"{self.charm.unit.name} is raft candidate")
                 data_flags["raft_candidate"] = "True"
             self.charm.unit_peer_data.update(data_flags)
 
@@ -867,7 +869,7 @@ def remove_raft_member(self, member_ip: str) -> None:
             result = syncobj_util.executeCommand(raft_host, ["remove", f"{member_ip}:2222"])
         except UtilityException:
             logger.debug("Remove raft member: Remove call failed")
-            raise RemoveRaftMemberFailedError()
+            raise RemoveRaftMemberFailedError() from None
 
         if not result.startswith("SUCCESS"):
             raise RemoveRaftMemberFailedError()
diff --git a/tests/integration/ha_tests/helpers.py b/tests/integration/ha_tests/helpers.py
index 8ba998aefa..3666a40440 100644
--- a/tests/integration/ha_tests/helpers.py
+++ b/tests/integration/ha_tests/helpers.py
@@ -624,7 +624,7 @@ async def is_replica(ops_test: OpsTest, unit_name: str, use_ip_from_inside: bool
 
 async def get_cluster_roles(
     ops_test: OpsTest, unit_name: str, use_ip_from_inside: bool = False
-) -> Dict[str, Union[Optional[str], list[str]]]:
+) -> dict[str, str | list[str] | None]:
     """Returns whether the unit a replica in the cluster."""
     unit_ip = await (
         get_ip_from_inside_the_unit(ops_test, unit_name)

From 7d49b47fb57a38b63066c5e46cdf07477edb6a64 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Thu, 23 Jan 2025 18:28:24 +0200
Subject: [PATCH 84/85] Restore psutil

---
 poetry.lock    | 39 +++++++++++++++++++++++++++++++++++----
 pyproject.toml |  1 +
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index 7cb15b0e08..a2c5bd0fc6 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 2.0.0 and should not be changed by hand.
+# This file is automatically @generated by Poetry 2.0.1 and should not be changed by hand.
 
 [[package]]
 name = "allure-pytest"
@@ -579,7 +579,6 @@ files = [
     {file = "cryptography-44.0.0-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:761817a3377ef15ac23cd7834715081791d4ec77f9297ee694ca1ee9c2c7e5eb"},
     {file = "cryptography-44.0.0-cp37-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:3c672a53c0fb4725a29c303be906d3c1fa99c32f58abe008a82705f9ee96f40b"},
     {file = "cryptography-44.0.0-cp37-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:4ac4c9f37eba52cb6fbeaf5b59c152ea976726b865bd4cf87883a7e7006cc543"},
-    {file = "cryptography-44.0.0-cp37-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:60eb32934076fa07e4316b7b2742fa52cbb190b42c2df2863dbc4230a0a9b385"},
     {file = "cryptography-44.0.0-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:ed3534eb1090483c96178fcb0f8893719d96d5274dfde98aa6add34614e97c8e"},
     {file = "cryptography-44.0.0-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:f3f6fdfa89ee2d9d496e2c087cebef9d4fcbb0ad63c40e821b39f74bf48d9c5e"},
     {file = "cryptography-44.0.0-cp37-abi3-win32.whl", hash = "sha256:eb33480f1bad5b78233b0ad3e1b0be21e8ef1da745d8d2aecbb20671658b9053"},
@@ -590,7 +589,6 @@ files = [
     {file = "cryptography-44.0.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:c5eb858beed7835e5ad1faba59e865109f3e52b3783b9ac21e7e47dc5554e289"},
     {file = "cryptography-44.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:f53c2c87e0fb4b0c00fa9571082a057e37690a8f12233306161c8f4b819960b7"},
     {file = "cryptography-44.0.0-cp39-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:9e6fc8a08e116fb7c7dd1f040074c9d7b51d74a8ea40d4df2fc7aa08b76b9e6c"},
-    {file = "cryptography-44.0.0-cp39-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:9abcc2e083cbe8dde89124a47e5e53ec38751f0d7dfd36801008f316a127d7ba"},
     {file = "cryptography-44.0.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:d2436114e46b36d00f8b72ff57e598978b37399d2786fd39793c36c6d5cb1c64"},
     {file = "cryptography-44.0.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:a01956ddfa0a6790d594f5b34fc1bfa6098aca434696a03cfdbe469b8ed79285"},
     {file = "cryptography-44.0.0-cp39-abi3-win32.whl", hash = "sha256:eca27345e1214d1b9f9490d200f9db5a874479be914199194e746c893788d417"},
@@ -1525,6 +1523,37 @@ files = [
     {file = "protobuf-4.25.5.tar.gz", hash = "sha256:7f8249476b4a9473645db7f8ab42b02fe1488cbe5fb72fddd445e0665afd8584"},
 ]
 
+[[package]]
+name = "psutil"
+version = "6.1.1"
+description = "Cross-platform lib for process and system monitoring in Python."
+optional = false
+python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,>=2.7"
+groups = ["main"]
+files = [
+    {file = "psutil-6.1.1-cp27-cp27m-macosx_10_9_x86_64.whl", hash = "sha256:9ccc4316f24409159897799b83004cb1e24f9819b0dcf9c0b68bdcb6cefee6a8"},
+    {file = "psutil-6.1.1-cp27-cp27m-manylinux2010_i686.whl", hash = "sha256:ca9609c77ea3b8481ab005da74ed894035936223422dc591d6772b147421f777"},
+    {file = "psutil-6.1.1-cp27-cp27m-manylinux2010_x86_64.whl", hash = "sha256:8df0178ba8a9e5bc84fed9cfa61d54601b371fbec5c8eebad27575f1e105c0d4"},
+    {file = "psutil-6.1.1-cp27-cp27mu-manylinux2010_i686.whl", hash = "sha256:1924e659d6c19c647e763e78670a05dbb7feaf44a0e9c94bf9e14dfc6ba50468"},
+    {file = "psutil-6.1.1-cp27-cp27mu-manylinux2010_x86_64.whl", hash = "sha256:018aeae2af92d943fdf1da6b58665124897cfc94faa2ca92098838f83e1b1bca"},
+    {file = "psutil-6.1.1-cp27-none-win32.whl", hash = "sha256:6d4281f5bbca041e2292be3380ec56a9413b790579b8e593b1784499d0005dac"},
+    {file = "psutil-6.1.1-cp27-none-win_amd64.whl", hash = "sha256:c777eb75bb33c47377c9af68f30e9f11bc78e0f07fbf907be4a5d70b2fe5f030"},
+    {file = "psutil-6.1.1-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:fc0ed7fe2231a444fc219b9c42d0376e0a9a1a72f16c5cfa0f68d19f1a0663e8"},
+    {file = "psutil-6.1.1-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:0bdd4eab935276290ad3cb718e9809412895ca6b5b334f5a9111ee6d9aff9377"},
+    {file = "psutil-6.1.1-cp36-abi3-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b6e06c20c05fe95a3d7302d74e7097756d4ba1247975ad6905441ae1b5b66003"},
+    {file = "psutil-6.1.1-cp36-abi3-manylinux_2_12_x86_64.manylinux2010_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:97f7cb9921fbec4904f522d972f0c0e1f4fabbdd4e0287813b21215074a0f160"},
+    {file = "psutil-6.1.1-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:33431e84fee02bc84ea36d9e2c4a6d395d479c9dd9bba2376c1f6ee8f3a4e0b3"},
+    {file = "psutil-6.1.1-cp36-cp36m-win32.whl", hash = "sha256:384636b1a64b47814437d1173be1427a7c83681b17a450bfc309a1953e329603"},
+    {file = "psutil-6.1.1-cp36-cp36m-win_amd64.whl", hash = "sha256:8be07491f6ebe1a693f17d4f11e69d0dc1811fa082736500f649f79df7735303"},
+    {file = "psutil-6.1.1-cp37-abi3-win32.whl", hash = "sha256:eaa912e0b11848c4d9279a93d7e2783df352b082f40111e078388701fd479e53"},
+    {file = "psutil-6.1.1-cp37-abi3-win_amd64.whl", hash = "sha256:f35cfccb065fff93529d2afb4a2e89e363fe63ca1e4a5da22b603a85833c2649"},
+    {file = "psutil-6.1.1.tar.gz", hash = "sha256:cf8496728c18f2d0b45198f06895be52f36611711746b7f30c464b422b50e2f5"},
+]
+
+[package.extras]
+dev = ["abi3audit", "black", "check-manifest", "coverage", "packaging", "pylint", "pyperf", "pypinfo", "pytest-cov", "requests", "rstcheck", "ruff", "sphinx", "sphinx_rtd_theme", "toml-sort", "twine", "virtualenv", "vulture", "wheel"]
+test = ["pytest", "pytest-xdist", "setuptools"]
+
 [[package]]
 name = "psycopg2"
 version = "2.9.10"
@@ -1539,6 +1568,7 @@ files = [
     {file = "psycopg2-2.9.10-cp311-cp311-win_amd64.whl", hash = "sha256:0435034157049f6846e95103bd8f5a668788dd913a7c30162ca9503fdf542cb4"},
     {file = "psycopg2-2.9.10-cp312-cp312-win32.whl", hash = "sha256:65a63d7ab0e067e2cdb3cf266de39663203d38d6a8ed97f5ca0cb315c73fe067"},
     {file = "psycopg2-2.9.10-cp312-cp312-win_amd64.whl", hash = "sha256:4a579d6243da40a7b3182e0430493dbd55950c493d8c68f4eec0b302f6bbf20e"},
+    {file = "psycopg2-2.9.10-cp313-cp313-win_amd64.whl", hash = "sha256:91fd603a2155da8d0cfcdbf8ab24a2d54bca72795b90d2a3ed2b6da8d979dee2"},
     {file = "psycopg2-2.9.10-cp39-cp39-win32.whl", hash = "sha256:9d5b3b94b79a844a986d029eee38998232451119ad653aea42bb9220a8c5066b"},
     {file = "psycopg2-2.9.10-cp39-cp39-win_amd64.whl", hash = "sha256:88138c8dedcbfa96408023ea2b0c369eda40fe5d75002c0964c78f46f11fa442"},
     {file = "psycopg2-2.9.10.tar.gz", hash = "sha256:12ec0b40b0273f95296233e8750441339298e6a572f7039da5b260e3c8b60e11"},
@@ -1599,6 +1629,7 @@ files = [
     {file = "psycopg2_binary-2.9.10-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:bb89f0a835bcfc1d42ccd5f41f04870c1b936d8507c6df12b7737febc40f0909"},
     {file = "psycopg2_binary-2.9.10-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:f0c2d907a1e102526dd2986df638343388b94c33860ff3bbe1384130828714b1"},
     {file = "psycopg2_binary-2.9.10-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:f8157bed2f51db683f31306aa497311b560f2265998122abe1dce6428bd86567"},
+    {file = "psycopg2_binary-2.9.10-cp313-cp313-win_amd64.whl", hash = "sha256:27422aa5f11fbcd9b18da48373eb67081243662f9b46e6fd07c3eb46e4535142"},
     {file = "psycopg2_binary-2.9.10-cp38-cp38-macosx_12_0_x86_64.whl", hash = "sha256:eb09aa7f9cecb45027683bb55aebaaf45a0df8bf6de68801a6afdc7947bb09d4"},
     {file = "psycopg2_binary-2.9.10-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b73d6d7f0ccdad7bc43e6d34273f70d587ef62f824d7261c4ae9b8b1b6af90e8"},
     {file = "psycopg2_binary-2.9.10-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ce5ab4bf46a211a8e924d307c1b1fcda82368586a19d0a24f8ae166f5c784864"},
@@ -2656,4 +2687,4 @@ type = ["pytest-mypy"]
 [metadata]
 lock-version = "2.1"
 python-versions = "^3.10"
-content-hash = "c0bc686ebd559e3dad53fd7c4c1d96957d716760719981d3513e410f25f34aa8"
+content-hash = "e04a64cc429c59eb73a477073b910a8cc5e0e6c7934630dea3f83f8a0a3b522b"
diff --git a/pyproject.toml b/pyproject.toml
index 62166ca481..6dd00bc099 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,6 +16,7 @@ psycopg2 = "^2.9.10"
 pydantic = "^1.10.19"
 jinja2 = "^3.1.4"
 pysyncobj = "^0.3.13"
+psutil = "^6.1.1"
 
 [tool.poetry.group.charm-libs.dependencies]
 # data_platform_libs/v0/data_interfaces.py

From d20c568a6b5e82cbef490898a1c54e4b494d80ce Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Fri, 31 Jan 2025 17:24:16 +0200
Subject: [PATCH 85/85] Squashed commit of the following:

commit 72ddcac88e70637502af1efb67e1b8a68ecd7cd9
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Wed Jan 29 16:31:18 2025 +0200

    Timeout check

commit adb8ba9b40fdde4441388fe2291cd0f39a4a131e
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Wed Jan 29 15:08:52 2025 +0200

    Missed auto recovery

commit 7d2070a0f3c9d048581a4135330a993ef720f3d3
Merge: d08402610 16e36d0bf
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Wed Jan 29 11:16:53 2025 +0200

    Merge branch 'dpe-3684-reinitialise-raft' into dpe-3684-three-units

commit d0840261071dd1a000f5f2e13a225415b4864698
Merge: 9c00278d5 5d49c72d9
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Mon Jan 27 16:03:19 2025 +0200

    Merge branch 'dpe-3684-reinitialise-raft' into dpe-3684-three-units

commit 9c00278d59f5f6235310f9ad146d6ed3967b99c9
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Thu Jan 23 23:24:39 2025 +0200

    Remove dead code

commit 84c66018ce5c65ef7b399aae8290e9c04948339c
Merge: caf076f58 7d49b47fb
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Thu Jan 23 18:30:01 2025 +0200

    Merge branch 'dpe-3684-reinitialise-raft' into dpe-3684-three-units

commit caf076f5812de805957f7e64d68d872249c6ee8d
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Thu Jan 23 18:02:12 2025 +0200

    Correct timeout exception

commit ecb70d5a1e3c89ed9e8bc2d370b3aa253ce83dde
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Thu Jan 23 16:18:45 2025 +0200

    Check for Patroni self-healing

commit b3f44bcdc986d2c67f540de21ad6a2405f4c92fc
Merge: 7d4bb84bf 3681434f5
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Thu Jan 23 16:04:49 2025 +0200

    Merge branch 'dpe-3684-reinitialise-raft' into dpe-3684-three-units

commit 7d4bb84bfea9299875b267fb9f8ad73300fef150
Merge: f7cef49fe de521225b
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Wed Jan 22 16:57:45 2025 +0200

    Merge branch 'dpe-3684-reinitialise-raft' into dpe-3684-three-units

commit f7cef49fe4f97553c4bffb9df0c27ad3e1480059
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Wed Jan 22 16:52:44 2025 +0200

    Unit tests

commit e1de0c4a97d5cb2d6d5f3d628b92bc194412eb0b
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Wed Jan 22 16:22:19 2025 +0200

    Remove replica test

commit b76412e2ace187112ea0c0d914af56fa1c1ea22d
Merge: 5f6ad5ede 6ae7e52cc
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Wed Jan 22 16:19:02 2025 +0200

    Merge branch 'main' into dpe-3684-three-units

commit 5f6ad5edebae366f92cf8e9b14a6f12eebdbce0d
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Tue Jan 21 22:32:39 2025 +0200

    Switch test order

commit a39bae9360575b95d6cb1e9fb4a75639b9f793ec
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Tue Jan 21 17:48:55 2025 +0200

    Regular promote

commit 66088cb95bad7f5a0a13375f26ee504758aeeb6f
Merge: 720acf004 6917be373
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Mon Jan 20 01:47:28 2025 +0200

    Merge branch 'dpe-3684-reinitialise-raft' into dpe-3684-three-units

commit 720acf0044ac5bbbe206a5d9e3703edb1c210695
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Wed Jan 8 16:55:48 2025 +0200

    Use a sync unit to run the restore action

commit 9b3bcfd3cc91e974b49b4bfe92b7048abf98f6b6
Merge: 86d14b9dd 6ffaf174f
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Wed Jan 8 15:49:31 2025 +0200

    Merge branch 'dpe-3684-reinitialise-raft' into dpe-3684-three-units

commit 86d14b9dd29ca55e3ab6779d00c1a4d5c484acab
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Mon Jan 6 10:55:02 2025 +0200

    GH build job workaround

commit c062fc0b1a575d8b0aee9fed810a349cf5a854ec
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Sun Jan 5 00:42:32 2025 +0200

    Fix tests

commit df65bd0646d8c14082c43b60cd9cdc0af9aea05b
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Mon Dec 30 22:22:13 2024 +0200

    Wrong call

commit 5f3609f9510eef8abaf14d868f64a541de608863
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Mon Dec 30 20:04:23 2024 +0200

    Raise flag on read only primary

commit e459d058c19feb1258afec3ef0f85853180941b0
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Thu Dec 26 16:23:45 2024 +0200

    Tests

commit ffd6d582155314d49363684b904471a1db96c3ac
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Tue Dec 24 17:41:26 2024 +0200

    Bump timeout

commit 74b4d481a23785f1432569a7c89cecb437200364
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Tue Dec 24 17:01:35 2024 +0200

    Switch action

commit dd8819806d41e66b77fa50d4501377570712463e
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Tue Dec 24 16:22:34 2024 +0200

    Promote via existing action

commit 55a7a7ed0c7178501d400631bb1e55efe028812b
Merge: 2308e8caa a20cf72bd
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Tue Dec 24 15:51:57 2024 +0200

    Merge branch 'dpe-3684-three-units' into dpe-3684-reinit-action

commit a20cf72bd5c4cd5466bec2fbc5efab449e019a18
Merge: 1e17cedd0 d155abf7b
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Tue Dec 24 15:51:23 2024 +0200

    Merge branch 'dpe-3684-reinitialise-raft' into dpe-3684-three-units

commit 2308e8caa17c571131e83634a0807f26d041e95c
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Tue Nov 26 19:21:16 2024 +0200

    Run reinit as part of the action for leader

commit e224cafa4958de04687b359558c596c35a1e1ed3
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Tue Nov 26 18:13:13 2024 +0200

    Add experimental action

commit 1e17cedd0e9507b32f2aa09254c36811cad4f34b
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Tue Nov 26 16:31:08 2024 +0200

    Wrong parameter

commit 5e18530c6f6687e6805e8036c2c0e6bd67dc07e7
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Tue Nov 26 16:28:14 2024 +0200

    Try to reinit on update status

commit 07fd96917bd88e9494d18caa25b9fb305e3d3b48
Merge: ebf136147 a307bd85b
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Mon Nov 25 18:59:58 2024 +0200

    Merge branch 'dpe-3684-reinitialise-raft' into dpe-3684-three-units

commit ebf136147aadfe5e1eae467a74bf5479cc334182
Merge: 782979b7f 553589674
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Fri Nov 22 16:22:16 2024 +0200

    Merge branch 'dpe-3684-reinitialise-raft' into dpe-3684-three-units

commit 782979b7f8b5789f7d04cbf01bfaa3e06ad2051e
Merge: da33c75cb 089624548
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Tue Nov 19 12:50:18 2024 +0200

    Merge branch 'dpe-3684-reinitialise-raft' into dpe-3684-three-units

commit da33c75cb94e01261f0d15a1c7c8c32e24f5a6d2
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Wed Oct 30 00:05:04 2024 +0200

    Linting

commit ae5b560e634adfe32af0c42a18909fa784d01566
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Tue Oct 29 23:54:44 2024 +0200

    Degraded status message

commit cc16d5fde1b191f9765728e2caf7a43147fd63ec
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Mon Oct 28 22:25:12 2024 +0200

    Set maintanance status during reinit

commit 5088114b164f26d6c95401237b21f038ac5e633b
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Sun Oct 27 23:11:21 2024 +0200

    Unit tests

commit 634a9bb652c243285ee8a082d5b73fc010647cb0
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Sun Oct 27 22:58:32 2024 +0200

    Track candidate by cluster status

commit c0fece549975c7b08f2382ce172d6460ba80a293
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Sun Oct 27 21:40:41 2024 +0200

    Log roles removed

commit e06c0770744fd97fed03e01a395d7db6c814e132
Merge: 7123659dd c769b8e9d
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Sun Oct 27 20:27:38 2024 +0200

    Merge branch 'dpe-3684-reinitialise-raft' into dpe-3684-three-units

commit 7123659dd3112a7d9409b5d2ff65a660289643d9
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Sun Oct 27 19:45:42 2024 +0200

    Parameterised test

commit 5d972aca654169f1009181c6d8d24530616e3218
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Sun Oct 27 18:37:41 2024 +0200

    Parameterised test

commit 42899c17a10f62e8ceeabd0336ca7668e48d500d
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Sun Oct 27 15:13:29 2024 +0200

    Pass the machine name when stopping

commit e0a09ae17c9d29baecc017fda1f7832ac050f01c
Author: Dragomir Penev <dragomir.penev@canonical.com>
Date:   Sun Oct 27 13:46:39 2024 +0200

    Three units scaling tests
---
 actions.yaml                                  |   7 +-
 src/charm.py                                  |  45 +++++-
 src/cluster.py                                |  47 ++++--
 src/relations/async_replication.py            |   5 +-
 tests/integration/ha_tests/test_scaling.py    |  38 +++++
 .../ha_tests/test_scaling_three_units.py      | 134 ++++++++++++++++++
 tests/integration/test_charm.py               |  19 +++
 tests/unit/test_charm.py                      |  58 +++++++-
 tests/unit/test_cluster.py                    |  28 +++-
 9 files changed, 352 insertions(+), 29 deletions(-)
 create mode 100644 tests/integration/ha_tests/test_scaling_three_units.py

diff --git a/actions.yaml b/actions.yaml
index fa9ed87da3..7069c4668e 100644
--- a/actions.yaml
+++ b/actions.yaml
@@ -33,8 +33,13 @@ list-backups:
 pre-upgrade-check:
   description: Run necessary pre-upgrade checks and preparations before executing a charm refresh.
 promote-to-primary:
-  description: Promotes the cluster of choice to a primary cluster. Must be ran against the leader unit.
+  description: Promotes the cluster of choice to a primary cluster. Must be ran against the leader unit when promoting a cluster
+   or against the unit to be promoted within the cluster.
   params:
+    scope:
+      type: string
+      default: cluster
+      description: Whether to promote a unit or a cluster. Must be set to either unit or cluster.
     force:
       type: boolean
       description: Force the promotion of a cluster when there is already a primary cluster.
diff --git a/src/charm.py b/src/charm.py
index b738e19490..db938ff4c6 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -185,6 +185,7 @@ def __init__(self, *args):
         self.framework.observe(self.on.start, self._on_start)
         self.framework.observe(self.on.get_password_action, self._on_get_password)
         self.framework.observe(self.on.set_password_action, self._on_set_password)
+        self.framework.observe(self.on.promote_to_primary_action, self._on_promote_to_primary)
         self.framework.observe(self.on.update_status, self._on_update_status)
         self.cluster_name = self.app.name
         self._member_name = self.unit.name.replace("/", "-")
@@ -631,6 +632,7 @@ def _raft_reinitialisation(self) -> None:
                 and "raft_primary" not in self.unit_peer_data
                 and "raft_followers_stopped" in self.app_peer_data
             ):
+                self.unit.status = MaintenanceStatus("Reinitialising raft")
                 logger.info(f"Reinitialising {self.unit.name} as primary")
                 self._patroni.reinitialise_raft_data()
                 self.unit_peer_data["raft_primary"] = "True"
@@ -644,6 +646,7 @@ def _raft_reinitialisation(self) -> None:
             self.unit_peer_data.pop("raft_stopped", None)
             self.update_config()
             self._patroni.start_patroni()
+            self._set_primary_status_message()
 
             if self.unit.is_leader():
                 self._stuck_raft_cluster_cleanup()
@@ -1506,6 +1509,32 @@ def _on_set_password(self, event: ActionEvent) -> None:
 
         event.set_results({"password": password})
 
+    def _on_promote_to_primary(self, event: ActionEvent) -> None:
+        if event.params.get("scope") == "cluster":
+            return self.async_replication.promote_to_primary(event)
+        elif event.params.get("scope") == "unit":
+            return self.promote_primary_unit(event)
+        else:
+            event.fail("Scope should be either cluster or unit")
+
+    def promote_primary_unit(self, event: ActionEvent) -> None:
+        """Handles promote to primary for unit scope."""
+        if event.params.get("force"):
+            if self.has_raft_keys():
+                self.unit_peer_data.update({"raft_candidate": "True"})
+                if self.unit.is_leader():
+                    self._raft_reinitialisation()
+                return
+            event.fail("Raft is not stuck")
+        else:
+            if self.has_raft_keys():
+                event.fail("Raft is stuck. Set force to reinitialise with new primary")
+                return
+            try:
+                self._patroni.switchover(self._member_name)
+            except SwitchoverFailedError:
+                event.fail("Unit is not sync standby")
+
     def _on_update_status(self, _) -> None:
         """Update the unit status message and users list in the database."""
         if not self._can_run_on_update_status():
@@ -1675,10 +1704,18 @@ def _set_primary_status_message(self) -> None:
                     self.app_peer_data["s3-initialization-block-message"]
                 )
                 return
-            if self._patroni.get_primary(unit_name_pattern=True) == self.unit.name:
-                self.unit.status = ActiveStatus("Primary")
-            elif self.is_standby_leader:
-                self.unit.status = ActiveStatus("Standby")
+            if (
+                self._patroni.get_primary(unit_name_pattern=True) == self.unit.name
+                or self.is_standby_leader
+            ):
+                danger_state = ""
+                if not self._patroni.has_raft_quorum():
+                    danger_state = " (read-only)"
+                elif len(self._patroni.get_running_cluster_members()) < self.app.planned_units():
+                    danger_state = " (degraded)"
+                self.unit.status = ActiveStatus(
+                    f"{'Standby' if self.is_standby_leader else 'Primary'}{danger_state}"
+                )
             elif self._patroni.member_started:
                 self.unit.status = ActiveStatus()
         except (RetryError, ConnectionError) as e:
diff --git a/src/cluster.py b/src/cluster.py
index a4f891ba79..cf7801b95d 100644
--- a/src/cluster.py
+++ b/src/cluster.py
@@ -18,6 +18,7 @@
 import requests
 from charms.operator_libs_linux.v2 import snap
 from jinja2 import Template
+from ops import BlockedStatus
 from pysyncobj.utility import TcpUtility, UtilityException
 from tenacity import (
     AttemptManager,
@@ -746,15 +747,18 @@ def stop_patroni(self) -> bool:
             logger.exception(error_message, exc_info=e)
             return False
 
-    def switchover(self) -> None:
+    def switchover(self, candidate: str | None = None) -> None:
         """Trigger a switchover."""
         # Try to trigger the switchover.
         for attempt in Retrying(stop=stop_after_delay(60), wait=wait_fixed(3)):
             with attempt:
                 current_primary = self.get_primary()
+                body = {"leader": current_primary}
+                if candidate:
+                    body["candidate"] = candidate
                 r = requests.post(
                     f"{self._patroni_url}/switchover",
-                    json={"leader": current_primary},
+                    json=body,
                     verify=self.verify,
                     auth=self._patroni_auth,
                     timeout=PATRONI_TIMEOUT,
@@ -774,6 +778,19 @@ def primary_changed(self, old_primary: str) -> bool:
         primary = self.get_primary()
         return primary != old_primary
 
+    def has_raft_quorum(self) -> bool:
+        """Check if raft cluster has quorum."""
+        # Get the status of the raft cluster.
+        syncobj_util = TcpUtility(password=self.raft_password, timeout=3)
+
+        raft_host = "127.0.0.1:2222"
+        try:
+            raft_status = syncobj_util.executeCommand(raft_host, ["status"])
+        except UtilityException:
+            logger.warning("Has raft quorum: Cannot connect to raft cluster")
+            return False
+        return raft_status["has_quorum"]
+
     def remove_raft_data(self) -> None:
         """Stops Patroni and removes the raft journals."""
         logger.info("Stopping patroni")
@@ -827,6 +844,21 @@ def reinitialise_raft_data(self) -> None:
                     raise RaftPostgresqlNotUpError()
         logger.info("Raft should be unstuck")
 
+    def get_running_cluster_members(self) -> list[str]:
+        """List running patroni members."""
+        try:
+            members = requests.get(
+                f"{self._patroni_url}/{PATRONI_CLUSTER_STATUS_ENDPOINT}",
+                verify=self.verify,
+                timeout=API_REQUEST_TIMEOUT,
+                auth=self._patroni_auth,
+            ).json()["members"]
+            return [
+                member["name"] for member in members if member["state"] in ("streaming", "running")
+            ]
+        except Exception:
+            return []
+
     def remove_raft_member(self, member_ip: str) -> None:
         """Remove a member from the raft cluster.
 
@@ -860,18 +892,9 @@ def remove_raft_member(self, member_ip: str) -> None:
         if not raft_status["has_quorum"] and (
             not raft_status["leader"] or raft_status["leader"].host == member_ip
         ):
+            self.charm.unit.status = BlockedStatus("Raft majority loss, run: promote-to-primary")
             logger.warning("Remove raft member: Stuck raft cluster detected")
             data_flags = {"raft_stuck": "True"}
-            try:
-                health_status = self.get_patroni_health()
-            except Exception:
-                logger.warning("Remove raft member: Unable to get health status")
-                health_status = {}
-            if health_status.get("role") in ("leader", "master") or health_status.get(
-                "sync_standby"
-            ):
-                logger.info(f"{self.charm.unit.name} is raft candidate")
-                data_flags["raft_candidate"] = "True"
             self.charm.unit_peer_data.update(data_flags)
 
             # Leader doesn't always trigger when changing it's own peer data.
diff --git a/src/relations/async_replication.py b/src/relations/async_replication.py
index ebfa06752d..db40602cea 100644
--- a/src/relations/async_replication.py
+++ b/src/relations/async_replication.py
@@ -106,9 +106,6 @@ def __init__(self, charm):
         self.framework.observe(
             self.charm.on.create_replication_action, self._on_create_replication
         )
-        self.framework.observe(
-            self.charm.on.promote_to_primary_action, self._on_promote_to_primary
-        )
 
         self.framework.observe(self.charm.on.secret_changed, self._on_secret_changed)
 
@@ -583,7 +580,7 @@ def _on_create_replication(self, event: ActionEvent) -> None:
         # Set the status.
         self.charm.unit.status = MaintenanceStatus("Creating replication...")
 
-    def _on_promote_to_primary(self, event: ActionEvent) -> None:
+    def promote_to_primary(self, event: ActionEvent) -> None:
         """Promote this cluster to the primary cluster."""
         if (
             self.charm.app.status.message != READ_ONLY_MODE_BLOCKING_MESSAGE
diff --git a/tests/integration/ha_tests/test_scaling.py b/tests/integration/ha_tests/test_scaling.py
index 3d92cd9b2c..f3105d38cd 100644
--- a/tests/integration/ha_tests/test_scaling.py
+++ b/tests/integration/ha_tests/test_scaling.py
@@ -69,6 +69,24 @@ async def test_removing_stereo_primary(ops_test: OpsTest, continuous_writes) ->
     primary = await get_primary(ops_test, app)
     await ops_test.model.destroy_unit(primary, force=True, destroy_storage=False, max_wait=1500)
 
+    left_unit = ops_test.model.units[original_roles["sync_standbys"][0]]
+    for left_unit in ops_test.model.applications[DATABASE_APP_NAME].units:
+        if left_unit.name not in original_roles["primaries"]:
+            break
+
+    await ops_test.model.block_until(
+        lambda: left_unit.workload_status == "blocked"
+        and left_unit.workload_status_message == "Raft majority loss, run: promote-to-primary",
+        timeout=600,
+    )
+
+    run_action = (
+        await ops_test.model.applications[DATABASE_APP_NAME]
+        .units[0]
+        .run_action("promote-to-primary", scope="unit", force=True)
+    )
+    await run_action.wait()
+
     await ops_test.model.wait_for_idle(status="active", timeout=600, idle_period=45)
 
     await are_writes_increasing(ops_test, primary)
@@ -154,6 +172,16 @@ async def test_removing_raft_majority(ops_test: OpsTest, continuous_writes) -> N
         ),
     )
 
+    left_unit = ops_test.model.units[original_roles["sync_standbys"][1]]
+    await ops_test.model.block_until(
+        lambda: left_unit.workload_status == "blocked"
+        and left_unit.workload_status_message == "Raft majority loss, run: promote-to-primary",
+        timeout=600,
+    )
+
+    run_action = await left_unit.run_action("promote-to-primary", scope="unit", force=True)
+    await run_action.wait()
+
     await ops_test.model.wait_for_idle(status="active", timeout=900, idle_period=45)
 
     await are_writes_increasing(
@@ -202,6 +230,16 @@ async def test_removing_raft_majority_async(ops_test: OpsTest, continuous_writes
         ),
     )
 
+    left_unit = ops_test.model.units[original_roles["sync_standbys"][0]]
+    await ops_test.model.block_until(
+        lambda: left_unit.workload_status == "blocked"
+        and left_unit.workload_status_message == "Raft majority loss, run: promote-to-primary",
+        timeout=600,
+    )
+
+    run_action = await left_unit.run_action("promote-to-primary", scope="unit", force=True)
+    await run_action.wait()
+
     await ops_test.model.wait_for_idle(status="active", timeout=900, idle_period=45)
 
     await are_writes_increasing(
diff --git a/tests/integration/ha_tests/test_scaling_three_units.py b/tests/integration/ha_tests/test_scaling_three_units.py
new file mode 100644
index 0000000000..6817cd238a
--- /dev/null
+++ b/tests/integration/ha_tests/test_scaling_three_units.py
@@ -0,0 +1,134 @@
+#!/usr/bin/env python3
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+import logging
+from asyncio import exceptions, gather, sleep
+
+import pytest
+from pytest_operator.plugin import OpsTest
+
+from .. import markers
+from ..helpers import (
+    CHARM_BASE,
+    DATABASE_APP_NAME,
+    get_machine_from_unit,
+    stop_machine,
+)
+from .conftest import APPLICATION_NAME
+from .helpers import (
+    app_name,
+    are_writes_increasing,
+    check_writes,
+    get_cluster_roles,
+    start_continuous_writes,
+)
+
+logger = logging.getLogger(__name__)
+
+charm = None
+
+
+@pytest.mark.group(1)
+@markers.juju3
+@pytest.mark.abort_on_fail
+async def test_build_and_deploy(ops_test: OpsTest) -> None:
+    """Build and deploy two PostgreSQL clusters."""
+    # This is a potentially destructive test, so it shouldn't be run against existing clusters
+    charm = await ops_test.build_charm(".")
+    async with ops_test.fast_forward():
+        # Deploy the first cluster with reusable storage
+        await gather(
+            ops_test.model.deploy(
+                charm,
+                application_name=DATABASE_APP_NAME,
+                num_units=3,
+                base=CHARM_BASE,
+                config={"profile": "testing"},
+            ),
+            ops_test.model.deploy(
+                APPLICATION_NAME,
+                application_name=APPLICATION_NAME,
+                base=CHARM_BASE,
+                channel="edge",
+            ),
+        )
+
+        await ops_test.model.wait_for_idle(status="active", timeout=1500)
+
+
+@pytest.mark.group(1)
+@markers.juju3
+@pytest.mark.parametrize(
+    "roles",
+    [
+        ["primaries"],
+        ["sync_standbys"],
+        ["replicas"],
+        ["primaries", "replicas"],
+        ["sync_standbys", "replicas"],
+    ],
+)
+@pytest.mark.abort_on_fail
+async def test_removing_unit(ops_test: OpsTest, roles: list[str], continuous_writes) -> None:
+    logger.info(f"removing {', '.join(roles)}")
+    # Start an application that continuously writes data to the database.
+    app = await app_name(ops_test)
+    original_roles = await get_cluster_roles(
+        ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
+    )
+    await start_continuous_writes(ops_test, app)
+    units = [original_roles[role][0] for role in roles]
+    for unit in units:
+        logger.info(f"Stopping unit {unit}")
+        await stop_machine(ops_test, await get_machine_from_unit(ops_test, unit))
+    await sleep(15)
+    for unit in units:
+        logger.info(f"Deleting unit {unit}")
+        await ops_test.model.destroy_unit(unit, force=True, destroy_storage=False, max_wait=1500)
+
+    if len(roles) > 1:
+        for left_unit in ops_test.model.applications[DATABASE_APP_NAME].units:
+            if left_unit.name not in units:
+                break
+        try:
+            await ops_test.model.block_until(
+                lambda: left_unit.workload_status == "blocked"
+                and left_unit.workload_status_message
+                == "Raft majority loss, run: promote-to-primary",
+                timeout=600,
+            )
+
+            run_action = (
+                await ops_test.model.applications[DATABASE_APP_NAME]
+                .units[0]
+                .run_action("promote-to-primary", scope="unit", force=True)
+            )
+            await run_action.wait()
+        except exceptions.TimeoutError:
+            # Check if Patroni self healed
+            assert (
+                left_unit.workload_status == "active"
+                and left_unit.workload_status_message == "Primary"
+            )
+            logger.warning(f"Patroni self-healed without raft reinitialisation for roles {roles}")
+
+    await ops_test.model.wait_for_idle(status="active", timeout=600, idle_period=45)
+
+    await are_writes_increasing(ops_test, units)
+
+    logger.info("Scaling back up")
+    await ops_test.model.applications[DATABASE_APP_NAME].add_unit(count=len(roles))
+    await ops_test.model.wait_for_idle(status="active", timeout=1500)
+
+    new_roles = await get_cluster_roles(
+        ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
+    )
+    assert len(new_roles["primaries"]) == 1
+    assert len(new_roles["sync_standbys"]) == 1
+    assert len(new_roles["replicas"]) == 1
+    if "primaries" in roles:
+        assert new_roles["primaries"][0] == original_roles["sync_standbys"][0]
+    else:
+        assert new_roles["primaries"][0] == original_roles["primaries"][0]
+
+    await check_writes(ops_test)
diff --git a/tests/integration/test_charm.py b/tests/integration/test_charm.py
index 943e521cef..53ac76e4b2 100644
--- a/tests/integration/test_charm.py
+++ b/tests/integration/test_charm.py
@@ -15,6 +15,7 @@
 
 from locales import SNAP_LOCALES
 
+from .ha_tests.helpers import get_cluster_roles
 from .helpers import (
     CHARM_BASE,
     DATABASE_APP_NAME,
@@ -323,6 +324,24 @@ async def test_scale_down_and_up(ops_test: OpsTest):
     await scale_application(ops_test, DATABASE_APP_NAME, initial_scale)
 
 
+@pytest.mark.group(1)
+async def test_switchover_sync_standby(ops_test: OpsTest):
+    original_roles = await get_cluster_roles(
+        ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
+    )
+    run_action = await ops_test.model.units[original_roles["sync_standbys"][0]].run_action(
+        "promote-to-primary", scope="unit"
+    )
+    await run_action.wait()
+
+    await ops_test.model.wait_for_idle(status="active", timeout=200)
+
+    new_roles = await get_cluster_roles(
+        ops_test, ops_test.model.applications[DATABASE_APP_NAME].units[0].name
+    )
+    assert new_roles["primaries"][0] == original_roles["sync_standbys"][0]
+
+
 @pytest.mark.group(1)
 async def test_persist_data_through_primary_deletion(ops_test: OpsTest):
     """Test data persists through a primary deletion."""
diff --git a/tests/unit/test_charm.py b/tests/unit/test_charm.py
index 1ce381e8a4..77b62a08dc 100644
--- a/tests/unit/test_charm.py
+++ b/tests/unit/test_charm.py
@@ -35,7 +35,7 @@
     PRIMARY_NOT_REACHABLE_MESSAGE,
     PostgresqlOperatorCharm,
 )
-from cluster import NotReadyError, RemoveRaftMemberFailedError
+from cluster import NotReadyError, RemoveRaftMemberFailedError, SwitchoverFailedError
 from constants import PEER, POSTGRESQL_SNAP_NAME, SECRET_INTERNAL_LABEL, SNAP_PACKAGES
 
 CREATE_CLUSTER_CONF_PATH = "/etc/postgresql-common/createcluster.d/pgcharm.conf"
@@ -1989,6 +1989,7 @@ def test_raft_reinitialisation(harness):
         patch("charm.Patroni.remove_raft_data") as _remove_raft_data,
         patch("charm.Patroni.reinitialise_raft_data") as _reinitialise_raft_data,
         patch("charm.PostgresqlOperatorCharm.update_config") as _update_config,
+        patch("charm.PostgresqlOperatorCharm._set_primary_status_message"),
     ):
         # No data
         harness.charm._raft_reinitialisation()
@@ -2569,6 +2570,8 @@ def test_update_new_unit_status(harness):
 @pytest.mark.parametrize("is_leader", [True, False])
 def test_set_primary_status_message(harness, is_leader):
     with (
+        patch("charm.Patroni.has_raft_quorum", return_value=True),
+        patch("charm.Patroni.get_running_cluster_members", return_value=["test"]),
         patch("charm.Patroni.member_started", new_callable=PropertyMock) as _member_started,
         patch(
             "charm.PostgresqlOperatorCharm.is_standby_leader", new_callable=PropertyMock
@@ -2602,7 +2605,9 @@ def test_set_primary_status_message(harness, is_leader):
                     assert isinstance(harness.charm.unit.status, MaintenanceStatus)
                 else:
                     _is_standby_leader.side_effect = None
-                    _is_standby_leader.return_value = values[1]
+                    _is_standby_leader.return_value = (
+                        values[0] != harness.charm.unit.name and values[1]
+                    )
                     harness.charm._set_primary_status_message()
                     assert isinstance(
                         harness.charm.unit.status,
@@ -2754,3 +2759,52 @@ def test_get_plugins(harness):
             "insert_username",
             "moddatetime",
         ]
+
+
+def test_on_promote_to_primary(harness):
+    with (
+        patch("charm.PostgresqlOperatorCharm._raft_reinitialisation") as _raft_reinitialisation,
+        patch("charm.PostgreSQLAsyncReplication.promote_to_primary") as _promote_to_primary,
+        patch("charm.Patroni.switchover") as _switchover,
+    ):
+        event = Mock()
+        event.params = {"scope": "cluster"}
+
+        # Cluster
+        harness.charm._on_promote_to_primary(event)
+        _promote_to_primary.assert_called_once_with(event)
+
+        # Unit, no force, regular promotion
+        event.params = {"scope": "unit"}
+
+        harness.charm._on_promote_to_primary(event)
+
+        _switchover.assert_called_once_with("postgresql-0")
+
+        # Unit, no force, switchover failed
+        event.params = {"scope": "unit"}
+        _switchover.side_effect = SwitchoverFailedError
+
+        harness.charm._on_promote_to_primary(event)
+
+        event.fail.assert_called_once_with("Unit is not sync standby")
+        event.fail.reset_mock()
+
+        # Unit, no force, raft stuck
+        event.params = {"scope": "unit"}
+        rel_id = harness.model.get_relation(PEER).id
+        with harness.hooks_disabled():
+            harness.update_relation_data(rel_id, harness.charm.unit.name, {"raft_stuck": "True"})
+
+        harness.charm._on_promote_to_primary(event)
+        event.fail.assert_called_once_with(
+            "Raft is stuck. Set force to reinitialise with new primary"
+        )
+
+        # Unit, raft reinit
+        event.params = {"scope": "unit", "force": "true"}
+        with harness.hooks_disabled():
+            harness.set_leader()
+        harness.charm._on_promote_to_primary(event)
+        _raft_reinitialisation.assert_called_once_with()
+        assert harness.charm.unit_peer_data["raft_candidate"] == "True"
diff --git a/tests/unit/test_cluster.py b/tests/unit/test_cluster.py
index 3eb42242fc..07f01eed47 100644
--- a/tests/unit/test_cluster.py
+++ b/tests/unit/test_cluster.py
@@ -460,6 +460,18 @@ def test_switchover(peers_ips, patroni):
             auth=patroni._patroni_auth,
             timeout=PATRONI_TIMEOUT,
         )
+        _post.reset_mock()
+
+        # Test candidate
+        patroni.switchover("candidate")
+
+        _post.assert_called_once_with(
+            "http://1.1.1.1:8008/switchover",
+            json={"leader": "primary", "candidate": "candidate"},
+            verify=True,
+            auth=patroni._patroni_auth,
+            timeout=PATRONI_TIMEOUT,
+        )
 
 
 def test_update_synchronous_node_count(peers_ips, patroni):
@@ -735,7 +747,7 @@ def test_remove_raft_member(patroni):
 def test_remove_raft_member_no_quorum(patroni, harness):
     with (
         patch("cluster.TcpUtility") as _tcp_utility,
-        patch("cluster.Patroni.get_patroni_health") as _get_patroni_health,
+        patch("cluster.requests.get") as _get,
         patch(
             "charm.PostgresqlOperatorCharm.unit_peer_data", new_callable=PropertyMock
         ) as _unit_peer_data,
@@ -747,7 +759,9 @@ def test_remove_raft_member_no_quorum(patroni, harness):
             "has_quorum": False,
             "leader": None,
         }
-        _get_patroni_health.return_value = {"role": "replica", "sync_standby": False}
+        _get.return_value.json.return_value = {
+            "members": [{"role": "async_replica", "name": "postgresql-0"}]
+        }
 
         patroni.remove_raft_member("1.2.3.4")
         assert harness.charm.unit_peer_data == {"raft_stuck": "True"}
@@ -759,7 +773,7 @@ def test_remove_raft_member_no_quorum(patroni, harness):
             "has_quorum": False,
             "leader": None,
         }
-        _get_patroni_health.side_effect = Exception
+        _get.side_effect = Exception
 
         patroni.remove_raft_member("1.2.3.4")
 
@@ -774,12 +788,14 @@ def test_remove_raft_member_no_quorum(patroni, harness):
             "has_quorum": False,
             "leader": leader_mock,
         }
-        _get_patroni_health.side_effect = None
-        _get_patroni_health.return_value = {"role": "replica", "sync_standby": True}
+        _get.side_effect = None
+        _get.return_value.json.return_value = {
+            "members": [{"role": "sync_standby", "name": "postgresql-0"}]
+        }
 
         patroni.remove_raft_member("1.2.3.4")
 
-        assert harness.charm.unit_peer_data == {"raft_candidate": "True", "raft_stuck": "True"}
+        assert harness.charm.unit_peer_data == {"raft_stuck": "True"}
 
 
 def test_remove_raft_data(patroni):