mongo 1.3.1 → 1.4.0

data/lib/mongo/util/logging.rb

@@ -0,0 +1,42 @@
+module Mongo
+  module Logging
+
+    # Log a message with the given level.
+    def log(level, msg)
+      return unless @logger
+      case level
+        when :debug then
+          @logger.debug "MONGODB [DEBUG] #{msg}"
+        when :warn then
+          @logger.warn "MONGODB [WARNING] #{msg}"
+        when :error then
+          @logger.error "MONGODB [ERROR] #{msg}"
+        when :fatal then
+          @logger.fatal "MONGODB [FATAL] #{msg}"
+        else
+          @logger.info "MONGODB [INFO] #{msg}"
+      end
+    end
+
+    # Execute the block and log the operation described by name and payload.
+    def instrument(name, payload = {}, &blk)
+      res = yield
+      log_operation(name, payload)
+      res
+    end
+
+    protected
+
+    def log_operation(name, payload)
+      @logger ||= nil
+      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
+
+  end
+end

data/lib/mongo/util/node.rb

@@ -0,0 +1,183 @@
+module Mongo
+  class Node
+
+    attr_accessor :host, :port, :address, :config, :connection, :socket
+
+    def initialize(connection, data)
+      @connection = connection
+      if data.is_a?(String)
+        @host, @port = split_nodes(data)
+      else
+        @host = data[0]
+        @port = data[1].nil? ? Connection::DEFAULT_PORT : data[1].to_i
+      end
+      @address = "#{host}:#{port}"
+      @config = nil
+    end
+
+    def eql?(other)
+      other.is_a?(Node) && host == other.host && port == other.port
+    end
+    alias :== :eql?
+
+    def host_string
+      address
+    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
+      begin
+        socket = nil
+        if @connection.connect_timeout
+          Mongo::TimeoutHandler.timeout(@connection.connect_timeout, OperationTimeout) do
+            socket = @connection.socket_class.new(@host, @port)
+          end
+        else
+          socket = @connection.socket_class.new(@host, @port)
+        end
+
+        if socket.nil?
+          return nil
+        else
+          socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+        end
+      rescue OperationTimeout, OperationFailure, SocketError, SystemCallError, IOError => ex
+        @connection.log(:debug, "Failed connection to #{host_string} with #{ex.class}, #{ex.message}.")
+        socket.close if socket
+        return nil
+      end
+
+      @socket = socket
+    end
+
+    def close
+      if @socket
+        @socket.close
+        @socket = nil
+        @config = nil
+      end
+    end
+
+    def connected?
+      @socket != nil
+    end
+
+    def active?
+      begin
+        result = @connection['admin'].command({:ping => 1}, :socket => @socket)
+        return result['ok'] == 1
+      rescue OperationFailure, SocketError, SystemCallError, IOError => ex
+        return nil
+      end
+    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
+      begin
+        @config = @connection['admin'].command({:ismaster => 1}, :socket => @socket)
+
+        if @config['msg'] && @logger
+          @connection.log(:warn, "#{config['msg']}")
+        end
+
+        check_set_membership(config)
+        check_set_name(config)
+      rescue ConnectionFailure, OperationFailure, SocketError, SystemCallError, IOError => ex
+        @connection.log(:warn, "Attempted connection to node #{host_string} raised " +
+                            "#{ex.class}: #{ex.message}")
+        return nil
+      end
+
+      @config
+    end
+
+    # Return a list of replica set nodes from the config.
+    # Note: this excludes arbiters.
+    def node_list
+      connect unless connected?
+      set_config unless @config
+
+      return [] unless config
+
+      nodes = []
+      nodes += config['hosts'] if config['hosts']
+      nodes += config['passives'] if config['passives']
+      nodes
+    end
+
+    def arbiters
+      connect unless connected?
+      set_config unless @config
+      return [] unless config['arbiters']
+
+      config['arbiters'].map do |arbiter|
+        split_nodes(arbiter)
+      end
+    end
+
+    def tags
+      connect unless connected?
+      set_config unless @config
+      return {} unless config['tags'] && !config['tags'].empty?
+
+      config['tags']
+    end
+
+    def primary?
+      @config['ismaster'] == true || @config['ismaster'] == 1
+    end
+
+    def secondary?
+      @config['secondary'] == true || @config['secondary'] == 1
+    end
+
+    def host_port
+      [@host, @port]
+    end
+
+    def hash
+      address.hash
+    end
+
+    private
+
+    def split_nodes(host_string)
+      data = host_string.split(":")
+      host = data[0]
+      port = data[1].nil? ? Connection::DEFAULT_PORT : data[1].to_i
+
+      [host, port]
+    end
+
+    # Ensure that this node is a member of a replica set.
+    def check_set_membership(config)
+      if !config['hosts']
+        message = "Will not connect to #{host_string} because it's not a member " +
+          "of a 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 @connection.replica_set_name
+        if !config['setName']
+          @connection.log(:warn, "Could not verify replica set name for member #{host_string} " +
+            "because ismaster does not return name in this version of MongoDB")
+        elsif @connection.replica_set_name != config['setName']
+          message = "Attempting to connect to replica set '#{config['setName']}' on member #{host_string} " +
+            "but expected '#{@connection.replica_set_name}'"
+          raise ReplicaSetConnectionError, message
+        end
+      end
+    end
+  end
+end

