RUBY-475 refactor read preference logic

Move read_pool into Mongo::ReadPreference
Add #read_preference to produce read preference hash

Author: TylerBrock

lib/mongo/cursor.rb

@@ -6,12 +6,13 @@ class Cursor
     include Mongo::Constants
     include Mongo::Conversions
     include Mongo::Logging
+    include Mongo::ReadPreference
 
     attr_reader :collection, :selector, :fields,
       :order, :hint, :snapshot, :timeout,
       :full_collection_name, :transformer,
       :options, :cursor_id, :show_disk_loc,
-      :comment, :read, :tag_sets
+      :comment, :read, :tag_sets, :acceptable_latency
 
     # Create a new cursor.
     #
@@ -530,9 +531,9 @@ def checkout_socket_from_connection
         if @pool
           socket = @pool.checkout
         elsif @command && !Mongo::Support::secondary_ok?(@selector)
-          socket = @connection.checkout_reader(:primary)
+          socket = @connection.checkout_reader({:mode => :primary})
         else
-          socket = @connection.checkout_reader(@read, @tag_sets, @acceptable_latency)
+          socket = @connection.checkout_reader(read_preference)
         end
       rescue SystemStackError, NoMemoryError, SystemCallError => ex
         @connection.close

lib/mongo/mongo_client.rb

@@ -351,10 +351,14 @@ def [](db_name)
     def refresh
     end
 
-    def pin_pool(pool)
+    def pinned_pool
+      @primary_pool
+    end
+
+    def pin_pool(pool, read_prefs)
     end
 
-    def unpin_pool(pool)
+    def unpin_pool
     end
 
     # Drop a database.
@@ -521,7 +525,7 @@ def max_message_size
 
     # Checkout a socket for reading (i.e., a secondary node).
     # Note: this is overridden in MongoReplicaSetClient.
-    def checkout_reader(mode=:primary, tag_sets={}, acceptable_latency=15)
+    def checkout_reader(read_preference)
       connect unless connected?
       @primary_pool.checkout
     end

lib/mongo/mongo_replica_set_client.rb

@@ -2,6 +2,7 @@ module Mongo
 
   # Instantiates and manages connections to a MongoDB replica set.
   class MongoReplicaSetClient < MongoClient
