feat: fastest response strategy plugin

Author: joyc-bq

aws_advanced_python_wrapper/driver_dialect.py

@@ -152,3 +152,13 @@ def unwrap_connection(self, conn_obj: object) -> Any:
 
     def transfer_session_state(self, from_conn: Connection, to_conn: Connection):
         return
+
+    def ping(self, conn: Connection) -> bool:
+        try:
+            with conn.cursor() as cursor:
+                query = "SELECT 1"
+                self.execute("Cursor.execute", lambda: cursor.execute(query), query, exec_timeout=10)
+                cursor.fetchone()
+                return True
+        except Exception:
+            return False

aws_advanced_python_wrapper/fastest_response_strategy_plugin.py

@@ -0,0 +1,330 @@
+#  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, time
+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_millis = 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_millis)
+
+        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 HostResponseTimeMonitor:
+
+    _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._host_id: Optional[str] = self._host_info.host_id
+        if self._host_id is None or self._host_id == "":
+            self._host_id = self._host_info.host
+
+        self._executor: Executor = ThreadPoolExecutor(thread_name_prefix="HostResponseTimeMonitorExecutor")
+
+    @property
+    def response_time(self):
+        return self._response_time
+
+    @response_time.setter
+    def response_time(self, response_time: int):
+        self._response_time = response_time
+
+    @property
+    def check_timestamp(self):
+        return self._check_timestamp
+
+    @check_timestamp.setter
+    def check_timestamp(self, check_timestamp: datetime):
+        self._check_timestamp = 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("HostResponseTimeMonitor.Stopped", self._host_info.host)
+
+    def _get_current_time(self):
+        return time()
+
+    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 range(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):
+                                calculated_response_time = self._get_current_time() - start_time
+                                response_time_sum = response_time_sum + calculated_response_time
+                                count = count + 1
+
+                            if count > 0:
+                                self.response_time = response_time_sum / count / 1000
+                            else:
+                                self.response_time = 2 ^ 31
+                            self.check_timestamp = self._get_current_time()
+                            logger.debug("HostResponseTimeMonitor.ResponseTime", self._host_info.host, self._response_time)
+
+                        sleep(self._interval_ms/1000)
+
+                except InterruptedError:
+                    # exit thread
+                    logger.debug("HostResponseTimeMonitor.InterruptedExceptionDuringMonitoring", self._host_info.host)
+        except Exception as e:
+            # this should not be reached; log and exit thread
+            logger.debug("HostResponseTimeMonitor.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("HostResponseTimeMonitor.OpeningConnection", self._host_info.url)
+                self._monitoring_conn = self._plugin_service.force_connect(self._host_info, monitoring_conn_properties, None)
+                logger.debug("HostResponseTimeMonitor.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:
+
+    @abstractmethod
+    def get_response_time(self, host_info: HostInfo) -> int:
+        """
+        Return a response time in milliseconds to the host.
+        Return 2 ^ 31 if response time is not available.
+
+        @param hostSpec the host details
+        @return response time in milliseconds for a desired host. It should return 2 ^ 31
+        if response time couldn't be measured.
+        """
+        ...
+
+    @abstractmethod
+    def set_hosts(self, hosts: Tuple[HostInfo, ...]) -> None:
+        ...
+
+
+class HostResponseTimeServiceImpl(HostResponseTimeService):
+    _CACHE_EXPIRATION_MILLIS: int = 6 * 10 ^ 5
+    _CACHE_CLEANUP_NANO: int = 6 * 10 ^ 4
+    _lock: Lock = Lock()
+    _monitoring_nodes: ClassVar[SlidingExpirationCacheWithCleanupThread[str, HostResponseTimeMonitor]] = \
+        SlidingExpirationCacheWithCleanupThread(_CACHE_CLEANUP_NANO,
+                                                should_dispose_func=lambda monitor: True,
+                                                item_disposal_func=lambda monitor: monitor.close())
+
+    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[HostResponseTimeMonitor] = HostResponseTimeServiceImpl._monitoring_nodes.get(host_info.url)
+        if monitor is None:
+            return 2 ^ 31
+        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 _: HostResponseTimeMonitor(
+                                                                 self._plugin_service,
+                                                                 new_host,
+                                                                 self._properties,
+                                                                 self._interval_ms), self._CACHE_EXPIRATION_MILLIS)

aws_advanced_python_wrapper/mysql_driver_dialect.py

@@ -148,6 +148,9 @@ def transfer_session_state(self, from_conn: Connection, to_conn: Connection):
                 isinstance(to_conn, CMySQLConnection) or isinstance(to_conn, MySQLConnection)):
             to_conn.autocommit = from_conn.autocommit
 
+    def ping(self, conn: Connection) -> bool:
+        return not self.is_closed(conn)
+
     def prepare_connect_info(self, host_info: HostInfo, original_props: Properties) -> Properties:
         driver_props: Properties = Properties(original_props.copy())
         PropertiesUtils.remove_wrapper_props(driver_props)

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,