aws_advanced_ruby_driver_wrapper 1.0.0

data/lib/aws_advanced_ruby_driver_wrapper/plugins/default_plugin.rb

@@ -0,0 +1,108 @@
+# 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 '../host/host_info'
+require_relative '../host/host_availability'
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    class DefaultPlugin
+      SUBSCRIBED_METHODS = Set['*'].freeze
+      HOST_PORT_KEYS = %i[host port].freeze
+
+      def initialize(service_container, _props)
+        @service_container = service_container
+      end
+
+      def subscribed_methods
+        SUBSCRIBED_METHODS
+      end
+
+      def connect(host_info, driver_props, is_initial_connection, _pipeline_callable)
+        connection_service = @service_container.connection_service
+        driver_dialect = @service_container.dialect_service.driver_dialect
+        target_host_info = if is_initial_connection && connection_service.multi_host_url?
+                             # If the user specified a multi-host URL, we should always pass the same hosts/ports they
+                             # specified. Note that host_info may have a different value than initial_host_info.
+                             connection_service.initial_host_info
+                           else
+                             host_info
+                           end
+
+        conn = driver_dialect.connect(target_host_info, driver_props)
+
+        if conn.nil?
+          raise Errors::AwsError,
+                "Failed to connect to #{target_host_info&.host}: driver returned nil connection"
+        end
+
+        # If host was not specified (Unix socket / localhost), fill in from the live connection.
+        if host_info.nil? || !host_info.host_specified?
+          host_info = Host::HostInfo.new(
+            host: conn.respond_to?(:host) && conn.host ? conn.host : Host::HostInfo::NO_HOST,
+            port: conn.respond_to?(:port) && conn.port ? conn.port.to_s : Host::HostInfo::NO_PORT
+          )
+        else
+          @service_container.host_service.set_availability(host_info, Host::HostAvailability::AVAILABLE)
+        end
+
+        connection_service.update_current_connection(conn, host_info)
+
+        if is_initial_connection
+          @service_container.dialect_service.update_dialect(conn)
+          @service_container.host_service.refresh_host_list
+
+          if connection_service.pg? && connection_service.multi_host_url?
+            connection_service.config.initial_host_info = Host::HostInfo.new(
+              host: conn.host,
+              port: conn.port.to_i
+            )
+          end
+        end
+
+        init_func = AwsAdvancedRubyDriverWrapper.config.connection_init_func
+        init_func&.call(conn, host_info)
+
+        conn
+      end
+
+      def internal_connect(host_info, driver_props, _wrapper_props, _is_initial_connection, _pipeline_callable)
+        driver_dialect = @service_container.dialect_service.driver_dialect
+        driver_dialect.connect(host_info, driver_props)
+      end
+
+      def execute(target_method_name, target_callable, *args, **, &)
+        session = @service_container.session_state_service
+        autocommit_before = session&.autocommit?
+        succeeded = false
+
+        begin
+          result = target_callable.call(*args, **, &)
+          succeeded = true
+          result
+        ensure
+          if session && (conn = @service_container.connection_service.current_connection)
+            dialect = @service_container.dialect_service.driver_dialect
+            session.update_transaction_state(target_method_name, args, autocommit_before,
+                                             dialect, conn, succeeded: succeeded)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/plugins/failover_mode.rb

@@ -0,0 +1,43 @@
+# 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 FailoverMode
+      STRICT_WRITER = :strict_writer
+      STRICT_READER = :strict_reader
+      READER_OR_WRITER = :reader_or_writer
+
+      def self.from_value(value)
+        return nil if value.nil?
+
+        normalized = value.to_s.strip.downcase
+        return nil if normalized.empty?
+
+        case normalized
+        when 'strict_writer', 'strict-writer', 'strictwriter'
+          STRICT_WRITER
+        when 'strict_reader', 'strict-reader', 'strictreader'
+          STRICT_READER
+        when 'reader_or_writer', 'reader-or-writer', 'readerorwriter'
+          READER_OR_WRITER
+        else
+          raise ArgumentError, "Invalid failover mode: '#{value}'"
+        end
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/plugins/failover_plugin.rb

