aws_advanced_ruby_driver_wrapper 1.0.0

data/lib/aws_advanced_ruby_driver_wrapper/plugins/custom_endpoint/custom_endpoint_monitor.rb

@@ -0,0 +1,266 @@
+# frozen_string_literal: true
+
+#
+#  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.
+
+require_relative '../../monitoring/monitor'
+require_relative '../../logging'
+require_relative 'info'
+require_relative 'member_list_type'
+require 'concurrent'
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    module CustomEndpoint
+      class CustomEndpointMonitor < Monitoring::Monitor
+        include Logging
+
+        TERMINATION_TIMEOUT_SEC = 30.0
+        ENDPOINT_INFO_EXPIRATION_SEC = 300.0
+        UNAUTHORIZED_SLEEP_SEC = 300.0
+        ENDPOINT_INFO_CACHE_NAME = :custom_endpoint
+        ALLOWED_BLOCKED_CACHE_NAME = :custom_endpoint_allowed_blocked
+
+        def initialize(
+          service_container,
+          custom_endpoint_host,
+          endpoint_id,
+          region,
+          refresh_rate_sec,
+          refresh_rate_backoff_factor,
+          max_refresh_rate_sec,
+          rds_client_func: ->(_, region) { Aws::RDS::Client.new(region: region) }
+        )
+          super(termination_timeout_sec: TERMINATION_TIMEOUT_SEC)
+
+          @service_container = service_container
+          @custom_endpoint_host = custom_endpoint_host
+          @endpoint_id = endpoint_id
+          @min_refresh_rate_sec = refresh_rate_sec
+          @refresh_rate_sec = refresh_rate_sec
+          @refresh_rate_backoff_factor = refresh_rate_backoff_factor
+          @max_refresh_rate_sec = max_refresh_rate_sec
+          @rds_client = rds_client_func.call(custom_endpoint_host, region)
+
+          @refresh_mutex = Mutex.new
+          @refresh_cv = ConditionVariable.new
+          @refresh_required = false
+          @connection_issue = false
+        end
+
+        def endpoint_info?
+          info = cached_info
+          request_refresh if info.nil? && @refresh_mutex.synchronize { !@refresh_required && !@connection_issue }
+          !info.nil?
+        end
+
+        # Waits up to timeout_sec for the monitor to place endpoint info in the cache. Polls rather than
+        # waiting on a one-shot signal: #endpoint_info? re-checks the cache and re-requests a refresh
+        # (waking the monitor) on each pass, so a transient empty cache - e.g. an entry that aged out
+        # between monitor iterations - is repopulated within the window instead of failing immediately.
+        def wait_for_info?(timeout_sec)
+          deadline = monotonic_time + timeout_sec
+          loop do
+            return true if endpoint_info?
+            return false if monotonic_time >= deadline
+
+            sleep(0.1)
+          end
+        end
+
+        def request_endpoint_info_update
+          return if @refresh_mutex.synchronize { @connection_issue }
+
+          request_refresh
+        end
+
+        def close
+          remove_cached_info
+        end
+
+        def self.clear_cache(storage_service)
+          storage_service.clear(ENDPOINT_INFO_CACHE_NAME)
+          storage_service.clear(ALLOWED_BLOCKED_CACHE_NAME)
+        end
+
+        private
+
+        def monitor
+          logger.debug("[#{@endpoint_id}] Started custom endpoint monitor for #{@custom_endpoint_host.url}")
+
+          until stopped?
+            update_activity
+            run_monitor_iteration
+          end
+        rescue StandardError => e
+          logger.error("[#{@endpoint_id}] Unexpected error in custom endpoint monitor: #{e.message}")
+        ensure
+          remove_cached_info
+          logger.debug("[#{@endpoint_id}] Stopped custom endpoint monitor for #{@custom_endpoint_host.url}")
+        end
+
+        def run_monitor_iteration
+          start = monotonic_time
+
+          endpoints = fetch_endpoints
+          return unless endpoints
+
+          @refresh_mutex.synchronize do
+            @connection_issue = false
+            @refresh_required = false
+          end
+
+          unless valid_endpoint_count?(endpoints)
+            sleep_ignoring_refresh_requests(@refresh_rate_sec)
+            return
+          end
+
+          endpoint_info = Info.from_db_cluster_endpoint(endpoints.first)
+          process_endpoint_info(endpoint_info, start)
+        rescue Aws::RDS::Errors::ServiceError => e
+          handle_rds_error(e)
+        rescue StandardError => e
+          logger.error("[#{@endpoint_id}] Exception monitoring #{@custom_endpoint_host.url}: #{e.message}")
+          sleep_ignoring_refresh_requests(@refresh_rate_sec)
+        end
+
+        def fetch_endpoints
+          response = @rds_client.describe_db_cluster_endpoints(
+            db_cluster_endpoint_identifier: @endpoint_id,
+            filters: [{ name: 'db-cluster-endpoint-type', values: ['custom'] }]
+          )
+          response.db_cluster_endpoints
+        end
+
+        def valid_endpoint_count?(endpoints)
+          return true if endpoints.size == 1
+
+          logger.warn("[#{@endpoint_id}] Expected 1 endpoint, got #{endpoints.size}: #{endpoints.map(&:endpoint)}")
+          false
+        end
+
+        def process_endpoint_info(endpoint_info, start)
+          if cached_info == endpoint_info
+            elapsed = monotonic_time - start
+            interruptible_sleep([0, @refresh_rate_sec - elapsed].max)
+            return
+          end
+
+          logger.debug("[#{@endpoint_id}] Custom endpoint info changed: #{endpoint_info}")
+          cache_allowed_blocked(endpoint_info)
+          cache_info(endpoint_info)
+          speedup_refresh_rate
+
+          elapsed = monotonic_time - start
+          interruptible_sleep([0, @refresh_rate_sec - elapsed].max)
+        end
+
+        def cache_allowed_blocked(endpoint_info)
+          value = if endpoint_info.member_list_type == MemberListType::STATIC_LIST
+                    { allowed: endpoint_info.static_members, blocked: nil, required_role: endpoint_info.required_role }
+                  else
+                    { allowed: nil, blocked: endpoint_info.excluded_members, required_role: endpoint_info.required_role }
+                  end
+          storage_service.set(ALLOWED_BLOCKED_CACHE_NAME, @custom_endpoint_host.url, value)
+        end
+
+        def handle_rds_error(error)
+          logger.error("[#{@endpoint_id}] RDS error for #{@custom_endpoint_host.url}: #{error.message}")
+
+          if throttling_error?(error)
+            slowdown_refresh_rate
+            sleep_ignoring_refresh_requests(@refresh_rate_sec)
+          elsif unauthorized_error?(error)
+            sleep_ignoring_refresh_requests(UNAUTHORIZED_SLEEP_SEC)
+          else
+            sleep_ignoring_refresh_requests(@refresh_rate_sec)
+          end
+        end
+
+        def throttling_error?(error)
+          error.context&.http_response&.status_code == 429 ||
+            error.code == 'ThrottlingException' ||
+            error.code == 'Throttling'
+        end
+
+        def unauthorized_error?(error)
+          [401, 403].include?(error.context&.http_response&.status_code)
+        end
+
+        def speedup_refresh_rate
+          return unless @refresh_rate_sec > @min_refresh_rate_sec
+
+          @refresh_rate_sec = [@refresh_rate_sec / @refresh_rate_backoff_factor, @min_refresh_rate_sec].max
+        end
+
+        def slowdown_refresh_rate
+          return unless @refresh_rate_sec < @max_refresh_rate_sec
+
+          @refresh_rate_sec = [@refresh_rate_sec * @refresh_rate_backoff_factor, @max_refresh_rate_sec].min
+        end
+
+        def request_refresh
+          @refresh_mutex.synchronize do
+            @refresh_required = true
+            @refresh_cv.broadcast
+          end
+        end
+
+        # Wakes early if a refresh is requested or the monitor is stopped.
+        def interruptible_sleep(duration_sec)
+          end_time = monotonic_time + duration_sec
+          wait_sec = [0.5, duration_sec].min
+
+          @refresh_mutex.synchronize do
+            @refresh_cv.wait(@refresh_mutex, wait_sec) until @refresh_required || monotonic_time >= end_time || stopped?
+          end
+        end
+
+        # Sleeps for the full duration, ignoring refresh requests. Used on error/backoff paths
+        # to prevent connections from bypassing throttling backoff.
+        def sleep_ignoring_refresh_requests(duration_sec)
+          end_time = monotonic_time + duration_sec
+          until stopped?
+            remaining = end_time - monotonic_time
+            break if remaining <= 0
+
+            sleep([0.5, remaining].min)
+          end
+        end
+
+        def cached_info
+          storage_service.get(ENDPOINT_INFO_CACHE_NAME, @custom_endpoint_host.url, register_access: false)
+        end
+
+        def cache_info(info)
+          storage_service.set(ENDPOINT_INFO_CACHE_NAME, @custom_endpoint_host.url, info)
+        end
+
+        def remove_cached_info
+          storage_service.remove(ENDPOINT_INFO_CACHE_NAME, @custom_endpoint_host.url)
+          storage_service.remove(ALLOWED_BLOCKED_CACHE_NAME, @custom_endpoint_host.url)
+        end
+
+        def storage_service
+          @service_container.storage_service
+        end
+
+        def monotonic_time
+          Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        end
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/plugins/custom_endpoint/custom_endpoint_plugin.rb

