feat: fastest response strategy plugin

joyc-bq · joyc-bq · commit 36a895f2885e · 2024-02-21T14:07:15.000-08:00
diff --git a/aws_advanced_python_wrapper/fastest_response_strategy_plugin.py b/aws_advanced_python_wrapper/fastest_response_strategy_plugin.py
@@ -0,0 +1,322 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from concurrent.futures import Executor, ThreadPoolExecutor
+from copy import copy
+from dataclasses import dataclass
+from datetime import datetime
+from logging import Logger
+from threading import Event, Lock
+from time import sleep
+from typing import (TYPE_CHECKING, Callable, ClassVar, Dict, Optional, Set,
+                    Tuple)
+
+from aws_advanced_python_wrapper.errors import AwsWrapperError
+from aws_advanced_python_wrapper.hostselector import RandomHostSelector
+from aws_advanced_python_wrapper.plugin import Plugin
+from aws_advanced_python_wrapper.utils.cache_map import CacheMap
+from aws_advanced_python_wrapper.utils.messages import Messages
+from aws_advanced_python_wrapper.utils.properties import (Properties,
+                                                          WrapperProperties)
+from aws_advanced_python_wrapper.utils.sliding_expiration_cache import \
+    SlidingExpirationCacheWithCleanupThread
+
+if TYPE_CHECKING:
+    from aws_advanced_python_wrapper.driver_dialect import DriverDialect
+    from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole
+    from aws_advanced_python_wrapper.pep249 import Connection
+    from aws_advanced_python_wrapper.plugin_service import PluginService
+    from aws_advanced_python_wrapper.utils.notifications import HostEvent
+
+logger = Logger(__name__)
+
+
+class FastestResponseStrategyPlugin(Plugin):
+    _FASTEST_RESPONSE_STRATEGY_NAME = "fastest_response"
+    _SUBSCRIBED_METHODS: Set[str] = {"accepts_strategy",
+                                     "get_host_info_by_strategy",
+                                     "notify_host_list_changed"}
+
+    def __init__(self, plugin_service: PluginService, props: Properties, host_response_time_service: Optional[HostResponseTimeService] = None):
+        self._plugin_service = plugin_service
+        self._properties = props
+        self._host_response_time_service: Optional[HostResponseTimeService] = host_response_time_service
+        self._cache_expiration_nano = WrapperProperties.RESPONSE_MEASUREMENT_INTERVAL_MILLIS.get_int(props)
+        self._cached_fastest_response_host_by_role: CacheMap[str, HostInfo] = CacheMap()
+        self._random_host_selector = RandomHostSelector()
+        self._hosts: Tuple[HostInfo, ...] = ()
+
+    @property
+    def subscribed_methods(self) -> Set[str]:
+        return self._SUBSCRIBED_METHODS
+
+    def connect(
+            self,
+            target_driver_func: Callable,
+            driver_dialect: DriverDialect,
+            host_info: HostInfo,
+            props: Properties,
+            is_initial_connection: bool,
+            connect_func: Callable) -> Connection:
+        return self._connect(host_info, props, is_initial_connection, connect_func)
+
+    def force_connect(
+            self,
+            target_driver_func: Callable,
+            driver_dialect: DriverDialect,
+            host_info: HostInfo,
+            props: Properties,
+            is_initial_connection: bool,
+            force_connect_func: Callable) -> Connection:
+        return self._connect(host_info, props, is_initial_connection, force_connect_func)
+
+    def _connect(
+            self,
+            host: HostInfo,
+            properties: Properties,
+            is_initial_connection: bool,
+            connect_func: Callable) -> Connection:
+        conn = connect_func()
+
+        if is_initial_connection:
+            self._plugin_service.refresh_host_list(conn)
+
+        return conn
+
+    def accepts_strategy(self, role: HostRole, strategy: str) -> bool:
+        return strategy == FastestResponseStrategyPlugin._FASTEST_RESPONSE_STRATEGY_NAME
+
+    def get_host_info_by_strategy(self, role: HostRole, strategy: str) -> HostInfo:
+        if not self.accepts_strategy(role, strategy):
+            raise AwsWrapperError(Messages.get_formatted("DriverConnectionProvider.UnsupportedStrategy", strategy))
+
+        fastest_response_host: Optional[HostInfo] = self._cached_fastest_response_host_by_role.get(role.name)
+        if fastest_response_host is not None:
+
+            # Found a fastest host. Let's find it in the the latest topology.
+            for host in self._plugin_service.hosts:
+                if host == fastest_response_host:
+                    # found the fastest host in the topology
+                    return host
+                # It seems that the fastest cached host isn't in the latest topology.
+                # Let's ignore cached results and find the fastest host.
+
+            # Cached result isn't available. Need to find the fastest response time host.
+            host_dict = {}
+            for host in self._plugin_service.hosts:
+                if role == host.role and self._host_response_time_service is not None:
+                    response_time_tuple = FastestResponseStrategyPlugin.ResponseTimeTuple(host,
+                                                                                          self._host_response_time_service.get_response_time(host))
+                    host_dict[response_time_tuple.host_info] = response_time_tuple.response_time
+            # sort by response time then retrieve the first host
+            sorted_host_dict = dict(sorted(host_dict.items(), key=lambda item: item[1]))
+            calculated_fastest_response_host: HostInfo = next(iter(sorted_host_dict.keys()))
+            if calculated_fastest_response_host is None:
+                return self._random_host_selector.get_host(self._plugin_service.hosts, role, self._properties)
+
+            self._cached_fastest_response_host_by_role.put(role.name, calculated_fastest_response_host, self._cache_expiration_nano)
+
+        return calculated_fastest_response_host
+
+    def notify_host_list_changed(self, changes: Dict[str, Set[HostEvent]]):
+        self._hosts = self._plugin_service.hosts
+        if self._host_response_time_service:
+            self._host_response_time_service.set_hosts(self._hosts)
+
+    @dataclass
+    class ResponseTimeTuple:
+        host_info: HostInfo
+        response_time: int
+
+
+class FastestResponseStrategyPluginFactory:
+
+    def get_instance(self, plugin_service: PluginService, props: Properties) -> Plugin:
+        return FastestResponseStrategyPlugin(plugin_service, props)
+
+
+class NodeResponseTimeMonitor:
+
+    _MONITORING_PROPERTY_PREFIX: str = "frt-"
+    _NUM_OF_MEASURES: int = 5
+
+    def __init__(self, plugin_service: PluginService, host_info: HostInfo, props: Properties, interval_ms: int):
+        self._plugin_service = plugin_service
+        self._host_info = host_info
+        self._properties = props
+        self._interval_ms = interval_ms
+
+        self._response_time: int = 0
+        self._check_timestamp = datetime.now()
+        self._lock: Lock = Lock()
+        self._monitoring_conn: Optional[Connection] = None
+        self._is_stopped: Event = Event()
+
+        self._node_id: Optional[str] = self._host_info.host_id
+        if self._node_id is None or self._node_id == "":
+            self._node_id = self._host_info.host
+
+        self._executor: Executor = ThreadPoolExecutor(thread_name_prefix="NodeResponseTimeMonitorExecutor")
+
+    @property
+    def response_time(self):
+        return self._response_time
+
+    @property
+    def check_timestamp(self):
+        return self._check_timestamp
+
+    @property
+    def host_info(self):
+        return self._host_info
+
+    @property
+    def is_stopped(self):
+        return self._is_stopped.is_set()
+
+    def close(self):
+        self._is_stopped.set()
+
+        self._executor.shutdown(wait=True)
+        logger.debug("NodeResponseTimeMonitor.Stopped", self._host_info.host)
+
+    def _get_current_time(self):
+        return datetime.now()
+
+    def run(self):
+        try:
+            while not self.is_stopped:
+                try:
+                    self._open_connection()
+
+                    if self._monitoring_conn is not None:
+
+                        response_time_sum = 0
+                        count = 0
+                        for i in self._NUM_OF_MEASURES:
+                            if self._is_stopped:
+                                break
+                            start_time = self._get_current_time()
+                            if self._plugin_service.driver_dialect.ping(self._monitoring_conn):
+                                response_time = self._get_current_time() - start_time
+                                response_time_sum = response_time_sum + response_time
+                                count = count + 1
+                            if count > 0:
+                                self.response_time = response_time_sum / count / 1000
+                            else:
+                                self.response_time = 0
+                            self.check_timestamp = self._get_current_time()
+                            logger.debug("NodeResponseTimeMonitor.ResponseTime", self._host_info.host, self._response_time)
+
+                        sleep(self._interval_ms/1000)
+
+                except InterruptedError:
+                    # exit thread
+                    logger.debug("NodeResponseTimeMonitor.InterruptedExceptionDuringMonitoring", self._host_info.host)
+        except Exception as e:
+            # this should not be reached; log and exit thread
+            logger.debug("NodeResponseTimeMonitor.ExceptionDuringMonitoringStop",
+                         self._host_info.host,
+                         e)  # print full trace stack of the exception.
+        finally:
+            self._is_stopped.set()
+            if self._monitoring_conn is not None:
+                try:
+                    self._monitoring_conn.close()
+                except Exception:
+                    # Do nothing
+                    pass
+
+    def _open_connection(self):
+        try:
+            driver_dialect = self._plugin_service.driver_dialect
+            if self._monitoring_conn is None or driver_dialect.is_closed(self._monitoring_conn):
+                monitoring_conn_properties: Properties = copy(self._properties)
+                for key, value in self._properties.items():
+                    if key.startswith(self._MONITORING_PROPERTY_PREFIX):
+                        monitoring_conn_properties[key[len(self._MONITORING_PROPERTY_PREFIX):len(key)]] = value
+                        monitoring_conn_properties.pop(key, None)
+
+                logger.debug("NodeResponseTimeMonitor.OpeningConnection", self._host_info.url)
+                self._monitoring_conn = self._plugin_service.force_connect(self._host_info, monitoring_conn_properties, None)
+                logger.debug("NodeResponseTimeMonitor.OpenedConnection", self._host_info.url)
+
+        except Exception:
+            if self._monitoring_conn is not None:
+                try:
+                    self._monitoring_conn.close()
+                except Exception:
+                    pass  # ignore
+
+                self._monitoring_conn = None
+
+
+class HostResponseTimeService:
+    """
+    Return a response time in milliseconds to the host.
+    Return HostResponseTimeService._MAX_VALUE if response time is not available.
+
+    @param hostSpec the host details
+    @return response time in milliseconds for a desired host. It should return HostResponseTimeService._MAX_VALUE
+    if response time couldn't be measured.
+   """
+
+    @abstractmethod
+    def get_response_time(self, host_info: HostInfo) -> int:
+        ...
+
+    @abstractmethod
+    def set_hosts(self, hosts: Tuple[HostInfo, ...]) -> None:
+        ...
+
+
+class HostResponseTimeServiceImpl(HostResponseTimeService):
+    _MAX_VALUE = 2 ^ 31 - 1
+    _CACHE_EXPIRATION_NANO: int = 10 * 10 ^ 9
+    _CACHE_CLEANUP_NANO: int = 1 * 10 ^ 9
+    _lock: Lock = Lock()
+    _monitoring_nodes: ClassVar[SlidingExpirationCacheWithCleanupThread[str, NodeResponseTimeMonitor]] = \
+        SlidingExpirationCacheWithCleanupThread(_CACHE_CLEANUP_NANO,
+                                                should_dispose_func=lambda monitor: True,
+                                                item_disposal_func=lambda monitor: monitor.dispose())
+
+    def __init__(self, plugin_service: PluginService, props: Properties, interval_ms: int):
+        self._plugin_service = plugin_service
+        self._properties = props
+        self._interval_ms = interval_ms
+        self._hosts: Tuple[HostInfo, ...] = ()
+
+    def get_response_time(self, host_info: HostInfo) -> int:
+        monitor: Optional[NodeResponseTimeMonitor] = HostResponseTimeServiceImpl._monitoring_nodes.get(host_info.url)
+        if monitor is None:
+            return HostResponseTimeServiceImpl._MAX_VALUE
+        return monitor.response_time
+
+    def set_hosts(self, hosts: Tuple[HostInfo, ...]) -> None:
+        old_hosts_dict = {x.url: x for x in hosts}
+        self._hosts = hosts
+
+        for host in self._hosts:
+            new_host = host if host.url not in old_hosts_dict else None
+            if new_host:
+                with self._lock:
+                    self._monitoring_nodes.compute_if_absent(host.url,
+                                                             lambda _: NodeResponseTimeMonitor(
+                                                                 self._plugin_service,
+                                                                 new_host,
+                                                                 self._properties,
+                                                                 self._interval_ms), self._CACHE_EXPIRATION_NANO)
diff --git a/aws_advanced_python_wrapper/plugin_service.py b/aws_advanced_python_wrapper/plugin_service.py
@@ -16,6 +16,8 @@
 
 from typing import TYPE_CHECKING, ClassVar, List, Type
 