@@ -0,0 +1,467 @@
+# 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 '../host/host_info'
+require_relative '../host/host_role'
+require_relative '../host/host_availability'
+require_relative '../logging'
+require_relative '../utils/rds_utils'
+require_relative '../utils/rds_url_type'
+require_relative '../utils/retry_util'
+require_relative '../wrapper_property'
+require_relative 'failover_mode'
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    class FailoverPlugin
+      include Logging
+
+      ReaderFailoverResult = Data.define(:connection, :host_info)
+
+      def initialize(service_container, props = ::Concurrent::Map.new)
+        @service_container = service_container
+        @wrapper_props = props
+
+        @retry_util = Utils::RetryUtil.new(service_container)
+        @failover_timeout = PropertyDefinition::FAILOVER_TIMEOUT_SEC.get_float(props)
+        @reader_selector_strategy = PropertyDefinition::FAILOVER_READER_HOST_SELECTOR_STRATEGY.get(props)
+        @failover_mode = nil
+        @rds_url_type = nil
+        @closed_explicitly = false
+
+        network_methods = @service_container.dialect_service.driver_dialect.network_bound_methods
+        @subscribed_methods = (Set['connect'] | network_methods).freeze
+      end
+
+      attr_reader :subscribed_methods
+
+      def connect(host_info, driver_props, is_initial_connection, pipeline_callable)
+        init_failover_mode
+
+        unless PropertyDefinition::ENABLE_CONNECT_FAILOVER.get_bool(@wrapper_props)
+          return verified_connection(is_initial_connection, host_info, driver_props, pipeline_callable)
+        end
+
+        topology_host = host_service.hosts.find { |h| h.host_and_port == host_info&.host_and_port }
+
+        if !topology_host.nil? && topology_host.availability == Host::HostAvailability::UNAVAILABLE
+          host_service.refresh_host_list
+          return connect_via_failover(is_initial_connection)
+        end
+
+        begin
+          conn = verified_connection(is_initial_connection, host_info, driver_props, pipeline_callable)
+          host_service.refresh_host_list if is_initial_connection
+          conn
+        rescue StandardError => e
+          raise unless trigger_failover?(e)
+
+          host_service.set_availability(host_info, Host::HostAvailability::UNAVAILABLE)
+          connect_via_failover(is_initial_connection)
+        end
+      end
+
+      def execute(method_name, pipeline_callable, *, **, &)
+        if can_direct_execute?(method_name)
+          @closed_explicitly = true if method_name == RubyMethod::CONNECTION_CLOSE.name
+          return pipeline_callable.call
+        end
+
+        conn = connection_service.current_connection
+        if conn.nil?
+          logger.warn do
+            "[execute] current_connection is nil for #{method_name} — " \
+              "host_info=#{connection_service.current_host_info&.host}, " \
+              "closed_explicitly=#{@closed_explicitly}, failover_mode=#{@failover_mode.inspect}"
+          end
+        elsif !@closed_explicitly && driver_dialect.closed?(conn)
+          logger.warn("#{method_name} was called on closed connection #{conn} to #{connection_service.current_host_info}. " \
+                      'The driver will attempt to failover and then execute.')
+          failover
+        end
+
+        begin
+          pipeline_callable.call
+        rescue StandardError => e
+          handle_error(e)
+        end
+      end
+
+      private
+
+      def connect_via_failover(is_initial_connection)
+        begin
+          failover
+        rescue Errors::FailoverSuccessError => _e
+          conn = connection_service.current_connection
+        end
+
+        host_service.refresh_host_list if is_initial_connection
+        conn
+      end
+
+      def connection_service
+        @service_container.connection_service
+      end
+
+      def host_service
+        @service_container.host_service
+      end
+
+      def dialect_service
+        @service_container.dialect_service
+      end
+
+      def driver_dialect
+        dialect_service.driver_dialect
+      end
+
+      def db_dialect
+        dialect_service.db_dialect
+      end
+
+      def init_failover_mode
+        return unless @rds_url_type.nil?
+
+        @failover_mode = FailoverMode.from_value(PropertyDefinition::FAILOVER_MODE.get(@wrapper_props))
+        initial_host = connection_service.initial_host_info
+
+        @rds_url_type = Utils::RdsUtils.identify_rds_type(initial_host&.host)
+        reject_rds_proxy_endpoint
+
+        if @failover_mode.nil?
+          @failover_mode = if @rds_url_type == Utils::RdsUrlType::RDS_READER_CLUSTER
+                             FailoverMode::READER_OR_WRITER
+                           else
+                             FailoverMode::STRICT_WRITER
+                           end
+        end
+
+        logger.debug { "failover_mode=#{@failover_mode}" }
+      end
+
+      # RDS Proxy handles failover internally, so a failover plugin must not be layered on top of it.
+      #
+      # @raise [Errors::AwsError] if the initial endpoint is an RDS Proxy endpoint
+      def reject_rds_proxy_endpoint
+        return unless [Utils::RdsUrlType::RDS_PROXY, Utils::RdsUrlType::RDS_PROXY_ENDPOINT].include?(@rds_url_type)
+
+        raise Errors::AwsError,
+              'The failover plugin is not compatible with RDS Proxy endpoints. ' \
+              'RDS Proxy handles failover internally - please remove the failover plugin from your configuration.'
+      end
+
+      def can_direct_execute?(method_name)
+        method_name == RubyMethod::CONNECTION_CLOSE.name ||
+          method_name == RubyMethod::CONNECTION_PING.name
+      end
+
+      def handle_error(error)
+        logger.debug { "Detected error: #{error.message}" }
+        raise error unless trigger_failover?(error)
+
+        invalidate_current_connection
+        host_service.set_availability(
+          connection_service.current_host_info,
+          Host::HostAvailability::UNAVAILABLE
+        )
+
+        failover
+        # Failover will be skipped if the connection was explicitly closed. In this case we raise the error here.
+        raise error
+      end
+
+      def trigger_failover?(error)
+        if host_service.all_hosts.empty?
+          logger.debug do
+            "Skipping failover for #{connection_service.current_host_info}: no topology available."
+          end
+          return false
+        end
+
+        return true if dialect_service.network_error?(error)
+
+        dialect_service.read_only_error?(error) && failover_on_read_only_error?
+      end
+
+      # Whether a read-only error should trigger failover. It should only do so when failover would
+      # target a writer, since a read-only error means the current connection is not one.
+      #
+      # @return [Boolean]
+      def failover_on_read_only_error?
+        @failover_mode == FailoverMode::STRICT_WRITER
+      end
+
+      # Rolls back an open transaction on the current connection so it does not linger server-side.
+      #
+      # The broken connection is intentionally left open rather than closed here. A real pg/mysql2
+      # connection is left broken-but-open on network failure; its owner (ActiveRecord via
+      # disconnect!/discard!, or a raw-driver caller) is responsible for closing it. Closing it eagerly
+      # mid-error-handling would leave the facade pointing at a freed object on a failed failover, so
+      # introspection calls AR makes during its error path (e.g. transaction_status, ping) would raise
+      # instead of reporting a broken state. On a successful failover the old connection is closed when
+      # the new one is swapped in via ConnectionService#update_current_connection.
+      def invalidate_current_connection
+        conn = connection_service.current_connection
+        return if conn.nil?
+        return unless @service_container.session_state_service.in_transaction?
+
+        begin
+          driver_dialect.execute(conn, 'ROLLBACK')
+        rescue StandardError
+          nil
+        end
+      end
+
+      def failover
+        if @closed_explicitly
+          logger.debug { 'Connection was explicitly closed, skipping failover' }
+          return
+        end
+
+        if @failover_mode == FailoverMode::STRICT_WRITER
+          failover_writer
+        else
+          failover_reader
+        end
+      end
+
+      def failover_reader
+        failover_start = Time.now
+        failover_deadline = failover_start + @failover_timeout
+
+        logger.info { 'Starting reader failover' }
+
+        unless host_service.force_refresh_host_list?(verify_writer: false, timeout_sec: 0)
+          raise Errors::FailoverFailedError, 'The request to discover the new topology was unsuccessful'
+        end
+
+        begin
+          was_in_transaction = @service_container.session_state_service.in_transaction?
+          result = reader_failover_connection(failover_deadline)
+          raise Errors::FailoverFailedError, 'Unable to connect to a reader instance' unless result&.connection
+
+          connection_service.update_current_connection(result.connection, result.host_info)
+        rescue Timeout::Error
+          raise Errors::FailoverFailedError, 'Unable to connect to a reader instance'
+        end
+
+        raise_failover_success_error(was_in_transaction)
+      ensure
+        duration_ms = ((Time.now - failover_start) * 1000).round
+        logger.debug { "Reader failover duration: #{duration_ms}ms" }
+      end
+
+      def failover_writer
+        failover_start = Time.now
+        failover_deadline = failover_start + @failover_timeout
+        result = nil
+
+        logger.info { 'Starting writer failover' }
+
+        begin
+          unless host_service.force_refresh_host_list?(verify_writer: true, timeout_sec: @failover_timeout)
+            raise Errors::FailoverFailedError, 'The request to discover the new topology timed out or was unsuccessful'
+          end
+
+          was_in_transaction = @service_container.session_state_service.in_transaction?
+          result = @retry_util.connect_to_writer(self, @service_container.plugin_manager, deadline: failover_deadline)
+          raise Errors::FailoverFailedError, 'Unable to connect to a writer instance' unless result&.connection && result.host_info
+
+          connection_service.update_current_connection(result.connection, result.host_info)
+        rescue Timeout::Error
+          close_quietly(result&.connection)
+          raise Errors::FailoverFailedError,
+                "Writer failover timed out after #{@failover_timeout}s. Unable to connect to a new writer instance."
+        ensure
+          duration_ms = ((Time.now - failover_start) * 1000).round
+          logger.debug { "Writer failover duration: #{duration_ms}ms" }
+        end
+
+        raise_failover_success_error(was_in_transaction)
+      end
+
+      def raise_failover_success_error(was_in_transaction)
+        logger.debug { "Established connection to: #{connection_service.current_host_info}" }
+        raise Errors::FailoverSuccessError unless was_in_transaction
+
+        @service_container.session_state_service.in_transaction = false
+        raise Errors::TransactionStateUnknownError
+      end
+
+      def reader_failover_connection(deadline)
+        original_writer = nil
+        original_writer_still_writer = false
+
+        loop do
+          break if Time.now >= deadline
+
+          hosts = host_service.hosts
+          reader_candidates = hosts.select { |h| h.role == Host::HostRole::READER }
+          original_writer ||= hosts.find { |h| h.role == Host::HostRole::WRITER }
+
+          result = try_reader_candidates(reader_candidates, deadline)
+          return result if result
+
+          result = try_original_writer(original_writer, original_writer_still_writer)
+          case result
+          when ReaderFailoverResult
+            return result
+          when :still_writer
+            original_writer_still_writer = true
+          end
+
+          sleep(0.1)
+        end
+
+        raise Timeout::Error, 'The reader failover process was not able to establish a connection before timing out.'
+      end
+
+      def try_reader_candidates(reader_candidates, deadline)
+        remaining = reader_candidates.dup
+
+        while !remaining.empty? && Time.now < deadline
+          candidate = select_reader_candidate(remaining)
+          if candidate.nil?
+            # Unable to find available candidate in the host list. Let's try assuming all hosts are available.
+            available = remaining.map do |h|
+              h.deep_dup.tap { |dup| dup.availability = Host::HostAvailability::AVAILABLE }
+            end
+            candidate = select_reader_candidate(available)
+          end
+
+          if candidate.nil?
+            logger.debug { 'Unable to find reader in the updated host list.' }
+            break
+          end
+
+          outcome, result = attempt_reader_connection(candidate)
+          case outcome
+          when :success
+            return result
+          when :writer
+            reader_candidates.delete(candidate)
+            remaining.delete(candidate)
+          else
+            remaining.delete(candidate)
+          end
+        end
+
+        nil
+      end
+
+      def try_original_writer(original_writer, original_writer_still_writer)
+        return nil if original_writer.nil?
+        return nil if @failover_mode == FailoverMode::STRICT_READER && original_writer_still_writer
+
+        outcome, result = attempt_reader_connection(original_writer)
+        case outcome
+        when :success
+          result
+        when :writer
+          :still_writer
+        else
+          logger.debug { "Failed to connect to host: #{original_writer.url}" }
+          nil
+        end
+      end
+
+      def attempt_reader_connection(host_info)
+        conn = @service_container.plugin_manager.connect(host_info, connection_service.driver_props, false, plugin_to_skip: self)
+        # Since the roles in the host list might not be accurate, we execute a query to check the instance's role.
+        role = db_dialect.host_role(conn)
+
+        if role == Host::HostRole::READER || @failover_mode != FailoverMode::STRICT_READER
+          updated_host = host_info.deep_dup(role: role)
+          return [:success, ReaderFailoverResult.new(conn, updated_host)]
+        end
+
+        # The role is WRITER or UNKNOWN, and we are in STRICT_READER mode, so the connection is not valid.
+        close_quietly(conn)
+        if role == Host::HostRole::WRITER
+          [:writer, nil]
+        else
+          logger.debug do
+            "Unable to determine host role for #{host_info.url}. " \
+              'Since failover mode is set to STRICT_READER and the host may be a writer, ' \
+              'it will not be selected for reader failover.'
+          end
+          [:unknown, nil]
+        end
+      rescue StandardError
+        close_quietly(conn)
+        [:failed, nil]
+      end
+
+      def select_reader_candidate(hosts)
+        host_service.select_host(
+          hosts,
+          Host::HostRole::READER,
+          @reader_selector_strategy
+        )
+      rescue StandardError
+        nil
+      end
+
+      def verified_connection(is_initial_connection, host_info, driver_props, connect_func)
+        url_type = Utils::RdsUtils.identify_rds_type(host_info&.host)
+        if url_type != Utils::RdsUrlType::RDS_WRITER_CLUSTER
+          # We are not using a writer cluster endpoint. No verification needed - continue with the regular workflow.
+          return connect_func.call
+        end
+
+        conn = connect_func.call
+        if db_dialect.host_role(conn) == Host::HostRole::WRITER
+          host_service.refresh_host_list
+          return conn
+        end
+
+        # The writer cluster URL resolved to a reader. We will try to redirect to the writer instance.
+        host_service.force_refresh_host_list?(verify_writer: false, timeout_sec: 5.0)
+        writer = host_service.all_hosts.find { |h| h.role == Host::HostRole::WRITER }
+        if writer.nil? || Utils::RdsUtils.rds_cluster_dns?(writer.host)
+          # Writer instance endpoint not found - unable to redirect.
+          close_quietly(conn)
+          raise Errors::AwsError, 'Stale DNS detected - a writer was requested, but the writer cluster endpoint resolved to a reader'
+        end
+
+        allowed_hosts = host_service.hosts
+        unless allowed_hosts.any? { |h| h.host_and_port == writer.host_and_port }
+          raise Errors::AwsError, "Current writer #{writer.host_and_port} is not in allowed hosts"
+        end
+
+        # Attempt to correct the stale DNS problem by connecting to the writer instance.
+        logger.debug { "Stale DNS data detected. Opening a connection to #{writer.host}" }
+        writer_conn = @service_container.plugin_manager.connect(writer, driver_props, false, plugin_to_skip: self)
+        connection_service.initial_host_info = writer if is_initial_connection
+
+        # Close the incorrect reader connection.
+        close_quietly(conn)
+        writer_conn
+      end
+
+      def close_quietly(conn)
+        return if conn.nil?
+
+        driver_dialect.close_connection(conn)
+      rescue StandardError
+        # ignore
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/plugins/gdb/gdb_failover_mode.rb

