mongo 1.3.1 → 1.4.0

data/lib/mongo/connection.rb

@@ -24,18 +24,18 @@ module Mongo
 
   # Instantiates and manages connections to MongoDB.
   class Connection
+    include Mongo::Logging
+
     TCPSocket = ::TCPSocket
     Mutex = ::Mutex
     ConditionVariable = ::ConditionVariable
 
-    # Abort connections if a ConnectionError is raised.
-    Thread.abort_on_exception = true
-
     DEFAULT_PORT = 27017
     STANDARD_HEADER_SIZE = 16
     RESPONSE_HEADER_SIZE = 20
 
-    attr_reader :logger, :size, :auths, :primary, :safe, :primary_pool, :host_to_try, :pool_size
+    attr_reader :logger, :size, :auths, :primary, :safe, :host_to_try,
+      :pool_size, :connect_timeout, :primary_pool, :socket_class
 
     # Counter for generating unique request ids.
     @@current_request_id = 0
@@ -65,11 +65,14 @@ module Mongo
     #   logging negatively impacts performance; therefore, it should not be used for high-performance apps.
     # @option opts [Integer] :pool_size (1) The maximum number of socket connections allowed per
     #   connection pool. Note: this setting is relevant only for multi-threaded applications.
-    # @option opts [Float] :timeout (5.0) When all of the connections a pool are checked out,
+    # @option opts [Float] :pool_timeout (5.0) When all of the connections a pool are checked out,
     #   this is the number of seconds to wait for a new connection to be released before throwing an exception.
     #   Note: this setting is relevant only for multi-threaded applications (which in Ruby are rare).
     # @option opts [Float] :op_timeout (nil) The number of seconds to wait for a read operation to time out.
     #   Disabled by default.
+    # @option opts [Float] :connect_timeout (nil) The number of seconds to wait before timing out a
+    #   connection attempt.
+    # @option opts [Boolean] :ssl (false) If true, create the connection to the server using SSL.
     #
     # @example localhost, 27017
     #   Connection.new
@@ -128,7 +131,7 @@ module Mongo
     #
     # @deprecated
     def self.multi(nodes, opts={})
-      warn "Connection.multi is now deprecated. Please use ReplSetConnection.new instead."
+      warn "Connection.multi is now deprecated and will be removed in v2.0. Please use ReplSetConnection.new instead."
 
       nodes << opts
       ReplSetConnection.new(*nodes)
@@ -306,7 +309,7 @@ module Mongo
     #
     # @core databases []-instance_method
     def [](db_name)
-      DB.new(db_name, self, :safe => @safe)
+      DB.new(db_name, self)
     end
 
     # Drop a database.
@@ -406,11 +409,7 @@ module Mongo
 
         send_message_on_socket(packed_message, socket)
       ensure
-        if connection == :writer
-          checkin_writer(socket)
-        else
-          checkin_reader(socket)
-        end
+        checkin(socket)
       end
     end
 
@@ -442,13 +441,13 @@ module Mongo
           docs, num_received, cursor_id = receive(sock, last_error_id)
         end
       ensure
-        checkin_writer(sock)
+        checkin(sock)
       end
 
       if num_received == 1 && (error = docs[0]['err'] || docs[0]['errmsg'])
         close if error == "not master"
         error = "wtimeout" if error == "timeout"
-        raise Mongo::OperationFailure, docs[0]['code'].to_s + ': ' + error
+        raise OperationFailure.new(docs[0]['code'].to_s + ': ' + error, docs[0]['code'], docs[0])
       end
 
       docs[0]
@@ -462,30 +461,40 @@ module Mongo
     # @param [Socket] socket a socket to use in lieu of checking out a new one.
     # @param [Boolean] command (false) indicate whether this is a command. If this is a command,
     #   the message will be sent to the primary node.
+    # @param [Boolean] command (false) indicate whether the cursor should be exhausted. Set
+    #   this to true only when the OP_QUERY_EXHAUST flag is set.
     #
     # @return [Array]
     #   An array whose indexes include [0] documents returned, [1] number of document received,
     #   and [3] a cursor_id.