+from aws_advanced_python_wrapper.fastest_response_strategy_plugin import \
+    FastestResponseStrategyPluginFactory
 from aws_advanced_python_wrapper.federated_plugin import \
     FederatedAuthPluginFactory
 
@@ -571,6 +573,7 @@ class PluginManager(CanReleaseResources):
         "host_monitoring": HostMonitoringPluginFactory,
         "failover": FailoverPluginFactory,
         "read_write_splitting": ReadWriteSplittingPluginFactory,
+        "fastest_response_strategy": FastestResponseStrategyPluginFactory,
         "stale_dns": StaleDnsPluginFactory,
         "connect_time": ConnectTimePluginFactory,
         "execute_time": ExecuteTimePluginFactory,
@@ -589,8 +592,9 @@ class PluginManager(CanReleaseResources):
         ReadWriteSplittingPluginFactory: 300,
         FailoverPluginFactory: 400,
         HostMonitoringPluginFactory: 500,
-        IamAuthPluginFactory: 600,
-        AwsSecretsManagerPluginFactory: 700,
+        FastestResponseStrategyPluginFactory: 600,
+        IamAuthPluginFactory: 700,
+        AwsSecretsManagerPluginFactory: 800,
         ConnectTimePluginFactory: WEIGHT_RELATIVE_TO_PRIOR_PLUGIN,
         ExecuteTimePluginFactory: WEIGHT_RELATIVE_TO_PRIOR_PLUGIN,
         DeveloperPluginFactory: WEIGHT_RELATIVE_TO_PRIOR_PLUGIN,