data/lib/mongo/util/pool.rb

@@ -17,8 +17,9 @@
 
 module Mongo
   class Pool
+    PING_ATTEMPTS = 6
 
-    attr_accessor :host, :port, :size, :timeout, :safe, :checked_out
+    attr_accessor :host, :port, :size, :timeout, :safe, :checked_out, :connection
 
     # Create a new pool of connections.
     #
@@ -27,6 +28,9 @@ module Mongo
 
       @host, @port = host, port
 
+      # A Mongo::Node object.
+      @node = opts[:node]
+
       # Pool size and timeout.
       @size      = opts[:size] || 1
       @timeout   = opts[:timeout]   || 5.0
@@ -43,20 +47,77 @@ module Mongo
       @sockets      = []
       @pids         = {}
       @checked_out  = []
+      @ping_time    = nil
+      @last_ping    = nil
     end
 
     def close
-      @sockets.each do |sock|
-        begin
-          sock.close
-        rescue IOError => ex
-          warn "IOError when attempting to close socket connected to #{@host}:#{@port}: #{ex.inspect}"
+      @connection_mutex.synchronize do
+        @sockets.each do |sock|
+          begin
+            sock.close
+          rescue IOError => ex
+            warn "IOError when attempting to close socket connected to #{@host}:#{@port}: #{ex.inspect}"
+          end
         end
+        @host = @port = nil
+        @sockets.clear
+        @pids.clear
+        @checked_out.clear
       end
-      @host = @port = nil
-      @sockets.clear
-      @pids.clear
-      @checked_out.clear
+    end
+
+    def inspect
+      "#<Mongo::Pool:0x#{self.object_id.to_s(16)} @host=#{@host} @port=#{port} " +
+        "@ping_time=#{@ping_time} #{@checked_out.size}/#{@size} sockets available.>"
+    end
+
+    def host_string
+      "#{@host}:#{@port}"
+    end
+
+    def host_port
+      [@host, @port]
+    end
+
+    # Refresh ping time only if we haven't
+    # checked within the last five minutes.
+    def ping_time
+      if !@last_ping
+        @last_ping = Time.now
+        @ping_time = refresh_ping_time
+      elsif Time.now - @last_ping > 300
+        @last_ping = Time.now
+        @ping_time = refresh_ping_time
+      else
+        @ping_time
+      end
+    end
+
+    # Return the time it takes on average
+    # to do a round-trip against this node.
+    def refresh_ping_time
+      trials = []
+      begin
+         PING_ATTEMPTS.times do
+           t1 = Time.now
+           self.connection['admin'].command({:ping => 1}, :socket => @node.socket)
+           trials << (Time.now - t1) * 1000
+         end
+       rescue OperationFailure, SocketError, SystemCallError, IOError => ex
+         return nil
+      end
+
+      trials.sort!
+
+      # Delete shortest and longest times
+      trials.delete_at(trials.length-1)
+      trials.delete_at(0)
+
+      total = 0.0
+      trials.each { |t| total += t }
+
+      (total / trials.length).ceil
     end
 
     # Return a socket to the pool.
@@ -74,10 +135,12 @@ module Mongo
     # therefore, it runs within a mutex.
     def checkout_new_socket
       begin
-      socket = TCPSocket.new(@host, @port)
-      socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+        socket = self.connection.socket_class.new(@host, @port)
+        socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
       rescue => ex
+        socket.close if socket
         raise ConnectionFailure, "Failed to connect to host #{@host} and port #{@port}: #{ex}"
+        @node.close if @node
       end
 
       # If any saved authentications exist, we want to apply those
@@ -128,7 +191,7 @@ module Mongo
       if @pids[socket] != Process.pid
          @pids[socket] = nil
          @sockets.delete(socket)
-         socket.close
+         socket.close if socket
          checkout_new_socket
       else
         @checked_out << socket

data/lib/mongo/util/pool_manager.rb