-    def receive_message(operation, message, log_message=nil, socket=nil, command=false)
+    def receive_message(operation, message, log_message=nil, socket=nil, command=false, read=:primary, exhaust=false)
       request_id = add_message_headers(message, operation)
       packed_message = message.to_s
       begin
         if socket
           sock = socket
-          checkin = false
+          should_checkin = false
         else
-          sock = (command ? checkout_writer : checkout_reader)
-          checkin = true
+          if command
+            sock = checkout_writer
+          elsif read == :primary
+            sock = checkout_writer
+          elsif read == :secondary
+            sock = checkout_reader
+          else
+            sock = checkout_tagged(read)
+          end
+          should_checkin = true
         end
 
         result = ''
         @safe_mutexes[sock].synchronize do
           send_message_on_socket(packed_message, sock)
-          result = receive(sock, request_id)
+          result = receive(sock, request_id, exhaust)
         end
       ensure
-        if checkin
-          command ? checkin_writer(sock) : checkin_reader(sock)
+        if should_checkin
+          checkin(sock)
         end
       end
       result
@@ -509,12 +518,11 @@ module Mongo
           @read_primary = false
         end
 
+        @max_bson_size = config['maxBsonObjectSize'] || Mongo::DEFAULT_MAX_BSON_SIZE
         set_primary(@host_to_try)
       end
 
-      if connected?
-        BSON::BSON_CODER.update_max_bson_size(self)
-      else
+      if !connected?
         raise ConnectionFailure, "Failed to connect to a master node at #{@host_to_try[0]}:#{@host_to_try[1]}"
       end
     end
@@ -556,6 +564,14 @@ module Mongo
     end
     alias :primary? :read_primary?
 
+    # The value of the read preference. Because
+    # this is a single-node connection, the value
+    # is +:primary+, and the connection will read
+    # from whichever type of node it's connected to.
+    def read_preference
+      :primary
+    end
+
     # Close the connection to the database.
     def close
       @primary_pool.close if @primary_pool
@@ -568,8 +584,7 @@ module Mongo
     #
     # @return [Integer]
     def max_bson_size
-      config = self['admin'].command({:ismaster => 1})
-      config['maxBsonObjectSize'] || Mongo::DEFAULT_MAX_BSON_SIZE
+      @max_bson_size
     end
 
     # Checkout a socket for reading (i.e., a secondary node).
@@ -589,32 +604,41 @@ module Mongo
     # Checkin a socket used for reading.
     # Note: this is overridden in ReplSetConnection.
     def checkin_reader(socket)
-      if @primary_pool
-        @primary_pool.checkin(socket)
-      end
+      warn "Connection#checkin_writer is not deprecated and will be removed " +
+        "in driver v2.0. Use Connection#checkin instead."
+      checkin(socket)
     end
 
     # Checkin a socket used for writing.
     # Note: this is overridden in ReplSetConnection.
     def checkin_writer(socket)
+      warn "Connection#checkin_writer is not deprecated and will be removed " +
+        "in driver v2.0. Use Connection#checkin instead."
+      checkin(socket)
+    end
+
+    # Check a socket back into its pool.
+    def checkin(socket)
       if @primary_pool
         @primary_pool.checkin(socket)
       end
     end
 
-    # Execute the block and log the operation described by name
-    # and payload.
-    # TODO: Not sure if this should take a block.
-    def instrument(name, payload = {}, &blk)
-      res = yield
-      log_operation(name, payload)
-      res
-    end
-
     protected
 
     # Generic initialization code.
     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
+        @socket_class = ::TCPSocket
+      end
+
       # Authentication objects
       @auths = opts.fetch(:auths, [])
 
@@ -623,18 +647,26 @@ module Mongo
 
       # Pool size and timeout.
       @pool_size = opts[:pool_size] || 1
