[MISC] Parallel patroni calls

pyproject.toml

@@ -18,6 +18,7 @@ jinja2 = "^3.1.6"
 pysyncobj = "^0.3.14"
 psutil = "^7.0.0"
 charm-refresh = "^3.0.0.3"
+httpx = "^0.28.1"
 
 [tool.poetry.group.charm-libs.dependencies]
 # data_platform_libs/v0/data_interfaces.py

scripts/cluster_topology_observer.py

@@ -6,8 +6,10 @@
 import json
 import subprocess
 import sys
+from asyncio import as_completed, get_running_loop, run, wait
+from contextlib import suppress
 from os import environ
-from ssl import CERT_NONE, create_default_context
+from ssl import create_default_context
 from time import sleep
 from urllib.parse import urljoin
 from urllib.request import urlopen
@@ -16,6 +18,10 @@
 
 API_REQUEST_TIMEOUT = 5
 PATRONI_CLUSTER_STATUS_ENDPOINT = "cluster"
+TLS_CA_BUNDLE_FILE = "peer_ca_bundle.pem"
+SNAP_CURRENT_PATH = "/var/snap/charmed-postgresql/current"
+SNAP_CONF_PATH = f"{SNAP_CURRENT_PATH}/etc"
+PATRONI_CONF_PATH = f"{SNAP_CONF_PATH}/patroni"
 
 # File path for the spawned cluster topology observer process to write logs.
 LOG_FILE_PATH = "/var/log/cluster_topology_observer.log"
@@ -25,6 +31,20 @@ class UnreachableUnitsError(Exception):
     """Cannot reach any known cluster member."""
 
 
+def call_url(url, context):
+    """Task handler for calling an url."""
+    try:
+        # Scheme is generated by the charm
+        resp = urlopen(  # noqa: S310
+            url,
+            timeout=API_REQUEST_TIMEOUT,
+            context=context,
+        )
+        return json.loads(resp.read())
+    except Exception as e:
+        print(f"Failed to contact {url} with {e}")
+
+
 def check_for_authorisation_rules_changes(run_cmd, unit, charm_dir, previous_authorisation_rules):
     """Check for changes in the authorisation rules.
 
@@ -120,7 +140,7 @@ def dispatch(run_cmd, unit, charm_dir, custom_event):
     subprocess.run([run_cmd, "-u", unit, dispatch_sub_cmd.format(custom_event, charm_dir)])  # noqa: S603
 
 
-def main():
+async def main():
     """Main watch and dispatch loop.
 
     Watch the Patroni API cluster info. When changes are detected, dispatch the change event.
@@ -135,23 +155,19 @@ def main():
     while True:
         # Disable TLS chain verification
         context = create_default_context()
-        context.check_hostname = False
-        context.verify_mode = CERT_NONE
+        with suppress(FileNotFoundError):
+            context.load_verify_locations(cafile=f"{PATRONI_CONF_PATH}/{TLS_CA_BUNDLE_FILE}")
 
         cluster_status = None
-        for url in urls:
-            try:
-                # Scheme is generated by the charm
-                resp = urlopen(  # noqa: S310
-                    url,
-                    timeout=API_REQUEST_TIMEOUT,
-                    context=context,
-                )
-                cluster_status = json.loads(resp.read())
+        loop = get_running_loop()
+        tasks = [loop.run_in_executor(None, call_url, url, context) for url in urls]
+        for task in as_completed(tasks):
+            if result := await task:
+                for task in tasks:
+                    task.cancel()
+                await wait(tasks)
+                cluster_status = result
                 break
-            except Exception as e:
-                print(f"Failed to contact {url} with {e}")
-                continue
         if not cluster_status:
             raise UnreachableUnitsError("Unable to reach cluster members")
         current_cluster_topology = {}
@@ -186,4 +202,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    run(main())

src/charm.py

@@ -130,6 +130,7 @@
 logger = logging.getLogger(__name__)
 logging.getLogger("httpx").setLevel(logging.WARNING)
 logging.getLogger("httpcore").setLevel(logging.WARNING)
