refactor: core services

Author: sophia-bq

aws_advanced_python_wrapper/aws_secrets_manager_plugin.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+from datetime import timedelta
 from json import JSONDecodeError, loads
 from re import search
 from types import SimpleNamespace
@@ -23,7 +24,7 @@
 
 from aws_advanced_python_wrapper.aws_credentials_manager import \
     AwsCredentialsManager
-from aws_advanced_python_wrapper.utils.cache_map import CacheMap
+from aws_advanced_python_wrapper.utils import core_services
 
 if TYPE_CHECKING:
     from boto3 import Session
@@ -46,14 +47,20 @@
 logger = Logger(__name__)
 
 
+class Secret:
+    """Wrapper type for secrets, used as StorageService type key."""
+
+    def __init__(self, value: SimpleNamespace):
+        self.value = value
+
+
 class AwsSecretsManagerPlugin(Plugin):
     _SUBSCRIBED_METHODS: Set[str] = {DbApiMethod.CONNECT.method_name, DbApiMethod.FORCE_CONNECT.method_name}
 
     _SECRETS_ARN_PATTERN = r"^arn:aws:secretsmanager:(?P<region>[^:\n]*):[^:\n]*:([^:/\n]*[:/])?(.*)$"
     _ONE_YEAR_IN_SECONDS = 60 * 60 * 24 * 365
 
     _secret: Optional[SimpleNamespace] = None
-    _secrets_cache: CacheMap[Tuple, SimpleNamespace] = CacheMap()
     _secret_key: Tuple = ()
 
     @property
@@ -63,6 +70,8 @@ def subscribed_methods(self) -> Set[str]:
     def __init__(self, plugin_service: PluginService, props: Properties, session: Optional[Session] = None):
         self._plugin_service = plugin_service
         self._session = session
+        self._storage_service = core_services.get_storage_service()
+        self._storage_service.register(Secret, item_expiration_time=timedelta(minutes=30))
 
         secret_id = WrapperProperties.SECRETS_MANAGER_SECRET_ID.get(props)
         if not secret_id:
@@ -100,13 +109,13 @@ def force_connect(
         return self._connect(host_info, props, force_connect_func)
 
     def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callable) -> Connection:
-        token_expiration_sec: int = WrapperProperties.SECRETS_MANAGER_EXPIRATION.get_int(props)
+        token_expiration_sec = WrapperProperties.SECRETS_MANAGER_EXPIRATION.get_int(props)
         # if value is less than 0, default to one year
         if token_expiration_sec < 0:
             token_expiration_sec = AwsSecretsManagerPlugin._ONE_YEAR_IN_SECONDS
         token_expiration_ns = token_expiration_sec * 1_000_000_000
 
-        secret_fetched: bool = self._update_secret(host_info, props, token_expiration_ns=token_expiration_ns)
+        secret_fetched: bool = self._update_secret(host_info, props, token_expiration_ns)
 
         try:
             self._apply_secret_to_properties(props)
@@ -120,7 +129,7 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
                 raise AwsWrapperError(
                     Messages.get_formatted("AwsSecretsManagerPlugin.ConnectException", e), e) from e
 
-            secret_fetched = self._update_secret(host_info, props, token_expiration_ns=token_expiration_ns, force_refetch=True)
+            secret_fetched = self._update_secret(host_info, props, token_expiration_ns, force_refetch=True)
 
             if secret_fetched:
                 try:
