mongo 1.2.0 → 1.2.1

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

data/test/unit/collection_test.rb

@@ -5,6 +5,7 @@ class CollectionTest < Test::Unit::TestCase
   context "Basic operations: " do
     setup do
       @logger = mock()
+      @logger.expects(:debug)
     end
 
     should "send update message" do
@@ -14,7 +15,7 @@ class CollectionTest < Test::Unit::TestCase
       @conn.expects(:send_message).with do |op, msg, log|
         op == 2001
       end
-      @logger.stubs(:debug)
+      @conn.stubs(:log_operation)
       @coll.update({}, {:title => 'Moby Dick'})
     end
 
@@ -25,8 +26,8 @@ class CollectionTest < Test::Unit::TestCase
       @conn.expects(:send_message).with do |op, msg, log|
         op == 2002
       end
-      @logger.expects(:debug).with do |msg|
-        msg.include?("Moby")
+      @conn.expects(:log_operation).with do |name, payload|
+        (name == :insert) && payload[:documents][0][:title].include?('Moby')
       end
       @coll.insert({:title => 'Moby Dick'})
     end
@@ -38,8 +39,8 @@ class CollectionTest < Test::Unit::TestCase
       @conn.expects(:receive_message).with do |op, msg, log, sock|
         op == 2004
       end.returns([[], 0, 0])
-      @logger.expects(:debug).with do |msg|
-        msg.include?('Moby')
+      @conn.expects(:log_operation).with do |name, payload|
+        (name == :find) && payload[:selector][:title].include?('Moby')
       end
       @coll.find({:title => 'Moby Dick'}).sort([['title', 1], ['author', 1]]).next_document
     end
@@ -52,8 +53,8 @@ class CollectionTest < Test::Unit::TestCase
       @conn.expects(:send_message).with do |op, msg, log|
         op == 2002
       end
-      @logger.expects(:debug).with do |msg|
-        msg.include?("Binary")
+      @conn.expects(:log_operation).with do |name, payload|
+        (name == :insert) && payload[:documents][0][:data].inspect.include?('Binary')
       end
       @coll.insert({:data => data})
     end
@@ -65,8 +66,8 @@ class CollectionTest < Test::Unit::TestCase
       @conn.expects(:send_message_with_safe_check).with do |op, msg, db_name, log|
         op == 2001
       end
-      @logger.expects(:debug).with do |msg|
-        msg.include?("testing['books'].update")
+      @conn.expects(:log_operation).with do |name, payload|
+        (name == :update) && payload[:document][:title].include?('Moby')
       end
       @coll.update({}, {:title => 'Moby Dick'}, :safe => true)
     end
@@ -78,43 +79,42 @@ class CollectionTest < Test::Unit::TestCase
       @conn.expects(:send_message_with_safe_check).with do |op, msg, db_name, log|
         op == 2001
       end
-      @logger.stubs(:debug)
+      @conn.stubs(:log_operation)
       @coll.update({}, {:title => 'Moby Dick'}, :safe => true)
     end
-    
+
     should "not call insert for each ensure_index call" do
       @conn = Connection.new('localhost', 27017, :logger => @logger, :connect => false)
       @db   = @conn['testing']
       @coll = @db.collection('books')
       @coll.expects(:generate_indexes).once
-      
+
       @coll.ensure_index [["x", Mongo::DESCENDING]]
       @coll.ensure_index [["x", Mongo::DESCENDING]]
-      
     end
+
     should "call generate_indexes for a new direction on the same field for ensure_index" do
       @conn = Connection.new('localhost', 27017, :logger => @logger, :connect => false)
       @db   = @conn['testing']
       @coll = @db.collection('books')
       @coll.expects(:generate_indexes).twice
-      
+
       @coll.ensure_index [["x", Mongo::DESCENDING]]
       @coll.ensure_index [["x", Mongo::ASCENDING]]
-      
+
     end
-    
+
     should "call generate_indexes twice because the cache time is 0 seconds" do
       @conn = Connection.new('localhost', 27017, :logger => @logger, :connect => false)
       @db   = @conn['testing']
       @db.cache_time = 0
       @coll = @db.collection('books')
       @coll.expects(:generate_indexes).twice
-      
 
       @coll.ensure_index [["x", Mongo::DESCENDING]]
       @coll.ensure_index [["x", Mongo::DESCENDING]]
     end
-    
+
     should "call generate_indexes for each key when calling ensure_indexes" do
       @conn = Connection.new('localhost', 27017, :logger => @logger, :connect => false)
       @db   = @conn['testing']
@@ -123,11 +123,8 @@ class CollectionTest < Test::Unit::TestCase
       @coll.expects(:generate_indexes).once.with do |a, b, c|
         a == {"x"=>-1, "y"=>-1}
       end
-      
+
       @coll.ensure_index [["x", Mongo::DESCENDING], ["y", Mongo::DESCENDING]]
     end
-
-    
-    
   end
 end