-      @timeout   = opts[:timeout]   || 5.0
+      if opts[:timeout]
+        warn "The :timeout option has been deprecated " +
+          "and will be removed in the 2.0 release. Use :pool_timeout instead."
+      end
+      @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
+
       # Mutex for synchronizing pool access
+      # TODO: remove this.
       @connection_mutex = Mutex.new
 
       # Global safe option. This is false by default.
       @safe = opts[:safe] || false
 
-      # Create a mutex when a new key, in this case a socket,
+            # Create a mutex when a new key, in this case a socket,
       # is added to the hash.
       @safe_mutexes = Hash.new { |h, k| h[k] = Mutex.new }
 
@@ -672,18 +704,6 @@ module Mongo
       end
     end
 
-    ## Logging methods
-
-    def log_operation(name, payload)
-      return unless @logger
-      msg = "#{payload[:database]}['#{payload[:collection]}'].#{name}("
-      msg += payload.values_at(:selector, :document, :documents, :fields ).compact.map(&:inspect).join(', ') + ")"
-      msg += ".skip(#{payload[:skip]})"  if payload[:skip]
-      msg += ".limit(#{payload[:limit]})"  if payload[:limit]
-      msg += ".sort(#{payload[:order]})"  if payload[:order]
-      @logger.debug "MONGODB #{msg}"
-    end
-
     private
 
     ## Methods for establishing a connection:
@@ -698,8 +718,16 @@ module Mongo
     def check_is_master(node)
       begin
         host, port = *node
-        socket = TCPSocket.new(host, port)
-        socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+
+        if @connect_timeout
+          Mongo::TimeoutHandler.timeout(@connect_timeout, OperationTimeout) do
+            socket = @socket_class.new(host, port)
+            socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+          end
+        else
+          socket = @socket_class.new(host, port)
+          socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+        end
 
         config = self['admin'].command({:ismaster => 1}, :socket => socket)
       rescue OperationFailure, SocketError, SystemCallError, IOError => ex
@@ -720,21 +748,38 @@ module Mongo
 
     ## Low-level connection methods.
 
-    def receive(sock, expected_response)
+    def receive(sock, cursor_id, exhaust=false)
       begin
-      receive_header(sock, expected_response)
-      number_received, cursor_id = receive_response_header(sock)
-      read_documents(number_received, cursor_id, sock)
+        if exhaust
+          docs = []
+          num_received = 0
+
+          while(cursor_id != 0) do
+            receive_header(sock, cursor_id, exhaust)
+            number_received, cursor_id = receive_response_header(sock)
+            new_docs, n = read_documents(number_received, sock)
+            docs += new_docs
+            num_received += n
+          end
+
+          return [docs, num_received, cursor_id]
+        else
+          receive_header(sock, cursor_id, exhaust)
+          number_received, cursor_id = receive_response_header(sock)
+          docs, num_received = read_documents(number_received, sock)
+
+          return [docs, num_received, cursor_id]
+        end
       rescue Mongo::ConnectionFailure => ex
         close
         raise ex
       end
     end
 
-    def receive_header(sock, expected_response)
+    def receive_header(sock, expected_response, exhaust=false)
       header = receive_message_on_socket(16, sock)
       size, request_id, response_to = header.unpack('VVV')
-      if expected_response != response_to
+      if !exhaust && expected_response != response_to
         raise Mongo::ConnectionFailure, "Expected response #{expected_response} but got #{response_to}"
       end
 
@@ -766,7 +811,7 @@ module Mongo
       end
     end
 
-    def read_documents(number_received, cursor_id, sock)
+    def read_documents(number_received, sock)
       docs = []
       number_remaining = number_received
       while number_remaining > 0 do
@@ -776,7 +821,7 @@ module Mongo
         number_remaining -= 1
         docs << BSON::BSON_CODER.deserialize(buf)
       end
-      [docs, number_received, cursor_id]
+      [docs, number_received]
     end
 
     # Constructs a getlasterror message. This method is used exclusively by