+logging.getLogger("asyncio").setLevel(logging.WARNING)
 
 PRIMARY_NOT_REACHABLE_MESSAGE = "waiting for primary to be reachable from this unit"
 EXTENSIONS_DEPENDENCY_MESSAGE = "Unsatisfied plugin dependencies. Please check the logs"

src/cluster.py

@@ -11,18 +11,21 @@
 import re
 import shutil
 import subprocess
+from asyncio import as_completed, create_task, run, wait
+from contextlib import suppress
 from pathlib import Path
+from ssl import CERT_NONE, create_default_context
 from typing import TYPE_CHECKING, Any, TypedDict
 
 import charm_refresh
 import psutil
 import requests
 from charms.operator_libs_linux.v2 import snap
+from httpx import AsyncClient, BasicAuth, HTTPError
 from jinja2 import Template
 from ops import BlockedStatus
 from pysyncobj.utility import TcpUtility, UtilityException
 from tenacity import (
-    AttemptManager,
     RetryError,
     Retrying,
     retry,
@@ -172,6 +175,10 @@ def __init__(
     def _patroni_auth(self) -> requests.auth.HTTPBasicAuth:
         return requests.auth.HTTPBasicAuth("patroni", self.patroni_password)
 
+    @property
+    def _patroni_async_auth(self) -> BasicAuth:
+        return BasicAuth("patroni", password=self.patroni_password)
+
     @property
     def _patroni_url(self) -> str:
         """Patroni REST API URL."""
@@ -249,28 +256,14 @@ def get_postgresql_version(self) -> str:
             if snp["name"] == charm_refresh.snap_name():
                 return snp["version"]
 
-    def cluster_status(
-        self, alternative_endpoints: list | None = None
-    ) -> list[ClusterMember] | None:
+    def cluster_status(self, alternative_endpoints: list | None = None) -> list[ClusterMember]:
         """Query the cluster status."""
         # Request info from cluster endpoint (which returns all members of the cluster).
-        # TODO we don't know the other cluster's ca
-        verify = self.verify if not alternative_endpoints else False
-        for attempt in Retrying(
-            stop=stop_after_attempt(
-                len(alternative_endpoints) if alternative_endpoints else len(self.peers_ips)
-            )
+        if response := self.parallel_patroni_get_request(
+            f"/{PATRONI_CLUSTER_STATUS_ENDPOINT}", alternative_endpoints
         ):
-            with attempt:
-                request_url = self._get_alternative_patroni_url(attempt, alternative_endpoints)
-
-                cluster_status = requests.get(
-                    f"{request_url}/{PATRONI_CLUSTER_STATUS_ENDPOINT}",
-                    verify=verify,
-                    timeout=API_REQUEST_TIMEOUT,
-                    auth=self._patroni_auth,
-                )
-                return cluster_status.json()["members"]
+            return response["members"]
+        raise RetryError(last_attempt=Exception("Unable to reach any units"))
 
     def get_member_ip(self, member_name: str) -> str | None:
         """Get cluster member IP address.
@@ -281,13 +274,14 @@ def get_member_ip(self, member_name: str) -> str | None:
         Returns:
             IP address of the cluster member.
         """
-        cluster_status = self.cluster_status()
-        if not cluster_status:
-            return
+        try:
+            cluster_status = self.cluster_status()
 
-        for member in cluster_status:
-            if member["name"] == member_name:
-                return member["host"]
+            for member in cluster_status:
+                if member["name"] == member_name:
+                    return member["host"]
+        except RetryError:
+            logger.debug("Unable to get IP. Cluster status unreachable")
 
     def get_member_status(self, member_name: str) -> str:
         """Get cluster member status.
@@ -307,6 +301,44 @@ def get_member_status(self, member_name: str) -> str:
                     return member["state"]
         return ""
 
+    async def _httpx_get_request(self, url: str, verify: bool = True):
+        ssl_ctx = create_default_context()
+        if verify:
+            with suppress(FileNotFoundError):
+                ssl_ctx.load_verify_locations(cafile=f"{PATRONI_CONF_PATH}/{TLS_CA_BUNDLE_FILE}")
+        else:
+            ssl_ctx.check_hostname = False
+            ssl_ctx.verify_mode = CERT_NONE
+        async with AsyncClient(
+            auth=self._patroni_async_auth, timeout=API_REQUEST_TIMEOUT, verify=ssl_ctx
+        ) as client:
+            try:
+                return (await client.get(url)).json()
+            except (HTTPError, ValueError):
+                return None
+
+    async def _async_get_request(self, uri: str, endpoints: list[str], verify: bool = True):
+        tasks = [
+            create_task(self._httpx_get_request(f"https://{ip}:8008{uri}", verify))
+            for ip in endpoints
+        ]
+        for task in as_completed(tasks):
+            if result := await task:
+                for task in tasks:
+                    task.cancel()
+                await wait(tasks)
+                return result
+
+    def parallel_patroni_get_request(self, uri: str, endpoints: list[str] | None = None) -> dict:
+        """Call all possible patroni endpoints in parallel."""
+        if not endpoints:
+            endpoints = (self.unit_ip, *self.peers_ips)
+            verify = True
+        else:
+            # TODO we don't know the other cluster's ca
+            verify = False
+        return run(self._async_get_request(uri, endpoints, verify))
+
     def get_primary(
         self, unit_name_pattern=False, alternative_endpoints: list[str] | None = None
     ) -> str | None:
@@ -320,14 +352,17 @@ def get_primary(
             primary pod or unit name.
         """
         # Request info from cluster endpoint (which returns all members of the cluster).
-        if cluster_status := self.cluster_status(alternative_endpoints):
+        try:
+            cluster_status = self.cluster_status(alternative_endpoints)
             for member in cluster_status:
                 if member["role"] == "leader":
                     primary = member["name"]
                     if unit_name_pattern:
                         # Change the last dash to / in order to match unit name pattern.
                         primary = label2name(primary)
                     return primary
+        except RetryError:
+            logger.debug("Unable to get primary. Cluster status unreachable")
 
     def get_standby_leader(
         self, unit_name_pattern=False, check_whether_is_running: bool = False
@@ -366,31 +401,6 @@ def get_sync_standby_names(self) -> list[str]:
                     sync_standbys.append(label2name(member["name"]))
         return sync_standbys
 
-    def _get_alternative_patroni_url(
-        self, attempt: AttemptManager, alternative_endpoints: list[str] | None = None
-    ) -> str:
-        """Get an alternative REST API URL from another member each time.
-
-        When the Patroni process is not running in the current unit it's needed
-        to use a URL from another cluster member REST API to do some operations.
-        """
-        if alternative_endpoints is not None:
-            return self._patroni_url.replace(
-                self.unit_ip, alternative_endpoints[attempt.retry_state.attempt_number - 1]
-            )
-        attempt_number = attempt.retry_state.attempt_number
-        if attempt_number > 1:
-            url = self._patroni_url
-            if (attempt_number - 1) <= len(self.peers_ips):
-                unit_number = attempt_number - 2
-            else:
-                unit_number = attempt_number - 2 - len(self.peers_ips)
-            other_unit_ip = list(self.peers_ips)[unit_number]
-            url = url.replace(self.unit_ip, other_unit_ip)
-        else:
-            url = self._patroni_url
-        return url
-
     def are_all_members_ready(self) -> bool:
         """Check if all members are correctly running Patroni and PostgreSQL.
 

tests/unit/test_charm.py

@@ -602,7 +602,7 @@ def test_on_start(harness):
         patch(
             "charm.PostgresqlOperatorCharm._is_storage_attached",
             side_effect=[False, True, True, True, True, True],
-        ) as _is_storage_attached,
+        ),
         patch(
             "charm.PostgresqlOperatorCharm._can_connect_to_postgresql",
             new_callable=PropertyMock,