lib/mongo/server_selector/base.rb in mongo-2.19.3 vs lib/mongo/server_selector/base.rb in mongo-2.20.0

@@ -162,30 +162,34 @@
       #   Deprecated and ignored.
       # @param [ Session | nil ] session Optional session to take into account
       #   for mongos pinning. Added in version 2.10.0.
       # @param [ true | false ] write_aggregation Whether we need a server that
       #   supports writing aggregations (e.g. with $merge/$out) on secondaries.
+      # @param [ Array<Server> ] deprioritized A list of servers that should
+      #   be selected from only if no other servers are available. This is
+      #   used to avoid selecting the same server twice in a row when
+      #   retrying a command.
       #
       # @return [ Mongo::Server ] A server matching the server preference.
       #
       # @raise [ Error::NoServerAvailable ] No server was found matching the
       #   specified preference / pinning requirement in the server selection
       #   timeout.
       # @raise [ Error::LintError ] An unexpected condition was detected, and
       #   lint mode is enabled.
       #
       # @since 2.0.0
-      def select_server(cluster, ping = nil, session = nil, write_aggregation: false)
-        select_server_impl(cluster, ping, session, write_aggregation).tap do |server|
+      def select_server(cluster, ping = nil, session = nil, write_aggregation: false, deprioritized: [])
+        select_server_impl(cluster, ping, session, write_aggregation, deprioritized).tap do |server|
           if Lint.enabled? && !server.pool.ready?
             raise Error::LintError, 'Server selector returning a server with a pool which is not ready'
           end
         end
       end
 
       # Parameters and return values are the same as for select_server.
-      private def select_server_impl(cluster, ping, session, write_aggregation)
+      private def select_server_impl(cluster, ping, session, write_aggregation, deprioritized)
         if cluster.topology.is_a?(Cluster::Topology::LoadBalanced)
           return cluster.servers.first
         end
 
         server_selection_timeout = cluster.options[:server_selection_timeout] || SERVER_SELECTION_TIMEOUT
@@ -264,11 +268,11 @@
                 raise Error::LintError, "Server #{server.summary} is known but has non-ready pool"
               end
             end
           end
 
-          server = try_select_server(cluster, write_aggregation: write_aggregation)
+          server = try_select_server(cluster, write_aggregation: write_aggregation, deprioritized: deprioritized)
 
           if server
             unless cluster.topology.compatible?
               raise Error::UnsupportedFeatures, cluster.topology.compatibility_error.to_s
             end
@@ -319,15 +323,19 @@
       #
       # @param [ Mongo::Cluster ] cluster The cluster from which to select
       #   an eligible server.
       # @param [ true | false ] write_aggregation Whether we need a server that
       #   supports writing aggregations (e.g. with $merge/$out) on secondaries.
+      # @param [ Array<Server> ] deprioritized A list of servers that should
+      #   be selected from only if no other servers are available. This is
+      #   used to avoid selecting the same server twice in a row when
+      #   retrying a command.
       #
       # @return [ Server | nil ] A suitable server, if one exists.
       #
       # @api private
-      def try_select_server(cluster, write_aggregation: false)
+      def try_select_server(cluster, write_aggregation: false, deprioritized: [])
         servers = if write_aggregation && cluster.replica_set?
           # 1. Check if ALL servers in cluster support secondary writes.
           is_write_supported = cluster.servers.reduce(true) do |res, server|
             res && server.features.merge_out_on_secondary_enabled?
           end
@@ -345,11 +353,11 @@
 
         # This list of servers may be ordered in a specific way
         # by the selector (e.g. for secondary preferred, the first
         # server may be a secondary and the second server may be primary)
         # and we should take the first server here respecting the order
-        server = servers.first
+        server = suitable_server(servers, deprioritized)
 
         if server
           if Lint.enabled?
             # It is possible for a server to have a nil average RTT here
             # because the ARTT comes from description which may be updated
@@ -415,9 +423,27 @@
           []
         end
       end
 
       private
+
+      # Returns a server from the list of servers that is suitable for
+      # executing the operation.
+      #
+      # @param [ Array<Server> ] servers The candidate servers.
+      # @param [ Array<Server> ] deprioritized A list of servers that should
+      #  be selected from only if no other servers are available.
+      #
+      # @return [ Server | nil ] The suitable server or nil if no suitable
+      #  server is available.
+      def suitable_server(servers, deprioritized)
+        preferred = servers - deprioritized
+        if preferred.empty?
+          servers.first
+        else
+          preferred.first
+        end
+      end
 
       # Convert this server preference definition into a format appropriate
       #   for sending to a MongoDB server (i.e., as a command field).
       #
       # @return [ Hash ] The server preference formatted as a command field value.