data/lib/mongo/cursor.rb

@@ -18,12 +18,15 @@ module Mongo
 
   # A cursor over query results. Returned objects are hashes.
   class Cursor
-    include Mongo::Conversions
     include Enumerable
+    include Mongo::Constants
+    include Mongo::Conversions
+    include Mongo::Logging
 
     attr_reader :collection, :selector, :fields,
       :order, :hint, :snapshot, :timeout,
-      :full_collection_name, :transformer
+      :full_collection_name, :transformer,
+      :options, :cursor_id
 
     # Create a new cursor.
     #
@@ -59,6 +62,7 @@ module Mongo
       @limit      = opts[:limit]    || 0
       @tailable   = opts[:tailable] || false
       @timeout    = opts.fetch(:timeout, true)
+      @options    = 0
 
       # Use this socket for the query
       @socket     = opts[:socket]
@@ -67,12 +71,28 @@ module Mongo
       @query_run    = false
 
       @transformer = opts[:transformer]
+      if value = opts[:read]
+        Mongo::Support.validate_read_preference(value)
+      else
+        value = collection.read_preference
+      end
+      @read_preference = value.is_a?(Hash) ? value.dup : value
       batch_size(opts[:batch_size] || 0)
 
       @full_collection_name = "#{@collection.db.name}.#{@collection.name}"
       @cache        = []
       @returned     = 0
 
+      if(!@timeout)
+        add_option(OP_QUERY_NO_CURSOR_TIMEOUT)
+      end
+      if(@connection.slave_ok?)
+        add_option(OP_QUERY_SLAVE_OK)
+      end
+      if(@tailable)
+        add_option(OP_QUERY_TAILABLE)
+      end
+
       if @collection.name =~ /^\$cmd/ || @collection.name =~ /^system/
         @command = true
       else
@@ -80,11 +100,30 @@ module Mongo
       end
     end
 
+    # Guess whether the cursor is alive on the server.
+    #
+    # Note that this method only checks whether we have
+    # a cursor id. The cursor may still have timed out
+    # on the server. This will be indicated in the next
+    # call to Cursor#next.
+    #
+    # @return [Boolean]
+    def alive?
+      @cursor_id && @cursor_id != 0
+    end
+
     # Get the next document specified the cursor options.
     #
     # @return [Hash, Nil] the next document or Nil if no documents remain.
-    def next_document
-      refresh if @cache.length == 0
+    def next
+      if @cache.length == 0
+        if @query_run && (@options & OP_QUERY_EXHAUST != 0)
+          close
+          return nil
+        else
+          refresh
+        end
+      end
       doc = @cache.shift
 
       if doc && doc['$err']
@@ -95,10 +134,10 @@ module Mongo
         # connection. The next request will re-open on master server.
         if err == "not master"
           @connection.close
-          raise ConnectionFailure, err
+          raise ConnectionFailure.new(err, doc['code'], doc)
         end
 
-        raise OperationFailure, err
+        raise OperationFailure.new(err, doc['code'], doc)
       end
 
       if @transformer.nil?
@@ -107,7 +146,7 @@ module Mongo
         @transformer.call(doc) if doc
       end
     end
-    alias :next :next_document
+    alias :next_document :next
 
     # Reset this cursor on the server. Cursor options, such as the
     # query string and the values for skip and limit, are preserved.
@@ -148,7 +187,7 @@ module Mongo
       response = @db.command(command)
       return response['n'].to_i if Mongo::Support.ok?(response)
       return 0 if response['errmsg'] == "ns missing"
-      raise OperationFailure, "Count failed: #{response['errmsg']}"
+      raise OperationFailure.new("Count failed: #{response['errmsg']}", response['code'], response)
     end
 
     # Sort this cursor's results.
@@ -220,12 +259,13 @@ module Mongo
     #   the server will determine the batch size.
     #
     # @return [Cursor]
