Mongoid / Mongo bench with Celluloid::IO + nio4r

mongoid_bench.rb

=begin
# Gemfile:
source :rubygems
gem 'celluloid',      :git => 'git://github.com/tarcieri/celluloid.git'
gem 'celluloid-io',   :git => 'git://github.com/tarcieri/celluloid-io.git'
gem 'mongoid'
gem 'bson_ext'
=end

require 'bundler'
Bundler.setup
# require 'ruby-debug'

require 'celluloid/io'
require 'mongoid'
require 'benchmark'

NUM_DOCUMENTS = 100 # Adjust here...

# Patch Mongo driver to use Celluloid's TCPSocket class...
# from version mongo-ruby-driver 1.6.0
module Mongo
  class Connection
    def setup(opts)
      # Default maximum BSON object size
      @max_bson_size = Mongo::DEFAULT_MAX_BSON_SIZE

      # Determine whether to use SSL.
      @ssl = opts.fetch(:ssl, false)
      if @ssl
        @socket_class = Mongo::SSLSocket
      else
        # Hax...
        # @socket_class = ::TCPSocket
        @socket_class = Celluloid::IO::TCPSocket
      end

      # Authentication objects
      @auths = opts.fetch(:auths, [])

      # Lock for request ids.
      @id_lock = Mutex.new

      # Pool size and timeout.
      @pool_size = opts[:pool_size] || 1
      if opts[:timeout]
        warn "The :timeout option has been deprecated " +
          "and will be removed in the 2.0 release. Use :pool_timeout instead."
      end
      @pool_timeout = opts[:pool_timeout] || opts[:timeout] || 5.0

      # Timeout on socket read operation.
      @op_timeout = opts[:op_timeout] || nil

      # Timeout on socket connect.
      @connect_timeout = opts[:connect_timeout] || nil

      # Global safe option. This is false by default.
      @safe = opts[:safe] || false
      
      # Connection pool for primay node
      @primary      = nil
      @primary_pool = nil

      @logger = opts.fetch(:logger, nil)

      if @logger
        write_logging_startup_message
      end

      should_connect = opts.fetch(:connect, true)
      connect if should_connect
    end
  end
end
#--- end-of-hack

class Person
  include Mongoid::Document

  field :uid, :type => Integer
  field :first_name, :type => String
  field :last_name, :type => String

  index :uid
end

class MongoidBench
  # Unfortunattely in this case, since the mongo driver references
  # ::TCPSocket, I had to patch the Mongo::Connection class above..
  # This is just for testing purposes, a patch can be submitted
  # to the mongo ruby driver maintainers later if this works out..
  include Celluloid::IO

  def initialize
    # Conncet to the db..
    Mongoid.configure do |config|
      config.logger = false
      config.from_hash({
        host: "localhost",
        database: "bench",
        allow_dynamic_fields: false,
        autocreate_indexes: true,
        identity_map_enabled: false,
        include_root_in_json: false,
        max_retries_on_connection_failure: 3,
        parameterize_keys: true,
        persist_in_safe_mode: true,
        preload_models: false,
        raise_not_found_error: false,
        skip_version_check: true
      }.with_indifferent_access)
    end

    bench!
  end

  def bench!
    Person.destroy_all

    # Create a bunch of records first..
    Benchmark.bm do |bm|
      bm.report("Creating #{NUM_DOCUMENTS} documents") do
        NUM_DOCUMENTS.times do |i|
          p = Person.new
          p.uid = i
          p.first_name = rand_str
          p.last_name = rand_str
          p.save
        end
      end
    end

    # Print records...
    Benchmark.bm do |bm|
      bm.report("Querying #{NUM_DOCUMENTS} documents") do
        NUM_DOCUMENTS.times do |i|
          p = Person.where(:uid => i).first
          # puts p.inspect
        end
      end
    end

  end

  def rand_str
    (0...12).map{65.+(rand(25)).chr}.join
  end

end

MongoidBench.new

__END__

$ ruby mongoid_bench.rb
E, [2012-02-23T17:32:34.840107 #6107] ERROR -- : MongoidBench crashed!
NoMethodError: undefined method `setsockopt' for #<Celluloid::IO::TCPSocket:0x007faa913cd8e8>
/Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongo-1.6.0/lib/mongo/connection.rb:599:in `check_is_master'
/Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongo-1.6.0/lib/mongo/connection.rb:393:in `connect'
mongoid_bench.rb:72:in `setup'
/Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongo-1.6.0/lib/mongo/connection.rb:105:in `initialize'
/Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongo-1.6.0/lib/mongo/connection.rb:156:in `new'
/Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongo-1.6.0/lib/mongo/connection.rb:156:in `from_uri'
/Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongoid-2.4.5/lib/mongoid/config/database.rb:89:in `master'
/Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongoid-2.4.5/lib/mongoid/config/database.rb:22:in `configure'
/Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongoid-2.4.5/lib/mongoid/config.rb:273:in `configure_databases'
/Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongoid-2.4.5/lib/mongoid/config.rb:94:in `from_hash'
mongoid_bench.rb:99:in `block in initialize'
/Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongoid-2.4.5/lib/mongoid.rb:115:in `configure'
mongoid_bench.rb:97:in `initialize'
/Users/peter/Development/other/celluloid/lib/celluloid/calls.rb:47:in `dispatch'
/Users/peter/Development/other/celluloid/lib/celluloid/actor.rb:202:in `block in handle_message'
/Users/peter/Development/other/celluloid/lib/celluloid/task.rb:45:in `block in initialize'
/Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongo-1.6.0/lib/mongo/connection.rb:599:in `check_is_master': undefined method `setsockopt' for #<Celluloid::IO::TCPSocket:0x007faa913cd8e8> (NoMethodError)
	from /Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongo-1.6.0/lib/mongo/connection.rb:393:in `connect'
	from mongoid_bench.rb:72:in `setup'
	from /Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongo-1.6.0/lib/mongo/connection.rb:105:in `initialize'
	from /Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongo-1.6.0/lib/mongo/connection.rb:156:in `new'
	from /Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongo-1.6.0/lib/mongo/connection.rb:156:in `from_uri'
	from /Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongoid-2.4.5/lib/mongoid/config/database.rb:89:in `master'
	from /Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongoid-2.4.5/lib/mongoid/config/database.rb:22:in `configure'
	from /Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongoid-2.4.5/lib/mongoid/config.rb:273:in `configure_databases'
	from /Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongoid-2.4.5/lib/mongoid/config.rb:94:in `from_hash'
	from mongoid_bench.rb:99:in `block in initialize'
	from /Users/peter/.rvm/gems/ruby-1.9.3-p125/gems/mongoid-2.4.5/lib/mongoid.rb:115:in `configure'
	from mongoid_bench.rb:97:in `initialize'
	from /Users/peter/Development/other/celluloid/lib/celluloid/calls.rb:47:in `dispatch'
	from /Users/peter/Development/other/celluloid/lib/celluloid/actor.rb:202:in `block in handle_message'
	from /Users/peter/Development/other/celluloid/lib/celluloid/task.rb:45:in `block in initialize'