mongo 1.1.5 → 1.3.0

data/test/replica_sets/connect_test.rb

@@ -16,8 +16,8 @@ class ConnectTest < Test::Unit::TestCase
 
   def test_connect_with_deprecated_multi
     @conn = Connection.multi([[RS.host, RS.ports[0]], [RS.host, RS.ports[1]]], :name => RS.name)
-    assert @conn.connected?
     assert @conn.is_a?(ReplSetConnection)
+    assert @conn.connected?
   end
 
   def test_connect_bad_name
@@ -31,12 +31,22 @@ class ConnectTest < Test::Unit::TestCase
     @conn = ReplSetConnection.new([RS.host, RS.ports[0]], [RS.host, RS.ports[1]],
       [RS.host, RS.ports[2]], :name => RS.name)
     assert @conn.connected?
+    assert @conn.read_primary?
+    assert @conn.primary?
 
     assert_equal RS.primary, @conn.primary
     assert_equal RS.secondaries.sort, @conn.secondaries.sort
     assert_equal RS.arbiters.sort, @conn.arbiters.sort
   end
 
+  def test_host_port_accessors
+    @conn = ReplSetConnection.new([RS.host, RS.ports[0]], [RS.host, RS.ports[1]],
+      [RS.host, RS.ports[2]], :name => RS.name)
+
+    assert_equal @conn.host, RS.primary[0]
+    assert_equal @conn.port, RS.primary[1]
+  end
+
   def test_connect_with_primary_node_killed
     node = RS.kill_primary
 

data/test/replica_sets/connection_string_test.rb

@@ -0,0 +1,32 @@
+$:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
+require './test/replica_sets/rs_test_helper'
+
+# NOTE: This test expects a replica set of three nodes to be running on RS.host,
+# on ports TEST_PORT, RS.ports[1], and TEST + 2.
+class ConnectionStringTest < Test::Unit::TestCase
+  include Mongo
+
+  def setup
+    RS.restart_killed_nodes
+  end
+
+  def teardown
+    RS.restart_killed_nodes
+  end
+
+  def test_connect_with_connection_string
+    @conn = Connection.from_uri("mongodb://#{RS.host}:#{RS.ports[0]},#{RS.host}:#{RS.ports[1]}?replicaset=#{RS.name}")
+    assert @conn.is_a?(ReplSetConnection)
+    assert @conn.connected?
+  end
+
+  def test_connect_with_full_connection_string
+    @conn = Connection.from_uri("mongodb://#{RS.host}:#{RS.ports[0]},#{RS.host}:#{RS.ports[1]}?replicaset=#{RS.name};safe=true;w=2;fsync=true;slaveok=true")
+    assert @conn.is_a?(ReplSetConnection)
+    assert @conn.connected?
+    assert_equal 2, @conn.safe[:w]
+    assert @conn.safe[:fsync]
+    assert @conn.read_pool
+  end
+
+end

data/test/replica_sets/query_secondaries.rb

@@ -16,6 +16,13 @@ class ReplicaSetQuerySecondariesTest < Test::Unit::TestCase
     RS.restart_killed_nodes
   end
 
+  def test_read_primary
+    rescue_connection_failure do
+      assert !@conn.read_primary?
+      assert !@conn.primary?
+    end
+  end
+
   def test_con
     assert @conn.primary_pool, "No primary pool!"
     assert @conn.read_pool, "No read pool!"
@@ -54,6 +61,12 @@ class ReplicaSetQuerySecondariesTest < Test::Unit::TestCase
     # Should still be able to read immediately after killing master node
     RS.kill_primary
     assert_equal 2, @coll.find.to_a.length
+    rescue_connection_failure do
+      @coll.save({:a => 50}, :safe => {:w => 2, :wtimeout => 10000})
+    end
+    RS.restart_killed_nodes
+    @coll.save({:a => 50}, :safe => {:w => 2, :wtimeout => 10000})
+    assert_equal 4, @coll.find.to_a.length
   end
 
   def test_kill_secondary
@@ -66,6 +79,7 @@ class ReplicaSetQuerySecondariesTest < Test::Unit::TestCase
     RS.kill(read_node)
 
     # Should fail immediately on next read
+    old_read_pool_port = @conn.read_pool.port
     assert_raise ConnectionFailure do
       @coll.find.to_a.length
     end
