[DPE-3473] Stabilise SST and network cut tests (canonical#385)

* Stabilise network cut test

Signed-off-by: Marcelo Henrique Neppel <marcelo.neppel@canonical.com>

* Fix forceful restart without data and transaction logs test

Signed-off-by: Marcelo Henrique Neppel <marcelo.neppel@canonical.com>

* Minor fixes + add unit tests

Signed-off-by: Marcelo Henrique Neppel <marcelo.neppel@canonical.com>

---------

Signed-off-by: Marcelo Henrique Neppel <marcelo.neppel@canonical.com>

Author: marceloneppel

poetry.lock

@@ -398,7 +398,7 @@ def _on_peer_relation_changed(self, event: HookEvent) -> None:
         ):
             self._patroni.reinitialize_postgresql()
             logger.debug("Deferring on_peer_relation_changed: reinitialising replica")
-            self.unit.status = WaitingStatus("reinitialising replica")
+            self.unit.status = MaintenanceStatus("reinitialising replica")
             event.defer()
             return
 
@@ -1079,6 +1079,22 @@ def _handle_processes_failures(self) -> bool:
                 return False
             return True
 
+        if (
+            not self.is_primary
+            and self._patroni.member_started
+            and not self._patroni.member_streaming
+        ):
+            try:
+                self._patroni.reinitialize_postgresql()
+                logger.info("restarted the replica because it was not streaming from primary")
+                self.unit.status = MaintenanceStatus("reinitialising replica")
+            except RetryError:
+                logger.error(
+                    "failed to reinitialise replica after checking that it was not streaming from primary"
+                )
+                return False
+            return True
+
         return False
 
     def _set_primary_status_message(self) -> None:

src/charm.py

@@ -268,6 +268,23 @@ def member_started(self) -> bool:
 
         return r.json()["state"] in RUNNING_STATES
 
+    @property
+    def member_streaming(self) -> bool:
+        """Has the member started to stream data from primary.
+
+        Returns:
+            True if it's streaming False otherwise. Retries over a period of 60 seconds times to
+            allow server time to start up.
+        """
+        try:
+            for attempt in Retrying(stop=stop_after_delay(60), wait=wait_fixed(3)):
+                with attempt:
+                    r = requests.get(f"{self._patroni_url}/health", verify=self._verify)
+        except RetryError:
+            return False
+
+        return r.json().get("replication_state") == "streaming"
+
     @property
     def is_database_running(self) -> bool:
         """Returns whether the PostgreSQL database process is running (and isn't frozen)."""

src/patroni.py

@@ -155,7 +155,7 @@ async def is_cluster_updated(ops_test: OpsTest, primary_name: str) -> None:
     # Verify that old primary is up-to-date.
     assert await is_secondary_up_to_date(
         ops_test, primary_name, total_expected_writes
-    ), "secondary not up to date with the cluster after restarting."
+    ), f"secondary ({primary_name}) not up to date with the cluster after restarting."
 
 
 def get_member_lag(cluster: Dict, member_name: str) -> int:
@@ -191,7 +191,7 @@ async def check_writes(ops_test) -> int:
     for member, count in actual_writes.items():
         assert (
             count == max_number_written[member]
-        ), f"{member}: writes to the db were missed: count of actual writes different from the max number written."
+        ), f"{member}: writes to the db were missed: count of actual writes ({count}) on {member} different from the max number written ({max_number_written[member]})."
         assert total_expected_writes == count, f"{member}: writes to the db were missed."
     return total_expected_writes
 
@@ -200,10 +200,12 @@ async def are_writes_increasing(ops_test, down_unit: str = None) -> None:
     """Verify new writes are continuing by counting the number of writes."""
     writes, _ = await count_writes(ops_test, down_unit=down_unit)
     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, down_unit=down_unit)
