From b70c9ce1526e7221226485b6f65df5484b4f3ee0 Mon Sep 17 00:00:00 2001
From: Tyler Brock <tyler.brock@gmail.com>
Date: Sat, 18 Feb 2012 16:51:57 -0500
Subject: [PATCH] RUBY-406 enhancements to :secondary_only read preference

Improved implementation
-- read_preference :secondary_only is now communicated via invocation
of ReplSetConnection#checkout_secondary

Better tests
-- ensures reads go to secondaries
-- ensures reads do not go to primaries
---
 lib/mongo/cursor.rb                       |  2 ++
 lib/mongo/repl_set_connection.rb          | 26 +++++++++++++++++++++-
 lib/mongo/util/pool_manager.rb            | 13 ++++++-----
 test/replica_sets/read_preference_test.rb | 27 ++++++++++++++++++-----
 4 files changed, 57 insertions(+), 11 deletions(-)

diff --git a/lib/mongo/cursor.rb b/lib/mongo/cursor.rb
index ddbce9d..dc6fea9 100644
--- a/lib/mongo/cursor.rb
+++ b/lib/mongo/cursor.rb
@@ -522,6 +522,8 @@ module Mongo
         @checkin_connection = true
         if @command || @read_preference == :primary
           socket = @connection.checkout_writer
+        elsif @read_preference == :secondary_only
+          socket = @connection.checkout_secondary
         else
           @read_pool = @connection.read_pool
           socket = @connection.checkout_reader
diff --git a/lib/mongo/repl_set_connection.rb b/lib/mongo/repl_set_connection.rb
index 1443cab..5591024 100644
--- a/lib/mongo/repl_set_connection.rb
+++ b/lib/mongo/repl_set_connection.rb
@@ -321,7 +321,7 @@ module Mongo
       end
       begin
         socket = get_socket_from_pool(self.read_pool)
-        if !socket && @read != :secondary_only
+        if !socket
           connect
           socket = get_socket_from_pool(self.primary_pool)
         end
@@ -336,6 +336,26 @@ module Mongo
         raise ConnectionFailure.new("Could not connect to a node for reading.")
       end
     end
+    
+    def checkout_secondary
+      if connected?
+        sync_refresh
+      else
+        connect
+      end
+      begin
+        socket = get_socket_from_pool(self.secondary_pool)
+      rescue => ex
+        checkin(socket) if socket
+        raise ex
+      end
+
+      if socket
+        socket
+      else
+        raise ConnectionFailure.new("Could not connect to a secondary for reading.")
+      end
+    end
 
     # Checkout a socket for writing (i.e., a primary node).
     def checkout_writer
@@ -424,6 +444,10 @@ module Mongo
     def read_pool
       @manager ? @manager.read_pool : nil
     end
+    
+    def secondary_pool
+      @manager ? @manager.secondary_pool : nil
+    end
 
     def secondary_pools
       @manager ? @manager.secondary_pools : []
diff --git a/lib/mongo/util/pool_manager.rb b/lib/mongo/util/pool_manager.rb
index cb26172..0525ec4 100644
--- a/lib/mongo/util/pool_manager.rb
+++ b/lib/mongo/util/pool_manager.rb
@@ -2,8 +2,8 @@ module Mongo
   class PoolManager
 
     attr_reader :connection, :arbiters, :primary, :secondaries, :primary_pool,
-      :read_pool, :secondary_pools, :hosts, :nodes, :max_bson_size,
-      :tags_to_pools, :tag_map, :members
+      :read_pool, :secondary_pool, :secondary_pools, :hosts, :nodes,
+      :max_bson_size, :tags_to_pools, :tag_map, :members
 
     # Create a new set of connection pools.
     #
@@ -140,6 +140,7 @@ module Mongo
       @read_pool = nil
       @arbiters = []
       @secondaries = []
+      @secondary_pool = nil
       @secondary_pools = []
       @hosts = Set.new
       @members = Set.new
@@ -232,12 +233,14 @@ module Mongo
     # If more than one node is available, use the ping
     # time to figure out which nodes to choose from.
     def set_read_pool
-      if @secondary_pools.empty? && @connection.read_preference != :secondary_only
-         @read_pool = @primary_pool
+      if @secondary_pools.empty?
+        @read_pool = @primary_pool
       elsif @secondary_pools.size == 1
-         @read_pool = @secondary_pools[0]
+        @read_pool = @secondary_pools[0]
+        @secondary_pool = @read_pool
       else
         @read_pool = nearby_pool_from_set(@secondary_pools)
+        @secondary_pool = @read_pool
       end
     end
 
diff --git a/test/replica_sets/read_preference_test.rb b/test/replica_sets/read_preference_test.rb
index 1981ceb..bff4874 100644
--- a/test/replica_sets/read_preference_test.rb
+++ b/test/replica_sets/read_preference_test.rb
@@ -34,17 +34,34 @@ class ReadPreferenceTest < Test::Unit::TestCase
   end
 
   def test_read_secondary_only
-    @conn = ReplSetConnection.new([@rs.host, @rs.ports[0]], [@rs.host, @rs.ports[1]], :read => :secondary_only)
-    assert_equal @conn.read_preference, :secondary_only
+    @rs.add_arbiter
+    @rs.remove_secondary_node
     
+    @conn = ReplSetConnection.new(["#{@rs.host}:#{@rs.ports[0]}","#{@rs.host}:#{@rs.ports[1]}"],
+      :read => :secondary_only)
+
     @db = @conn.db(MONGO_TEST_DB)
     @coll = @db.collection("test-sets")
-    @coll.save({:a => 20})
-    @rs.kill_all_secondaries
     
-    assert_raise ConnectionFailure do
+    @coll.save({:a => 20}, :safe => {:w => 2})
+
+    # Test that reads are going to secondary on ReplSetConnection
+    @secondary = Connection.new(@rs.host, @conn.read_pool.port, :slave_ok => true)
+    queries_before = @secondary['admin'].command({:serverStatus => 1})['opcounters']['query']
+    @coll.find_one
+    queries_after = @secondary['admin'].command({:serverStatus => 1})['opcounters']['query']
+    assert_equal 1, queries_after - queries_before
+
+    @rs.kill_secondary
+    @conn.refresh
+    
+    # Test that reads are only allowed from secondaries
+    assert_raise ConnectionFailure.new("Could not connect to a secondary for reading.") do
       @coll.find_one
     end
+        
+    @rs = ReplSetManager.new
+    @rs.start_set
   end
 
   def test_query_secondaries