diff --git a/aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties b/aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties
@@ -148,6 +148,13 @@ MonitorService.ErrorPopulatingAliases=[MonitorService] An error occurred while p
 
 MultiAzHostListProvider.UnableToParseInstanceName=[MultiAzHostListProvider] The MultiAzHostListProvider was unable to parse the instance name from the endpoint returned by the topology query.
 
+NodeResponseTimeMonitor.ExceptionDuringMonitoringStop=[NodeResponseTimeMonitor] Stopping thread after unhandled exception was thrown in Response time thread for node {}.
+NodeResponseTimeMonitor.InterruptedExceptionDuringMonitoring=[NodeResponseTimeMonitor] Response time thread for node {} was interrupted.
+NodeResponseTimeMonitor.OpenedConnection=[NodeResponseTimeMonitor] Opened Response time connection: {}.
+NodeResponseTimeMonitor.OpeningConnection=[NodeResponseTimeMonitor] Opening a Response time connection to ''{}''.
+NodeResponseTimeMonitor.ResponseTime=[NodeResponseTimeMonitor] Response time for ''{}'': {} ms
+NodeResponseTimeMonitor.Stopped=[NodeResponseTimeMonitor] Stopped Response time thread for node ''{}''.
+
 OpenedConnectionTracker.OpenedConnectionsTracked=[OpenedConnectionTracker] Opened Connections Tracked: {}
 OpenedConnectionTracker.InvalidatingConnections=[OpenedConnectionTracker] Invalidating opened connections to host: {}
 OpenedConnectionTracker.UnableToPopulateOpenedConnectionSet=[OpenedConnectionTracker] The driver is unable to track this opened connection because the instance endpoint is unknown.
