RUBY-558 batch size on initial query, RUBY-505 test intervening refresh

lib/mongo/cursor.rb

@@ -552,7 +552,7 @@ def construct_query_message
       message.put_int(@options)
       BSON::BSON_RUBY.serialize_cstr(message, "#{@db.name}.#{@collection.name}")
       message.put_int(@skip)
-      message.put_int(@limit)
+      @batch_size > 1 ? message.put_int(@batch_size) : message.put_int(@limit)
       spec = query_contains_special_fields? ? construct_query_spec : @selector
       message.put_binary(BSON::BSON_CODER.serialize(spec, false, false, @connection.max_bson_size).to_s)
       message.put_binary(BSON::BSON_CODER.serialize(@fields, false, false, @connection.max_bson_size).to_s) if @fields

test/functional/collection_test.rb

@@ -902,6 +902,66 @@ def test_save_symbol_find_string
     # assert_equal :mike, @@test.find_one("foo" => "mike")["foo"]
   end
 
+  def test_batch_size
+    n_docs = 6
+    batch_size = n_docs/2
+    n_docs.times do |i|
+      @@test.save(:foo => i)
+    end
+
+    doc_count = 0
+    cursor = @@test.find({}, :batch_size => batch_size)
+    cursor.next
+    assert_equal batch_size, cursor.instance_variable_get(:@returned)
+    doc_count += batch_size
+    batch_size.times { cursor.next }
+    assert_equal doc_count + batch_size, cursor.instance_variable_get(:@returned)
+    doc_count += batch_size
+    assert_equal n_docs, doc_count
+  end
+
+  def test_batch_size_with_smaller_limit
+    n_docs = 6
+    batch_size = n_docs/2
+    n_docs.times do |i|
+      @@test.insert(:foo => i)
+    end
+
+    cursor = @@test.find({}, :batch_size => batch_size, :limit => 2)
+    cursor.next
+    assert_equal 2, cursor.instance_variable_get(:@returned)
+  end
+
+  def test_batch_size_with_larger_limit
+    n_docs = 6
+    batch_size = n_docs/2
+    n_docs.times do |i|
+      @@test.insert(:foo => i)
+    end
+
+    doc_count = 0
+    cursor = @@test.find({}, :batch_size => batch_size, :limit => n_docs + 5)
+    cursor.next
+    assert_equal batch_size, cursor.instance_variable_get(:@returned)
+    doc_count += batch_size
+    batch_size.times { cursor.next }
+    assert_equal doc_count + batch_size, cursor.instance_variable_get(:@returned)
+    doc_count += batch_size
+    assert_equal n_docs, doc_count
+end
+
+  def test_batch_size_with_negative_limit
+    n_docs = 6
+    batch_size = n_docs/2
+    n_docs.times do |i|
+      @@test.insert(:foo => i)
+    end
+
+    cursor = @@test.find({}, :batch_size => batch_size, :limit => -7)
+    cursor.next
+    assert_equal n_docs, cursor.instance_variable_get(:@returned)
+  end
+
   def test_limit_and_skip
     10.times do |i|
       @@test.save(:foo => i)

test/replica_set/cursor_test.rb

@@ -41,20 +41,21 @@ def test_cursors_get_closed_secondary_query
     assert_cursors_on_members(:secondary)
   end
 
+  def test_intervening_query_secondary
+    setup_client(:primary)
+    refresh_while_iterating(:secondary)
+  end
+
   private
 
   def setup_client(read=:primary)
-    route_read ||= read
     # Setup ReplicaSet Connection
     @client = MongoReplicaSetClient.new(@rs.repl_set_seeds, :read => read)
 
     @db = @client.db(MONGO_TEST_DB)
     @db.drop_collection("cursor_tests")
     @coll = @db.collection("cursor_tests")
     insert_docs
-
-    # Setup Direct Connections
-    @primary = Mongo::MongoClient.new(*@client.manager.primary)
   end
 
   def insert_docs
@@ -169,4 +170,25 @@ def assert_cursors_on_members(read=:primary)
       cursor_clone.next
     end
   end
+
+  def refresh_while_iterating(read)
+    set_read_client_and_tag(read)
+
+    read_opts = {:read => read}
+    read_opts[:tag_sets] = [{:node => @tag}]
+    read_opts[:batch_size] = 2
+    cursor = @coll.find({}, read_opts)
+
+    2.times { cursor.next }
+    port = cursor.instance_variable_get(:@pool).port
+    host = cursor.instance_variable_get(:@pool).host
+    # Refresh connection
+    @client.refresh
+    assert_nothing_raised do
+      cursor.next
+    end
+
+    assert_equal port, cursor.instance_variable_get(:@pool).port
+    assert_equal host, cursor.instance_variable_get(:@pool).host
+  end
 end

test/replica_set/refresh_test.rb

@@ -1,104 +1,81 @@
 require 'test_helper'
