mongo 1.3.0.rc0 → 1.3.0

data/test/async/worker_pool_test.rb

@@ -0,0 +1,99 @@
+require 'test/test_helper'
+include Mongo
+
+class WorkerPoolTest < Test::Unit::TestCase
+  context "Initialization: " do
+
+    def wait_for_async
+      sleep 0.2
+    end
+
+    setup do
+      def new_mock_queue
+        stub_everything('queue')
+      end
+
+      def new_mock_thread
+        stub_everything('thread')
+      end
+    end
+
+    context "given a size" do
+      setup do
+        @size = 5
+      end
+
+      should "allocate a Thread 'size' times" do
+        Queue.stubs(:new).returns(new_mock_queue)
+        Thread.expects(:new).times(@size).returns(new_mock_thread)
+        Async::WorkerPool.new @size
+      end
+
+      should "set 'abort_on_exception' for each current thread" do
+        Queue.stubs(:new).returns(new_mock_queue)
+        thread = new_mock_thread
+        Thread.stubs(:new).returns(thread)
+
+        thread.expects(:abort_on_exception=).with(true).times(@size)
+
+        Async::WorkerPool.new @size
+      end
+
+      should "save each thread into the workers queue" do
+        assert_equal @size, Async::WorkerPool.new(@size).workers.size
+      end
+
+    end # context 'given a size'
+
+
+    context "given a job" do
+      setup do
+        @pool = Async::WorkerPool.new 1
+        @command = stub_everything('command')
+        @cmd_args = stub_everything('command args')
+        @callback = stub_everything('callback')
+      end
+
+      should "remove nils from the command args array and pass the results to the callback" do
+        args = [nil, @cmd_args]
+        @command.expects(:call).with(@cmd_args).returns(2)
+        @callback.expects(:call).with(nil, 2)
+
+        @pool.enqueue @command, args, @callback
+        wait_for_async
+      end
+
+      should "execute the original command with args and pass the results to the callback" do
+        @cmd_args.expects(:compact).returns(@cmd_args)
+        @command.expects(:call).with(@cmd_args).returns(2)
+        @callback.expects(:call).with(nil, 2)
+
+        @pool.enqueue @command, @cmd_args, @callback
+        wait_for_async
+      end
+
+      should "capture any exceptions and pass them to the callback" do
+        args = [@cmd_args]
+        error = StandardError.new
+        @command.expects(:call).with(@cmd_args).raises(error)
+        @callback.expects(:call).with(error, nil)
+
+        @pool.enqueue @command, args, @callback
+        wait_for_async
+      end
+
+      should "abort the thread when the callback raises an exception" do
+        args = [@cmd_args]
+        error = StandardError.new
+        @callback.expects(:call).raises(error)
+
+        assert_raises(StandardError) do
+          @pool.enqueue @command, args, @callback
+          wait_for_async
+        end
+      end
+    end # context 'given a job'
+
+
+  end
+end

data/test/bson/byte_buffer_test.rb

@@ -21,6 +21,13 @@ class ByteBufferTest < Test::Unit::TestCase
     assert_equal 1, @buf.get
   end
 
+  def test_unpack
+    @buf.put_array([17, 2, 3, 4])
+    assert_equal [17, 2, 3, 4], @buf.to_a
+    assert_equal ["11020304"], @buf.unpack("H*")
+    assert_equal ["11020304"], @buf.to_a("H*")
+  end
+
   def test_one_get_returns_array_length_one
     @buf.put_array([1, 2, 3, 4])
     @buf.rewind

data/test/load/resque/load.rb

@@ -0,0 +1,21 @@
+require File.join(File.dirname(__FILE__), '..', '..', '..', 'lib', 'mongo')
+require 'logger'
+require 'rubygems'
+require 'resque'
+require 'sinatra'
+require File.join(File.dirname(__FILE__), 'processor')
+
+$con = Mongo::Connection.new
+$db = $con['foo']
+
+
+configure do
+  LOGGER = Logger.new("sinatra.log")
+  enable :logging, :dump_errors
+  set :raise_errors, true
+end
+
+get '/' do
+  Processor.perform(1)
+  true
+end

data/test/load/resque/processor.rb

@@ -0,0 +1,26 @@
+require 'logger'
+
+class Processor
+  @queue = :processor
+
+  def self.connection
+    @log ||= Logger.new(STDOUT)
+    @con ||= Mongo::Connection.new("localhost", 27017)
+  end
+
+  def self.perform(n)
+    begin
+    100.times do |n|
+      self.connection['resque']['docs'].insert({:n => n, :data => "0" * 1000}, :safe => true)
+    end
+
+    5.times do |n|
+      num = rand(100)
+      self.connection['resque']['docs'].find({:n => {"$gt" => num}}).limit(1).to_a
+    end
+    rescue => e
+      @log.warn(e.inspect)
+    end
+  end
+
+end

data/test/load/unicorn/unicorn.rb

@@ -0,0 +1,29 @@
+# set path to app that will be used to configure unicorn, 
+# # note the trailing slash in this example
+@dir = "/home/kyle/work/10gen/ruby-driver/test/load/"
+
+worker_processes 10
+working_directory @dir
+
+preload_app true
+
+timeout 30
+
+# Specify path to socket unicorn listens to,
+# we will use this in our nginx.conf later
+listen "#{@dir}tmp/sockets/unicorn.sock", :backlog => 64
+
+# Set process id path
+pid "#{@dir}tmp/pids/unicorn.pid"
+
+# # Set log file paths
+stderr_path "#{@dir}log/unicorn.stderr.log"
+stdout_path "#{@dir}log/unicorn.stdout.log" 
+
+# NOTE: You need this when using forking web servers!
+after_fork do |server, worker|
+  $con.close if $con
+  $con = Mongo::Connection.new
+  $db = $con['foo']
+  STDERR << "FORKED #{server} #{worker}"
+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/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