@@ -0,0 +1,158 @@
+# frozen_string_literal: true
+
+#  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.
+
+require 'concurrent'
+require_relative '../../errors'
+require_relative '../../logging'
+require_relative '../../property_definition'
+require_relative '../../ruby_method'
+require_relative '../../utils/rds_utils'
+require_relative 'custom_endpoint_monitor'
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    module CustomEndpoint
+      class CustomEndpointPlugin
+        include Logging
+
+        MONITOR_TYPE = :custom_endpoint
+
+        attr_reader :subscribed_methods
+
+        def initialize(service_container, props = ::Concurrent::Map.new)
+          ensure_aws_sdk!
+          @service_container = service_container
+          @props = props
+          @should_wait_for_info = PropertyDefinition::WAIT_FOR_CUSTOM_ENDPOINT_INFO.get_bool(props)
+          @wait_timeout_sec = PropertyDefinition::WAIT_FOR_CUSTOM_ENDPOINT_INFO_TIMEOUT_SEC.get_float(props)
+
+          @custom_endpoint_host = nil
+          @endpoint_id = nil
+          @region = nil
+
+          monitor_expiration_sec = PropertyDefinition::CUSTOM_ENDPOINT_MONITOR_EXPIRATION_SEC.get_float(props)
+          service_container.monitor_service.register_type(
+            MONITOR_TYPE,
+            expiration_timeout_sec: monitor_expiration_sec,
+            produced_data_type: CustomEndpointMonitor::ENDPOINT_INFO_CACHE_NAME
+          )
+          service_container.storage_service.register(
+            CustomEndpointMonitor::ENDPOINT_INFO_CACHE_NAME,
+            ttl: CustomEndpointMonitor::ENDPOINT_INFO_EXPIRATION_SEC
+          )
+          service_container.storage_service.register(
+            CustomEndpointMonitor::ALLOWED_BLOCKED_CACHE_NAME,
+            ttl: CustomEndpointMonitor::ENDPOINT_INFO_EXPIRATION_SEC
+          )
+
+          network_methods = service_container.dialect_service.driver_dialect.network_bound_methods
+          @subscribed_methods = (Set['connect'] | network_methods).freeze
+        end
+
+        def connect(host_info, _, _, pipeline_callable)
+          return pipeline_callable.call unless Utils::RdsUtils.rds_custom_cluster_dns?(host_info.host)
+
+          logger.debug("CustomEndpointPlugin: connection request to custom endpoint '#{host_info.url}'")
+          init_endpoint_state!(host_info) if @custom_endpoint_host.nil?
+          monitor = create_monitor_if_absent(@props)
+          wait_for_endpoint_info(monitor) if @should_wait_for_info
+          pipeline_callable.call
+        end
+
+        CLOSE_METHODS = Set[
+          RubyMethod::CONNECTION_CLOSE.name,
+          RubyMethod::CONNECTION_RESET.name
+        ].freeze
+
+        def execute(method_name, pipeline_callable, *_, **_)
+          return pipeline_callable.call if @custom_endpoint_host.nil?
+          return pipeline_callable.call if CLOSE_METHODS.include?(method_name.to_s)
+
+          monitor = create_monitor_if_absent(@props)
+          wait_for_endpoint_info(monitor) if @should_wait_for_info
+          pipeline_callable.call
+        end
+
+        def self.clear_cache(storage_service)
+          CustomEndpointMonitor.clear_cache(storage_service)
+        end
+
+        private
+
+        def init_endpoint_state!(host_info)
+          @custom_endpoint_host = host_info
+
+          endpoint_id = Utils::RdsUtils.rds_cluster_id(host_info.host)
+          if endpoint_id.nil? || endpoint_id.empty?
+            raise Errors::AwsError, "CustomEndpointPlugin: unable to parse endpoint identifier from '#{host_info.host}'"
+          end
+
+          @endpoint_id = endpoint_id
+          @region = resolve_region!(host_info.host)
+        end
+
+        def resolve_region!(host)
+          region = PropertyDefinition::CUSTOM_ENDPOINT_REGION.get(@props) || Utils::RdsUtils.rds_region(host)
+          unless region
+            raise Errors::AwsError,
+                  "CustomEndpointPlugin: unable to determine region for '#{host}'. " \
+                  "Set the '#{PropertyDefinition::CUSTOM_ENDPOINT_REGION.name}' property explicitly."
+          end
+
+          region
+        end
+
+        def create_monitor_if_absent(props)
+          @service_container.monitor_service.run_if_absent(
+            MONITOR_TYPE,
+            @custom_endpoint_host.url,
+            @service_container
+          ) do |service_container|
+            CustomEndpointMonitor.new(
+              service_container,
+              @custom_endpoint_host,
+              @endpoint_id,
+              @region,
+              PropertyDefinition::CUSTOM_ENDPOINT_INFO_REFRESH_RATE_SEC.get_float(props),
+              PropertyDefinition::CUSTOM_ENDPOINT_INFO_REFRESH_RATE_BACKOFF_FACTOR.get_int(props),
+              PropertyDefinition::CUSTOM_ENDPOINT_INFO_MAX_REFRESH_RATE_SEC.get_float(props)
+            )
+          end
+        end
+
+        def wait_for_endpoint_info(monitor)
+          return if monitor.endpoint_info?
+
+          monitor.request_endpoint_info_update
+          logger.debug("CustomEndpointPlugin: waiting up to #{@wait_timeout_sec}s for endpoint info on #{@custom_endpoint_host.url}")
+
+          return if monitor.wait_for_info?(@wait_timeout_sec)
+
+          raise Errors::AwsError,
+                "CustomEndpointPlugin: timed out after #{@wait_timeout_sec}s " \
+                "waiting for custom endpoint info for '#{@custom_endpoint_host.url}'"
+        end
+
+        def ensure_aws_sdk!
+          require 'aws-sdk-rds'
+        rescue LoadError
+          raise LoadError,
+                "The custom endpoint plugin requires 'aws-sdk-rds'. Add it to your Gemfile: gem 'aws-sdk-rds'"
+        end
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/plugins/custom_endpoint/info.rb