@@ -243,6 +250,8 @@ RoundRobinHostSelector.ClusterInfoNone=[RoundRobinHostSelector] The round robin
 RoundRobinHostSelector.RoundRobinInvalidDefaultWeight=[RoundRobinHostSelector] The provided default weight value is not valid. Weight values must be an integer greater than or equal to 1.
 RoundRobinHostSelector.RoundRobinInvalidHostWeightPairs= [RoundRobinHostSelector] The provided host weight pairs have not been configured correctly. Please ensure the provided host weight pairs is a comma separated list of pairs, each pair in the format of <host>:<weight>. Weight values must be an integer greater than or equal to the default weight value of 1.
 
+SlidingExpirationCache.CleaningUp=[SlidingExpirationCache] Cleaning up...
+
 SqlAlchemyPooledConnectionProvider.PoolNone=[SqlAlchemyPooledConnectionProvider] Attempted to find or create a pool for '{}' but the result of the attempt evaluated to None.
 SqlAlchemyPooledConnectionProvider.UnableToCreateDefaultKey=[SqlAlchemyPooledConnectionProvider] Unable to create a default key for internal connection pools. By default, the user parameter is used, but the given user evaluated to None or the empty string (""). Please ensure you have passed a valid user in the connection properties.
 
diff --git a/aws_advanced_python_wrapper/utils/properties.py b/aws_advanced_python_wrapper/utils/properties.py
@@ -282,6 +282,11 @@ class WrapperProperties:
     DB_USER = WrapperProperty("db_user",
                               "The database user used to access the database",
                               None)
+    # Fastest Response Strategy
+
+    RESPONSE_MEASUREMENT_INTERVAL_MILLIS = WrapperProperty("response_measurement_interval_ms",
+                                                           "Interval in millis between measuring response time to a database node",
+                                                           30000)
 
     # Telemetry
     ENABLE_TELEMETRY = WrapperProperty(
diff --git a/aws_advanced_python_wrapper/utils/sliding_expiration_cache.py b/aws_advanced_python_wrapper/utils/sliding_expiration_cache.py