+    include ReadPreference
     include ThreadLocalVariableManager
 
     REPL_SET_OPTS = [
@@ -268,7 +269,7 @@ def nodes
     #
     # @return [Boolean]
     def read_primary?
-      @manager.read_pool == @manager.primary_pool
+      read_pool == primary_pool
     end
     alias :primary? :read_primary?
 
@@ -335,9 +336,9 @@ def checkout
       end
     end
 
-    def checkout_reader(mode=@read, tag_sets=@tag_sets, acceptable_latency=@acceptable_latency)
+    def checkout_reader(read_pref={})
       checkout do
-        pool = read_pool(mode, tag_sets, acceptable_latency)
+        pool = read_pool(read_pref)
         get_socket_from_pool(pool)
       end
     end
@@ -361,11 +362,20 @@ def ensure_manager
       thread_local[:managers][self] = @manager
     end
 
-    def pin_pool(pool)
-      thread_local[:pinned_pools][@manager.object_id] = pool if @manager
+    def pinned_pool
+      thread_local[:pinned_pools][@manager.object_id]
     end
 
-    def unpin_pool(pool)
+    def pin_pool(pool, read_preference)
+      if @manager
+        thread_local[:pinned_pools][@manager.object_id] = {
+          :pool => pool,
+          :read_preference => read_preference
+        }
+      end
+    end
+
+    def unpin_pool
       thread_local[:pinned_pools].delete @manager.object_id if @manager
     end
 
@@ -402,10 +412,6 @@ def primary_pool
       local_manager ? local_manager.primary_pool : nil
     end
 
-    def read_pool(mode=@read, tags=@tag_sets, acceptable_latency=@acceptable_latency)
-      local_manager ? local_manager.read_pool(mode, tags, acceptable_latency) : nil
-    end
-
     def secondary_pool
       local_manager ? local_manager.secondary_pool : nil
     end

lib/mongo/util/pool.rb

@@ -85,25 +85,6 @@ def up?
       !@closed
     end
 
-    def matches_mode(mode)
-      if mode == :primary && @node.secondary? ||
-        mode == :secondary && @node.primary?
-        false
-      else
-        true
-      end
-    end
-
-    def matches_tag_set(tag_set)
-      tag_set.all? do |tag, value|
-        tags.has_key?(tag) && tags[tag] == value
-      end
-    end
-
-    def matches_tag_sets(tag_sets)
-      tag_sets.all? {|set| matches_tag_set(set)}
-    end
-
     def inspect
       "#<Mongo::Pool:0x#{self.object_id.to_s(16)} @host=#{@host} @port=#{port} " +
         "@ping_time=#{@ping_time} #{@checked_out.size}/#{@size} sockets available " +

lib/mongo/util/pool_manager.rb

@@ -1,6 +1,5 @@
 module Mongo
   class PoolManager
-    include Mongo::ReadPreference
     include ThreadLocalVariableManager
 
     attr_reader :client,
@@ -114,26 +113,6 @@ def read
       read_pool.host_port
     end
 
-    def read_pool(mode=@client.read,
-                  tags=@client.tag_sets,
-                  acceptable_latency=@client.acceptable_latency)
-
-      pinned = thread_local[:pinned_pools][self.object_id]
-
-      if pinned && pinned.matches_mode(mode) && pinned.matches_tag_sets(tags) && pinned.up?
-        pool = pinned
-      else
-        pool = select_pool(mode, tags, acceptable_latency)
-      end
-
-      unless pool
-        raise ConnectionFailure, "No replica set member available for query " +
-          "with read preference matching mode #{mode} and tags matching #{tags}."
-      end
-
-      pool
-    end
-
     def max_bson_size
       @max_bson_size ||= config_min('maxBsonObjectSize', DEFAULT_MAX_BSON_SIZE)
     end

lib/mongo/util/read_preference.rb

@@ -33,29 +33,56 @@ def self.validate(value)
       end
     end
 
-    def select_pool(mode, tags, latency)
-      return primary_pool if @client.mongos?
+    def read_preference
+      {
+        :mode => @read,
+        :tags => @tag_sets,
+        :latency => @acceptable_latency
+      }
+    end
+
+    def read_pool(read_preference_override={})
+      return primary_pool if mongos?
+
+      read_pref = read_preference.merge read_preference_override
+
+      if pinned_pool && pinned_pool[:read_preference] == read_pref
+        pool = pinned_pool[:pool]
+      else
+        unpin_pool
+        pool = select_pool(read_pref)
+      end
+
+      unless pool
+        raise ConnectionFailure, "No replica set member available for query " +
+          "with read preference matching mode #{read_pref[:mode]} and tags " +
+          "matching #{read_pref[:tags]}."
+      end
+
+      pool
+    end
 
-      if mode == :primary && !tags.empty?
+    def select_pool(read_pref)
+      if read_pref[:mode] == :primary && !read_pref[:tags].empty?
         raise MongoArgumentError, "Read preference :primary cannot be combined with tags"
       end
 
-      case mode
+      case read_pref[:mode]
         when :primary
           primary_pool
         when :primary_preferred
-          primary_pool || select_secondary_pool(secondary_pools, tags, latency)
+          primary_pool || select_secondary_pool(secondary_pools, read_pref)
         when :secondary
-          select_secondary_pool(secondary_pools, tags, latency)
+          select_secondary_pool(secondary_pools, read_pref)
         when :secondary_preferred
-          select_secondary_pool(secondary_pools, tags, latency) || primary_pool
+          select_secondary_pool(secondary_pools, read_pref) || primary_pool
         when :nearest
-          select_secondary_pool(pools, tags, latency)
+          select_secondary_pool(pools, read_pref)
       end
     end
 
-    def select_secondary_pool(candidates, tag_sets, latency)
-      tag_sets = [tag_sets] unless tag_sets.is_a?(Array)
+    def select_secondary_pool(candidates, read_pref)
+      tag_sets = read_pref[:tags]
 
       if !tag_sets.empty?
         matches = []
@@ -70,10 +97,11 @@ def select_secondary_pool(candidates, tag_sets, latency)
         matches = candidates
       end
 
-      matches.empty? ? nil : select_near_pool(matches, latency)
+      matches.empty? ? nil : select_near_pool(matches, read_pref)
     end
 
-    def select_near_pool(candidates, latency)
+    def select_near_pool(candidates, read_pref)
+      latency = read_pref[:latency]
       nearest_pool = candidates.min_by { |candidate| candidate.ping_time }
       near_pools = candidates.select do |candidate|
         (candidate.ping_time - nearest_pool.ping_time) <= latency

test/replica_set/basic_test.rb

@@ -68,7 +68,7 @@ def test_accessors
     assert_equal 2, client.secondaries.length
     assert_equal 2, client.secondary_pools.length
     assert_equal @rs.repl_set_name, client.replica_set_name
-    assert client.secondary_pools.include?(client.read_pool(:secondary))
+    assert client.secondary_pools.include?(client.read_pool({:mode => :secondary}))
     assert_equal 90, client.refresh_interval
     assert_equal client.refresh_mode, false
     client.close