@@ -75,6 +89,8 @@ class ReplicaSetQuerySecondariesTest < Test::Unit::TestCase
       length = @coll.find.to_a.length
       assert_equal 2, length
     end
+    new_read_pool_port = @conn.read_pool.port
+    assert old_read_pool != new_read_pool
   end
 
 end

data/test/replica_sets/query_test.rb

@@ -27,6 +27,8 @@ class ReplicaSetQueryTest < Test::Unit::TestCase
       assert results.any? {|r| r['a'] == a}, "Could not find record for a => #{a}"
     end
 
+    puts "Benchmark before failover: #{benchmark_queries}"
+
     RS.kill_primary
 
     results = []
@@ -35,7 +37,15 @@ class ReplicaSetQueryTest < Test::Unit::TestCase
       [20, 30, 40].each do |a|
         assert results.any? {|r| r['a'] == a}, "Could not find record for a => #{a}"
       end
+
+    puts "Benchmark after failover: #{benchmark_queries}"
     end
   end
 
+  def benchmark_queries
+    t1 = Time.now
+    10000.times { @coll.find_one }
+    Time.now - t1
+  end
+
 end

data/test/replica_sets/replication_ack_test.rb

@@ -13,6 +13,8 @@ class ReplicaSetAckTest < Test::Unit::TestCase
     @slave1 = Connection.new(@conn.secondary_pools[0].host,
       @conn.secondary_pools[0].port, :slave_ok => true)
 
+    assert !@slave1.read_primary?
+
     @db = @conn.db(MONGO_TEST_DB)
     @db.drop_collection("test-sets")
     @col = @db.collection("test-sets")

data/test/replica_sets/rs_test_helper.rb

@@ -12,17 +12,15 @@ class Test::Unit::TestCase
   # Generic code for rescuing connection failures and retrying operations.
   # This could be combined with some timeout functionality.
   def rescue_connection_failure(max_retries=60)
-    success = false
-    tries   = 0
-    while !success && tries < max_retries
-      begin
-        yield
-        success = true
-      rescue Mongo::ConnectionFailure
-        puts "Rescue attempt #{tries}\n"
-        tries += 1
-        sleep(1)
-      end
+    retries = 0
+    begin
+      yield
+    rescue Mongo::ConnectionFailure => ex
+      puts "Rescue attempt #{retries}: from #{ex}"
+      retries += 1
+      raise ex if retries > max_retries
+      sleep(1)
+      retry
     end
   end
 

data/test/support/hash_with_indifferent_access.rb

@@ -38,19 +38,6 @@ class Hash
 
   alias_method :to_options,  :symbolize_keys
   #alias_method :to_options!, :symbolize_keys!
-
-  # Validate all keys in a hash match *valid keys, raising ArgumentError on a mismatch.
-  # Note that keys are NOT treated indifferently, meaning if you use strings for keys but assert symbols
-  # as keys, this will fail.
-  #
-  # ==== Examples
-  #   { :name => "Rob", :years => "28" }.assert_valid_keys(:name, :age) # => raises "ArgumentError: Unknown key(s): years"
-  #   { :name => "Rob", :age => "28" }.assert_valid_keys("name", "age") # => raises "ArgumentError: Unknown key(s): name, age"
-  #   { :name => "Rob", :age => "28" }.assert_valid_keys(:name, :age) # => passes, raises nothing
-  def assert_valid_keys(*valid_keys)
-    unknown_keys = keys - [valid_keys].flatten
-    raise(ArgumentError, "Unknown key(s): #{unknown_keys.join(", ")}") unless unknown_keys.empty?
-  end
 end
 
 module ActiveSupport

data/test/support_test.rb

@@ -1,7 +1,6 @@
 require './test/test_helper'
 
 class SupportTest < Test::Unit::TestCase
-  include Mongo
 
   def test_command_response_succeeds
     assert Support.ok?('ok' => 1)

data/test/test_helper.rb

@@ -1,14 +1,22 @@
 $:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
-require 'rubygems' if ENV['C_EXT']
+require 'rubygems' if RUBY_VERSION < '1.9.0' && ENV['C_EXT']
 require 'mongo'
 require 'test/unit'
 
+def silently
+  warn_level = $VERBOSE
+  $VERBOSE = nil
+  result = yield
+  $VERBOSE = warn_level
+  result
+end
+
 begin
