RUBY-475 pin/unpin logic added to cursor

Author: TylerBrock

lib/mongo/cursor.rb

@@ -472,6 +472,7 @@ def send_initial_query
             nil, @options & OP_QUERY_EXHAUST != 0)
         rescue ConnectionFailure => ex
           socket.close if socket
+          @connection.unpin_pool
           @connection.refresh
           if tries < 3 && !@socket && (!@command || Mongo::Support::secondary_ok?(@selector))
             tries += 1
@@ -484,6 +485,9 @@ def send_initial_query
         ensure
           socket.checkin unless @socket || socket.nil?
         end
+        if !@socket && !@command
+          @connection.pin_pool(socket.pool, read_preference)
+        end
         @returned += @n_received
         @cache += results
         @query_run = true

lib/mongo/mongo_replica_set_client.rb

@@ -283,6 +283,7 @@ def close(opts={})
 
       # Clear the reference to this object.
       thread_local[:managers].delete(self)
+      unpin_pool
 
       @connected = false
     end
@@ -363,7 +364,7 @@ def ensure_manager
     end
 
     def pinned_pool
-      thread_local[:pinned_pools][@manager.object_id]
+      thread_local[:pinned_pools][@manager.object_id] if @manager
     end
 
     def pin_pool(pool, read_preference)

test/replica_set/pinning_test.rb

@@ -0,0 +1,41 @@
+require 'test_helper'
+
+class ReplicaSetPinningTest < Test::Unit::TestCase
+  def setup
+    ensure_cluster(:rs)
+    @client = MongoReplicaSetClient.new(@rs.repl_set_seeds, :name => @rs.repl_set_name)
+    @db = @client.db(MONGO_TEST_DB)
+    @coll = @db.collection("test-sets")
+    @coll.insert({:a => 1})
+  end
+
+  def test_unpinning
+    # pin primary
+    @coll.find_one
+    assert_equal @client.pinned_pool[:pool], @client.primary_pool
+
+    # pin secondary
+    @coll.find_one({}, :read => :secondary_preferred)
+    assert @client.secondary_pools.include? @client.pinned_pool[:pool]
+
+    # repin primary
+    @coll.find_one({}, :read => :primary_preferred)
+    assert_equal @client.pinned_pool[:pool], @client.primary_pool
+  end
+
+  def test_pinned_pool_is_local_to_thread
+    threads = []
+    30.times do |i|
+      threads << Thread.new do
+        if i % 2 == 0
+          @coll.find_one({}, :read => :secondary_preferred)
+          assert @client.secondary_pools.include? @client.pinned_pool[:pool]
+        else
+          @coll.find_one({}, :read => :primary_preferred)
+          assert_equal @client.pinned_pool[:pool], @client.primary_pool
+        end
+      end
+    end
+    threads.each(&:join)
+  end
+end