-    def batch_size(size=0)
+    def batch_size(size=nil)
+      return @batch_size unless size
       check_modifiable
       if size < 0 || size == 1
         raise ArgumentError, "Invalid value for batch_size #{size}; must be 0 or > 1."
       else
-        @batch_size = size > @limit ? @limit : size
+        @batch_size = @limit != 0 && size > @limit ? @limit : size
       end
 
       self
@@ -243,11 +283,8 @@ module Mongo
     #     puts doc['user']
     #   end
     def each
-      #num_returned = 0
-      #while has_next? && (@limit <= 0 || num_returned < @limit)
-      while doc = next_document
-        yield doc #next_document
-        #num_returned += 1
+      while doc = self.next
+        yield doc
       end
     end
 
@@ -273,7 +310,8 @@ module Mongo
     #
     # @core explain explain-instance_method
     def explain
-      c = Cursor.new(@collection, query_options_hash.merge(:limit => -@limit.abs, :explain => true))
+      c = Cursor.new(@collection,
+        query_options_hash.merge(:limit => -@limit.abs, :explain => true))
       explanation = c.next_document
       c.close
 
@@ -295,7 +333,7 @@ module Mongo
         message = BSON::ByteBuffer.new([0, 0, 0, 0])
         message.put_int(1)
         message.put_long(@cursor_id)
-        @logger.debug("MONGODB cursor.close #{@cursor_id}") if @logger
+        log(:debug, "Cursor#close #{@cursor_id}")
         @connection.send_message(Mongo::Constants::OP_KILL_CURSORS, message, :connection => :reader)
       end
       @cursor_id = 0
@@ -305,7 +343,9 @@ module Mongo
     # Is this cursor closed?
     #
     # @return [Boolean]
-    def closed?; @closed; end
+    def closed?
+      @closed
+    end
 
     # Returns an integer indicating which query options have been selected.
     #
@@ -314,11 +354,43 @@ module Mongo
     # @see http://www.mongodb.org/display/DOCS/Mongo+Wire+Protocol#MongoWireProtocol-Mongo::Constants::OPQUERY
     # The MongoDB wire protocol.
     def query_opts
-      opts     = 0
-      opts    |= Mongo::Constants::OP_QUERY_NO_CURSOR_TIMEOUT unless @timeout
-      opts    |= Mongo::Constants::OP_QUERY_SLAVE_OK if @connection.slave_ok?
-      opts    |= Mongo::Constants::OP_QUERY_TAILABLE if @tailable
-      opts
+      warn "The method Cursor#query_opts has been deprecated " +
+        "and will removed in v2.0. Use Cursor#options instead."
+      @options
+    end
+
+    # Add an option to the query options bitfield.
+    #
+    # @param opt a valid query option
+    #
+    # @raise InvalidOperation if this method is run after the cursor has bee
+    #   iterated for the first time.
+    #
+    # @return [Integer] the current value of the options bitfield for this cursor.
+    #
+    # @see http://www.mongodb.org/display/DOCS/Mongo+Wire+Protocol#MongoWireProtocol-Mongo::Constants::OPQUERY
+    def add_option(opt)
+      check_modifiable
+
+      @options |= opt
+      @options
+    end
+
+    # Remove an option from the query options bitfield.
+    #
+    # @param opt a valid query option
+    #
+    # @raise InvalidOperation if this method is run after the cursor has bee
+    #   iterated for the first time.
+    #
+    # @return [Integer] the current value of the options bitfield for this cursor.
+    #
+    # @see http://www.mongodb.org/display/DOCS/Mongo+Wire+Protocol#MongoWireProtocol-Mongo::Constants::OPQUERY
+    def remove_option(opt)
+      check_modifiable
+
+      @options &= ~opt
+      @options
     end
 
     # Get the query options for this Cursor.
@@ -341,7 +413,7 @@ module Mongo
     # Clean output for inspect.
     def inspect
       "<Mongo::Cursor:0x#{object_id.to_s(16)} namespace='#{@db.name}.#{@collection.name}' " +