-  require 'rubygems'
-  require 'shoulda'
-  require 'mocha'
-  rescue LoadError
-    puts <<MSG
+  require 'rubygems' if RUBY_VERSION < "1.9.0" && !ENV['C_EXT']
+  silently { require 'shoulda' }
+  silently { require 'mocha' }
+rescue LoadError
+  puts <<MSG
 
 This test suite requires shoulda and mocha.
 You can install them as follows:
@@ -16,7 +24,8 @@ You can install them as follows:
   gem install mocha
 
 MSG
-    exit
+
+  exit
 end
 
 require 'bson_ext/cbson' if !(RUBY_PLATFORM =~ /java/) && ENV['C_EXT']
@@ -69,7 +78,17 @@ class Test::Unit::TestCase
     self.class.mongo_port
   end
 
-  
+  def new_mock_socket(host='localhost', port=27017)
+    socket = Object.new
+    socket.stubs(:setsockopt).with(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+    socket.stubs(:close)
+    socket
+  end
+
+  def new_mock_db
+    db = Object.new
+  end
+
   def assert_raise_error(klass, message)
     begin
       yield

data/test/tools/auth_repl_set_manager.rb

@@ -0,0 +1,14 @@
+require File.join((File.expand_path(File.dirname(__FILE__))), 'repl_set_manager')
+
+class AuthReplSetManager < ReplSetManager
+  def initialize(opts={})
+    super(opts)
+
+    @key_path = opts[:key_path] || File.join(File.expand_path(File.dirname(__FILE__)), "keyfile.txt")
+    system("chmod 600 #{@key_path}")
+  end
+
+  def start_cmd(n)
+    super + " --keyFile #{@key_path}"
+  end
+end

data/test/tools/load.rb

@@ -0,0 +1,58 @@
+require 'rubygems'
+require 'mongo'
+require 'sharding_manager'
+
+class MongoLoader
+
+  def initialize
+    @mongo = Mongo::Connection.new("localhost", 50000)
+    @data  = BSON::Binary.new(File.open("tools.gz").read)
+    @count = 0
+    @manager = ShardingManager.new(:config_count => 3)
+    @manager.start_cluster
+  end
+
+  def kill
+    @manager.kill_random
+  end
+
+  def restart
+    @manager.restart_killed_nodes
+  end
+
+  def run
+    Thread.new do
+      ("a".."z").each do |p|
+        seed(p)
+      end
+    end
+  end
+
+  def seed(prefix)
+    @queue = []
+    1000.times do |n|
+      id = BSON::OrderedHash.new
+      id[:p] = prefix
+      id[:c] = n
+      @queue << {:tid => id, :data => @data}
+    end
+
+    while @queue.length > 0 do
+      begin
+        doc = @queue.pop
+        @mongo['app']['photos'].insert(doc, :safe => {:w => 3})
+        @count += 1
+        p @count
+        rescue StandardError => e
+          p e
+          p @count
+          @queue.push(doc)
+          @count -= 1
+          sleep(10)
+          retry
+      end
+    end
+  end
+end
+
+@m = MongoLoader.new

data/test/tools/repl_set_manager.rb

@@ -1,5 +1,7 @@
 #!/usr/bin/ruby
 
+require 'thread'
+
 STDOUT.sync = true
 
 unless defined? Mongo
@@ -17,6 +19,7 @@ class ReplSetManager
     @host       = opts[:host]  || 'localhost'
     @retries    = opts[:retries] || 60
     @config     = {"_id" => @name, "members" => []}
+    @durable    = opts.fetch(:durable, false)
     @path       = File.join(File.expand_path(File.dirname(__FILE__)), "data")
 
     @arbiter_count   = opts[:arbiter_count]   || 2
@@ -61,6 +64,13 @@ class ReplSetManager
     ensure_up
   end
 
+  def cleanup_set
+    system("killall mongod")
+    @count.times do |n|
+      system("rm -rf #{@mongods[n]['db_path']}")
+    end
+  end
+
   def init_node(n)
     @mongods[n] ||= {}
     port = @start_port + n
@@ -71,9 +81,7 @@ class ReplSetManager
     system("rm -rf #{@mongods[n]['db_path']}")
     system("mkdir -p #{@mongods[n]['db_path']}")
 
-    @mongods[n]['start'] = "mongod --replSet #{@name} --logpath '#{@mongods[n]['log_path']}' " +
-     " --dbpath #{@mongods[n]['db_path']} --port #{@mongods[n]['port']} --fork"
-
+    @mongods[n]['start'] = start_cmd(n)
     start(n)
 
     member = {'_id' => n, 'host' => "#{@host}:#{@mongods[n]['port']}"}
@@ -88,17 +96,24 @@ class ReplSetManager
     @config['members'] << member
   end
 
-  def kill(node)
+  def start_cmd(n)
+    @mongods[n]['start'] = "mongod --replSet #{@name} --logpath '#{@mongods[n]['log_path']}' " +
+     " --dbpath #{@mongods[n]['db_path']} --port #{@mongods[n]['port']} --fork"
+    @mongods[n]['start'] += " --dur" if @durable
+    @mongods[n]['start']
+  end
+
+  def kill(node, signal=2)
     pid = @mongods[node]['pid']
     puts "** Killing node with pid #{pid} at port #{@mongods[node]['port']}"
-    system("kill -2 #{@mongods[node]['pid']}")
+    system("kill -#{signal} #{@mongods[node]['pid']}")
     @mongods[node]['up'] = false
     sleep(1)
   end
 
-  def kill_primary
+  def kill_primary(signal=2)
     node = get_node_with_state(1)
-    kill(node)
+    kill(node, signal)
     return node
   end
 
@@ -174,6 +189,16 @@ class ReplSetManager
     get_all_host_pairs_with_state(7)
   end
 
+  # String used for adding a shard via mongos
+  # using the addshard command.
+  def shard_string
+    str = "#{@name}/"
+    str << @mongods.map do |k, mongod|
+      "#{@host}:#{mongod['port']}"
+    end.join(',')
+    str
+  end
+
   private
 
   def initiate
@@ -229,13 +254,13 @@ class ReplSetManager
     while count < @retries do
       begin
         return yield
-        rescue Mongo::OperationFailure, Mongo::ConnectionFailure
+        rescue Mongo::OperationFailure, Mongo::ConnectionFailure => ex
           sleep(1)
           count += 1
       end
     end
 
-    raise exception
+    raise ex
   end
 
 end

data/test/tools/sharding_manager.rb

@@ -0,0 +1,202 @@
+require 'repl_set_manager'
+require 'thread'
+
+class ShardingManager
+
+  attr_accessor :shards
+
+  def initialize(opts={})
+    @durable = opts.fetch(:durable, true)
+    @host    = "localhost"
+
+    @mongos_port = opts[:mongos_port] || 50000
+    @config_port = opts[:config_port] || 40000
+    @shard_start_port = opts[:start_shard_port] || 30000
+    @path       = File.join(File.expand_path(File.dirname(__FILE__)), "data")
+    system("rm -rf #{@path}")
+
+    @shard_count  = 2
+    @mongos_count = 1
+    @config_count = opts.fetch(:config_count, 1)
+    if ![1, 3].include?(@config_count)
+      raise ArgumentError, "Must specify 1 or 3 config servers."
+    end
+
+    @config_servers = {}
+    @mongos_servers = {}
+    @shards = []
+    @ports  = []
+  end
+
+  def kill_random
+    shard_to_kill = rand(@shard_count)
+    @shards[shard_to_kill].kill_primary
+  end
+
+  def restart_killed
+    threads = []
+    @shards.each do |k, shard|
+      threads << Thread.new do
+        shard.restart_killed_nodes
+      end
+    end
+  end
+
+  def start_cluster
+    start_sharding_components
+    start_mongos_servers
+    configure_cluster
+  end
+
+  def configure_cluster
+    add_shards
+    enable_sharding
+    shard_collection
+  end
+
+  def enable_sharding
+    mongos['admin'].command({:enablesharding => "app"})
+  end
+
+  def shard_collection
+    cmd = BSON::OrderedHash.new
+    cmd[:shardcollection] = "app.photos"
+    cmd[:key] = {:tid => 1}
+    p mongos['admin'].command(cmd)
+  end
+
+  def add_shards
+    @shards.each do |shard|
+      cmd = {:addshard => shard.shard_string}
+      p cmd
+      p mongos['admin'].command(cmd)
+    end
+    p mongos['admin'].command({:listshards => 1})
+  end
+
+  def mongos
+    attempt do
+      @mongos ||= Mongo::Connection.new(@host, @mongos_servers[0]['port'])
+    end
+  end
+
+  private
+
+  def start_sharding_components
+    system("killall mongos")
+
+    threads = []
+    threads << Thread.new do
+      start_shards
+    end
+
+    threads << Thread.new do
+      start_config_servers
+    end
+    threads.each {|t| t.join}
+    puts "\nShards and config servers up!"
+  end
+
+  def start_shards
+    threads = []
+    @shard_count.times do |n|
+        threads << Thread.new do
+        port = @shard_start_port + n * 100
+        shard = ReplSetManager.new(:arbiter_count => 0, :secondary_count => 2,
+                       :passive_count => 0, :start_port => port, :durable => @durable,
+                       :name => "shard-#{n}")
+        shard.start_set
+        shard.ensure_up
+        @shards << shard
+      end
+    end
+    threads.each {|t| t.join}
+  end
+
+  def start_config_servers
+    @config_count.times do |n|
+      @config_servers[n] ||= {}
+      port = @config_port + n
+      @ports << port
+      @config_servers[n]['port'] = port
+      @config_servers[n]['db_path'] = get_path("config-#{port}")
+      @config_servers[n]['log_path'] = get_path("log-config-#{port}")
+      system("rm -rf #{@config_servers[n]['db_path']}")
+      system("mkdir -p #{@config_servers[n]['db_path']}")
+
+      @config_servers[n]['start'] = start_config_cmd(n)
+
+      start(@config_servers, n)
+    end
+  end
+
+  def start_mongos_servers
+    @mongos_count.times do |n|
+      @mongos_servers[n] ||= {}
+      port = @mongos_port + n
+      @ports << port
+      @mongos_servers[n]['port'] = port
+      @mongos_servers[n]['db_path'] = get_path("mongos-#{port}")
+      @mongos_servers[n]['pidfile_path'] = File.join(@mongos_servers[n]['db_path'], "mongod.lock")
+      @mongos_servers[n]['log_path'] = get_path("log-mongos-#{port}")
+      system("rm -rf #{@mongos_servers[n]['db_path']}")
+      system("mkdir -p #{@mongos_servers[n]['db_path']}")
+
+      @mongos_servers[n]['start'] = start_mongos_cmd(n)
+
+      start(@mongos_servers, n)
+    end
+  end
+
+  def start_config_cmd(n)
+    cmd = "mongod --configsvr --logpath '#{@config_servers[n]['log_path']}' " +
+     " --dbpath #{@config_servers[n]['db_path']} --port #{@config_servers[n]['port']} --fork"
+    cmd += " --dur" if @durable
+    cmd
+  end
+
+  def start_mongos_cmd(n)
+    "mongos --configdb #{config_db_string} --logpath '#{@mongos_servers[n]['log_path']}' " +
+      "--pidfilepath #{@mongos_servers[n]['pidfile_path']} --port #{@mongos_servers[n]['port']} --fork"
+  end
+
+  def config_db_string
+    @config_servers.map do |k, v|
+      "#{@host}:#{v['port']}"
+    end.join(',')
+  end
+
+  def start(set, node)
+    system(set[node]['start'])
+    set[node]['up'] = true
+    sleep(0.5)
+    set[node]['pid'] = File.open(File.join(set[node]['db_path'], 'mongod.lock')).read.strip
+  end
+  alias :restart :start
+
+  private
+
+  def cleanup_config
+  end
+
+  def get_path(name)
+    File.join(@path, name)
+  end
+
+  # TODO: put this into a shared module
+  def attempt
+    raise "No block given!" unless block_given?
+    count = 0
+
+    while count < 50 do
+      begin
+        return yield
+        rescue Mongo::OperationFailure, Mongo::ConnectionFailure
+          sleep(1)
+          count += 1
+      end
+    end
+
+    raise exception
+  end
+end

data/test/tools/test.rb

@@ -1,13 +1,4 @@
-require 'repl_set_manager'
+require 'sharding_manager'
 
-m = ReplSetManager.new
-m.start_set
-
-node = m.kill_secondary
-m.ensure_up
-
-puts "Pausing..."
-gets
-
-m.start(node)
-m.ensure_up
+m = ShardingManager.new(:config_count => 3)
+m.start_cluster