@@ -146,13 +155,14 @@ def _update_secret(self, host_info: HostInfo, props: Properties, token_expiratio
 
         try:
             fetched: bool = False
-            self._secret: Optional[SimpleNamespace] = AwsSecretsManagerPlugin._secrets_cache.get(self._secret_key)
+            cached_secret = self._storage_service.get(Secret, self._secret_key)
+            self._secret = cached_secret.value if cached_secret is not None else None
             endpoint = self._secret_key[2]
             if not self._secret or force_refetch:
                 try:
                     self._secret = self._fetch_latest_credentials(host_info, props)
                     if self._secret:
-                        AwsSecretsManagerPlugin._secrets_cache.put(self._secret_key, self._secret, token_expiration_ns)
+                        self._storage_service.put(Secret, self._secret_key, Secret(self._secret), item_expiration_ns=token_expiration_ns)
                         fetched = True
                 except (ClientError, AttributeError) as e:
                     logger.debug("AwsSecretsManagerPlugin.FailedToFetchDbCredentials", e)

aws_advanced_python_wrapper/blue_green_plugin.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import socket
-from datetime import datetime
+from datetime import datetime, timedelta
 from time import perf_counter_ns
 from typing import TYPE_CHECKING, FrozenSet, List, cast
 
@@ -44,9 +44,11 @@
 from aws_advanced_python_wrapper.iam_plugin import IamAuthPlugin
 from aws_advanced_python_wrapper.pep249_methods import DbApiMethod
 from aws_advanced_python_wrapper.plugin import Plugin, PluginFactory
+from aws_advanced_python_wrapper.utils import core_services
 from aws_advanced_python_wrapper.utils.atomic import AtomicInt
 from aws_advanced_python_wrapper.utils.concurrent import (ConcurrentDict,
                                                           ConcurrentSet)
+from aws_advanced_python_wrapper.utils.events import MonitorResetEvent
 from aws_advanced_python_wrapper.utils.log import Logger
 from aws_advanced_python_wrapper.utils.messages import Messages
 from aws_advanced_python_wrapper.utils.properties import (Properties,
@@ -901,6 +903,9 @@ def _run(self):
 
         finally:
             self._close_connection()
+            if self._host_list_provider is not None:
+                self._host_list_provider.stop_monitor()
+                self._host_list_provider = None
             logger.debug("BlueGreenStatusMonitor.ThreadCompleted", self._bg_role)
 
     def _open_connection(self):
@@ -1237,6 +1242,8 @@ def __init__(self, plugin_service: PluginService, props: Properties, bg_id: str)
         self._green_dns_removed = False
         self._green_topology_changed = False
         self._all_green_hosts_changed_name = False
+        self._monitor_reset_on_in_progress_completed = False
+        self._monitor_reset_on_topology_completed = False
         self._post_status_end_time_ns = 0
         self._process_status_lock = RLock()
         self._status_check_intervals_ms: Dict[BlueGreenIntervalRate, int] = {}
@@ -1258,6 +1265,8 @@ def __init__(self, plugin_service: PluginService, props: Properties, bg_id: str)
                 Messages.get_formatted(
                     "BlueGreenStatusProvider.UnsupportedDialect", self._bg_id, dialect.__class__.__name__))
 
+        core_services.get_storage_service().register(BlueGreenStatus, item_expiration_time=timedelta(hours=1))
+
         current_host_info = self._plugin_service.current_host_info
         blue_monitor = BlueGreenStatusMonitor(
             BlueGreenRole.SOURCE,
@@ -1476,6 +1485,7 @@ def _update_summary_status(self, bg_role: BlueGreenRole, interim_status: BlueGre
         elif self._latest_phase == BlueGreenPhase.IN_PROGRESS:
             self._update_dns_flags(bg_role, interim_status)
             self._summary_status = self._get_status_of_in_progress()
+            self._reset_monitors("_monitor_reset_on_in_progress_completed", "- start")
 
         elif self._latest_phase == BlueGreenPhase.POST:
             self._update_dns_flags(bg_role, interim_status)
@@ -1503,6 +1513,7 @@ def _update_dns_flags(self, bg_role: BlueGreenRole, interim_status: BlueGreenInt
             logger.debug("BlueGreenStatusProvider.GreenTopologyChanged", self._bg_id)
             self._green_topology_changed = True
             self._store_event_phase_time("Green topology changed")
+            self._reset_monitors("_monitor_reset_on_topology_completed", "- green topology")
 
     def _store_event_phase_time(self, key_prefix: str, phase: Optional[BlueGreenPhase] = None):
         rollback_str = " (rollback)" if self._rollback else ""
@@ -1846,6 +1857,20 @@ def _update_status_cache(self):
             with latest_status.cv:
                 latest_status.cv.notify_all()
 
+    def _reset_monitors(self, completed_flag_attr: str, event_name: str):
+        if getattr(self, completed_flag_attr):
+            return
+        setattr(self, completed_flag_attr, True)
+
+        blue_endpoints = frozenset(
+            host for host, role in self._roles_by_host.items()
+            if role == BlueGreenRole.SOURCE)
+
+        cluster_id = self._plugin_service.host_list_provider.get_cluster_id()
+        core_services.get_event_publisher().publish(
+            MonitorResetEvent(cluster_id=cluster_id, endpoints=blue_endpoints))
+        self._store_event_phase_time(f"Monitor reset {event_name}")
+
     def _log_current_context(self):
         logger.debug(f"[bg_id: '{self._bg_id}'] Summary status: \n{self._summary_status}")
         hosts_str = "\n".join(
@@ -1913,6 +1938,8 @@ def _reset_context_when_completed(self):
         self._green_dns_removed = False
         self._green_topology_changed = False
         self._all_green_hosts_changed_name = False
+        self._monitor_reset_on_in_progress_completed = False
+        self._monitor_reset_on_topology_completed = False
         self._post_status_end_time_ns = 0
         self._interim_status_hashes = [0, 0]
         self._latest_context_hash = 0

aws_advanced_python_wrapper/cleanup.py

@@ -16,18 +16,11 @@
     OpenedConnectionTracker
 from aws_advanced_python_wrapper.aws_credentials_manager import \
     AwsCredentialsManager
-from aws_advanced_python_wrapper.host_monitoring_plugin import \
-    MonitoringThreadContainer
-from aws_advanced_python_wrapper.thread_pool_container import \
-    ThreadPoolContainer
-from aws_advanced_python_wrapper.utils.sliding_expiration_cache_container import \
-    SlidingExpirationCacheContainer
+from aws_advanced_python_wrapper.utils import core_services
 
 
 def release_resources() -> None:
     """Release all global resources used by the wrapper."""
-    MonitoringThreadContainer.clean_up()
-    ThreadPoolContainer.release_resources()
+    core_services.release_resources()
     AwsCredentialsManager.release_resources()
     OpenedConnectionTracker.release_resources()
-    SlidingExpirationCacheContainer.release_resources()

aws_advanced_python_wrapper/cluster_topology_monitor.py

@@ -18,16 +18,18 @@
 import time
 from abc import ABC, abstractmethod
 from concurrent.futures import ThreadPoolExecutor
+from time import perf_counter_ns
 from typing import TYPE_CHECKING, Dict, Optional
 
 from aws_advanced_python_wrapper.errors import AwsWrapperError
 from aws_advanced_python_wrapper.host_availability import HostAvailability
-from aws_advanced_python_wrapper.hostinfo import HostInfo
+from aws_advanced_python_wrapper.hostinfo import HostInfo, Topology
+from aws_advanced_python_wrapper.utils import core_services
 from aws_advanced_python_wrapper.utils.atomic import AtomicReference
+from aws_advanced_python_wrapper.utils.events import (EventBase,
+                                                      MonitorResetEvent)
 from aws_advanced_python_wrapper.utils.messages import Messages
 from aws_advanced_python_wrapper.utils.rds_utils import RdsUtils
-from aws_advanced_python_wrapper.utils.storage.storage_service import (
-    StorageService, Topology)
 from aws_advanced_python_wrapper.utils.thread_safe_connection_holder import \
     ThreadSafeConnectionHolder
 from aws_advanced_python_wrapper.utils.utils import LogUtils
@@ -55,10 +57,20 @@ def force_refresh(self, should_verify_writer: bool, timeout_sec: int) -> Topolog
     def force_refresh_with_connection(self, connection: Connection, timeout_sec: int) -> Topology:
         pass
 
+    @property
     @abstractmethod
     def can_dispose(self) -> bool:
         pass
 
+    @abstractmethod
+    def stop(self) -> None:
+        pass
+
+    @property
+    @abstractmethod
+    def last_activity_ns(self) -> int:
+        pass
+
     @abstractmethod
     def close(self) -> None:
         pass
@@ -109,6 +121,7 @@ def __init__(self, plugin_service: PluginService, topology_utils: TopologyUtils,
         self._high_refresh_rate_end_time_nano = 0
         self._stop = threading.Event()
         self._monitor_thread: Optional[threading.Thread] = None
+        self._last_activity_ns: int = perf_counter_ns()
 
         self._monitoring_properties = PropertiesUtils.create_topology_monitoring_properties(properties)
         if WrapperProperties.SOCKET_TIMEOUT_SEC.get(self._monitoring_properties) is None:
@@ -163,17 +176,41 @@ def _wait_till_topology_gets_updated(self, timeout_sec: int) -> Topology:
                         self._cluster_id, timeout_sec * 1000))
 
     def _get_stored_hosts(self) -> Topology:
-        hosts = StorageService.get(Topology, self._cluster_id)
+        hosts = core_services.get_storage_service().get(Topology, self._cluster_id)
         if hosts is None:
             return ()
         return hosts
 
+    def stop(self) -> None:
+        self._stop.set()
+        self.close()
+
+    @property
     def can_dispose(self) -> bool:
         return self._stop.is_set()
 
+    @property
+    def last_activity_ns(self) -> int:
+        return self._last_activity_ns
+
+    def process_event(self, event: EventBase) -> None:
+        if isinstance(event, MonitorResetEvent) and event.cluster_id == self._cluster_id:
+            logger.debug("ClusterTopologyMonitorImpl.ResetEventReceived", self._cluster_id)
+            self._host_threads_stop.set()
+            self._close_host_monitors()
+            self._close_connection_from_ref(self._host_threads_writer_connection)
+            self._close_connection_from_ref(self._host_threads_reader_connection)
+            self._host_threads_stop.clear()
+            self._submitted_hosts.clear()
+            self._host_threads_writer_host_info.set(None)
+            self._host_threads_latest_topology.set(None)
+            self._monitoring_connection.clear()
+            self._is_verified_writer_connection = False
+            self._writer_host_info.set(None)
+            self._high_refresh_rate_end_time_nano = 0
+
     def close(self) -> None:
         logger.debug("ClusterTopologyMonitorImpl.ClosingMonitor", self._cluster_id)
-        self._stop.set()
         self._request_to_update_topology.set()
 
         self._close_host_monitors()
@@ -195,6 +232,7 @@ def _monitor(self) -> None:
             logger.debug("ClusterTopologyMonitor.StartMonitoringThread", self._cluster_id, self._initial_host_info.host)
 
             while not self._stop.is_set():
+                self._last_activity_ns = perf_counter_ns()
                 if self._is_in_panic_mode():
                     if not self._submitted_hosts:
                         self._close_host_monitors()
@@ -444,7 +482,7 @@ def _get_instance_template(self, instance_id: str, connection: Connection) -> Ho
         return self._instance_template
 
     def _update_topology_cache(self, hosts: Topology) -> None:
-        StorageService.set(self._cluster_id, hosts, Topology)
+        core_services.get_storage_service().put(Topology, self._cluster_id, hosts)
         # Notify waiting threads
         self._request_to_update_topology.clear()
         self._topology_updated.set()