@@ -0,0 +1,111 @@
+# frozen_string_literal: true
+
+#  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.
+
+require_relative 'member_list_type'
+require_relative 'role'
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    module CustomEndpoint
+      # Represents custom endpoint information for a given custom endpoint.
+      class Info
+        attr_reader :endpoint_identifier, :cluster_identifier, :url, :role, :member_list_type, :members
+
+        # @param endpoint_identifier [String] The endpoint identifier for the custom endpoint. For example, if the
+        #   custom endpoint URL is "my-custom-endpoint.cluster-custom-XYZ.us-east-1.rds.amazonaws.com", the endpoint
+        #   identifier is "my-custom-endpoint".
+        # @param cluster_identifier [String] The cluster identifier for the cluster that the custom endpoint belongs to.
+        # @param url [String] The URL for the custom endpoint.
+        # @param role [Symbol] The role of the custom endpoint.
+        # @param members [Enumerable<String>] The instance IDs for the hosts in the custom endpoint.
+        # @param member_list_type [Symbol] The list type for +members+.
+        def initialize(endpoint_identifier:,
+                       cluster_identifier:,
+                       url:,
+                       role:,
+                       members:,
+                       member_list_type:)
+          @endpoint_identifier = endpoint_identifier
+          @cluster_identifier = cluster_identifier
+          @url = url
+          @role = role
+          @members = members.to_set.freeze
+          @member_list_type = member_list_type
+        end
+
+        def self.from_db_cluster_endpoint(response)
+          static = response.static_members&.any?
+          members = static ? response.static_members : response.excluded_members
+          member_list_type = static ? MemberListType::STATIC_LIST : MemberListType::EXCLUSION_LIST
+
+          new(
+            endpoint_identifier: response.db_cluster_endpoint_identifier,
+            cluster_identifier: response.db_cluster_identifier,
+            url: response.endpoint,
+            role: Role.parse(response.custom_endpoint_type),
+            members: members,
+            member_list_type: member_list_type
+          )
+        end
+
+        # Evaluates whether instances in the custom endpoint must match a particular role according to the custom
+        # endpoint properties. Note that custom clusters with static member lists always route to all static members,
+        # even if the member is a writer and the custom endpoint is of type READER, so there are never role
+        # requirements for static list custom clusters.
+        # @return [Symbol, nil] the required role of instances in the custom endpoint, or nil if there is no strict
+        #   role requirement.
+        def required_role
+          :reader if @member_list_type == MemberListType::EXCLUSION_LIST && @role == Role::READER
+        end
+
+        # Gets the static members of the custom endpoint. If the custom endpoint member list type is an exclusion
+        # list, returns nil.
+        # @return [Set<String>, nil]
+        def static_members
+          @members if @member_list_type == MemberListType::STATIC_LIST
+        end
+
+        # Gets the excluded members of the custom endpoint. If the custom endpoint member list type is a static
+        # list, returns nil.
+        # @return [Set<String>, nil]
+        def excluded_members
+          @members if @member_list_type == MemberListType::EXCLUSION_LIST
+        end
+
+        def ==(other)
+          other.is_a?(Info) &&
+            endpoint_identifier == other.endpoint_identifier &&
+            cluster_identifier == other.cluster_identifier &&
+            url == other.url &&
+            role == other.role &&
+            member_list_type == other.member_list_type &&
+            members == other.members
+        end
+
+        alias eql? ==
+
+        def hash
+          [@endpoint_identifier, @cluster_identifier, @url, @role, @member_list_type, @members].hash
+        end
+
+        def to_s
+          "Info[url=#{@url}, cluster=#{@cluster_identifier}, role=#{@role}, " \
+            "member_list_type=#{@member_list_type}, members=#{@members}]"
+        end
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/plugins/custom_endpoint/member_list_type.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+#  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.
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    module CustomEndpoint
+      # Represents the member list type of a custom endpoint.
+      # Used with a member list to determine which instances are included or excluded.
+      module MemberListType
+        # Only the listed instances are included. New cluster instances are NOT auto-added.
+        STATIC_LIST = :static_list
+
+        # The listed instances are excluded. New cluster instances ARE auto-added.
+        EXCLUSION_LIST = :exclusion_list
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/plugins/custom_endpoint/role.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+#  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.
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    module CustomEndpoint
+      # Represents the possible roles of instances specified by a custom endpoint.
+      module Role
+        # Instances may be either a writer or a reader.
+        ANY = :any
+
+        # Instance is always the writer.
+        WRITER = :writer
+
+        # Instances are always readers.
+        READER = :reader
+
+        ROLE_MAPPING = {
+          'ANY' => ANY,
+          'WRITER' => WRITER,
+          'READER' => READER
+        }.freeze
+
+        def self.parse(value)
+          raise ArgumentError, 'Role value is blank' if value.nil? || value.strip.empty?
+
+          ROLE_MAPPING.fetch(value.upcase) { raise ArgumentError, "Unknown role: #{value}" }
+        end
+      end
+    end
+  end
+end