mongo 1.3.0 → 1.12.5

data/lib/mongo/collection_writer.rb

@@ -0,0 +1,364 @@
+# Copyright (C) 2009-2013 MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+module Mongo
+
+  class CollectionWriter
+    include Mongo::Logging
+    include Mongo::WriteConcern
+
+    OPCODE = {
+      :insert => Mongo::Constants::OP_INSERT,
+      :update => Mongo::Constants::OP_UPDATE,
+      :delete => Mongo::Constants::OP_DELETE
+    }
+    WRITE_COMMAND_ARG_KEY = {
+      :insert => :documents,
+      :update => :updates,
+      :delete => :deletes
+    }
+
+    def initialize(collection)
+      @collection = collection
+      @name = @collection.name
+      @db = @collection.db
+      @connection = @db.connection
+      @logger     = @connection.logger
+      @max_write_batch_size = Mongo::MongoClient::DEFAULT_MAX_WRITE_BATCH_SIZE
+    end
+
+    # common implementation only for new batch write commands (insert, update, delete) and old batch insert
+    def batch_write_incremental(op_type, documents, check_keys=true, opts={})
+      raise Mongo::OperationFailure, "Request contains no documents" if documents.empty?
+      write_concern = get_write_concern(opts, @collection)
+      max_message_size, max_append_size, max_serialize_size = batch_write_max_sizes(write_concern)
+      ordered = opts[:ordered]
+      continue_on_error = !!opts[:continue_on_error] || ordered == false
+      collect_on_error = !!opts[:collect_on_error] || ordered == false
+      error_docs = [] # docs with serialization errors
+      errors = []
+      write_concern_errors = []
+      exchanges = []
+      serialized_doc = nil
+      message = BSON::ByteBuffer.new("", max_message_size)
+      @max_write_batch_size = @collection.db.connection.max_write_batch_size
+      docs = documents.dup
+      catch(:error) do
+        until docs.empty? || (!errors.empty? && !collect_on_error && !continue_on_error) # process documents a batch at a time
+          batch_docs = []
+          batch_message_initialize(message, op_type, continue_on_error, write_concern)
+          while !docs.empty? && batch_docs.size < @max_write_batch_size
+            begin
+              doc = docs.first
+              doc = doc[:d] if op_type == :insert && !ordered.nil? #check_keys for :update outside of serialize
+              serialized_doc ||= BSON::BSON_CODER.serialize(doc, check_keys, true, max_serialize_size)
+            rescue BSON::InvalidDocument, BSON::InvalidKeyName, BSON::InvalidStringEncoding => ex
+              bulk_message = "Bulk write error - #{ex.message} - examine result for complete information"
+              ex = BulkWriteError.new(bulk_message, Mongo::ErrorCode::INVALID_BSON,
+                                      {:op_type => op_type, :serialize => doc, :ord => docs.first[:ord], :error => ex}) unless ordered.nil?
+              error_docs << docs.shift
+              errors << ex
+              next if collect_on_error
+              throw(:error) if batch_docs.empty?
+              break # defer exit and send batch
+            end
+            break if message.size + serialized_doc.size > max_append_size
+            batch_docs << docs.shift
+            batch_message_append(message, serialized_doc, write_concern)
+            serialized_doc = nil
+          end
+          begin
+            response = batch_message_send(message, op_type, batch_docs, write_concern, continue_on_error) if batch_docs.size > 0
+            exchanges << {:op_type => op_type, :batch => batch_docs, :opts => opts, :response => response}
+          rescue Mongo::WriteConcernError => ex
+            write_concern_errors << ex
+            exchanges << {:op_type => op_type, :batch => batch_docs, :opts => opts, :response => ex.result}
+          rescue Mongo::OperationFailure => ex
+            errors << ex
+            exchanges << {:op_type => op_type, :batch => batch_docs, :opts => opts, :response => ex.result}
+            throw(:error) unless continue_on_error
+          end
+        end
+      end
+      [error_docs, errors, write_concern_errors, exchanges]
+    end
+
+    def batch_write_partition(op_type, documents, check_keys, opts)
+      raise Mongo::OperationFailure, "Request contains no documents" if documents.empty?
+      write_concern = get_write_concern(opts, @collection)
+      ordered = opts[:ordered]
+      continue_on_error = !!opts[:continue_on_error] || ordered == false # continue_on_error default false
+      collect_on_error = !!opts[:collect_on_error] # collect_on_error default false
+      error_docs = [] # docs with serialization errors
+      errors = []
+      write_concern_errors = []
+      exchanges = []
+      @max_write_batch_size = @collection.db.connection.max_write_batch_size
+      @write_batch_size = [documents.size, @max_write_batch_size].min
+      docs = documents.dup
+      until docs.empty?
+        batch = docs.take(@write_batch_size)
+        begin
+          batch_to_send = batch #(op_type == :insert && !ordered.nil?) ? batch.collect{|doc|doc[:d]} : batch
+          if @collection.db.connection.use_write_command?(write_concern) # TODO - polymorphic send_write including legacy insert
+            response = send_bulk_write_command(op_type, batch_to_send, check_keys, opts)
+          else
+            response = send_write_operation(op_type, nil, batch_to_send, check_keys, opts, write_concern)
+          end
+          exchanges << {:op_type => op_type, :batch => batch, :opts => opts, :response => response}
+          docs = docs.drop(batch.size)
+          @write_batch_size = [(@write_batch_size*1097) >> 10, @write_batch_size+1].max unless docs.empty? # 2**(1/10) multiplicative increase
+          @write_batch_size = @max_write_batch_size if @write_batch_size > @max_write_batch_size
+        rescue BSON::InvalidDocument, BSON::InvalidKeyName, BSON::InvalidStringEncoding => ex
+          if @write_batch_size > 1 # decrease batch size
+            @write_batch_size = (@write_batch_size+1) >> 1 # 2**(-1) multiplicative decrease
+            next
+          end
+          # error on a single document
+          bulk_message = "Bulk write error - #{ex.message} - examine result for complete information"
+          ex = BulkWriteError.new(bulk_message, Mongo::ErrorCode::INVALID_BSON,
+                                  {:op_type => op_type, :batch => batch, :ord => batch.first[:ord], :opts => opts, :error => ex}) unless ordered.nil?
+          error_docs << docs.shift
+          next if collect_on_error
+          errors << ex
+          break unless continue_on_error
+        rescue Mongo::WriteConcernError => ex
+          write_concern_errors << ex
+          exchanges << {:op_type => op_type, :batch => batch_docs, :opts => opts, :response => ex.result}
+          docs = docs.drop(batch.size)
+        rescue Mongo::OperationFailure => ex
+          errors << ex
+          exchanges << {:op_type => op_type, :batch => batch, :opts => opts, :response => ex.result}
+          docs = docs.drop(batch.size)
+          break if !continue_on_error && !collect_on_error
+        end
+      end
+      [error_docs, errors, write_concern_errors, exchanges]
+    end
+
+    alias :batch_write :batch_write_incremental
+
+    def send_bulk_write_command(op_type, documents, check_keys, opts, collection_name=@name)
+      if op_type == :insert
+        documents = documents.collect{|doc| doc[:d]} if opts.key?(:ordered)
+        documents.each do |doc|
+          # TODO - @pk_factory.create_pk(doc)
+          if check_keys
+            doc.each_key do |key|
+              key = key.to_s
+              raise BSON::InvalidKeyName.new("key #{key} must not start with '$'") if key[0] == ?$
+              raise BSON::InvalidKeyName.new("key #{key} must not contain '.'") if key.include? ?.
+            end
+          end
+        end
+      #elsif op_type == :update # TODO - check keys
+      #elsif op_type == :delete
+      #else
+      #  raise ArgumentError, "Write operation type must be :insert, :update or :delete"
+      end
+      request = BSON::OrderedHash[op_type, collection_name].merge!(
+          Mongo::CollectionWriter::WRITE_COMMAND_ARG_KEY[op_type] => documents,
+          :writeConcern => get_write_concern(opts, @collection),
+          :ordered => opts[:ordered] || !opts[:continue_on_error]
+      )
+      @db.command(request)
+    end
+
+    private
+
+    def sort_by_first_sym(pairs)
+      pairs = pairs.collect{|first, rest| [first.to_s, rest]} #stringify_first
+      pairs = pairs.sort{|x,y| x.first <=> y.first }
+      pairs.collect{|first, rest| [first.to_sym, rest]} #symbolize_first
+    end
+
+    def ordered_group_by_first(pairs)
+      pairs.inject([[], nil]) do |memo, pair|
+        result, previous_value = memo
+        current_value = pair.first
+        result << [current_value, []] if previous_value != current_value
+        result.last.last << pair.last
+        [result, current_value]
+      end.first
+    end
+
+  end
+
+  class CollectionOperationWriter < CollectionWriter
+    def initialize(collection)
+      super(collection)
+    end
+
+    def send_write_operation(op_type, selector, doc_or_docs, check_keys, opts, write_concern, collection_name=@name)
+      message = BSON::ByteBuffer.new("", @connection.max_message_size)
+      message.put_int((op_type == :insert && !!opts[:continue_on_error]) ? 1 : 0)
+      BSON::BSON_RUBY.serialize_cstr(message, "#{@db.name}.#{collection_name}")
+      if op_type == :update
+        update_options  = 0
+        update_options += 1 if opts[:upsert]
+        update_options += 2 if opts[:multi]
+        message.put_int(update_options)
+      elsif op_type == :delete
+        delete_options = 0
+        delete_options += 1 if opts[:limit] && opts[:limit] != 0
+        message.put_int(delete_options)
+      end
+      message.put_binary(BSON::BSON_CODER.serialize(selector, false, true, @connection.max_bson_size).to_s) if selector
+      [doc_or_docs].flatten(1).compact.each do |document|
+        message.put_binary(BSON::BSON_CODER.serialize(document, check_keys, true, @connection.max_bson_size).to_s)
+        if message.size > @connection.max_message_size
+          raise BSON::InvalidDocument, "Message is too large. This message is limited to #{@connection.max_message_size} bytes."
+        end
+      end
+      instrument(op_type, :database => @db.name, :collection => collection_name, :selector => selector, :documents => doc_or_docs) do
+        op_code = OPCODE[op_type]
+        if Mongo::WriteConcern.gle?(write_concern)
+          @connection.send_message_with_gle(op_code, message, @db.name, nil, write_concern)
+        else
+          @connection.send_message(op_code, message)
+        end
+      end
+    end
+
+    def bulk_execute(ops, options, opts = {})
+      write_concern = get_write_concern(opts, @collection)
+      errors = []
+      write_concern_errors = []
+      exchanges = []
+      ops.each do |op_type, doc|
+        doc = {:d => @collection.pk_factory.create_pk(doc[:d]), :ord => doc[:ord]} if op_type == :insert
+        doc_opts = doc.merge(opts)
+        d = doc_opts.delete(:d)
+        q = doc_opts.delete(:q)
+        u = doc_opts.delete(:u)
+        begin  # use single and NOT batch inserts since there no index for an error
+          response = @collection.operation_writer.send_write_operation(op_type, q, d || u, check_keys = false, doc_opts, write_concern)
+          exchanges << {:op_type => op_type, :batch => [doc], :opts => opts, :response => response}
+        rescue BSON::InvalidDocument, BSON::InvalidKeyName, BSON::InvalidStringEncoding => ex
+          bulk_message = "Bulk write error - #{ex.message} - examine result for complete information"
+          ex = BulkWriteError.new(bulk_message, Mongo::ErrorCode::INVALID_BSON,
+                                  {:op_type => op_type, :serialize => doc, :ord => doc[:ord], :error => ex})
+          errors << ex
+          break if options[:ordered]
+        rescue Mongo::WriteConcernError => ex
+          write_concern_errors << ex
+          exchanges << {:op_type => op_type, :batch => [doc], :opts => opts, :response => ex.result}
+        rescue Mongo::OperationFailure => ex
+          errors << ex
+          exchanges << {:op_type => op_type, :batch => [doc], :opts => opts, :response => ex.result}
+          break if options[:ordered] && ex.result["err"] != "norepl"
+        end
+      end
+      [errors, write_concern_errors, exchanges]
+    end
+
+    private
+
+    def batch_message_initialize(message, op_type, continue_on_error, write_concern)
+      message.clear!.clear
+      message.put_int(continue_on_error ? 1 : 0)
+      BSON::BSON_RUBY.serialize_cstr(message, "#{@db.name}.#{@name}")
+    end
+
+    def batch_message_append(message, serialized_doc, write_concern)
+      message.put_binary(serialized_doc.to_s)
+    end
+
+    def batch_message_send(message, op_type, batch_docs, write_concern, continue_on_error)
+      instrument(:insert, :database => @db.name, :collection => @name, :documents => batch_docs) do
+        if Mongo::WriteConcern.gle?(write_concern)
+          @connection.send_message_with_gle(Mongo::Constants::OP_INSERT, message, @db.name, nil, write_concern)
+        else
+          @connection.send_message(Mongo::Constants::OP_INSERT, message)
+        end
+      end
+    end
+
+    def batch_write_max_sizes(write_concern)
+      [@connection.max_message_size, @connection.max_message_size, @connection.max_bson_size]
+    end
+
+  end
+
+  class CollectionCommandWriter < CollectionWriter
+    def initialize(collection)
+      super(collection)
+    end
+
+    def send_write_command(op_type, selector, doc_or_docs, check_keys, opts, write_concern, collection_name=@name)
+      if op_type == :insert
+        argument = [doc_or_docs].flatten(1).compact
+      elsif op_type == :update
+        argument = [{:q => selector, :u => doc_or_docs, :multi => !!opts[:multi]}]
+        argument.first.merge!(:upsert => opts[:upsert]) if opts[:upsert]
+      elsif op_type == :delete
+        argument = [{:q => selector, :limit => (opts[:limit] || 0)}]
+      else
+        raise ArgumentError, "Write operation type must be :insert, :update or :delete"
+      end
+      request = BSON::OrderedHash[op_type, collection_name, WRITE_COMMAND_ARG_KEY[op_type], argument]
+      request.merge!(:writeConcern => write_concern, :ordered => !opts[:continue_on_error])
+      request.merge!(opts)
+      instrument(op_type, :database => @db.name, :collection => collection_name, :selector => selector, :documents => doc_or_docs) do
+        @db.command(request)
+      end
+    end
+
+    def bulk_execute(ops, options, opts = {})
+      errors = []
+      write_concern_errors = []
+      exchanges = []
+      ops = (options[:ordered] == false) ? sort_by_first_sym(ops) : ops # sort by write-type
+      ordered_group_by_first(ops).each do |op_type, documents|
+        documents.collect! {|doc| {:d => @collection.pk_factory.create_pk(doc[:d]), :ord => doc[:ord]} } if op_type == :insert
+        error_docs, batch_errors, batch_write_concern_errors, batch_exchanges =
+            batch_write(op_type, documents, check_keys = false, opts.merge(:ordered => options[:ordered]))
+        errors += batch_errors
+        write_concern_errors += batch_write_concern_errors
+        exchanges += batch_exchanges
+        break if options[:ordered] && !batch_errors.empty?
+      end
+      [errors, write_concern_errors, exchanges]
+    end
+
+    private
+
+    def batch_message_initialize(message, op_type, continue_on_error, write_concern)
+      message.clear!.clear
+      @bson_empty ||= BSON::BSON_CODER.serialize({})
+      message.put_binary(@bson_empty.to_s)
+      message.unfinish!.array!(WRITE_COMMAND_ARG_KEY[op_type])
+    end
+
+    def batch_message_append(message, serialized_doc, write_concern)
+      message.push_doc!(serialized_doc)
+    end
+
+    def batch_message_send(message, op_type, batch_docs, write_concern, continue_on_error)
+      message.finish!
+      request = BSON::OrderedHash[op_type, @name, :bson, message]
+      request.merge!(:writeConcern => write_concern, :ordered => !continue_on_error)
+      instrument(:insert, :database => @db.name, :collection => @name, :documents => batch_docs) do
+        @db.command(request)
+      end
+    end
+
+    def batch_write_max_sizes(write_concern)
+      [MongoClient::COMMAND_HEADROOM, MongoClient::APPEND_HEADROOM, MongoClient::SERIALIZE_HEADROOM].collect{|h| @connection.max_bson_size + h}
+    end
+
+  end
+
+end
+