-require 'benchmark'
 
 class ReplicaSetRefreshTest < Test::Unit::TestCase
 
   def setup
     ensure_cluster(:rs)
   end
 
-  def test_connect_and_manual_refresh_with_secondaries_down
+  def test_connect_and_manual_refresh_with_secondary_down
     num_secondaries = @rs.secondaries.size
     client = MongoReplicaSetClient.new(@rs.repl_set_seeds, :refresh_mode => false)
 
     assert_equal num_secondaries, client.secondaries.size
     assert client.connected?
     assert_equal client.read_pool, client.primary_pool
+    old_refresh_version = client.refresh_version
 
-    @rs.secondaries.each {|s| s.stop}
+    @rs.stop_secondary
 
     client.refresh
-    assert client.secondaries.empty?
+    assert_equal num_secondaries - 1, client.secondaries.size
     assert client.connected?
     assert_equal client.read_pool, client.primary_pool
+    assert client.refresh_version > old_refresh_version
+    old_refresh_version = client.refresh_version
 
     # Test no changes after restart until manual refresh
     @rs.restart
-    assert client.secondaries.empty?
+    assert_equal num_secondaries - 1, client.secondaries.size
     assert client.connected?
     assert_equal client.read_pool, client.primary_pool
+    assert_equal client.refresh_version, old_refresh_version
 
     # Refresh and ensure state
     client.refresh
-    assert_equal client.read_pool, client.primary_pool
     assert_equal num_secondaries, client.secondaries.size
+    assert client.connected?
+    assert_equal client.read_pool, client.primary_pool
+    assert client.refresh_version > old_refresh_version
   end
 
-  def test_automated_refresh_with_secondaries_down
+  def test_automated_refresh_with_secondary_down
     num_secondaries = @rs.secondaries.size
     client = MongoReplicaSetClient.new(@rs.repl_set_seeds,
       :refresh_interval => 1, :refresh_mode => :sync, :read => :secondary_preferred)
 
-    # Ensure secondaries not available and read from primary
+    # Ensure secondaries are all recognized by client and client is connected
     assert_equal num_secondaries, client.secondaries.size
     assert client.connected?
     assert client.secondary_pools.include?(client.read_pool)
+    pool = client.read_pool
 
-    @rs.secondaries.each{|s| s.stop}
+    @rs.member_by_name(pool.host_string).stop
     sleep(2)
 
+    old_refresh_version = client.refresh_version
+    # Trigger synchronous refresh
     client['foo']['bar'].find_one
 
-    assert client.secondaries.empty?
     assert client.connected?
-    assert_equal client.read_pool, client.primary_pool
-
-    old_refresh_version = client.refresh_version
+    assert client.refresh_version > old_refresh_version
+    assert_equal num_secondaries - 1, client.secondaries.size
+    assert client.secondary_pools.include?(client.read_pool)
+    assert_not_equal pool, client.read_pool
 
     # Restart nodes and ensure refresh interval has passed
     @rs.restart
     sleep(2)
 
-    assert client.refresh_version == old_refresh_version,
-      "Refresh version has changed."
-
+    old_refresh_version = client.refresh_version
     # Trigger synchronous refresh
     client['foo']['bar'].find_one
 
+    assert client.connected?
     assert client.refresh_version > old_refresh_version,
       "Refresh version hasn't changed."
     assert_equal num_secondaries, client.secondaries.size
       "No secondaries have been added."
-    assert client.manager.read_pool != client.manager.primary,
-      "Read pool and primary pool are identical."
-  end
-
-  def test_automated_refresh_when_secondary_goes_down
-    client = MongoReplicaSetClient.new(@rs.repl_set_seeds,
-      :refresh_interval => 1, :refresh_mode => :sync)
-
-    num_secondaries = client.secondary_pools.size
-    old_refresh_version = client.refresh_version
-
-    @rs.stop_secondary
-    sleep(2)
-
-    assert client.refresh_version == old_refresh_version,
-      "Refresh version has changed."
-
-    client['foo']['bar'].find_one
-
-    assert client.refresh_version > old_refresh_version,
-      "Refresh version hasn't changed."
-    assert_equal num_secondaries - 1, client.secondaries.size
-    assert_equal num_secondaries - 1, client.secondary_pools.size
-
-    @rs.restart
-    sleep(2)
-
-    client['foo']['bar'].find_one
-
-    assert_equal num_secondaries, client.secondaries.size
     assert_equal num_secondaries, client.secondary_pools.size
   end
 =begin

test/tools/mongo_config.rb

@@ -430,10 +430,14 @@ def arbiter_names
 
       def members_by_name(names)
         names.collect do |name|
-          servers.find{|server| server.host_port == name}
+          member_by_name(name)
         end.compact
       end
 
+      def member_by_name(name)
+        servers.find{|server| server.host_port == name}
+      end
+
       def primary
         members_by_name([primary_name]).first
       end