@@ -0,0 +1,68 @@
+# 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 Gdb
+      # The host roles and regions that the GDB failover plugin may target during failover.
+      module GdbFailoverMode
+        STRICT_WRITER = :strict_writer
+        STRICT_HOME_READER = :strict_home_reader
+        STRICT_OUT_OF_HOME_READER = :strict_out_of_home_reader
+        STRICT_ANY_READER = :strict_any_reader
+        HOME_READER_OR_WRITER = :home_reader_or_writer
+        OUT_OF_HOME_READER_OR_WRITER = :out_of_home_reader_or_writer
+        ANY_READER_OR_WRITER = :any_reader_or_writer
+
+        ALL = [
+          STRICT_WRITER,
+          STRICT_HOME_READER,
+          STRICT_OUT_OF_HOME_READER,
+          STRICT_ANY_READER,
+          HOME_READER_OR_WRITER,
+          OUT_OF_HOME_READER_OR_WRITER,
+          ANY_READER_OR_WRITER
+        ].freeze
+
+        # Maps every accepted spelling of a mode to the mode itself. Kebab-case, snake_case and
+        # squashed spellings are all accepted, e.g. 'strict-home-reader', 'strict_home_reader'
+        # and 'stricthomereader'.
+        NAME_TO_VALUE = ALL.each_with_object({}) do |mode, mapping|
+          snake = mode.to_s
+          mapping[snake] = mode
+          mapping[snake.tr('_', '-')] = mode
+          mapping[snake.delete('_')] = mode
+        end.freeze
+
+        # Resolves a configured mode string into one of the mode constants.
+        #
+        # @param value [String, Symbol, nil] the configured value
+        # @return [Symbol, nil] the mode, or nil when no value was configured
+        # @raise [ArgumentError] if the value does not name a known mode
+        def self.from_value(value)
+          return nil if value.nil?
+
+          normalized = value.to_s.strip.downcase
+          return nil if normalized.empty?
+
+          NAME_TO_VALUE.fetch(normalized) do
+            raise ArgumentError, "Invalid global database failover mode: '#{value}'"
+          end
+        end
+      end
+    end
+  end
+end