data/lib/mongo/connection/node.rb

@@ -0,0 +1,249 @@
+# Copyright (C) 2009-2013 MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+module Mongo
+  class Node
+
+    attr_accessor :host, :port, :address, :client, :socket, :last_state
+
+    def initialize(client, host_port)
+      @client = client
+      @manager = @client.local_manager
+      @host, @port = Support.normalize_seeds(host_port)
+      @address = "#{@host}:#{@port}"
+      @config = nil
+      @socket = nil
+      @node_mutex = Monitor.new
+    end
+
+    def eql?(other)
+      (other.is_a?(Node) && @address == other.address)
+    end
+    alias :== :eql?
+
+    def =~(other)
+      if other.is_a?(String)
+        h, p = Support.normalize_seeds(other)
+        h == @host && p == @port
+      else
+        false
+      end
+    end
+
+    def host_string
+      address
+    end
+
+    def config
+      connect unless connected?
+      set_config unless @config || !connected?
+      @config
+    end
+
+    def inspect
+      "<Mongo::Node:0x#{self.object_id.to_s(16)} @host=#{@host} @port=#{@port}>"
+    end
+
+    # Create a connection to the provided node,
+    # and, if successful, return the socket. Otherwise,
+    # return nil.
+    def connect
+      @node_mutex.synchronize do
+        begin
+          @socket = @client.socket_class.new(@host, @port,
+                                             @client.op_timeout,
+                                             @client.connect_timeout,
+                                             @client.socket_opts)
+        rescue ConnectionTimeoutError, OperationTimeout, ConnectionFailure, OperationFailure,
+               SocketError, SystemCallError, IOError => ex
+          @client.log(:debug, "Failed connection to #{host_string} with #{ex.class}, #{ex.message}.")
+          close
+        end
+      end
+    end
+
+    # This should only be called within a mutex
+    def close
+      if @socket && !@socket.closed?
+        @socket.close
+      end
+      @socket = nil
+      @config = nil
+    end
+
+    def connected?
+      @socket != nil && !@socket.closed?
+    end
+
+    def active?
+      begin
+        result = @client['admin'].command({:ping => 1}, :socket => usable_socket)
+      rescue OperationFailure, SocketError, SystemCallError, IOError
+        return nil
+      end
+      result['ok'] == 1
+    end
+
+    # Get the configuration for the provided node as returned by the
+    # ismaster command. Additionally, check that the replica set name
+    # matches with the name provided.
+    def set_config
+      @node_mutex.synchronize do
+        begin
+          if @config
+            @last_state = @config['ismaster'] ? :primary : :other
+          end
+
+          if @client.connect_timeout
+            Timeout::timeout(@client.connect_timeout, OperationTimeout) do
+              @config = @client['admin'].command({:ismaster => 1}, :socket => usable_socket)
+            end
+          else
+            @config = @client['admin'].command({:ismaster => 1}, :socket => usable_socket)
+          end
+
+          update_max_sizes
+
+          if @config['msg']
+            @client.log(:warn, "#{config['msg']}")
+          end
+
+          unless @client.mongos?
+            check_set_membership(@config)
+            check_set_name(@config)
+          end
+        rescue ConnectionFailure, OperationFailure, OperationTimeout, SocketError, SystemCallError, IOError => ex
+          @client.log(:warn, "Attempted connection to node #{host_string} raised " +
+                              "#{ex.class}: #{ex.message}")
+          # Socket may already be nil from issuing command
+          close
+        end
+      end
+    end
+
+    # Return a list of replica set nodes from the config.
+    # Note: this excludes arbiters.
+    def node_list
+      nodes = []
+      nodes += config['hosts'] if config['hosts']
+      nodes += config['passives'] if config['passives']
+      nodes += ["#{@host}:#{@port}"] if @client.mongos?
+      nodes
+    end
+
+    def arbiters
+      return [] unless config['arbiters']
+      config['arbiters'].map do |arbiter|
+        Support.normalize_seeds(arbiter)
+      end
+    end
+
+    def primary?
+      config['ismaster'] == true || config['ismaster'] == 1
+    end
+
+    def secondary?
+      config['secondary'] == true || config['secondary'] == 1
+    end
+
+    def tags
+      config['tags'] || {}
+    end
+
+    def host_port
+      [@host, @port]
+    end
+
+    def hash
+      address.hash
+    end
+
+    def healthy?
+      connected? && config
+    end
+
+    def max_bson_size
+      @max_bson_size || DEFAULT_MAX_BSON_SIZE
+    end
+
+    def max_message_size
+      @max_message_size || max_bson_size * MESSAGE_SIZE_FACTOR
+    end
+
+    def max_wire_version
+      @max_wire_version || 0
+    end
+
+    def min_wire_version
+      @min_wire_version || 0
+    end
+
+    def wire_version_feature?(feature)
+      min_wire_version <= feature && feature <= max_wire_version
+    end
+
+    def max_write_batch_size
+      @max_write_batch_size || Mongo::MongoClient::DEFAULT_MAX_WRITE_BATCH_SIZE
+    end
+
+    protected
+
+    # Ensure that this node is a healthy member of a replica set.
+    def check_set_membership(config)
+      if !config.has_key?('hosts')
+        message = "Will not connect to #{host_string} because it's not a member " +
+          "of a replica set."
+        raise ConnectionFailure, message
+      elsif config['hosts'].length == 1 && !config['ismaster'] &&
+        !config['secondary']
+        message = "Attempting to connect to an unhealthy, single-node replica set."
+        raise ConnectionFailure, message
+      end
+    end
+
+    # Ensure that this node is part of a replica set of the expected name.
+    def check_set_name(config)
+      if @client.replica_set_name
+        if !config['setName']
+          @client.log(:warn, "Could not verify replica set name for member #{host_string} " +
+            "because ismaster does not return name in this version of MongoDB")
+        elsif @client.replica_set_name != config['setName']
+          message = "Attempting to connect to replica set '#{config['setName']}' on member #{host_string} " +
+            "but expected '#{@client.replica_set_name}'"
+          raise ReplicaSetConnectionError, message
+        end
+      end
+    end
+
+    private
+
+    def usable_socket
+      if @socket && @socket.pid != Process.pid
+        @socket.close
+        @socket = nil
+        connect
+      else
+        @socket
+      end
+    end
+
+    def update_max_sizes
+      @max_bson_size = config['maxBsonObjectSize'] || DEFAULT_MAX_BSON_SIZE
+      @max_message_size = config['maxMessageSizeBytes'] || @max_bson_size * MESSAGE_SIZE_FACTOR
+      @max_wire_version = config['maxWireVersion'] || 0
+      @min_wire_version = config['minWireVersion'] || 0
+      @max_write_batch_size = config['maxWriteBatchSize'] || Mongo::MongoClient::DEFAULT_MAX_WRITE_BATCH_SIZE
+    end
+  end
+end