@@ -0,0 +1,208 @@
+module Mongo
+  class PoolManager
+
+    attr_reader :connection, :seeds, :arbiters, :primary, :secondaries,
+      :primary_pool, :read_pool, :secondary_pools, :hosts, :nodes, :max_bson_size,
+      :tags_to_pools, :members
+
+    def initialize(connection, seeds)
+      @connection = connection
+      @seeds = seeds
+      @refresh_node = nil
+      @previously_connected = false
+    end
+
+    def inspect
+      "<Mongo::PoolManager:0x#{self.object_id.to_s(16)} @seeds=#{@seeds}>"
+    end
+
+    def connect
+      if @previously_connected
+        close
+      end
+
+      initialize_data
+      members = connect_to_members
+      initialize_pools(members)
+      update_seed_list(members)
+
+      @members = members
+      @previously_connected = true
+    end
+
+    def healthy?
+      if !@refresh_node || !refresh_node.set_config
+        return false
+      end
+
+      #if refresh_node.node_list
+    end
+
+    def close
+      begin
+        if @primary_pool
+          @primary_pool.close
+        end
+
+        if @secondary_pools
+          @secondary_pools.each do |pool|
+            pool.close
+          end
+        end
+
+        if @members
+          @members.each do |member|
+            member.close
+          end
+        end
+
+        rescue ConnectionFailure
+      end
+    end
+
+    private
+
+    def initialize_data
+      @primary = nil
+      @primary_pool = nil
+      @read_pool = nil
+      @arbiters = []
+      @secondaries = []
+      @secondary_pools = []
+      @hosts = Set.new
+      @members = Set.new
+      @tags_to_pools = {}
+    end
+
+    # Connect to each member of the replica set
+    # as reported by the given seed node, and return
+    # as a list of Mongo::Node objects.
+    def connect_to_members
+      members = []
+
+      seed = get_valid_seed_node
+
+      seed.node_list.each do |host|
+        node = Mongo::Node.new(self.connection, host)
+        if node.connect && node.set_config
+          members << node
+        end
+      end
+      seed.close
+
+      if members.empty?
+        raise ConnectionFailure, "Failed to connect to any given member."
+      end
+
+      members
+    end
+
+    def associate_tags_with_pool(tags, pool)
+      tags.each_key do |key|
+        @tags_to_pools[{key => tags[key]}] ||= []
+        @tags_to_pools[{key => tags[key]}] << pool
+      end
+    end
+
+    # Initialize the connection pools for the primary and secondary nodes.
+    def initialize_pools(members)
+      members.each do |member|
+        @hosts << member.host_string
+
+        if member.primary?
+          @primary = member.host_port
+          @primary_pool = Pool.new(self.connection, member.host, member.port,
+                                  :size => self.connection.pool_size,
+                                  :timeout => self.connection.connect_timeout,
+                                  :node => member)
+          associate_tags_with_pool(member.tags, @primary_pool)
+        elsif member.secondary? && !@secondaries.include?(member.host_port)
+          @secondaries << member.host_port
+          pool = Pool.new(self.connection, member.host, member.port,
+                                       :size => self.connection.pool_size,
+                                       :timeout => self.connection.connect_timeout,
+                                       :node => member)
+          @secondary_pools << pool
+          associate_tags_with_pool(member.tags, pool)
+        end
+      end
+
+
+      @max_bson_size = members.first.config['maxBsonObjectSize'] ||
+        Mongo::DEFAULT_MAX_BSON_SIZE
+      @arbiters = members.first.arbiters
+
+      set_read_pool
+      set_primary_tag_pools
+    end
+
+    # If there's more than one pool associated with
+    # a given tag, choose a close one using the bucket method.
+    def set_primary_tag_pools
+      @tags_to_pools.each do |k, pool_list|
+        if pool_list.length == 1
+          @tags_to_pools[k] = pool_list.first
+        else
+          @tags_to_pools[k] = nearby_pool_from_set(pool_list)
+        end
+      end
+    end
+
+    # Pick a node from the set of possible secondaries.
+    # If more than one node is available, use the ping
+    # time to figure out which nodes to choose from.
+    def set_read_pool
+      if @secondary_pools.empty?
+         @read_pool = @primary_pool
+      elsif @secondary_pools.size == 1
+         @read_pool = @secondary_pools[0]
+      else
+        @read_pool = nearby_pool_from_set(@secondary_pools)
+      end
+    end
+
+    def nearby_pool_from_set(pool_set)
+      ping_ranges = Array.new(3) { |i| Array.new }
+        pool_set.each do |pool|
+          case pool.ping_time
+            when 0..150
+              ping_ranges[0] << pool
+            when 150..1000
+              ping_ranges[1] << pool
+            else
+              ping_ranges[2] << pool
+          end
+        end
+
+        for list in ping_ranges do
+          break if !list.empty?
+        end
+
+      list[rand(list.length)]
+    end
+
+    # Iterate through the list of provided seed
+    # nodes until we've gotten a response from the
+    # replica set we're trying to connect to.
+    #
+    # If we don't get a response, raise an exception.
+    def get_valid_seed_node
+      @seeds.each do |seed|
+        node = Mongo::Node.new(self.connection, seed)
+        if node.connect && node.set_config
+          return node
+        else
+          node.close
+        end
+      end
+
+      raise ConnectionFailure, "Cannot connect to a replica set using seeds " +
+        "#{@seeds.map {|s| "#{s[0]}:#{s[1]}" }.join(', ')}"
+    end
+
+    def update_seed_list(members)
+      @seeds = members.map { |n| n.host_port }
+    end
+
+  end
+end