-        "@selector=#{@selector.inspect}>"
+        "@selector=#{@selector.inspect} @cursor_id=#{@cursor_id}>"
     end
 
     private
@@ -363,12 +435,43 @@ module Mongo
 
     # Return the number of documents remaining for this cursor.
     def num_remaining
-      refresh if @cache.length == 0
+      if @cache.length == 0
+        if @query_run && (@options & OP_QUERY_EXHAUST != 0)
+          close
+          return 0
+        else
+          refresh
+        end
+      end
+
       @cache.length
     end
 
+    # Refresh the documents in @cache. This means either
+    # sending the initial query or sending a GET_MORE operation.
     def refresh
-      return if send_initial_query || @cursor_id.zero?
+      if !@query_run
+        send_initial_query
+      elsif !@cursor_id.zero?
+        send_get_more
+      end
+    end
+
+    def send_initial_query
+      message = construct_query_message
+      payload = instrument_payload if @logger
+      instrument(:find, payload) do
+        results, @n_received, @cursor_id = @connection.receive_message(
+          Mongo::Constants::OP_QUERY, message, nil, @socket, @command,
+          @read_preference, @options & OP_QUERY_EXHAUST != 0)
+        @returned += @n_received
+        @cache += results
+        @query_run = true
+        close_cursor_if_query_complete
+      end
+    end
+
+    def send_get_more
       message = BSON::ByteBuffer.new([0, 0, 0, 0])
 
       # DB name.
@@ -387,37 +490,17 @@ module Mongo
 
       # Cursor id.
       message.put_long(@cursor_id)
-      @logger.debug("MONGODB cursor.refresh() for cursor #{@cursor_id}") if @logger
+      log(:debug, "cursor.refresh() for cursor #{@cursor_id}") if @logger
       results, @n_received, @cursor_id = @connection.receive_message(
-          Mongo::Constants::OP_GET_MORE, message, nil, @socket, @command)
+          Mongo::Constants::OP_GET_MORE, message, nil, @socket, @command, @read_preference)
       @returned += @n_received
       @cache += results
       close_cursor_if_query_complete
     end
 
-    # Run query the first time we request an object from the wire
-    # TODO: should we be calling instrument_payload even if logging
-    # is disabled?
-    def send_initial_query
-      if @query_run
-        false
-      else
-        message = construct_query_message
-        @connection.instrument(:find, instrument_payload) do
-          results, @n_received, @cursor_id = @connection.receive_message(
-            Mongo::Constants::OP_QUERY, message, nil, @socket, @command)
-          @returned += @n_received
-          @cache += results
-          @query_run = true
-          close_cursor_if_query_complete
-        end
-        true
-      end
-    end
-
     def construct_query_message
       message = BSON::ByteBuffer.new
-      message.put_int(query_opts)
+      message.put_int(@options)
       BSON::BSON_RUBY.serialize_cstr(message, "#{@db.name}.#{@collection.name}")
       message.put_int(@skip)
       message.put_int(@limit)
@@ -429,10 +512,10 @@ module Mongo
 
     def instrument_payload
       log = { :database => @db.name, :collection => @collection.name, :selector => selector }
-      log[:fields] = @fields  if @fields
-      log[:skip] = @skip  if @skip && (@skip > 0)
-      log[:limit] = @limit  if @limit && (@limit > 0)
-      log[:order] = @order  if @order
+      log[:fields] = @fields if @fields
+      log[:skip]   = @skip   if @skip && (@skip != 0)
+      log[:limit]  = @limit  if @limit && (@limit != 0)
+      log[:order]  = @order  if @order
       log
     end
 
@@ -455,10 +538,6 @@ module Mongo
       @order || @explain || @hint || @snapshot
     end
 
-    def to_s
-      "DBResponse(flags=#@result_flags, cursor_id=#@cursor_id, start=#@starting_from)"
-    end
-
     def close_cursor_if_query_complete
       if @limit > 0 && @returned >= @limit
         close