-                assert more_writes[member] > count, f"{member}: writes not continuing to DB"
+                assert (
+                    more_writes[member] > count
+                ), f"{member}: writes not continuing to DB (current writes: {more_writes[member]} - previous writes: {count})"
 
 
 def copy_file_into_pod(
@@ -422,14 +424,16 @@ async def is_connection_possible(ops_test: OpsTest, unit_name: str) -> bool:
     password = await get_password(ops_test, database_app_name=app, down_unit=unit_name)
     address = await get_unit_address(ops_test, unit_name)
     try:
-        with db_connect(
-            host=address, password=password
-        ) as connection, connection.cursor() as cursor:
-            cursor.execute("SELECT 1;")
-            success = cursor.fetchone()[0] == 1
-        connection.close()
-        return success
-    except psycopg2.Error:
+        for attempt in Retrying(stop=stop_after_delay(60), wait=wait_fixed(3)):
+            with attempt:
+                with db_connect(
+                    host=address, password=password
+                ) as connection, connection.cursor() as cursor:
+                    cursor.execute("SELECT 1;")
+                    success = cursor.fetchone()[0] == 1
+                connection.close()
+                return success
+    except (psycopg2.Error, RetryError):
         # Error raised when the connection is not possible.
         return False
 
@@ -614,14 +618,19 @@ async def is_secondary_up_to_date(ops_test: OpsTest, unit_name: str, expected_wr
     )
 
     try:
-        for attempt in Retrying(stop=stop_after_delay(60), wait=wait_fixed(3)):
+        for attempt in Retrying(stop=stop_after_delay(60 * 3), wait=wait_fixed(3)):
             with attempt:
                 with psycopg2.connect(
                     connection_string
                 ) as connection, connection.cursor() as cursor:
                     cursor.execute("SELECT COUNT(number), MAX(number) FROM continuous_writes;")
                     results = cursor.fetchone()
-                    assert results[0] == expected_writes and results[1] == expected_writes
+                    if results[0] != expected_writes or results[1] != expected_writes:
+                        async with ops_test.fast_forward(fast_interval="30s"):
+                            await ops_test.model.wait_for_idle(
+                                apps=[unit_name.split("/")[0]], idle_period=15, timeout=1000
+                            )
+                            raise Exception
     except RetryError:
         return False
     finally:

tests/integration/ha_tests/helpers.py

@@ -101,6 +101,7 @@ async def test_kill_db_process(
         await are_writes_increasing(ops_test, primary_name)
 
         # Verify that the database service got restarted and is ready in the old primary.
+        logger.info(f"waiting for the database service to restart on {primary_name}")
         assert await is_postgresql_ready(ops_test, primary_name)
 
     # Verify that a new primary gets elected (ie old primary is secondary).
@@ -147,6 +148,7 @@ async def test_freeze_db_process(
                         ops_test, primary_name, process, "SIGCONT", use_ssh
                     )
         # Verify that the database service got restarted and is ready in the old primary.
+        logger.info(f"waiting for the database service to restart on {primary_name}")
         assert await is_postgresql_ready(ops_test, primary_name)
 
     await is_cluster_updated(ops_test, primary_name)
@@ -174,6 +176,7 @@ async def test_restart_db_process(
         await are_writes_increasing(ops_test, primary_name)
 
         # Verify that the database service got restarted and is ready in the old primary.
+        logger.info(f"waiting for the database service to restart on {primary_name}")
         assert await is_postgresql_ready(ops_test, primary_name)
 
     # Verify that a new primary gets elected (ie old primary is secondary).
@@ -288,6 +291,7 @@ async def test_forceful_restart_without_data_and_transaction_logs(
     sleep(MEDIAN_ELECTION_TIME * 2)
 
     async with ops_test.fast_forward():
+        logger.info("checking whether writes are increasing")
         await are_writes_increasing(ops_test, primary_name)
 
         # Verify that a new primary gets elected (ie old primary is secondary).
@@ -305,6 +309,7 @@ async def test_forceful_restart_without_data_and_transaction_logs(
             await change_wal_settings(ops_test, unit.name, 32, 32, 1)
 
         # Rotate the WAL segments.
+        logger.info(f"rotating WAL segments on {new_primary_name}")
         files = await list_wal_files(ops_test, app)
         host = await get_unit_address(ops_test, new_primary_name)
         password = await get_password(ops_test, down_unit=primary_name)
@@ -317,17 +322,20 @@ async def test_forceful_restart_without_data_and_transaction_logs(
                 cursor.execute("SELECT pg_switch_wal();")
         connection.close()
         new_files = await list_wal_files(ops_test, app)
+
         # Check that the WAL was correctly rotated.
+
         for unit_name in files:
             assert not files[unit_name].intersection(
                 new_files
-            ), "WAL segments weren't correctly rotated"
+            ), f"WAL segments weren't correctly rotated on {unit_name}"
 
         # Start the systemd service in the old primary.
         logger.info(f"starting database on {primary_name}")
         await run_command_on_unit(ops_test, primary_name, "/charm/bin/pebble start postgresql")
 
         # Verify that the database service got restarted and is ready in the old primary.
+        logger.info(f"waiting for the database service to restart on {primary_name}")
         assert await is_postgresql_ready(ops_test, primary_name)
 
     await is_cluster_updated(ops_test, primary_name)
@@ -385,14 +393,14 @@ async def test_network_cut(
     remove_instance_isolation(ops_test)
 
     # Verify that the database service got restarted and is ready in the old primary.
-    logger.info("waiting for the database service to restart")
+    logger.info(f"waiting for the database service to restart on {primary_name}")
     assert await is_postgresql_ready(ops_test, primary_name)
 
     # Verify that connection is possible.
     logger.info("checking whether the connectivity to the database is working")
     assert await is_connection_possible(
         ops_test, primary_name
-    ), "Connection is not possible after network restore"
+    ), f"Connection is not possible to {primary_name} after network restore"
 
     await is_cluster_updated(ops_test, primary_name)