CVE-2015-4411

@@ -282,7 +282,7 @@ def from_time(time, options = {})
       #
       # @since 2.0.0
       def legal?(string)
-        string.to_s =~ /^[0-9a-f]{24}$/i ? true : false
+        string.to_s =~ /\A[0-9a-f]{24}\z/i ? true : false
       end
 
       # Executes the provided block only if the size of the provided object is

@@ -1,4 +1,7 @@
 # encoding: utf-8
+require "digest/md5"
+require "socket"
+
 module BSON
 
   # Represents object_id data.
@@ -7,25 +10,254 @@ module BSON
   #
   # @since 2.0.0
   class ObjectId
+    include Comparable
 
     # A object_id is type 0x07 in the BSON spec.
     #
     # @since 2.0.0
     BSON_TYPE = 7.chr.force_encoding(BINARY).freeze
 
-    # Encode the object_id type
+    # Check equality of the object id with another object.
+    #
+    # @example Check if the object id is equal to the other.
+    #   object_id == other
+    #
+    # @param [ Object ] other The object to check against.
+    #
+    # @return [ true, false ] If the objects are equal.
+    #
+    # @since 2.0.0
+    def ==(other)
+      return false unless other.is_a?(ObjectId)
+      to_bson == other.to_bson
+    end
+
+    # Check case equality on the object id.
+    #
+    # @example Check case equality.
+    #   object_id === other
+    #
+    # @param [ Object ] other The object to check against.
+    #
+    # @return [ true, false ] If the objects are equal in a case.
+    #
+    # @since 2.0.0
+    def ===(other)
+      return to_str === other.to_str if other.respond_to?(:to_str)
+      super
+    end
+
+    # Compare this object id with another object for use in sorting.
     #
-    # @example Encode the object_id.
+    # @example Compare the object id with the other object.
+    #   object <=> other
+    #
+    # @param [ Object ] other The object to compare to.
+    #
+    # @return [ Integer ] The result of the comparison.
+    #
+    # @since 2.0.0
+    def <=>(other)
+      to_bson <=> other.to_bson
+    end
+
+    # Return the UTC time at which this ObjectId was generated. This may
+    # be used instread of a created_at timestamp since this information
+    # is always encoded in the object id.
+    #
+    # @example Get the generation time.
+    #   object_id.generation_time
+    #
+    # @return [ Time ] The time the id was generated.
+    #
+    # @since 2.0.0
+    def generation_time
+      ::Time.at(to_bson.unpack("N")[0]).utc
+    end
+
+    # Get the object id as it's raw BSON data.
+    #
+    # @example Get the raw bson bytes.
     #   object_id.to_bson
     #
-    # @return [ String ] The encoded object_id.
+    # @note Since Moped's BSON and 10gen BSON before 2.0.0 have different
+    #   internal representations, we will attempt to repair the data for cases
+    #   where the object was instantiated in a non-standard way. (Like a
+    #   Marshal.load)
+    #
+    # @return [ String ] The raw bytes.
     #
     # @see http://bsonspec.org/#/specification
     #
     # @since 2.0.0
     def to_bson
+      repair!(@data) if defined?(@data)
+      @raw_data ||= @@generator.next
     end
 
+    # Get the string representation of the object id.
+    #
+    # @example Get the object id as a string.
+    #   object_id.to_s
+    #
+    # @return [ String ] The object id as a string.
+    #
+    # @since 2.0.0
+    def to_s
+      to_bson.unpack("H*")[0].force_encoding(UTF8)
+    end
+    alias :to_str :to_s
+
+    # Raised when trying to create an object id with invalid data.
+    #
+    # @since 2.0.0
+    class Invalid < RuntimeError; end
+
+    private
+
+    def data=(data)
+      @raw_data = data
+    end
+
+    class << self
+
+      # Create a new object id from raw bytes.
+      #
+      # @example Create an object id from raw bytes.
+      #   BSON::ObjectId.from_data(data)
+      #
+      # @param [ String ] data The raw bytes.
+      #
+      # @return [ ObjectId ] The new object id.
+      #
+      # @since 2.0.0
+      def from_data(data)
+        object_id = allocate
+        object_id.send(:data=, data)
+        object_id
+      end
+
+      # Create a new object id from a string.
+      #
+      # @example Create an object id from the string.
+      #   BSON::ObjectId.from_string(id)
+      #
+      # @param [ String ] string The string to create the id from.
+      #
+      # @raise [ BSON::ObjectId::Invalid ] If the provided string is invalid.
+      #
+      # @return [ BSON::ObjectId ] The new object id.
+      #
+      # @since 2.0.0
+      def from_string(string)
+        unless legal?(string)
+          raise Invalid.new("'#{string}' is an invalid ObjectId.")
+        end
+        from_data([ string ].pack("H*"))
+      end
+
+      # Create a new object id from a time.
+      #
+      # @example Create an object id from a time.
+      #   BSON::ObjectId.from_id(time)
+      #
+      # @example Create an object id from a time, ensuring uniqueness.
+      #   BSON::ObjectId.from_id(time, unique: true)
+      #
+      # @param [ Time ] time The time to generate from.
+      # @param [ Hash ] options The options.
+      #
+      # @option options [ true, false ] :unique Whether the id should be
+      #   unique.
+      #
+      # @return [ ObjectId ] The new object id.
+      #
+      # @since 2.0.0
+      def from_time(time, options = {})
+        from_data(options[:unique] ? @@generator.next(time.to_i) : [ time.to_i ].pack("Nx8"))
+      end
+
+      # Determine if the provided string is a legal object id.
+      #
+      # @example Is the string a legal object id?
+      #   BSON::ObjectId.legal?(string)
+      #
+      # @param [ String ] The string to check.
+      #
+      # @return [ true, false ] If the string is legal.
+      #
+      # @since 2.0.0
+      def legal?(string)
+        /\A\h{24}\Z/ === string.to_s
+      end
+    end
+
+    # Inner class that encapsulates the behaviour of actually generating each
+    # part of the ObjectId.
+    #
+    # @api private
+    #
+    # @since 2.0.0
+    class Generator
+
+      attr_reader :machine_id
+
+      # Instantiate the new object id generator. Will set the machine id once
+      # on the initial instantiation.
+      #
+      # @example Instantiate the generator.
+      #   BSON::ObjectId::Generator.new
+      #
+      # @since 2.0.0
+      def initialize
+        @counter = 0
+        @machine_id = Digest::MD5.digest(Socket.gethostname).unpack("N")[0]
+        @mutex = Mutex.new
+      end
+
+      # Return object id data based on the current time, incrementing the
+      # object id counter. Will use the provided time if not nil.
+      #
+      # @example Get the next object id data.
+      #   generator.next
+      #
+      # @param [ Time ] time The optional time to generate with.
+      #
+      # @return [ String ] The raw object id bytes.
+      #
+      # @since 2.0.0
+      def next(time = nil)
+        @mutex.lock
+        begin
+          count = @counter = (@counter + 1) % 0xFFFFFF
+        ensure
+          @mutex.unlock rescue nil
+        end
+        generate(time || ::Time.new.to_i, count)
+      end
+
+      # Generate object id data for a given time using the provided counter.
+      #
+      # @example Generate the object id bytes.
+      #   generator.generate(time)
+      #
+      # @param [ Integer ] time The time since epoch in seconds.
+      # @param [ Integer ] counter The optional counter.
+      #
+      # @return [ String ] The raw object id bytes.
+      #
+      # @since 2.0.0
+      def generate(time, counter = 0)
+        process_thread_id = "#{Process.pid}#{Thread.current.object_id}".hash % 0xFFFF
+        [ time, machine_id, process_thread_id, counter << 8 ].pack("N NX lXX NX")
+      end
+    end
+
+    # We keep one global generator for object ids.
+    #
+    # @since 2.0.0
+    @@generator = Generator.new
+
     # Register this type when the module is loaded.
     #
     # @since 2.0.0

@@ -19,6 +19,11 @@ module BSON
   #
   # @since 2.0.0
   NULL_BYTE = 0.chr.force_encoding(BINARY).freeze
+
+  # Constant for UTF-8 string encoding.
+  #
+  # @since 2.0.0
+  UTF8 = "UTF-8".freeze
 end
 
 require "bson/registry"

@@ -14,11 +14,6 @@ module String
     # @since 2.0.0
     BSON_TYPE = 2.chr.force_encoding(BINARY).freeze
 
-    # Constant for UTF-8 string encoding.
-    #
-    # @since 2.0.0
-    UTF8 = "UTF-8".freeze
-
     # Get the string as encoded BSON.
     #
     # @example Get the string as encoded BSON.

@@ -3,6 +3,203 @@
 
 describe BSON::ObjectId do
 
+  describe "#==" do
+
+    context "when data is identical" do
+
+      let(:time) do
+        Time.now
+      end
+
+      let(:object_id) do
+        described_class.from_time(time)
+      end
+
+      let(:other_id) do
+        described_class.from_time(time)
+      end
+
+      it "returns true" do
+        expect(object_id).to eq(other_id)
+      end
+    end
+
+    context "when the data is different" do
+
+      let(:time) do
+        Time.now
+      end
+
+      let(:object_id) do
+        described_class.from_time(time)
+      end
+
+      it "returns false" do
+        expect(object_id).to_not eq(described_class.new)
+      end
+    end
+
+    context "when other is not an object id" do
+
+      it "returns false" do
+        expect(described_class.new).to_not eq(nil)
+      end
+    end
+  end
+
+  describe "#===" do
+
+    let(:object_id) do
+      described_class.new
+    end
+
+    context "when comparing with another object id" do
+
+      context "when the data is equal" do
+
+        let(:other) do
+          described_class.from_string(object_id.to_s)
+        end
+
+        it "returns true" do
+          expect(object_id === other).to be_true
+        end
+      end
+
+      context "when the data is not equal" do
+
+        let(:other) do
+          described_class.new
+        end
+
+        it "returns false" do
+          expect(object_id === other).to be_false
+        end
+      end
+    end
+
+    context "when comparing to an object id class" do
+
+      it "returns false" do
+        expect(object_id === BSON::ObjectId).to be_false
+      end
+    end
+
+    context "when comparing with a string" do
+
+      context "when the data is equal" do
+
+        let(:other) do
+          object_id.to_s
+        end
+
+        it "returns true" do
+          expect(object_id === other).to be_true
+        end
+      end
+
+      context "when the data is not equal" do
+
+        let(:other) do
+          described_class.new.to_s
+        end
+
+        it "returns false" do
+          expect(object_id === other).to be_false
+        end
+      end
+    end
+
+    context "when comparing with a non string or object id" do
+
+      it "returns false" do
+        expect(object_id === "test").to be_false
+      end
+    end
+
+    context "when comparing with a non object id class" do
+
+      it "returns false" do
+        expect(object_id === String).to be_false
+      end
+    end
+  end
+
+  describe "#<" do
+
+    let(:object_id) do
+      described_class.from_time(Time.new(2012, 1, 1))
+    end
+
+    let(:other_id) do
+      described_class.from_time(Time.new(2012, 1, 30))
+    end
+
+    context "when the generation time before the other" do
+
+      it "returns true" do
+        expect(object_id < other_id).to be_true
+      end
+    end
+
+    context "when the generation time is after the other" do
+
+      it "returns false" do
+        expect(other_id < object_id).to be_false
+      end
+    end
+  end
+
+  describe "#>" do
+
+    let(:object_id) do
+      described_class.from_time(Time.new(2012, 1, 1))
+    end
+
+    let(:other_id) do
+      described_class.from_time(Time.new(2012, 1, 30))
+    end
+
+    context "when the generation time before the other" do
+
+      it "returns false" do
+        expect(object_id > other_id).to be_false
+      end
+    end
+
+    context "when the generation time is after the other" do
+
+      it "returns true" do
+        expect(other_id > object_id).to be_true
+      end
+    end
+  end
+
+  describe "#<=>" do
+
+    let(:object_id) do
+      described_class.from_time(Time.new(2012, 1, 1))
+    end
+
+    let(:other_id) do
+      described_class.from_time(Time.new(2012, 1, 30))
+    end
+
+    context "when the generation time before the other" do
+
+      it "returns -1" do
+        expect(object_id <=> other_id).to eq(-1)
+      end
+    end
+
+    context "when the generation time is after the other" do
+
+      it "returns false" do
+        expect(other_id <=> object_id).to eq(1)
+      end
+    end
+  end
+
   describe "::BSON_TYPE" do
 
     it "returns 0x07" do
@@ -21,7 +218,150 @@
     end
   end
 
-  pending "#to_bson"
+  describe ".from_string" do
+
+    context "when the string is valid" do
+
+      let(:string) do
+        "4e4d66343b39b68407000001"
+      end
+
+      let(:object_id) do
+        described_class.from_string(string)
+      end
+
+      it "initializes with the string's bytes" do
+        expect(object_id.to_s).to eq(string)
+      end
+    end
+
+    context "when the string is not valid" do
+
+      it "raises an error" do
+        expect {
+          described_class.from_string("asadsf")
+        }.to raise_error(BSON::ObjectId::Invalid)
+      end
+    end
+  end
+
+  describe ".from_time" do
+
+    context "when no unique option is provided" do
+
+      let(:time) do
+        Time.at((Time.now.utc - 64800).to_i).utc
+      end
+
+      let(:object_id) do
+        described_class.from_time(time)
+      end
+
+      it "sets the generation time" do
+        expect(object_id.generation_time).to eq(time)
+      end
+
+      it "does not include process or sequence information" do
+        expect(object_id.to_s =~ /\A\h{8}0{16}\Z/).to be_true
+      end
+    end
+
+    context "when a unique option is provided" do
+
+      let(:time) do
+        Time.at((Time.now.utc - 64800).to_i).utc
+      end
+
+      let(:object_id) do
+        described_class.from_time(time, unique: true)
+      end
+
+      let(:non_unique) do
+        described_class.from_time(time, unique: true)
+      end
+
+      it "creates a new unique object id" do
+        expect(object_id).to_not eq(non_unique)
+      end
+    end
+  end
+
+  describe ".legal?" do
+
+    context "when the string is too short to be an object id" do
+
+      it "returns false" do
+        expect(described_class).to_not be_legal("a" * 23)
+      end
+    end
+
+    context "when the string contains invalid hex characters" do
+
+      it "returns false" do
+        expect(described_class).to_not be_legal("y" + "a" * 23)
+      end
+    end
+
+    context "when the string is a valid object id" do
+
+      it "returns true" do
+        expect(described_class).to be_legal("a" * 24)
+      end
+    end
+
+    context "when checking against another object id" do
+
+      let(:object_id) do
+        described_class.new
+      end
+
+      it "returns true" do
+        expect(described_class).to be_legal(object_id)
+      end
+    end
+  end
+
+  describe "#to_bson" do
+
+    let(:object_id) do
+      described_class.new
+    end
+
+    let(:data) do
+      object_id.instance_variable_get(:@raw_data)
+    end
+
+    it "returns the raw bytes" do
+      expect(object_id.to_bson).to eq(data)
+    end
+  end
+
+  describe "#to_s" do
+
+    let(:time) do
+      Time.new(2013, 1, 1)
+    end
+
+    let(:expected) do
+      "50e218f00000000000000000"
+    end
+
+    let(:object_id) do
+      described_class.from_time(time)
+    end
+
+    it "returns a hex string representation of the id" do
+      expect(object_id.to_s).to eq(expected)
+    end
+
+    it "returns the string in UTF-8" do
+      expect(object_id.to_s.encoding).to eq(Encoding.find(BSON::UTF8))
+    end
+
+    it "converts to a readable yaml string" do
+      expect(YAML.dump(object_id.to_s)).to include(expected)
+    end
+  end
 
   context "when the class is loaded" do
 

@@ -2,5 +2,6 @@ Gemfile.lock
 doc
 .yardoc
 .rvmrc
+.env
 perf/results
 tmp/

@@ -8,31 +8,33 @@ class ObjectId
       # Formatting string for outputting an ObjectId.
       @@string_format = ("%02x" * 12).freeze
 
-      attr_reader :data
-
       class << self
         def from_string(string)
           raise Errors::InvalidObjectId.new(string) unless legal?(string)
-          data = []
+          data = ""
           12.times { |i| data << string[i*2, 2].to_i(16) }
-          new data
+          from_data data
+        end
+
+        def from_time(time)
+          from_data @@generator.generate(time.to_i)
         end
 
         def legal?(str)
-          !!str.match(/^[0-9a-f]{24}$/i)
+          !!str.match(/\A\h{24}\Z/i)
         end
-      end
 
-      def initialize(data = nil, time = nil)
-        if data
-          @data = data
-        elsif time
-          @data = @@generator.generate(time.to_i)
-        else
-          @data = @@generator.next
+        def from_data(data)
+          id = allocate
+          id.instance_variable_set :@data, data
+          id
         end
       end
 
+      def data
+        @data ||= @@generator.next
+      end
+
       def ==(other)
         BSON::ObjectId === other && data == other.data
       end
@@ -43,78 +45,56 @@ def hash
       end
 
       def to_s
-        @@string_format % data
+        @@string_format % data.unpack("C12")
       end
 
       # Return the UTC time at which this ObjectId was generated. This may
       # be used instread of a created_at timestamp since this information
       # is always encoded in the object id.
       def generation_time
-        Time.at(@data.pack("C4").unpack("N")[0]).utc
+        Time.at(data.unpack("N")[0]).utc
       end
 
       class << self
         def __bson_load__(io)
-          new io.read(12).unpack('C*')
+          from_data(io.read(12))
         end
-
       end
 
       def __bson_dump__(io, key)
         io << Types::OBJECT_ID
         io << key
         io << NULL_BYTE
-        io << data.pack('C12')
+        io << data
       end
 
       # @api private
       class Generator
         def initialize
           # Generate and cache 3 bytes of identifying information from the current
           # machine.
-          @machine_id = Digest::MD5.digest(Socket.gethostname).unpack("C3")
+          @machine_id = Digest::MD5.digest(Socket.gethostname).unpack("N")[0]
 
           @mutex = Mutex.new
-          @last_timestamp = nil
           @counter = 0
         end
 
-        # Return object id data based on the current time, incrementing a
-        # counter for object ids generated in the same second.
+        # Return object id data based on the current time, incrementing the
+        # object id counter.
         def next
-          now = Time.new.to_i
-
-          counter = @mutex.synchronize do
-            last_timestamp, @last_timestamp = @last_timestamp, now
-
-            if last_timestamp == now
-              @counter += 1
-            else
-              @counter = 0
-            end
+          @mutex.lock
+          begin
+            counter = @counter = (@counter + 1) % 0xFFFFFF
+          ensure
+            @mutex.unlock rescue nil
           end
 
-          generate(now, counter)
+          generate(Time.new.to_i, counter)
         end
 
-        # Generate object id data for a given time using the provided +inc+.
-        def generate(time, inc = 0)
-          pid = Process.pid % 0xFFFF
-
-          [
-            time >> 24 & 0xFF, # 4 bytes time (network order)
-            time >> 16 & 0xFF,
-            time >> 8  & 0xFF,
-            time       & 0xFF,
-            @machine_id[0],   # 3 bytes machine
-            @machine_id[1],
-            @machine_id[2],
-            pid  >> 8  & 0xFF, # 2 bytes process id
-            pid        & 0xFF,
-            inc  >> 16 & 0xFF, # 3 bytes increment
-            inc  >> 8  & 0xFF,
-            inc        & 0xFF,
-          ]
+        # Generate object id data for a given time using the provided +counter+.
+        def generate(time, counter = 0)
+          [time, @machine_id, Process.pid, counter << 8].pack("N NX lXX NX")
         end
       end
 

@@ -1,193 +1,143 @@
 module Moped
 
-  # @api private
-  #
-  # The internal class managing connections to both a single node and replica
-  # sets.
-  #
-  # @note Though the socket class itself *is* threadsafe, the cluster presently
-  #   is not. This means that in the course of normal operations sessions can be
-  #   shared across threads, but in failure modes (when a resync is required),
-  #   things can possibly go wrong.
   class Cluster
 
-    # @return [Array] the user supplied seeds
+    # @return [Array<String>] the seeds the replica set was initialized with
     attr_reader :seeds
 
-    # @return [Boolean] whether this is a direct connection
-    attr_reader :direct
-
-    # @return [Array] all available nodes
-    attr_reader :servers
-
-    # @return [Array] seeds gathered from cluster discovery
-    attr_reader :dynamic_seeds
-
-    # @param [Array] seeds an array of host:port pairs
-    # @param [Boolean] direct (false) whether to connect directly to the hosts
-    # provided or to find additional available nodes.
-    def initialize(seeds, direct = false)
-      @seeds  = seeds
-      @direct = direct
-
-      @servers = []
-      @dynamic_seeds = []
-    end
-
-    # @return [Array] available secondary nodes
-    def secondaries
-      servers.select(&:secondary?)
-    end
-
-    # @return [Array] available primary nodes
-    def primaries
-      servers.select(&:primary?)
-    end
-
-    # @return [Array] all known addresses from user supplied seeds, dynamically
-    # discovered seeds, and active servers.
-    def known_addresses
-      [].tap do |addresses|
-        addresses.concat seeds
-        addresses.concat dynamic_seeds
-        addresses.concat servers.map { |server| server.address }
-      end.uniq
-    end
-
-    def remove(server)
-      servers.delete(server)
+    # @option options :down_interval number of seconds to wait before attempting
+    # to reconnect to a down node. (30)
+    #
+    # @option options :refresh_interval number of seconds to cache information
+    # about a node. (300)
+    def initialize(hosts, options)
+      @options = {
+        down_interval: 30,
+        refresh_interval: 300
+      }.merge(options)
+
+      @seeds = hosts
+      @nodes = hosts.map { |host| Node.new(host) }
     end
 
-    def reconnect
-      @servers = servers.map { |server| Server.new(server.address) }
-    end
-
-    def sync
-      known = known_addresses.shuffle
-      seen  = {}
-
-      sync_seed = ->(seed) do
-        server = Server.new seed
-
-        unless seen[server.resolved_address]
-          seen[server.resolved_address] = true
-
-          hosts = sync_server(server)
-
-          hosts.each do |host|
-            sync_seed[host]
-          end
+    # Refreshes information for each of the nodes provided. The node list
+    # defaults to the list of all known nodes.
+    #
+    # If a node is successfully refreshed, any newly discovered peers will also
+    # be refreshed.
+    #
+    # @return [Array<Node>] the available nodes
+    def refresh(nodes_to_refresh = @nodes)
+      refreshed_nodes = []
+      seen = {}
+
+      # Set up a recursive lambda function for refreshing a node and it's peers.
+      refresh_node = ->(node) do
+        return if seen[node]
+        seen[node] = true
+
+        # Add the node to the global list of known nodes.
+        @nodes << node unless @nodes.include?(node)
+
+        begin
+          node.refresh
+
+          # This node is good, so add it to the list of nodes to return.
+          refreshed_nodes << node unless refreshed_nodes.include?(node)
+
+          # Now refresh any newly discovered peer nodes.
+          (node.peers - @nodes).each &refresh_node
+        rescue Errors::ConnectionFailure
+          # We couldn't connect to the node, so don't do anything with it.
         end
       end
 
-      known.each do |seed|
-        sync_seed[seed]
-      end
-
-      unless servers.empty?
-        @dynamic_seeds = servers.map(&:address)
-      end
-
-      true
+      nodes_to_refresh.each &refresh_node
+      refreshed_nodes.to_a
     end
 
-    def sync_server(server)
-      [].tap do |hosts|
-        socket = server.socket
-
-        if socket.connect
-          info = socket.simple_query Protocol::Command.new(:admin, ismaster: 1)
-
-          if info["ismaster"]
-            server.primary = true
-          end
+    # Returns the list of available nodes, refreshing 1) any nodes which were
+    # down and ready to be checked again and 2) any nodes whose information is
+    # out of date.
+    #
+    # @return [Array<Node>] the list of available nodes.
+    def nodes
+      # Find the nodes that were down but are ready to be refreshed, or those
+      # with stale connection information.
+      needs_refresh, available = @nodes.partition do |node|
+        (node.down? && node.down_at < (Time.new - @options[:down_interval])) ||
+          node.needs_refresh?(Time.new - @options[:refresh_interval])
+      end
 
-          if info["secondary"]
-            server.secondary = true
-          end
+      # Refresh those nodes.
+      available.concat refresh(needs_refresh)
 
-          if info["primary"]
-            hosts.push info["primary"]
-          end
-
-          if info["hosts"]
-            hosts.concat info["hosts"]
-          end
+      # Now return all the nodes that are available.
+      available.reject &:down?
+    end
 
-          if info["passives"]
-            hosts.concat info["passives"]
+    # Yields the replica set's primary node to the provided block. This method
+    # will retry the block in case of connection errors or replica set
+    # reconfiguration.
+    #
+    # @raises ConnectionFailure when no primary node can be found
+    def with_primary(retry_on_failure = true, &block)
+      if node = nodes.find(&:primary?)
+        begin
+          node.ensure_primary do
+            return yield node.apply_auth(auth)
           end
-
-          merge(server)
-
+        rescue Errors::ConnectionFailure, Errors::ReplicaSetReconfigured
+          # Fall through to the code below if our connection was dropped or the
+          # node is no longer the primary.
         end
-      end.uniq
-    end
-
-    def merge(server)
-      previous = servers.find { |other| other == server }
-      primary = server.primary?
-      secondary = server.secondary?
+      end
 
-      if previous
-        previous.merge(server)
+      if retry_on_failure
+        # We couldn't find a primary node, so refresh the list and try again.
+        refresh
+        with_primary(false, &block)
       else
-        servers << server
+        raise Errors::ConnectionFailure, "Could not connect to a primary node for replica set #{inspect}"
       end
     end
 
-    # @param [:read, :write] mode the type of socket to return
-    # @return [Socket] a socket valid for +mode+ operations
-    def socket_for(mode)
-      sync unless primaries.any? || (secondaries.any? && mode == :read)
-
-      server = nil
-      while primaries.any? || (secondaries.any? && mode == :read)
-        if mode == :write || secondaries.empty?
-          server = primaries.sample
-        else
-          server = secondaries.sample
-        end
-
-        if server
-          socket = server.socket
-          socket.connect unless socket.connection
-
-          if socket.alive?
-            break server
-          else
-            remove server
-          end
+    # Yields a secondary node if available, otherwise the primary node. This
+    # method will retry the block in case of connection errors.
+    #
+    # @raises ConnectionError when no secondary or primary node can be found
+    def with_secondary(retry_on_failure = true, &block)
+      available_nodes = nodes.shuffle!.partition(&:secondary?).flatten
+
+      while node = available_nodes.shift
+        begin
+          return yield node.apply_auth(auth)
+        rescue Errors::ConnectionFailure
+          # That node's no good, so let's try the next one.
+          next
         end
       end
 
-      unless server
-        raise Errors::ConnectionFailure.new("Could not connect to any primary or secondary servers")
+      if retry_on_failure
+        # We couldn't find a secondary or primary node, so refresh the list and
+        # try again.
+        refresh
+        with_secondary(false, &block)
+      else
+        raise Errors::ConnectionFailure, "Could not connect to any secondary or primary nodes for replica set #{inspect}"
       end
-
-      socket = server.socket
-      socket.apply_auth auth
-      socket
     end
 
     # @return [Hash] the cached authentication credentials for this cluster.
     def auth
       @auth ||= {}
     end
 
-    # Log in to +database+ with +username+ and +password+. Does not perform the
-    # actual log in, but saves the credentials for later authentication on a
-    # socket.
-    def login(database, username, password)
-      auth[database.to_s] = [username, password]
-    end
+    private
 
-    # Log out of +database+. Does not perform the actual log out, but will log
-    # out when the socket is used next.
-    def logout(database)
-      auth.delete(database.to_s)
+    def initialize_copy(_)
+      @nodes = @nodes.map &:dup
     end
 
   end
-
 end

@@ -56,12 +56,10 @@ def find(selector = {})
     #   @param [Array<Hash>] documents the documents to insert
     def insert(documents)
       documents = [documents] unless documents.is_a? Array
-      insert = Protocol::Insert.new(database.name, name, documents)
 
       database.session.with(consistency: :strong) do |session|
-        session.execute insert
+        session.context.insert(database.name, name, documents)
       end
-
     end
   end
 end

@@ -0,0 +1,95 @@
+require "timeout"
+
+module Moped
+  class Connection
+
+    class TCPSocket < ::TCPSocket
+      def self.connect(host, port, timeout)
+        Timeout::timeout(timeout) do
+          new(host, port).tap do |sock|
+            sock.set_encoding 'binary'
+          end
+        end
+      end
+
+      def alive?
+        if Kernel::select([self], nil, nil, 0)
+          !eof? rescue false
+        else
+          true
+        end
+      end
+
+      def write(*args)
+        raise Errors::ConnectionFailure, "Socket connection was closed by remote host" unless alive?
+        super
+      end
+    end
+
+    def initialize
+      @sock = nil
+      @request_id = 0
+    end
+
+    def connect(host, port, timeout)
+      @sock = TCPSocket.connect host, port, timeout
+    end
+
+    def alive?
+      connected? ? @sock.alive? : false
+    end
+
+    def connected?
+      !!@sock
+    end
+
+    def disconnect
+      @sock.close
+    rescue
+    ensure
+      @sock = nil
+    end
+
+    def write(operations)
+      buf = ""
+
+      operations.each do |operation|
+        operation.request_id = (@request_id += 1)
+        operation.serialize(buf)
+      end
+
+      @sock.write buf
+    end
+
+    def receive_replies(operations)
+      operations.map do |operation|
+        read if operation.is_a?(Protocol::Query) || operation.is_a?(Protocol::GetMore)
+      end
+    end
+
+    def read
+      reply = Protocol::Reply.allocate
+
+      reply.length,
+        reply.request_id,
+        reply.response_to,
+        reply.op_code,
+        reply.flags,
+        reply.cursor_id,
+        reply.offset,
+        reply.count = @sock.read(36).unpack('l5<q<l2<')
+
+      if reply.count == 0
+        reply.documents = []
+      else
+        buffer = StringIO.new(@sock.read(reply.length - 36))
+
+        reply.documents = reply.count.times.map do
+          BSON::Document.deserialize(buffer)
+        end
+      end
+
+      reply
+    end
+  end
+end

@@ -10,50 +10,69 @@ class Cursor
 
     def initialize(session, query_operation)
       @session = session
-      @query_op = query_operation.dup
 
-      @get_more_op = Protocol::GetMore.new(
-        @query_op.database,
-        @query_op.collection,
-        0,
-        @query_op.limit
-      )
+      @database    = query_operation.database
+      @collection  = query_operation.collection
+      @selector    = query_operation.selector
 
-      @kill_cursor_op = Protocol::KillCursors.new([0])
+      @cursor_id = 0
+      @limit = query_operation.limit
+      @limited = @limit > 0
+
+      @options = {
+        request_id: query_operation.request_id,
+        flags: query_operation.flags,
+        limit: query_operation.limit,
+        skip: query_operation.skip,
+        fields: query_operation.fields
+      }
     end
 
     def each
-      documents = query @query_op
+      documents = load
       documents.each { |doc| yield doc }
 
       while more?
-        return kill if limited? && @get_more_op.limit <= 0
+        return kill if limited? && @limit <= 0
 
-        documents = query @get_more_op
+        documents = get_more
         documents.each { |doc| yield doc }
       end
     end
 
-    def query(operation)
-      reply = session.query operation
+    def load
+      consistency = session.consistency
+      @options[:flags] |= [:slave_ok] if consistency == :eventual
+
+      reply, @node = session.context.with_node do |node|
+        [node.query(@database, @collection, @selector, @options), node]
+      end
 
-      @get_more_op.limit -= reply.count if limited?
-      @get_more_op.cursor_id = reply.cursor_id
-      @kill_cursor_op.cursor_ids = [reply.cursor_id]
+      @limit -= reply.count if limited?
+      @cursor_id = reply.cursor_id
 
       reply.documents
     end
 
     def limited?
-      @query_op.limit > 0
+      @limited
     end
 
     def more?
-      @get_more_op.cursor_id != 0
+      @cursor_id != 0
+    end
+
+    def get_more
+      reply = @node.get_more @database, @collection, @cursor_id, @limit
+
+      @limit -= reply.count if limited?
+      @cursor_id = reply.cursor_id
+
+      reply.documents
     end
 
     def kill
-      session.execute kill_cursor_op
+      @node.kill_cursors [@cursor_id]
     end
   end
 

@@ -29,20 +29,22 @@ def initialize(session, name)
 
     # Drop the database.
     def drop
-      command dropDatabase: 1
+      session.with(consistency: :strong) do |session|
+        session.context.command name, dropDatabase: 1
+      end
     end
 
     # Log in with +username+ and +password+ on the current database.
     #
     # @param [String] username the username
     # @param [String] password the password
     def login(username, password)
-      session.cluster.login(name, username, password)
+      session.context.login(name, username, password)
     end
 
     # Log out from the current database.
     def logout
-      session.cluster.logout(name)
+      session.context.logout(name)
     end
 
     # Run +command+ on the database.
@@ -54,17 +56,7 @@ def logout
     # @param [Hash] command the command to run
     # @return [Hash] the result of the command
     def command(command)
-      operation = Protocol::Command.new(name, command)
-
-      result = session.with(consistency: :strong) do |session|
-        session.simple_query(operation)
-      end
-
-      raise Errors::OperationFailure.new(
-        operation, result
-      ) unless result["ok"] == 1.0
-
-      result
+      session.context.command name, command
     end
 
     # @param [Symbol, String] collection the collection name

@@ -1,6 +1,4 @@
 module Moped
-
-  # The namespace for all errors generated by Moped.
   module Errors
 
     # Mongo's exceptions are sparsely documented, but this is the most accurate
@@ -10,21 +8,12 @@ module Errors
     # Generic error class for exceptions related to connection failures.
     class ConnectionFailure < StandardError; end
 
-    # Raised when providing an invalid string from an object id.
-    class InvalidObjectId < StandardError
-      def initialize(string)
-        super("'#{string}' is not a valid object id.")
-      end
-    end
+    # Tag applied to unhandled exceptions on a node.
+    module SocketError end
 
     # Generic error class for exceptions generated on the remote MongoDB
     # server.
-    class MongoError < StandardError; end
-
-    # Exception class for exceptions generated as a direct result of an
-    # operation, such as a failed insert or an invalid command.
-    class OperationFailure < MongoError
-
+    class MongoError < StandardError
       # @return the command that generated the error
       attr_reader :command
 
@@ -60,9 +49,28 @@ def error_message
       end
     end
 
-    # A special kind of OperationFailure, raised when Mongo sets the
-    # :query_failure flag on a query response.
-    class QueryFailure < OperationFailure; end
+    # Exception class for exceptions generated as a direct result of an
+    # operation, such as a failed insert or an invalid command.
+    class OperationFailure < MongoError; end
+
+    # Exception raised on invalid queries.
+    class QueryFailure < MongoError; end
+
+    # Exception raised when authentication fails.
+    class AuthenticationFailure < MongoError; end
+
+    # Raised when providing an invalid string from an object id.
+    class InvalidObjectId < StandardError
+      def initialize(string)
+        super("'#{string}' is not a valid object id.")
+      end
+    end
+
+    # @api private
+    #
+    # Internal exception raised by Node#ensure_primary and captured by
+    # Cluster#with_primary.
+    class ReplicaSetReconfigured < StandardError; end
 
   end
 end

@@ -0,0 +1,334 @@
+module Moped
+  class Node
+
+    attr_reader :address
+    attr_reader :resolved_address
+    attr_reader :ip_address
+    attr_reader :port
+
+    attr_reader :peers
+    attr_reader :timeout
+
+    def initialize(address)
+      @address = address
+
+      host, port = address.split(":")
+      @ip_address = ::Socket.getaddrinfo(host, nil, ::Socket::AF_INET, ::Socket::SOCK_STREAM).first[3]
+      @port = port.to_i
+      @resolved_address = "#{@ip_address}:#{@port}"
+
+      @timeout = 5
+    end
+
+    def command(database, cmd, options = {})
+      operation = Protocol::Command.new(database, cmd, options)
+
+      process(operation) do |reply|
+        result = reply.documents[0]
+
+        raise Errors::OperationFailure.new(
+          operation, result
+        ) if result["ok"] != 1 || result["err"] || result["errmsg"]
+
+        result
+      end
+    end
+
+    def kill_cursors(cursor_ids)
+      process Protocol::KillCursors.new(cursor_ids)
+    end
+
+    def get_more(database, collection, cursor_id, limit)
+      process Protocol::GetMore.new(database, collection, cursor_id, limit)
+    end
+
+    def remove(database, collection, selector, options = {})
+      process Protocol::Delete.new(database, collection, selector, options)
+    end
+
+    def update(database, collection, selector, change, options = {})
+      process Protocol::Update.new(database, collection, selector, change, options)
+    end
+
+    def insert(database, collection, documents)
+      process Protocol::Insert.new(database, collection, documents)
+    end
+
+    def query(database, collection, selector, options = {})
+      operation = Protocol::Query.new(database, collection, selector, options)
+
+      process operation do |reply|
+        if reply.flags.include? :query_failure
+          raise Errors::QueryFailure.new(operation, reply.documents.first)
+        end
+
+        reply
+      end
+    end
+
+    # @return [true/false] whether the node needs to be refreshed.
+    def needs_refresh?(time)
+      !@refreshed_at || @refreshed_at < time
+    end
+
+    def primary?
+      @primary
+    end
+
+    def secondary?
+      @secondary
+    end
+
+    # Refresh information about the node, such as it's status in the replica
+    # set and it's known peers.
+    #
+    # Returns nothing.
+    # Raises Errors::ConnectionFailure if the node cannot be reached
+    # Raises Errors::ReplicaSetReconfigured if the node is no longer a primary node and
+    #   refresh was called within an +#ensure_primary+ block.
+    def refresh
+      info = command "admin", ismaster: 1
+
+      @refreshed_at = Time.now
+      primary = true if info["ismaster"]
+      secondary = true if info["secondary"]
+
+      peers = []
+      peers.push info["primary"] if info["primary"]
+      peers.concat info["hosts"] if info["hosts"]
+      peers.concat info["passives"] if info["passives"]
+      peers.concat info["arbiters"] if info["arbiters"]
+
+      @peers = peers.map { |peer| Node.new(peer) }
+      @primary, @secondary = primary, secondary
+
+      if !primary && Threaded.executing?(:ensure_primary)
+        raise Errors::ReplicaSetReconfigured, "#{inspect} is no longer the primary node."
+      end
+    end
+
+    attr_reader :down_at
+
+    def down?
+      @down_at
+    end
+
+    # Set a flag on the node for the duration of provided block so that an
+    # exception is raised if the node is no longer the primary node.
+    #
+    # Returns nothing.
+    def ensure_primary
+      Threaded.begin :ensure_primary
+      yield
+    ensure
+      Threaded.end :ensure_primary
+    end
+
+    # Yields the block if a connection can be established, retrying when a
+    # connection error is raised.
+    #
+    # @raises ConnectionFailure when a connection cannot be established.
+    def ensure_connected
+      # Don't run the reconnection login if we're already inside an
+      # +ensure_connected+ block.
+      return yield if Threaded.executing? :connection
+      Threaded.begin :connection
+
+      retry_on_failure = true
+
+      begin
+        connect unless connected?
+        yield
+      rescue Errors::ReplicaSetReconfigured
+        # Someone else wrapped this in an #ensure_primary block, so let the
+        # reconfiguration exception bubble up.
+        raise
+      rescue Errors::OperationFailure, Errors::AuthenticationFailure, Errors::QueryFailure
+        # These exceptions are "expected" in the normal course of events, and
+        # don't necessitate disconnecting.
+        raise
+      rescue Errors::ConnectionFailure
+        disconnect
+
+        if retry_on_failure
+          # Maybe there was a hiccup -- try reconnecting one more time
+          retry_on_failure = false
+          retry
+        else
+          # Nope, we failed to connect twice. Flag the node as down and re-raise
+          # the exception.
+          down!
+          raise
+        end
+      rescue
+        # Looks like we got an unexpected error, so we'll clean up the connection
+        # and re-raise the exception.
+        disconnect
+        raise $!.extend(Errors::SocketError)
+      end
+    ensure
+      Threaded.end :connection
+    end
+
+    def pipeline
+      Threaded.begin :pipeline
+
+      begin
+        yield
+      ensure
+        Threaded.end :pipeline
+      end
+
+      flush unless Threaded.executing? :pipeline
+    end
+
+    def apply_auth(credentials)
+      unless auth == credentials
+        logouts = auth.keys - credentials.keys
+
+        logouts.each do |database|
+          logout database
+        end
+
+        credentials.each do |database, (username, password)|
+          login(database, username, password) unless auth[database] == [username, password]
+        end
+      end
+
+      self
+    end
+
+    def ==(other)
+      resolved_address == other.resolved_address
+    end
+    alias eql? ==
+
+    def hash
+      [ip_address, port].hash
+    end
+
+    private
+
+    def auth
+      @auth ||= {}
+    end
+
+    def login(database, username, password)
+      getnonce = Protocol::Command.new(database, getnonce: 1)
+      connection.write [getnonce]
+      result = connection.read.documents.first
+      raise Errors::OperationFailure.new(getnonce, result) unless result["ok"] == 1
+
+      authenticate = Protocol::Commands::Authenticate.new(database, username, password, result["nonce"])
+      connection.write [authenticate]
+      result = connection.read.documents.first
+      raise Errors::AuthenticationFailure.new(authenticate, result) unless result["ok"] == 1
+
+      auth[database] = [username, password]
+    end
+
+    def logout(database)
+      command = Protocol::Command.new(database, logout: 1)
+      connection.write [command]
+      result = connection.read.documents.first
+      raise Errors::OperationFailure.new(command, result) unless result["ok"] == 1
+
+      auth.delete(database)
+    end
+
+    def initialize_copy(_)
+      @connection = nil
+    end
+
+    def connection
+      @connection ||= Connection.new
+    end
+
+    def disconnect
+      auth.clear
+      connection.disconnect
+    end
+
+    def connected?
+      connection.connected?
+    end
+
+    # Mark the node as down.
+    #
+    # Returns nothing.
+    def down!
+      @down_at = Time.new
+
+      disconnect
+    end
+
+    # Connect to the node.
+    #
+    # Returns nothing.
+    # Raises Moped::ConnectionError if the connection times out.
+    # Raises Moped::ConnectionError if the server is unavailable.
+    def connect
+      connection.connect ip_address, port, timeout
+      @down_at = nil
+
+      refresh
+    rescue Timeout::Error
+      raise Errors::ConnectionFailure, "Timed out connection to Mongo on #{address}"
+    rescue Errno::ECONNREFUSED
+      raise Errors::ConnectionFailure, "Could not connect to Mongo on #{address}"
+    end
+
+    def process(operation, &callback)
+      if Threaded.executing? :pipeline
+        queue.push [operation, callback]
+      else
+        flush([[operation, callback]])
+      end
+    end
+
+    def queue
+      Threaded.stack(:pipelined_operations)
+    end
+
+    def flush(ops = queue)
+      operations, callbacks = ops.transpose
+
+      logging(operations) do
+        ensure_connected do
+          connection.write operations
+          replies = connection.receive_replies(operations)
+
+          replies.zip(callbacks).map do |reply, callback|
+            callback ? callback[reply] : reply
+          end.last
+        end
+      end
+    ensure
+      ops.clear
+    end
+
+    def logging(operations)
+      instrument_start = (logger = Moped.logger) && logger.debug? && Time.new
+      yield
+    ensure
+      log_operations(logger, operations, Time.new - instrument_start) if instrument_start && !$!
+    end
+
+    def log_operations(logger, ops, duration)
+      prefix  = "  MOPED: #{address} "
+      indent  = " "*prefix.length
+      runtime = (" (%.1fms)" % duration)
+
+      if ops.length == 1
+        logger.debug prefix + ops.first.log_inspect + runtime
+      else
+        first, *middle, last = ops
+
+        logger.debug prefix + first.log_inspect
+        middle.each { |m| logger.debug indent + m.log_inspect }
+        logger.debug indent + last.log_inspect + runtime
+      end
+    end
+
+  end
+end

@@ -11,8 +11,9 @@ class Command < Query
 
       # @param [String, Symbol] database the database to run this command on
       # @param [Hash] command the command to run
-      def initialize(database, command)
-        super database, :$cmd, command, limit: -1
+      # @param [Hash] additional query options
+      def initialize(database, command, options = {})
+        super database, :$cmd, command, options.merge(limit: -1)
       end
 
       def log_inspect

@@ -99,15 +99,15 @@ def select(select)
 
     # @return [Hash] the first document that matches the selector.
     def first
-      session.simple_query(operation)
+      limit(-1).each.first
     end
     alias one first
 
     # Iterate through documents matching the query's selector.
     #
     # @yieldparam [Hash] document each matching document
     def each
-      cursor = Cursor.new(session.with(retain_socket: true), operation)
+      cursor = Cursor.new(session, operation)
       cursor.to_enum.tap do |enum|
         enum.each do |document|
           yield document
@@ -129,6 +129,7 @@ def distinct(key)
         key: key.to_s,
         query: selector
       )
+
       result["values"]
     end
 
@@ -151,16 +152,12 @@ def count
     # @param [Array] flags an array of operation flags. Valid values are:
     #   +:multi+ and +:upsert+
     def update(change, flags = nil)
-      update = Protocol::Update.new(
-        operation.database,
-        operation.collection,
-        operation.selector,
-        change,
-        flags: flags
-      )
-
       session.with(consistency: :strong) do |session|
-        session.execute update
+        session.context.update operation.database,
+          operation.collection,
+          operation.selector,
+          change,
+          flags: flags
       end
     end
 
@@ -193,15 +190,11 @@ def upsert(change)
     # @example
     #   db[:people].find(name: "John").remove
     def remove
-      delete = Protocol::Delete.new(
-        operation.database,
-        operation.collection,
-        operation.selector,
-        flags: [:remove_first]
-      )
-
       session.with(consistency: :strong) do |session|
-        session.execute delete
+        session.context.remove operation.database,
+          operation.collection,
+          operation.selector,
+          flags: [:remove_first]
       end
     end
 
@@ -210,14 +203,10 @@ def remove
     # @example
     #   db[:people].find(name: "John").remove_all
     def remove_all
-      delete = Protocol::Delete.new(
-        operation.database,
-        operation.collection,
-        operation.selector
-      )
-
       session.with(consistency: :strong) do |session|
-        session.execute delete
+        session.context.remove operation.database,
+          operation.collection,
+          operation.selector
       end
     end
 

@@ -6,14 +6,16 @@
 require "moped/bson"
 require "moped/cluster"
 require "moped/collection"
+require "moped/connection"
 require "moped/cursor"
 require "moped/database"
 require "moped/errors"
 require "moped/indexes"
 require "moped/logging"
+require "moped/node"
 require "moped/protocol"
 require "moped/query"
-require "moped/server"
 require "moped/session"
-require "moped/socket"
+require "moped/session/context"
+require "moped/threaded"
 require "moped/version"

@@ -1,73 +0,0 @@
-module Moped
-
-  # @api private
-  #
-  # The internal class for storing information about a server.
-  class Server
-
-    # @return [String] the original host:port address provided
-    attr_reader :address
-
-    # @return [String] the resolved host:port address
-    attr_reader :resolved_address
-
-    # @return [String] the resolved ip address
-    attr_reader :ip_address
-
-    # @return [Integer] the resolved port
-    attr_reader :port
-
-    attr_writer :primary
-    attr_writer :secondary
-
-    def initialize(address)
-      @address = address
-
-      host, port = address.split(":")
-      port = port ? port.to_i : 27017
-
-      ip_address = ::Socket.getaddrinfo(host, nil, ::Socket::AF_INET, ::Socket::SOCK_STREAM).first[3]
-
-      @primary = @secondary = false
-      @ip_address = ip_address
-      @port = port
-      @resolved_address = "#{ip_address}:#{port}"
-    end
-
-    def primary?
-      !!@primary
-    end
-
-    def secondary?
-      !!@secondary
-    end
-
-    def merge(other)
-      @primary = other.primary?
-      @secondary = other.secondary?
-
-      other.close
-    end
-
-    def close
-      if @socket
-        @socket.close
-        @socket = nil
-      end
-    end
-
-    def socket
-      @socket ||= Socket.new(ip_address, port)
-    end
-
-    def ==(other)
-      self.class === other && hash == other.hash
-    end
-    alias eql? ==
-
-    def hash
-      [ip_address, port].hash
-    end
-
-  end
-end

@@ -0,0 +1,105 @@
+module Moped
+  class Session
+
+    # @api private
+    class Context
+      extend Forwardable
+
+      def initialize(session)
+        @session = session
+      end
+
+      delegate :safety => :@session
+      delegate :safe?  => :@session
+      delegate :consistency  => :@session
+      delegate :cluster  => :@session
+
+      def login(database, username, password)
+        cluster.auth[database.to_s] = [username, password]
+      end
+
+      def logout(database)
+        cluster.auth.delete database.to_s
+      end
+
+      def query(database, collection, selector, options = {})
+        if consistency == :eventual
+          options[:flags] ||= []
+          options[:flags] |= [:slave_ok]
+        end
+
+        with_node do |node|
+          node.query(database, collection, selector, options)
+        end
+      end
+
+      def command(database, command)
+        options = consistency == :eventual ? { :flags => [:slave_ok] } : {}
+        with_node do |node|
+          node.command(database, command, options)
+        end
+      end
+
+      def insert(database, collection, documents)
+        with_node do |node|
+          if safe?
+            node.pipeline do
+              node.insert(database, collection, documents)
+              node.command("admin", { getlasterror: 1 }.merge(safety))
+            end
+          else
+            node.insert(database, collection, documents)
+          end
+        end
+      end
+
+      def update(database, collection, selector, change, options = {})
+        with_node do |node|
+          if safe?
+            node.pipeline do
+              node.update(database, collection, selector, change, options)
+              node.command("admin", { getlasterror: 1 }.merge(safety))
+            end
+          else
+            node.update(database, collection, selector, change, options)
+          end
+        end
+      end
+
+      def remove(database, collection, selector, options = {})
+        with_node do |node|
+          if safe?
+            node.pipeline do
+              node.remove(database, collection, selector, options)
+              node.command("admin", { getlasterror: 1 }.merge(safety))
+            end
+          else
+            node.remove(database, collection, selector, options)
+          end
+        end
+      end
+
+      def get_more(*args)
+        raise NotImplementedError, "#get_more cannot be called on Context; it must be called directly on a node"
+      end
+
+      def kill_cursors(*args)
+        raise NotImplementedError, "#kill_cursors cannot be called on Context; it must be called directly on a node"
+      end
+
+      def with_node
+        if consistency == :eventual
+          cluster.with_secondary do |node|
+            yield node
+          end
+        else
+          cluster.with_primary do |node|
+            yield node
+          end
+        end
+      end
+
+    end
+
+  end
+end

@@ -37,6 +37,10 @@ class Session
     # @return [Cluster] this session's cluster
     attr_reader :cluster
 
+    # @private
+    # @return [Context] this session's context
+    attr_reader :context
+
     # @param [Array] seeds an of host:port pairs
     # @param [Hash] options
     # @option options [Boolean] :safe (false) ensure writes are persisted
@@ -45,7 +49,8 @@ class Session
     # @option options [Symbol, String] :database the database to use
     # @option options [:strong, :eventual] :consistency (:eventual)
     def initialize(seeds, options = {})
-      @cluster = Cluster.new(seeds)
+      @cluster = Cluster.new(seeds, {})
+      @context = Context.new(self)
       @options = options
       @options[:consistency] ||= :eventual
     end
@@ -55,11 +60,16 @@ def safe?
       !!safety
     end
 
+    # @return [:strong, :eventual] the session's consistency
+    def consistency
+      options[:consistency]
+    end
+
     # Switch the session's current database.
     #
     # @example
     #   session.use :moped
-    #   session[:people].     john, mary = session[:people].find.one # => { :name => "John" }
+    #   session[:people].find.one # => { :name => "John" }
     #
     # @param [String] database the database to use
     def use(database)
@@ -111,7 +121,7 @@ def with(options = {})
     # @return [Moped::Session] the new session
     def new(options = {})
       session = with(options)
-      session.cluster.reconnect
+      session.instance_variable_set(:@cluster, cluster.dup)
 
       if block_given?
         yield session
@@ -155,64 +165,6 @@ def new(options = {})
     # @raise (see Moped::Database#login)
     delegate :logout => :current_database
 
-    # @private
-    def current_database
-      return @current_database if defined? @current_database
-
-      if database = options[:database]
-        set_current_database(database)
-      else
-        raise "No database set for session. Call #use or #with before accessing the database"
-      end
-    end
-
-    # @private
-    def simple_query(query)
-      query.limit = -1
-
-      query(query).documents.first
-    end
-
-    # @private
-    def query(query)
-      if options[:consistency] == :eventual
-        query.flags |= [:slave_ok] if query.respond_to? :flags
-        mode = :read
-      else
-        mode = :write
-      end
-
-      reply = socket_for(mode).execute(query)
-
-      reply.tap do |reply|
-        if reply.flags.include?(:query_failure)
-          raise Errors::QueryFailure.new(query, reply.documents.first)
-        end
-      end
-    end
-
-    # @private
-    def execute(op)
-      mode = options[:consistency] == :eventual ? :read : :write
-      socket = socket_for(mode)
-
-      if safe?
-        last_error = Protocol::Command.new(
-          "admin", { getlasterror: 1 }.merge(safety)
-        )
-
-        socket.execute(op, last_error).documents.first.tap do |result|
-          raise Errors::OperationFailure.new(
-            op, result
-          ) if result["err"] || result["errmsg"]
-        end
-      else
-        socket.execute(op)
-      end
-    end
-
-    private
-
     # @return [Boolean, Hash] the safety level for this session
     def safety
       safe = options[:safe]
@@ -227,27 +179,29 @@ def safety
       end
     end
 
-    def socket_for(mode)
-      if options[:retain_socket]
-        @socket ||= cluster.socket_for(mode)
+    private
+
+    def current_database
+      return @current_database if defined? @current_database
+
+      if database = options[:database]
+        set_current_database(database)
       else
-        cluster.socket_for(mode)
+        raise "No database set for session. Call #use or #with before accessing the database"
       end
     end
 
     def set_current_database(database)
       @current_database = Database.new(self, database)
     end
 
-    def dup
-      session = super
-      session.instance_variable_set :@options, options.dup
+    def initialize_copy(_)
+      @context = Context.new(self)
+      @options = @options.dup
 
       if defined? @current_database
-        session.send(:remove_instance_variable, :@current_database)
+        remove_instance_variable :@current_database
       end
-
-      session
     end
   end
 end

@@ -1,201 +0,0 @@
-require "timeout"
-
-module Moped
-
-  # @api private
-  #
-  # The internal class wrapping a socket connection.
-  class Socket
-
-    # Thread-safe atomic integer.
-    class RequestId
-      def initialize
-        @mutex = Mutex.new
-        @id = 0
-      end
-
-      def next
-        @mutex.synchronize { @id += 1 }
-      end
-    end
-
-    attr_reader :connection
-
-    attr_reader :host
-    attr_reader :port
-
-    def initialize(host, port)
-      @host = host
-      @port = port
-
-      @mutex = Mutex.new
-      @request_id = RequestId.new
-    end
-
-    # @return [true, false] whether the connection was successful
-    # @note The connection timeout is currently just 0.5 seconds, which should
-    #   be sufficient, but may need to be raised or made configurable for
-    #   high-latency situations. That said, if connecting to the remote server
-    #   takes that long, we may not want to use the node any way.
-    def connect
-      return true if connection
-
-      Timeout::timeout 0.5 do
-        @connection = TCPSocket.new(host, port)
-      end
-    rescue Errno::ECONNREFUSED, Timeout::Error
-      return false
-    end
-
-    # @return [true, false] whether this socket connection is alive
-    def alive?
-      if connection
-        return false if connection.closed?
-
-        @mutex.synchronize do
-          if select([connection], nil, nil, 0)
-            !connection.eof? rescue false
-          else
-            true
-          end
-        end
-      else
-        false
-      end
-    end
-
-    # Execute the operations on the connection.
-    def execute(*ops)
-      instrument(ops) do
-        buf = ""
-
-        last = ops.each do |op|
-          op.request_id = @request_id.next
-          op.serialize buf
-        end.last
-
-        if Protocol::Query === last || Protocol::GetMore === last
-          length = nil
-
-          @mutex.synchronize do
-            connection.write buf
-
-            length, = connection.read(4).unpack('l<')
-
-            # Re-use the already allocated buffer used for writing the command.
-            connection.read(length - 4, buf)
-          end
-
-          parse_reply length, buf
-        else
-          @mutex.synchronize do
-            connection.write buf
-          end
-
-          nil
-        end
-      end
-    end
-
-    def parse_reply(length, data)
-      buffer = StringIO.new data
-
-      reply = Protocol::Reply.allocate
-
-      reply.length = length
-
-      reply.request_id,
-        reply.response_to,
-        reply.op_code,
-        reply.flags,
-        reply.cursor_id,
-        reply.offset,
-        reply.count = buffer.read(32).unpack('l4<q<l2<')
-
-      reply.documents = reply.count.times.map do
-        BSON::Document.deserialize(buffer)
-      end
-
-      reply
-    end
-
-    # Executes a simple (one result) query and returns the first document.
-    #
-    # @return [Hash] the first document in a result set.
-    def simple_query(query)
-      query = query.dup
-      query.limit = -1
-
-      execute(query).documents.first
-    end
-
-    # Manually closes the connection
-    def close
-      @mutex.synchronize do
-        connection.close if connection && !connection.closed?
-        @connection = nil
-      end
-    end
-
-    def auth
-      @auth ||= {}
-    end
-
-    def apply_auth(credentials)
-      return if auth == credentials
-      logouts = auth.keys - credentials.keys
-
-      logouts.each do |database|
-        logout database
-      end
-
-      credentials.each do |database, (username, password)|
-        login(database, username, password) unless auth[database] == [username, password]
-      end
-    end
-
-    def login(database, username, password)
-      getnonce = Protocol::Command.new(database, getnonce: 1)
-      result = simple_query getnonce
-
-      raise Errors::OperationFailure.new(getnonce, result) unless result["ok"] == 1
-
-      authenticate = Protocol::Commands::Authenticate.new(database, username, password, result["nonce"])
-      result = simple_query authenticate
-      raise Errors::OperationFailure.new(authenticate, result) unless result["ok"] == 1
-
-      auth[database.to_s] = [username, password]
-    end
-
-    def logout(database)
-      command = Protocol::Command.new(database, logout: 1)
-      result = simple_query command
-      raise Errors::OperationFailure.new(command, result) unless result["ok"] == 1
-      auth.delete(database.to_s)
-    end
-
-    def instrument(ops)
-      instrument_start = (logger = Moped.logger) && logger.debug? && Time.now
-      yield
-    ensure
-      log_operations(logger, ops, Time.now - instrument_start) if instrument_start && !$!
-    end
-
-    def log_operations(logger, ops, duration)
-      prefix  = "  MOPED: #{host}:#{port} "
-      indent  = " "*prefix.length
-      runtime = (" (%.1fms)" % duration)
-
-      if ops.length == 1
-        logger.debug prefix + ops.first.log_inspect + runtime
-      else
-        first, *middle, last = ops
-
-        logger.debug prefix + first.log_inspect
-        middle.each { |m| logger.debug indent + m.log_inspect }
-        logger.debug indent + last.log_inspect + runtime
-      end
-    end
-
-  end
-end

@@ -0,0 +1,32 @@
+module Moped
+
+  # This module contains logic for easy access to objects that have a lifecycle
+  # on the current thread.
+  #
+  # Extracted from Mongoid's +Threaded+ module.
+  #
+  # @api private
+  module Threaded
+    extend self
+
+    # Begin a thread-local stack for +name+.
+    def begin(name)
+      stack(name).push true
+    end
+
+    # @return [Boolean] whether the stack is being executed
+    def executing?(name)
+      !stack(name).empty?
+    end
+
+    # End the thread-local stack for +name+.
+    def end(name)
+      stack(name).pop
+    end
+
+    # @return [Array] a named, thread-local stack.
+    def stack(name)
+      Thread.current["[moped]:#{name}-stack"] ||= []
+    end
+  end
+end

@@ -57,9 +57,11 @@
 profile "Insert and find one (1000x, 2 threads)" do
   2.times.map do
     Thread.new do
-      1000.times do
-        session[:people].insert(name: "John")
-        session[:people].find.one
+      session.new do |session|
+        1000.times do
+          session[:people].insert(name: "John")
+          session[:people].find.one
+        end
       end
     end
   end.each &:join
@@ -68,9 +70,11 @@
 profile "Insert and find one (1000x, 5 threads)" do
   5.times.map do |i|
     Thread.new do
-      1000.times do
-        session[:people].insert(name: "John")
-        session[:people].find.one
+      session.new do |session|
+        1000.times do
+          session[:people].insert(name: "John")
+          session[:people].find.one
+        end
       end
     end
   end.each &:join

@@ -15,6 +15,21 @@ session[:artists].find(name: "Syd Vicious").
   )
 ```
 
+## Features
+
+* Automated replica set node discovery and failover.
+* No C or Java extensions
+* No external dependencies
+* Simple, stable, public API.
+
+### Unsupported Features
+
+* GridFS
+* Map/Reduce
+
+These features are possible to implement, but outside the scope of Moped's
+goals. Consider them perfect opportunities to write a companion gem!
+
 # Project Breakdown
 
 Moped is composed of three parts: an implementation of the [BSON
@@ -43,6 +58,31 @@ id.generation_time # => 2012-04-11 13:14:29 UTC
 id == Moped::BSON::ObjectId.from_string(id.to_s) # => true
 ```
 
+<table><tbody>
+
+<tr><th>new</th>
+<td>Creates a new object id.</td></tr>
+
+<tr><th>from_string</th>
+<td>Creates a new object id from an object id string.
+<br>
+<code>Moped::BSON::ObjectId.from_string("4f8d8c66e5a4e45396000009")</code>
+</td></tr>
+
+<tr><th>from_time</th>
+<td>Creates a new object id from a time.
+<br>
+<code>Moped::BSON::ObjectId.from_time(Time.new)</code>
+</td></tr>
+
+<tr><th>legal?</th>
+<td>Validates an object id string.
+<br>
+<code>Moped::BSON::ObjectId.legal?("4f8d8c66e5a4e45396000009")</code>
+</td></tr>
+
+</tbody></table>
+
 ### Moped::BSON::Code
 
 The `Code` class is used for working with javascript on the server.
@@ -299,6 +339,126 @@ scope.one # nil
 
 </tbody></table>
 
+# Exceptions
+
+Here's a list of the exceptions generated by Moped.
+
+<table><tbody>
+
+<tr><th>Moped::Errors::ConnectionFailure</th>
+<td>Raised when a node cannot be reached or a connection is lost.
+<br>
+<strong>Note:</strong> this exception is only raised if Moped could not
+reconnect, so you shouldn't attempt to rescue this.</td></tr>
+
+<tr><th>Moped::Errors::OperationFailure</th>
+<td>Raised when a command fails or is invalid, such as when an insert fails in
+safe mode.</td></tr>
+
+<tr><th>Moped::Errors::QueryFailure</th>
+<td>Raised when an invalid query was sent to the database.</td></tr>
+
+<tr><th>Moped::Errors::AuthenticationFailure</th>
+<td>Raised when invalid credentials were passed to `session.login`.</td></tr>
+
+<tr><th>Moped::Errors::SocketError</th>
+<td>Not a real exception, but a module used to tag unhandled exceptions inside
+of a node's networking code. Allows you to `rescue Moped::SocketError` which
+preserving the real exception.</td></tr>
+
+</tbody></table>
+
+Other exceptions are possible while running commands, such as IO Errors around
+failed connections. Moped tries to be smart about managing its connections,
+such as checking if they're dead before executing a command; but those checks
+aren't foolproof, and Moped is conservative about handling unexpected errors on
+its connections. Namely, Moped will *not* retry a command if an unexpected
+exception is raised. Why? Because it's impossible to know whether the command
+was actually received by the remote Mongo instance, and without domain
+knowledge it cannot be safely retried.
+
+Take for example this case:
+
+```ruby
+session.with(safe: true)["users"].insert(name: "John")
+```
+
+It's entirely possible that the insert command will be sent to Mongo, but the
+connection gets closed before we read the result for `getLastError`. In this
+case, there's no way to know whether the insert was actually successful!
+
+If, however, you want to gracefully handle this in your own application, you
+could do something like:
+
+```ruby
+document = { _id: Moped::BSON::ObjectId.new, name: "John" }
+
+begin
+  session["users"].insert(document)
+rescue Moped::Errors::SocketError
+  session["users"].find(_id: document[:_id]).upsert(document)
+end
+```
+
+# Replica Sets
+
+Moped has full support for replica sets including automatic failover and node
+discovery.
+
+## Automatic Failover
+
+Moped will automatically retry lost connections and attempt to detect dead
+connections before sending an operation. Note, that it will *not* retry
+individual operations! For example, these cases will work and not raise any
+exceptions:
+
+```ruby
+session[:users].insert(name: "John")
+# kill primary node and promote secondary
+session[:users].insert(name: "John")
+session[:users].find.count # => 2.0
+
+# primary node drops our connection
+session[:users].insert(name: "John")
+```
+
+However, you'll get an operation error in a case like:
+
+```ruby
+# primary node goes down while reading the reply
+session.with(safe: true)[:users].insert(name: "John")
+```
+
+And you'll get a connection error in a case like:
+
+```ruby
+# primary node goes down, no new primary available yet
+session[:users].insert(name: "John")
+```
+
+If your session is running with eventual consistency, read operations will
+never raise connection errors as long as any secondary or primary node is
+running. The only case where you'll see a connection failure is if a node goes
+down while attempting to retrieve more results from a cursor, because cursors
+are tied to individual nodes.
+
+When two attempts to connect to a node fail, it will be marked as down. This
+removes it from the list of available nodes for `:down_interval` (default 30
+seconds). Note that the `:down_interval` only applies to normal operations;
+that is, if you ask for a primary node and none is available, all nodes will be
+retried. Likewise, if you ask for a secondary node, and no secondary or primary
+node is available, all nodes will be retreied.
+
+## Node Discovery
+
+The addresses you pass into your session are used as seeds for setting up
+replica set connections. After connection, each seed node will return a list of
+other known nodes which will be added to the set.
+
+This information is cached according to the `:refresh_interval` option (default:
+5 minutes). That means, e.g., that if you add a new node to your replica set,
+it should be represented in Moped within 5 minutes.
+
 # Thread-Safety
 
 Moped is thread-safe -- depending on your definition of thread-safe. For Moped,

@@ -1,103 +0,0 @@
-require "spec_helper"
-
-describe Moped::Protocol do
-  Protocol = Moped::Protocol
-
-  let(:connection) do
-    TCPSocket.new("localhost", 27017)
-  end
-
-  after do
-    connection.close unless connection.closed?
-  end
-
-  describe "authentication" do
-    context "when nonce is invalid" do
-      let(:auth) do
-        Protocol::Commands::Authenticate.new :admin, "user", "pass", "fakenonce"
-      end
-
-      it "fails" do
-        connection.write auth
-        reply = Protocol::Reply.deserialize(connection).documents[0]
-        reply["ok"].should eq 0.0
-      end
-    end
-
-    context "when nonce is valid but user doesn't exist" do
-      let(:nonce) do
-        command = Protocol::Command.new :admin, getnonce: 1
-        connection.write command
-        Protocol::Reply.deserialize(connection).documents[0]["nonce"]
-      end
-
-      let(:auth) do
-        Protocol::Commands::Authenticate.new :admin, "user", "pass", nonce
-      end
-
-      it "fails" do
-        connection.write auth
-        reply = Protocol::Reply.deserialize(connection).documents[0]
-        reply["ok"].should eq 0.0
-      end
-    end
-
-    context "when nonce is valid but password is wrong" do
-      let(:nonce) do
-        command = Protocol::Command.new "moped-protocol-spec", getnonce: 1
-        connection.write command
-        Protocol::Reply.deserialize(connection).documents[0]["nonce"]
-      end
-
-      let(:auth) do
-        Protocol::Commands::Authenticate.new "moped-protocol-spec",
-          "moped",
-          "pass",
-          nonce
-      end
-
-      before do
-        connection.write Protocol::Insert.new(
-          "moped-protocol-spec",
-          "system.users",
-          [{ user: "moped", pwd: Digest::MD5.hexdigest("moped:mongo:password") }]
-        )
-      end
-
-      it "fails" do
-        connection.write auth
-        reply = Protocol::Reply.deserialize(connection).documents[0]
-        reply["ok"].should eq 0.0
-      end
-    end
-
-    context "when authentication is valid" do
-      let(:nonce) do
-        command = Protocol::Command.new "moped-protocol-spec", getnonce: 1
-        connection.write command
-        Protocol::Reply.deserialize(connection).documents[0]["nonce"]
-      end
-
-      let(:auth) do
-        Protocol::Commands::Authenticate.new "moped-protocol-spec",
-          "moped",
-          "password",
-          nonce
-      end
-
-      before do
-        connection.write Protocol::Insert.new(
-          "moped-protocol-spec",
-          "system.users",
-          [{ user: "moped", pwd: Digest::MD5.hexdigest("moped:mongo:password") }]
-        )
-      end
-
-      it "succeeds" do
-        connection.write auth
-        reply = Protocol::Reply.deserialize(connection).documents[0]
-        reply["ok"].should eq 1.0
-      end
-    end
-  end
-end

@@ -1,58 +0,0 @@
-require "spec_helper"
-
-describe Moped::Protocol do
-  let(:Protocol) { Moped::Protocol }
-
-  let(:connection) do
-    TCPSocket.new("localhost", 27017)
-  end
-
-  after do
-    connection.close unless connection.closed?
-  end
-
-  describe "reply response flags" do
-    let(:reply) do
-      Protocol::Reply.deserialize(connection)
-    end
-
-    context "when get more is called with an invalid cursor" do
-      let(:get_more) do
-        Protocol::GetMore.new("moped-protocol", "suite", 0, 0)
-      end
-
-      specify "the cursor not found flag is set" do
-        connection.write get_more
-        reply.flags.should include :cursor_not_found
-      end
-    end
-
-    context "when query generates an error" do
-      let(:query) do
-        Protocol::Query.new "moped-protocol", "people", { '$in' => 1 }, limit: -1
-      end
-
-      specify "the query failure flag is set" do
-        connection.write query
-        reply.flags.should eq [:query_failure]
-      end
-    end
-
-    context "when mongod supports await data query option" do
-      let(:query) do
-        Protocol::Query.new "admin", "$cmd", { buildinfo: 1 }, limit: -1
-      end
-
-      specify "the await capable flag is set" do
-        connection.write query
-
-        if reply.documents[0]["version"] >= "1.6"
-          reply.flags.should eq [:await_capable]
-        else
-          reply.flags.should eq []
-        end
-      end
-    end
-
-  end
-end

@@ -1,176 +0,0 @@
-# encoding: utf-8
-
-require "spec_helper"
-
-describe Moped::Session do
-  context "with a single master node" do
-    let(:session) { Moped::Session.new ["127.0.0.1:27017"], database: "moped_test" }
-
-    after do
-      session[:people].drop if session[:people].find.count > 0
-      session.cluster.servers.each(&:close)
-    end
-
-    it "inserts and queries a single document" do
-      id = Moped::BSON::ObjectId.new
-      session[:people].insert(_id: id, name: "John")
-      john = session[:people].find(_id: id).one
-      john["_id"].should eq id
-      john["name"].should eq "John"
-    end
-
-    it "inserts and queries on utf-8 data" do
-      id = Moped::BSON::ObjectId.new
-      doc = {
-        "_id" => id,
-        "gültig" => "1",
-        "1" => "gültig",
-        "2" => :"gültig",
-        "3" => ["gültig"],
-        "4" => /gültig/
-      }
-      session[:people].insert(doc)
-      session[:people].find(_id: id).one.should eq doc
-    end
-
-    it "can explain a query" do
-      id = Moped::BSON::ObjectId.new
-      session[:people].find(_id: id).explain["cursor"].should eq("BasicCursor")
-    end
-
-    it "can explain a query with a sort" do
-      id = Moped::BSON::ObjectId.new
-      query = session[:people].find(_id: id)
-      query.sort(_id: 1).explain["cursor"].should eq("BasicCursor")
-    end
-
-    it "drops a collection" do
-      session.command(count: :people)["n"].should eq 0
-      session[:people].insert(name: "John")
-      session.command(count: :people)["n"].should eq 1
-      session[:people].drop
-      session.command(count: :people)["n"].should eq 0
-    end
-
-    it "can be inserted into safely" do
-      session.with(safe: true) do |session|
-        session[:people].insert(name: "John")["ok"].should eq 1
-      end
-    end
-
-    it "raises an error on a failed insert in safe mode" do
-      session.with(safe: true) do |session|
-        lambda do
-          session[:people].insert("$invalid" => nil)
-        end.should raise_exception(Moped::Errors::OperationFailure)
-      end
-    end
-
-    it "can sort documents" do
-      session[:people].insert([{name: "John"}, {name: "Mary"}])
-      session[:people].find.sort(_id: -1).first["name"].should eq "Mary"
-    end
-
-    it "can update documents" do
-      id = Moped::BSON::ObjectId.new
-      session[:people].insert(_id: id, name: "John")
-      mary = session[:people].find(_id: id).one
-      mary["name"].should eq "John"
-      session[:people].find(_id: id).update(name: "Mary")
-      mary = session[:people].find(_id: id).one
-      mary["_id"].should eq id
-      mary["name"].should eq "Mary"
-    end
-
-    it "can update documents safely" do
-      id = Moped::BSON::ObjectId.new
-      session[:people].insert(_id: id, name: "John")
-      mary = session[:people].find(_id: id).one
-      mary["name"].should eq "John"
-      session.with(safe: true) do |session|
-        session[:people].find(_id: id).update(name: "Mary")["ok"].should eq 1
-      end
-    end
-
-    it "can update multiple documents" do
-      session[:people].insert([{name: "John"}, {name: "Mary"}])
-      session[:people].find.update_all("$set" => { "last_name" => "Unknown" })
-
-      session[:people].find.sort(_id: -1).first["last_name"].should eq "Unknown"
-      session[:people].find.sort(_id: 1).first["last_name"].should eq "Unknown"
-    end
-
-    it "can upsert documents" do
-      session[:people].find.upsert(name: "Mary")
-      mary = session[:people].find(name: "Mary").one
-      mary["name"].should eq "Mary"
-    end
-
-    it "can delete a single document" do
-      session[:people].insert([{name: "John"}, {name: "John"}])
-      session[:people].find(name: "John").remove
-      session[:people].find.count.should eq 1
-    end
-
-    it "can delete a single document safely" do
-      session[:people].insert([{name: "John"}, {name: "John"}])
-      session.with(safe: true) do |session|
-        session[:people].find(name: "John").remove["ok"].should eq 1
-      end
-      session[:people].find.count.should eq 1
-    end
-
-    it "can delete a multiple documents" do
-      session[:people].insert([{name: "John"}, {name: "John"}])
-      session[:people].find(name: "John").remove_all
-      session[:people].find.count.should eq 0
-    end
-
-    it "can retrieve multiple documents with fixed limit" do
-      session[:people].insert([{name: "John"}, {name: "Mary"}])
-      john, mary = session[:people].find.limit(-2).sort(name: 1).to_a
-      john["name"].should eq "John"
-      mary["name"].should eq "Mary"
-    end
-
-    it "can retrieve distinct values" do
-      session[:people].insert([{name: "John"}, {name: "Mary"}])
-      values = session[:people].find.distinct(:name)
-      values.should eq [ "John", "Mary" ]
-    end
-
-    it "can retrieve no documents" do
-      session[:people].find.limit(-2).sort(name: 1).to_a.should eq []
-    end
-
-    it "can limit a result set" do
-      documents = 100.times.map { { _id: Moped::BSON::ObjectId.new } }
-      session[:people].insert(documents)
-      session[:people].find.limit(20).to_a.length.should eq 20
-    end
-
-    it "does not leave open cursors" do
-      documents = 100.times.map { { _id: Moped::BSON::ObjectId.new } }
-      session[:people].insert(documents)
-      session[:people].find.limit(20).to_a.length.should eq 20
-      status = session.command serverStatus: 1
-      status["cursors"]["totalOpen"].should eq 0
-    end
-
-    it "can retrieve large result sets" do
-      documents = 1000.times.map do
-        { _id: Moped::BSON::ObjectId.new }
-      end
-      session[:people].insert(documents)
-      session[:people].find.to_a.length.should eq 1000
-    end
-
-    it "can have multiple connections" do
-      status = session.command serverStatus: 1
-      count = status["connections"]["current"]
-      new_session = session.new
-      status = new_session.command serverStatus: 1
-      status["connections"]["current"].should eq count + 1
-    end
-  end
-end

@@ -2,15 +2,15 @@
 
 describe Moped::BSON::ObjectId do
   let(:bytes) do
-    [78, 77, 102, 52, 59, 57, 182, 132, 7, 0, 0, 1]
+    [78, 77, 102, 52, 59, 57, 182, 132, 7, 0, 0, 1].pack("C12")
   end
 
   describe ".from_string" do
 
     context "when the string is valid" do
 
-      it "initializes with the strings bytes" do
-        Moped::BSON::ObjectId.should_receive(:new).with(bytes)
+      it "initializes with the string's bytes" do
+        Moped::BSON::ObjectId.should_receive(:from_data).with(bytes)
         Moped::BSON::ObjectId.from_string "4e4d66343b39b68407000001"
       end
     end
@@ -47,14 +47,14 @@
 
   end
 
-  describe "#initialize" do
-
-    context "with data" do
-      it "sets the object id's data" do
-        Moped::BSON::ObjectId.new(bytes).data.should == bytes
-      end
+  describe "#from_time" do
+    it "sets the generation time" do
+      time = Time.at((Time.now.utc - 64800).to_i).utc
+      Moped::BSON::ObjectId.from_time(time).generation_time.should == time
     end
+  end
 
+  describe "#initialize" do
     context "with no data" do
       it "increments the id on each call" do
         Moped::BSON::ObjectId.new.should_not eq Moped::BSON::ObjectId.new
@@ -65,21 +65,13 @@
         ids[0].value.should_not eq ids[1].value
       end
     end
-
-    context "with a time" do
-      it "sets the generation time" do
-        time = Time.at((Time.now.utc - 64800).to_i).utc
-        Moped::BSON::ObjectId.new(nil, time).generation_time.should == time
-      end
-    end
-
   end
 
   describe "#==" do
 
     context "when data is identical" do
       it "returns true" do
-        Moped::BSON::ObjectId.new(bytes).should == Moped::BSON::ObjectId.new(bytes)
+        Moped::BSON::ObjectId.from_data(bytes).should == Moped::BSON::ObjectId.from_data(bytes)
       end
     end
 
@@ -95,7 +87,7 @@
 
     context "when data is identical" do
       it "returns true" do
-        Moped::BSON::ObjectId.new(bytes).should eql Moped::BSON::ObjectId.new(bytes)
+        Moped::BSON::ObjectId.from_data(bytes).should eql Moped::BSON::ObjectId.from_data(bytes)
       end
     end
 
@@ -111,7 +103,7 @@
 
     context "when data is identical" do
       it "returns the same hash" do
-        Moped::BSON::ObjectId.new(bytes).hash.should eq Moped::BSON::ObjectId.new(bytes).hash
+        Moped::BSON::ObjectId.from_data(bytes).hash.should eq Moped::BSON::ObjectId.from_data(bytes).hash
       end
     end
 
@@ -126,7 +118,7 @@
   describe "#to_s" do
 
     it "returns a hex string representation of the id" do
-      Moped::BSON::ObjectId.new(bytes).to_s.should eq "4e4d66343b39b68407000001"
+      Moped::BSON::ObjectId.from_data(bytes).to_s.should eq "4e4d66343b39b68407000001"
     end
 
   end

@@ -1,358 +1,319 @@
 require "spec_helper"
 
-describe Moped::Cluster do
-
-  let(:master) do
-    TCPServer.new "127.0.0.1", 0
-  end
-
-  let(:secondary_1) do
-    TCPServer.new "127.0.0.1", 0
-  end
-
-  let(:secondary_2) do
-    TCPServer.new "127.0.0.1", 0
-  end
-
-  describe "initialize" do
-    let(:cluster) do
-      Moped::Cluster.new(["127.0.0.1:27017","127.0.0.1:27018"], true)
-    end
-
-    it "stores the list of seeds" do
-      cluster.seeds.should eq ["127.0.0.1:27017", "127.0.0.1:27018"]
-    end
-
-    it "stores whether the connection is direct" do
-      cluster.direct.should be_true
-    end
-
-    it "has an empty list of primaries" do
-      cluster.primaries.should be_empty
-    end
-
-    it "has an empty list of secondaries" do
-      cluster.secondaries.should be_empty
-    end
-
-    it "has an empty list of servers" do
-      cluster.servers.should be_empty
-    end
-
-    it "has an empty list of dynamic seeds" do
-      cluster.dynamic_seeds.should be_empty
-    end
+describe Moped::Cluster, replica_set: true do
+  let(:replica_set) do
+    Moped::Cluster.new(seeds, {})
   end
 
-  describe "#sync" do
-    let(:cluster) { Moped::Cluster.new(["127.0.0.1:27017"]) }
-
-    it "syncs each seed node" do
-      server = Moped::Server.allocate
-      Moped::Server.should_receive(:new).with("127.0.0.1:27017").and_return(server)
-
-      cluster.should_receive(:sync_server).with(server).and_return([])
-      cluster.sync
-    end
-  end
-
-  describe "#sync_server" do
-    let(:cluster) { Moped::Cluster.new [""], false }
-    let(:server) { Moped::Server.new("localhost:27017") }
-    let(:socket) { Moped::Socket.new "", 99999 }
-    let(:connection) { Support::MockConnection.new }
-
-    before do
-      socket.stub(connection: connection, alive?: true)
-      server.stub(socket: socket)
+  context "when the replica set hasn't connected yet" do
+    describe "#with_primary" do
+      it "connects and yields the primary node" do
+        replica_set.with_primary do |node|
+          node.address.should eq @primary.address
+        end
+      end
     end
 
-    context "when node is not running" do
-      it "returns nothing" do
-        socket.stub(connect: false)
-
-        cluster.sync_server(server).should be_empty
+    describe "#with_secondary" do
+      it "connects and yields a secondary node" do
+        replica_set.with_secondary do |node|
+          @secondaries.map(&:address).should include node.address
+        end
       end
     end
 
-    context "when talking to a single node" do
+    context "and the primary is down" do
       before do
-        connection.pending_replies << Hash[
-          "ismaster" => true,
-          "maxBsonObjectSize" => 16777216,
-          "ok" => 1.0
-        ]
-      end
-
-      it "adds the node to the master set" do
-        cluster.sync_server server
-        cluster.primaries.should include server
+        @primary.stop
       end
-    end
-
-    context "when talking to a replica set node" do
-
-      context "that is not configured" do
-        before do
-          connection.pending_replies << Hash[
-            "ismaster" => false,
-            "secondary" => false,
-            "info" => "can't get local.system.replset config from self or any seed (EMPTYCONFIG)",
-            "isreplicaset" => true,
-            "maxBsonObjectSize" => 16777216,
-            "ok" => 1.0
-          ]
-        end
 
-        it "returns nothing" do
-          cluster.sync_server(server).should be_empty
+      describe "#with_primary" do
+        it "raises a connection error" do
+          lambda do
+            replica_set.with_primary do |node|
+              node.command "admin", ping: 1
+            end
+          end.should raise_exception(Moped::Errors::ConnectionFailure)
         end
       end
 
-      context "that is being initiated" do
-        before do
-          connection.pending_replies << Hash[
-            "ismaster" => false,
-            "secondary" => false,
-            "info" => "Received replSetInitiate - should come online shortly.",
-            "isreplicaset" => true,
-            "maxBsonObjectSize" => 16777216,
-            "ok" => 1.0
-          ]
-        end
-
-        it "raises a connection failure exception" do
-          cluster.sync_server(server).should be_empty
+      describe "#with_secondary" do
+        it "connects and yields a secondary node" do
+          replica_set.with_secondary do |node|
+            @secondaries.map(&:address).should include node.address
+          end
         end
       end
+    end
 
-      context "that is ready but not elected" do
-        before do
-          connection.pending_replies << Hash[
-            "setName" => "3fef4842b608",
-            "ismaster" => false,
-            "secondary" => false,
-            "hosts" => ["localhost:61085", "localhost:61086", "localhost:61084"],
-            "primary" => "localhost:61084",
-            "me" => "localhost:61085",
-            "maxBsonObjectSize" => 16777216,
-            "ok" => 1.0
-          ]
-        end
-
-        it "raises no exception" do
-          lambda do
-            cluster.sync_server server
-          end.should_not raise_exception
-        end
+    context "and a single secondary is down" do
+      before do
+        @secondaries.first.stop
+      end
 
-        it "adds the server to the list" do
-          cluster.sync_server server
-          cluster.servers.should include server
+      describe "#with_primary" do
+        it "connects and yields the primary node" do
+          replica_set.with_primary do |node|
+            node.address.should eq @primary.address
+          end
         end
+      end
 
-        it "returns all other known hosts" do
-          cluster.sync_server(server).should =~ ["localhost:61085", "localhost:61086", "localhost:61084"]
+      describe "#with_secondary" do
+        it "connects and yields a secondary node" do
+          replica_set.with_secondary do |node|
+            node.address.should eq @secondaries.last.address
+          end
         end
       end
+    end
 
-      context "that is ready" do
-        before do
-          connection.pending_replies << Hash[
-            "setName" => "3ff029114780",
-            "ismaster" => true,
-            "secondary" => false,
-            "hosts" => ["localhost:59246", "localhost:59248", "localhost:59247"],
-            "primary" => "localhost:59246",
-            "me" => "localhost:59246",
-            "maxBsonObjectSize" => 16777216,
-            "ok" => 1.0
-          ]
-        end
+    context "and all secondaries are down" do
+      before do
+        @secondaries.each &:stop
+      end
 
-        it "adds the node to the master set" do
-          cluster.sync_server server
-          cluster.primaries.should include server
+      describe "#with_primary" do
+        it "connects and yields the primary node" do
+          replica_set.with_primary do |node|
+            node.address.should eq @primary.address
+          end
         end
+      end
 
-        it "returns all other known hosts" do
-          cluster.sync_server(server).should =~ ["localhost:59246", "localhost:59248", "localhost:59247"]
+      describe "#with_secondary" do
+        it "connects and yields the primary node" do
+          replica_set.with_secondary do |node|
+            node.address.should eq @primary.address
+          end
         end
       end
-
     end
   end
 
-  describe "#socket_for" do
-    let(:cluster) do
-      Moped::Cluster.new ""
+  context "when the replica set is connected" do
+    before do
+      replica_set.refresh
     end
 
-    let(:server) do
-      Moped::Server.new("localhost:27017").tap do |server|
-        server.stub(socket: socket)
+    describe "#with_primary" do
+      it "connects and yields the primary node" do
+        replica_set.with_primary do |node|
+          node.address.should eq @primary.address
+        end
       end
     end
 
-    let(:socket) do
-      Moped::Socket.new("127.0.0.1", 27017).tap do |socket|
-        socket.connect
+    describe "#with_secondary" do
+      it "connects and yields a secondary node" do
+        replica_set.with_secondary do |node|
+          @secondaries.map(&:address).should include node.address
+        end
       end
     end
 
-    context "when socket is dead" do
-      let(:dead_server) do
-        Moped::Server.allocate.tap do |server|
-          server.stub(socket: dead_socket)
+    context "and the primary is down" do
+      before do
+        @primary.stop
+      end
+
+      describe "#with_primary" do
+        it "raises a connection error" do
+          lambda do
+            replica_set.with_primary do |node|
+              node.command "admin", ping: 1
+            end
+          end.should raise_exception(Moped::Errors::ConnectionFailure)
         end
       end
 
-      let(:dead_socket) do
-        Moped::Socket.new("127.0.0.1", 27017).tap do |socket|
-          socket.stub(:alive? => false)
+      describe "#with_secondary" do
+        it "connects and yields a secondary node" do
+          replica_set.with_secondary do |node|
+            @secondaries.map(&:address).should include node.address
+          end
         end
       end
+    end
 
+    context "and a single secondary is down" do
       before do
-        primaries = [server, dead_server]
-        primaries.stub(:sample).and_return(dead_server, server)
-        cluster.stub(primaries: primaries)
+        @secondaries.first.stop
       end
 
-      it "removes the socket" do
-        cluster.should_receive(:remove).with(dead_server)
-        cluster.socket_for :write
+      describe "#with_primary" do
+        it "connects and yields the primary node" do
+          replica_set.with_primary do |node|
+            node.address.should eq @primary.address
+          end
+        end
       end
 
-      it "returns the living socket" do
-        cluster.socket_for(:write).should eq socket
+      describe "#with_secondary" do
+        it "connects and yields a secondary node" do
+          replica_set.with_secondary do |node|
+            node.command "admin", ping: 1
+            node.address.should eq @secondaries.last.address
+          end
+        end
       end
     end
 
-    context "when mode is write" do
+    context "and all secondaries are down" do
       before do
-        server.primary = true
+        @secondaries.each &:stop
       end
 
-      context "and the cluster is not synced" do
-        it "syncs the cluster" do
-          cluster.should_receive(:sync) do
-            cluster.servers << server
+      describe "#with_primary" do
+        it "connects and yields the primary node" do
+          replica_set.with_primary do |node|
+            node.address.should eq @primary.address
           end
-          cluster.socket_for :write
         end
+      end
 
-        it "returns the socket" do
-          cluster.stub(:sync) { cluster.servers << server }
-          cluster.socket_for(:write).should eq socket
+      describe "#with_secondary" do
+        it "connects and yields the primary node" do
+          replica_set.with_secondary do |node|
+            node.command "admin", ping: 1
+            node.address.should eq @primary.address
+          end
         end
+      end
+    end
+  end
 
-        it "applies the cached authentication" do
-          cluster.stub(:sync) { cluster.servers << server }
-          socket.should_receive(:apply_auth).with(cluster.auth)
-          cluster.socket_for(:write)
-        end
+  context "with down interval" do
+    let(:replica_set) do
+      Moped::Cluster.new(seeds, { down_interval: 5 })
+    end
+
+    context "and all secondaries are down" do
+      before do
+        replica_set.refresh
+        @secondaries.each &:stop
+        replica_set.refresh
       end
 
-      context "and the cluster is synced" do
-        before do
-          cluster.servers << server
+      describe "#with_secondary" do
+        it "connects and yields the primary node" do
+          replica_set.with_secondary do |node|
+            node.command "admin", ping: 1
+            node.address.should eq @primary.address
+          end
         end
+      end
 
-        it "does not re-sync the cluster" do
-          cluster.should_receive(:sync).never
-          cluster.socket_for :write
+      context "when a secondary node comes back up" do
+        before do
+          @secondaries.each &:restart
         end
 
-        it "returns the socket" do
-          cluster.socket_for(:write).should eq socket
+        describe "#with_secondary" do
+          it "connects and yields the primary node" do
+            replica_set.with_secondary do |node|
+              node.command "admin", ping: 1
+              node.address.should eq @primary.address
+            end
+          end
         end
 
-        it "applies the cached authentication" do
-          socket.should_receive(:apply_auth).with(cluster.auth)
-          cluster.socket_for(:write)
+        context "and the node is ready to be retried" do
+          it "connects and yields the secondary node" do
+            Time.stub(:new).and_return(Time.now + 10)
+            replica_set.with_secondary do |node|
+              node.command "admin", ping: 1
+              @secondaries.map(&:address).should include node.address
+            end
+          end
         end
       end
     end
+  end
 
-    context "when mode is read" do
-      context "and the cluster is not synced" do
-        before do
-          server.primary = true
-        end
+  context "with only primary provided as a seed" do
+    let(:replica_set) do
+      Moped::Cluster.new([@primary.address], {})
+    end
 
-        it "syncs the cluster" do
-          cluster.should_receive(:sync) do
-            cluster.servers << server
-          end
-          cluster.socket_for :read
+    describe "#with_primary" do
+      it "connects and yields the primary node" do
+        replica_set.with_primary do |node|
+          node.address.should eq @primary.address
         end
+      end
+    end
 
-        it "applies the cached authentication" do
-          cluster.stub(:sync) { cluster.servers << server }
-          socket.should_receive(:apply_auth).with(cluster.auth)
-          cluster.socket_for(:read)
+    describe "#with_secondary" do
+      it "connects and yields a secondary node" do
+        replica_set.with_secondary do |node|
+          @secondaries.map(&:address).should include node.address
         end
       end
+    end
+  end
 
-      context "and the cluster is synced" do
-        context "and no secondaries are found" do
-          before do
-            server.primary = true
-            cluster.servers << server
-          end
-
-          it "returns the master connection" do
-            cluster.socket_for(:read).should eq socket
-          end
+  context "with only primary provided as a seed" do
+    let(:replica_set) do
+      Moped::Cluster.new([@secondaries[0].address], {})
+    end
 
-          it "applies the cached authentication" do
-            socket.should_receive(:apply_auth).with(cluster.auth)
-            cluster.socket_for(:read)
-          end
+    describe "#with_primary" do
+      it "connects and yields the primary node" do
+        replica_set.with_primary do |node|
+          node.address.should eq @primary.address
         end
+      end
+    end
 
-        context "and a slave is found" do
-          it "returns a random slave connection" do
-            secondaries = [server]
-            cluster.stub(secondaries: secondaries)
-            secondaries.should_receive(:sample).and_return(server)
-            cluster.socket_for(:read).should eq socket
-          end
-
-          it "applies the cached authentication" do
-            cluster.stub(secondaries: [server])
-            socket.should_receive(:apply_auth).with(cluster.auth)
-            cluster.socket_for(:read)
-          end
+    describe "#with_secondary" do
+      it "connects and yields a secondary node" do
+        replica_set.with_secondary do |node|
+          @secondaries.map(&:address).should include node.address
         end
       end
     end
   end
+end
 
-  describe "#login" do
-    let(:cluster) do
-      Moped::Cluster.allocate
-    end
+describe Moped::Cluster, "authentication", mongohq: :auth do
+  let(:session) do
+    Support::MongoHQ.auth_session(false)
+  end
 
-    it "adds the credentials to the auth cache" do
-      cluster.login("admin", "username", "password")
-      cluster.auth.should eq("admin" => ["username", "password"])
+  describe "logging in with valid credentials" do
+    it "logs in and processes commands" do
+      session.login *Support::MongoHQ.auth_credentials
+      session.command(ping: 1).should eq("ok" => 1)
     end
   end
 
-  describe "#logout" do
-    let(:cluster) do
-      Moped::Cluster.allocate
+  describe "logging in with invalid credentials" do
+    it "raises an AuthenticationFailure exception" do
+      session.login "invalid-user", "invalid-password"
+
+      lambda do
+        session.command(ping: 1)
+      end.should raise_exception(Moped::Errors::AuthenticationFailure)
     end
+  end
 
+  describe "logging in with valid credentials and then logging out" do
     before do
-      cluster.login("admin", "username", "password")
+      session.login *Support::MongoHQ.auth_credentials
+      session.command(ping: 1).should eq("ok" => 1)
     end
 
-    it "removes the stored credentials" do
-      cluster.logout :admin
-      cluster.auth.should be_empty
+    it "logs out" do
+      lambda do
+        session.command dbStats: 1
+      end.should_not raise_exception
+
+      session.logout
+
+      lambda do
+        session.command dbStats: 1
+      end.should raise_exception(Moped::Errors::OperationFailure)
     end
   end
 end

@@ -1,92 +1,39 @@
 require "spec_helper"
 
 describe Moped::Collection do
-
   let(:session) do
-    mock(Moped::Session)
-  end
-
-  let(:database) do
-    mock(Moped::Database, session: session, name: "moped")
-  end
-
-  let(:collection) do
-    described_class.new database, :users
+    Moped::Session.new %w[127.0.0.1:27017], database: "moped_test"
   end
 
-  describe "#initialize" do
-
-    it "stores the database" do
-      collection.database.should eq database
-    end
-
-    it "stores the collection name" do
-      collection.name.should eq :users
-    end
-  end
-
-  describe "#indexes" do
-    it "returns a new indexes instance" do
-      collection.indexes.should be_an_instance_of Moped::Indexes
-    end
-  end
+  let(:scope) { object_id }
 
   describe "#drop" do
-
-    it "drops the collection" do
-      database.should_receive(:command).with(drop: :users)
-      collection.drop
-    end
-  end
-
-  describe "#find" do
-
-    let(:selector) do
-      Hash[ a: 1 ]
-    end
-
-    let(:query) do
-      mock(Moped::Query)
-    end
-
-    it "returns a new Query" do
-      Moped::Query.should_receive(:new).
-        with(collection, selector).and_return(query)
-      collection.find(selector).should eq query
+    before do
+      session.drop
+      session.command create: "users"
     end
 
-    it "defaults to an empty selector" do
-      Moped::Query.should_receive(:new).
-        with(collection, {}).and_return(query)
-      collection.find.should eq query
+    it "drops the collection" do
+      result = session[:users].drop
+      result["ns"].should eq "moped_test.users"
     end
   end
 
   describe "#insert" do
-
-    before do
-      session.should_receive(:with, :consistency => :strong).and_yield(session)
-      session.stub safe?: false
-    end
-
-    context "when passed a single document" do
-
-      it "inserts the document" do
-        session.should_receive(:execute).with do |insert|
-          insert.documents.should eq [{a: 1}]
-        end
-        collection.insert(a: 1)
-      end
+    it "inserts a single document" do
+      document = { "_id" => Moped::BSON::ObjectId.new, "scope" => scope }
+      session[:users].insert(document)
+      session[:users].find(document).one.should eq document
     end
 
-    context "when passed multiple documents" do
+    it "insert multiple documents" do
+      documents = [
+        { "_id" => Moped::BSON::ObjectId.new, "scope" => scope },
+        { "_id" => Moped::BSON::ObjectId.new, "scope" => scope }
+      ]
 
-      it "inserts the documents" do
-        session.should_receive(:execute).with do |insert|
-          insert.documents.should eq [{a: 1}, {b: 2}]
-        end
-        collection.insert([{a: 1}, {b: 2}])
-      end
+      session[:users].insert(documents)
+      session[:users].find(scope: scope).entries.should eq documents
     end
   end
 end

@@ -1,237 +0,0 @@
-require "spec_helper"
-
-describe Moped::Cursor do
-  let(:session) { mock Moped::Session }
-  let(:query_operation) { Moped::Protocol::Query.allocate }
-  let(:cursor) { Moped::Cursor.new(session, query_operation) }
-
-  describe "#initialize" do
-    it "stores the session" do
-      cursor.session.should eq session
-    end
-
-    it "stores a copy of the query operation" do
-      query_operation.should_receive(:dup).and_return(query_operation)
-      cursor.query_op.should eq query_operation
-    end
-
-    describe "the get_more operation" do
-      it "inherits the query's database" do
-        cursor.get_more_op.database.should eq query_operation.database
-      end
-
-      it "inherits the query's collection" do
-        cursor.get_more_op.collection.should eq query_operation.collection
-      end
-
-      it "inherits the query's limit" do
-        cursor.get_more_op.limit.should eq query_operation.limit
-      end
-    end
-  end
-
-  describe "#more?" do
-    context "when get more operation's cursor id is 0" do
-      it "returns false" do
-        cursor.get_more_op.cursor_id = 0
-        cursor.more?.should be_false
-      end
-    end
-    context "when get more operation's cursor id is not 0" do
-      it "returns true" do
-        cursor.get_more_op.cursor_id = 123
-        cursor.more?.should be_true
-      end
-    end
-  end
-
-  describe "#limited?" do
-    context "when original query's limit is greater than 0" do
-      before do
-        query_operation.limit = 20
-      end
-
-      it "returns true" do
-        cursor.should be_limited
-      end
-    end
-
-    context "when original query's limit is not greater than 0" do
-      before do
-        query_operation.limit = 0
-      end
-
-      it "returns true" do
-        cursor.should_not be_limited
-      end
-    end
-  end
-
-  describe "#query" do
-    let(:reply) do
-      Moped::Protocol::Reply.allocate.tap do |reply|
-        reply.cursor_id = 123
-        reply.count = 1
-        reply.documents = [{"a" => 1}]
-      end
-    end
-
-    before do
-      session.stub(query: reply)
-    end
-
-    context "when query is limited" do
-      before do
-        query_operation.limit = 21
-        cursor.query query_operation
-      end
-
-      it "updates the more operation's limit" do
-        cursor.get_more_op.limit.should eq 20
-      end
-
-      it "sets the kill cursor operation's cursor id" do
-        cursor.kill_cursor_op.cursor_ids.should eq [reply.cursor_id]
-      end
-
-      it "sets the more operation's cursor id" do
-        cursor.get_more_op.cursor_id.should eq reply.cursor_id
-      end
-    end
-
-    context "when query is limited" do
-      before do
-        query_operation.limit = 0
-        cursor.query query_operation
-      end
-
-      it "does not update the more operation's limit" do
-        cursor.get_more_op.limit.should eq query_operation.limit
-      end
-
-      it "sets the kill cursor operation's cursor id" do
-        cursor.kill_cursor_op.cursor_ids.should eq [reply.cursor_id]
-      end
-
-      it "sets the more operation's cursor id" do
-        cursor.get_more_op.cursor_id.should eq reply.cursor_id
-      end
-    end
-
-    it "returns the documents" do
-      cursor.query(query_operation).should eq reply.documents
-    end
-  end
-
-  describe "#each" do
-
-    context "when query returns all available documents" do
-      let(:reply) do
-        Moped::Protocol::Reply.allocate.tap do |reply|
-          reply.cursor_id = 0
-          reply.count = 21
-          reply.documents = [{"a" => 1}]
-        end
-      end
-
-      before do
-        session.stub(query: reply)
-      end
-
-      it "yields each document" do
-        results = []
-        cursor.each { |doc| results << doc }
-        results.should eq reply.documents
-      end
-
-      it "does not get more" do
-        session.should_receive(:query).once
-        cursor.each {}
-      end
-
-      it "does not kill the cursor" do
-        cursor.should_receive(:kill).never
-        cursor.each {}
-      end
-    end
-
-    context "when query is unlimited" do
-      let(:reply) do
-        Moped::Protocol::Reply.allocate.tap do |reply|
-          reply.cursor_id = 10
-          reply.count = 10
-          reply.documents = [{"a" => 1}]
-        end
-      end
-
-      let(:get_more_reply) do
-        Moped::Protocol::Reply.allocate.tap do |reply|
-          reply.cursor_id = 0
-          reply.count = 21
-          reply.documents = [{"a" => 1}]
-        end
-      end
-
-      before do
-        session.stub(:query).and_return(reply, get_more_reply)
-      end
-
-      it "yields each document" do
-        results = []
-        cursor.each { |doc| results << doc }
-        results.should eq reply.documents + get_more_reply.documents
-      end
-
-      it "gets more twice" do
-        session.should_receive(:query).twice
-        cursor.each {}
-      end
-
-      it "does not kill the cursor" do
-        cursor.should_receive(:kill).never
-        cursor.each {}
-      end
-    end
-
-    context "when query is limited" do
-      let(:reply) do
-        Moped::Protocol::Reply.allocate.tap do |reply|
-          reply.cursor_id = 10
-          reply.count = 10
-          reply.documents = [{"a" => 1}]
-        end
-      end
-
-      let(:get_more_reply) do
-        Moped::Protocol::Reply.allocate.tap do |reply|
-          reply.cursor_id = 10
-          reply.count = 10
-          reply.documents = [{"a" => 1}]
-        end
-      end
-
-      before do
-        query_operation.limit = 20
-        session.stub(:query).and_return(reply, get_more_reply)
-        session.stub(:execute)
-      end
-
-      it "yields each document" do
-        results = []
-        cursor.each { |doc| results << doc }
-        results.should eq reply.documents + get_more_reply.documents
-      end
-
-      it "gets more twice" do
-        session.should_receive(:query).at_least(2)
-        cursor.each {}
-      end
-
-      it "kills the cursor" do
-        cursor.should_receive(:kill).once
-        cursor.each {}
-      end
-    end
-
-  end
-end

@@ -1,86 +0,0 @@
-require "spec_helper"
-
-describe Moped::Database do
-
-  let(:session) do
-    Moped::Session.new ""
-  end
-
-  let(:database) do
-    Moped::Database.new(session, :admin)
-  end
-
-  describe "#initialize" do
-
-    it "stores the session" do
-      database.session.should eq session
-    end
-
-    it "stores the database name" do
-      database.name.should eq :admin
-    end
-  end
-
-  describe "#command" do
-
-    before do
-      session.stub(:with).and_yield(session)
-    end
-
-    it "runs the given command against the master connection" do
-      session.should_receive(:with, :consistency => :strong).
-        and_yield(session)
-      session.should_receive(:simple_query) do |query|
-        query.full_collection_name.should eq "admin.$cmd"
-        query.selector.should eq(ismaster: 1)
-
-        { "ok" => 1.0 }
-      end
-
-      database.command ismaster: 1
-    end
-
-    context "when the command fails" do
-
-      it "raises an exception" do
-        session.stub(simple_query: { "ok" => 0.0 })
-
-        expect {
-          database.command ismaster: 1
-        }.to raise_exception(Moped::Errors::OperationFailure)
-      end
-    end
-  end
-
-  describe "#drop" do
-
-    it "drops the database" do
-      database.should_receive(:command).with(dropDatabase: 1)
-      database.drop
-    end
-  end
-
-  describe "#[]" do
-
-    it "returns a collection with that name" do
-      Moped::Collection.should_receive(:new).with(database, :users)
-      database[:users]
-    end
-  end
-
-  describe "#login" do
-
-    it "logs in to the database with the username and password" do
-      session.cluster.should_receive(:login).with(:admin, "username", "password")
-      database.login("username", "password")
-    end
-  end
-
-  describe "#log out" do
-
-    it "logs out from the database" do
-      session.cluster.should_receive(:logout).with(:admin)
-      database.logout
-    end
-  end
-end

@@ -1,91 +0,0 @@
-require "spec_helper"
-
-describe Moped::Errors do
-
-  describe "OperationFailure" do
-    let(:command) do
-      Moped::Protocol::Query.allocate
-    end
-
-    let(:error_details) do
-      { "$err"=>"invalid query", "code"=>12580 }
-    end
-
-    let(:error) do
-      described_class::OperationFailure.new(command, error_details)
-    end
-
-    describe "#initialize" do
-      it "stores the command which generated the error" do
-        error.command.should eq command
-      end
-
-      it "stores the details about the error" do
-        error.details.should eq error_details
-      end
-    end
-
-    describe "#message" do
-      it "includes the command that generated the error" do
-        error.message.should include command.inspect
-      end
-
-      context "when code is included in error details" do
-        let(:error_details) do
-          { "err" => "invalid query", "code" => 12580 }
-        end
-
-        it "includes the code" do
-          error.message.should include error_details["code"].to_s
-        end
-
-        it "includes the error code reference site" do
-          error.message.should include Moped::Errors::ERROR_REFERENCE
-        end
-
-        it "includes the error message" do
-          error.message.should include error_details["err"].inspect
-        end
-      end
-
-      context "when err is in the error details" do
-        let(:error_details) do
-          { "err" => "invalid query" }
-        end
-
-        it "includes the error message" do
-          error.message.should include error_details["err"].inspect
-        end
-      end
-
-      context "when $err is in the error details" do
-        let(:error_details) do
-          { "$err" => "not master" }
-        end
-
-        it "includes the error message" do
-          error.message.should include error_details["$err"].inspect
-        end
-      end
-
-      context "when errmsg is in the error details" do
-        let(:error_details) do
-          { "errmsg" => "invalid query" }
-        end
-
-        it "includes the error message" do
-          error.message.should include error_details["errmsg"].inspect
-        end
-      end
-    end
-
-  end
-
-  describe "QueryFailure" do
-    it "is a kind of OperationFailure" do
-      Moped::Errors::QueryFailure.ancestors.should \
-        include Moped::Errors::OperationFailure
-    end
-  end
-
-end

@@ -1,126 +1,51 @@
 require "spec_helper"
 
 describe Moped::Indexes do
-  let(:session) { Moped::Session.new ["127.0.0.1:27017"], database: "moped_test" }
-  let(:indexes) do
-    described_class.new(session.current_database, :users)
-  end
-
-  after do
-    session.command(deleteIndexes: "users", index: "*")
+  let(:session) do
+    Moped::Session.new %w[127.0.0.1:27017], database: "moped_test"
   end
 
-  describe "#each" do
-    before do
-      session[:"system.indexes"].insert(ns: "moped_test.users", key: { name: 1 }, name: "name_1")
-    end
-
-    it "yields all indexes on the collection" do
-      indexes.to_a.should eq \
-        session[:"system.indexes"].find(ns: "moped_test.users").to_a
-    end
+  let(:indexes) do
+    session[:users].indexes
   end
 
-  describe "#[]" do
-    before do
-      session[:"system.indexes"].insert(ns: "moped_test.users", key: { name: 1 }, name: "name_1")
-    end
-
-    it "returns the index with the provided key" do
-      indexes[name: 1]["name"].should eq "name_1"
-    end
+  before do
+    indexes.drop
   end
 
   describe "#create" do
-    let(:key) do
-      Hash["location.latlong" => "2d", "name" => 1, "age" => -1]
-    end
-
-    context "with no options" do
-      it "creates an index with a generated name" do
-        indexes.create(key)
-        indexes[key]["name"].should eq "location.latlong_2d_name_1_age_-1"
-      end
-    end
-
-    context "with a name provided" do
-      it "creates an index with the provided name" do
-        indexes.create(key, name: "custom_index_name")
-        indexes[key]["name"].should eq "custom_index_name"
-      end
-    end
-
-    context "with background: true" do
-      it "creates an index" do
-        indexes.create(key, background: true)
-        indexes[key]["background"].should eq true
-      end
-    end
-
-    context "with dropDups: true" do
-      it "creates an index" do
-        indexes.create(key, dropDups: true)
-        indexes[key]["dropDups"].should eq true
+    context "when called without extra options" do
+      it "creates an index with no options" do
+        indexes.create name: 1
+        indexes[name: 1].should_not be_nil
       end
     end
 
-    context "with unique: true" do
-      it "creates an index" do
-        indexes.create(key, unique: true)
-        indexes[key]["unique"].should eq true
+    context "when called with extra options" do
+      it "creates an index with the extra options" do
+        indexes.create({name: 1}, {unique: true, dropDups: true})
+        index = indexes[name: 1]
+        index["unique"].should be_true
+        index["dropDups"].should be_true
       end
     end
-
-    context "with sparse: true" do
-      it "creates an index" do
-        indexes.create(key, sparse: true)
-        indexes[key]["sparse"].should eq true
-      end
-    end
-
-    context "with v: 0" do
-      it "creates an index" do
-        indexes.create(key, v: 0)
-        indexes[key]["v"].should eq 0
-      end
-    end
-
   end
 
   describe "#drop" do
-    before do
-      indexes.create name: 1
-      indexes.create age: -1
-    end
-
-    context "with no key" do
-      before do
-        indexes.drop
-      end
-
-      it "drops all indexes for the collection" do
-        indexes[name: 1].should be_nil
-        indexes[age: -1].should be_nil
+    context "when provided a key" do
+      it "drops the index" do
+        indexes.create name: 1
+        indexes.drop(name: 1).should be_true
       end
     end
 
-    context "with a key" do
-      before do
-        indexes.drop(name: 1)
-      end
-
-      it "drops the index that matches the key" do
+    context "when not provided a key" do
+      it "drops all indexes" do
+        indexes.create name: 1
+        indexes.create age: 1
+        indexes.drop
         indexes[name: 1].should be_nil
-      end
-
-      it "does not drop other indexes" do
-        indexes[age: -1].should_not be_nil
-      end
-    end
-
-    context "with a key that doesn't exist" do
-      it "returns false" do
-        indexes.drop(other: 1).should be_false
+        indexes[age: 1].should be_nil
       end
     end
   end

@@ -1,83 +0,0 @@
-require "spec_helper"
-
-describe Moped::Logging do
-  let(:config) do
-    Module.new { extend Moped::Logging }
-  end
-  let(:logger) { mock(Logger) }
-
-  describe ".rails_logger" do
-    context "when Rails is present" do
-      let(:rails) { Class.new }
-
-      before do
-        Object.const_set :Rails, rails
-      end
-
-      after do
-        Object.send(:remove_const, :Rails)
-      end
-
-      context "and it defines logger" do
-        before do
-          rails.stub(logger: logger)
-        end
-
-        it "returns the logger" do
-          config.rails_logger.should eq logger
-        end
-      end
-
-      context "but does not define logger" do
-        it "returns false" do
-          config.rails_logger.should be_false
-        end
-      end
-
-    end
-  end
-
-  describe ".default_logger" do
-    it "returns a new logger instance" do
-      config.default_logger.should be_a_kind_of Logger
-    end
-
-    it "sets the log level to info" do
-      config.default_logger.level.should eq Logger::INFO
-    end
-  end
-
-  describe ".logger" do
-    context "when a rails logger is available" do
-      before do
-        config.stub(rails_logger: logger)
-      end
-
-      it "returns the rails logger" do
-        config.logger.should eq logger
-      end
-    end
-
-    context "when a rails logger is not available" do
-      before do
-        config.stub(rails_logger: nil)
-        config.stub(default_logger: logger)
-      end
-
-      it "returns the default logger" do
-        config.logger.should eq logger
-      end
-    end
-
-    context "when the logger is set to nil" do
-      before do
-        config.logger = nil
-      end
-
-      it "returns nil" do
-        config.logger.should be_nil
-      end
-    end
-  end
-
-end

@@ -0,0 +1,57 @@
+require "spec_helper"
+
+describe Moped::Node, replica_set: true do
+  let(:replica_set_node) do
+    @replica_set.nodes.first
+  end
+
+  let(:node) do
+    Moped::Node.new(replica_set_node.address)
+  end
+
+  describe "#ensure_connected" do
+    context "when node is running" do
+      it "processes the block" do
+        node.ensure_connected do
+          node.command("admin", ping: 1)
+        end.should eq("ok" => 1)
+      end
+    end
+
+    context "when node is not running" do
+      before do
+        replica_set_node.stop
+      end
+
+      it "raises a connection error" do
+        lambda do
+          node.ensure_connected do
+            node.command("admin", ping: 1)
+          end
+        end.should raise_exception(Moped::Errors::ConnectionFailure)
+      end
+
+      it "marks the node as down" do
+        node.ensure_connected {} rescue nil
+        node.should be_down
+      end
+    end
+
+    context "when node is connected but connection is dropped" do
+      before do
+        node.ensure_connected do
+          node.command("admin", ping: 1)
+        end
+
+        replica_set_node.hiccup
+      end
+
+      it "reconnects without raising an exception" do
+        node.ensure_connected do
+          node.command("admin", ping: 1)
+        end.should eq("ok" => 1)
+      end
+    end
+  end
+
+end

@@ -1,288 +1,365 @@
 require "spec_helper"
 
 describe Moped::Query do
+  shared_examples_for "Query" do
+    let(:scope) do
+      object_id
+    end
 
-  let(:session) do
-    mock(Moped::Session)
-  end
-
-  let(:database) do
-    mock(
-      Moped::Database,
-      name: "moped",
-      session: session
-    )
-  end
-
-  let(:collection) do
-    mock(
-      Moped::Collection,
-      database: database,
-      name: "users"
-    )
-  end
-
-  let(:selector) do
-    Hash[ a: 1 ]
-  end
-
-  let(:query) do
-    described_class.new collection, selector
-  end
-
-  describe "#initialize" do
-
-    it "stores the collection" do
-      query.collection.should eq collection
+    before do
+      users.find.remove_all
     end
 
-    it "stores the selector" do
-      query.selector.should eq selector
+    let(:documents) do
+      [
+        { "_id" => Moped::BSON::ObjectId.new, "scope" => scope },
+        { "_id" => Moped::BSON::ObjectId.new, "scope" => scope }
+      ]
     end
-  end
 
-  describe "#limit" do
+    it "raises a query failure exception for invalid queries" do
+      lambda do
+        users.find("age" => { "$in" => nil }).first
+      end.should raise_exception(Moped::Errors::QueryFailure)
+    end
 
-    it "sets the query operation's limit field" do
-      query.limit(5)
-      query.operation.limit.should eq 5
+    describe "#limit" do
+      it "limits the query" do
+        users.insert(documents)
+        users.find(scope: scope).limit(1).to_a.should eq [documents.first]
+      end
     end
 
-    it "returns the query" do
-      query.limit(5).should eql query
+    describe "#skip" do
+      it "skips +n+ documents" do
+        users.insert(documents)
+        users.find(scope: scope).skip(1).to_a.should eq [documents.last]
+      end
     end
-  end
 
-  describe "#skip" do
+    describe "#sort" do
+      let(:documents) do
+        [
+          { "_id" => Moped::BSON::ObjectId.new, "scope" => scope, "n" => 0 },
+          { "_id" => Moped::BSON::ObjectId.new, "scope" => scope, "n" => 1 }
+        ]
+      end
 
-    it "sets the query operation's skip field" do
-      query.skip(5)
-      query.operation.skip.should eq 5
+      it "sorts the results" do
+        users.insert(documents)
+        users.find(scope: scope).sort(n: -1).to_a.should eq documents.reverse
+      end
     end
 
-    it "returns the query" do
-      query.skip(5).should eql query
+    describe "#distinct" do
+      let(:documents) do
+        [
+          { count: 0, scope: scope },
+          { count: 1, scope: scope },
+          { count: 1, scope: scope }
+        ]
+      end
+
+      it "returns distinct values for +key+" do
+        users.insert(documents)
+        users.find(scope: scope).distinct(:count).should =~ [0, 1]
+      end
     end
-  end
 
-  describe "#select" do
+    describe "#select" do
+      let(:documents) do
+        [
+          { "scope" => scope, "n" => 0 },
+          { "scope" => scope, "n" => 1 }
+        ]
+      end
 
-    it "sets the query operation's fields" do
-      query.select(a: 1)
-      query.operation.fields.should eq(a: 1)
+      it "changes the fields returned" do
+        users.insert(documents)
+        users.find(scope: scope).select(_id: 0).to_a.should eq documents
+      end
     end
 
-    it "returns the query" do
-      query.select(a: 1).should eql query
-    end
-  end
+    describe "#one" do
+      before do
+        users.insert(documents)
+      end
 
-  describe "#sort" do
+      it "returns the first matching document" do
+        users.find(scope: scope).one.should eq documents.first
+      end
 
-    context "when called for the first time" do
+      it "respects #skip" do
+        users.find(scope: scope).skip(1).one.should eq documents.last
+      end
 
-      it "updates the selector to mongo's advanced selector" do
-        query.sort(a: 1)
-        query.operation.selector.should eq(
-          "$query" => selector,
-          "$orderby" => { a: 1 }
-        )
+      it "respects #sort" do
+        users.find(scope: scope).sort(_id: -1).one.should eq documents.last
       end
     end
 
-    context "when called again" do
+    describe "#explain" do
+      context "when a sort exists" do
+        it "updates to a mongo advanced selector" do
+          stats = Support::Stats.collect do
+            users.find(scope: scope).sort(_id: 1).explain
+          end
+
+          operation = stats[node_for_reads].grep(Moped::Protocol::Query).last
+          operation.selector.should eq(
+            "$query" => { scope: scope },
+            "$explain" => true,
+            "$orderby" => { _id: 1 }
+          )
+        end
+      end
 
-      it "changes the $orderby" do
-        query.sort(a: 1)
-        query.sort(a: 2)
-        query.operation.selector.should eq(
-          "$query" => selector,
-          "$orderby" => { a: 2 }
-        )
+      context "when no sort exists" do
+        it "updates to a mongo advanced selector" do
+          stats = Support::Stats.collect do
+            users.find(scope: scope).explain
+          end
+
+          operation = stats[node_for_reads].grep(Moped::Protocol::Query).last
+          operation.selector.should eq(
+            "$query" => { scope: scope },
+            "$explain" => true,
+            "$orderby" => {}
+          )
+        end
       end
     end
 
-    it "returns the query" do
-      query.sort(a: 1).should eql query
-    end
-  end
+    describe "#each" do
+      it "yields each document" do
+        users.insert(documents)
+        users.find(scope: scope).each.with_index do |document, index|
+          document.should eq documents[index]
+        end
+      end
 
-  describe "#explain" do
+      context "with a limit" do
+        it "closes open cursors" do
+          users.insert(100.times.map { Hash["scope" => scope] })
 
-    before do
-      session.should_receive(:simple_query).with(query.operation)
-    end
+          stats = Support::Stats.collect do
+            users.find(scope: scope).limit(5).entries
+          end
 
-    context "when a sort exists" do
+          stats[node_for_reads].grep(Moped::Protocol::KillCursors).count.should eq 1
+        end
 
-      before do
-        query.sort(_id: 1)
       end
 
-      it "updates to a mongo advanced selector" do
-        query.explain
-        query.operation.selector.should eq(
-          "$query" => selector,
-          "$explain" => true,
-          "$orderby" => { _id: 1 }
-        )
-      end
-    end
+      context "without a limit" do
+        it "fetches more" do
+          users.insert(102.times.map { Hash["scope" => scope] })
 
-    context "when no sort exists" do
+          stats = Support::Stats.collect do
+            users.find(scope: scope).entries
+          end
 
-      it "updates to a mongo advanced selector" do
-        query.explain
-        query.operation.selector.should eq(
-          "$query" => selector,
-          "$explain" => true,
-          "$orderby" => {}
-        )
+          stats[node_for_reads].grep(Moped::Protocol::GetMore).count.should eq 1
+        end
       end
     end
-  end
 
-  describe "#one" do
+    describe "#count" do
+      let(:documents) do
+        [
+          { "_id" => Moped::BSON::ObjectId.new, "scope" => scope },
+          { "_id" => Moped::BSON::ObjectId.new, "scope" => scope },
+          { "_id" => Moped::BSON::ObjectId.new }
+        ]
+      end
 
-    it "executes a simple query" do
-      session.should_receive(:simple_query).with(query.operation)
-      query.one
+      it "returns the number of matching document" do
+        users.insert(documents)
+        users.find(scope: scope).count.should eq 2
+      end
     end
-  end
 
-  describe "#distinct" do
+    describe "#update" do
+      it "updates the first matching document" do
+        users.insert(documents)
+        users.find(scope: scope).update("$set" => { "updated" => true })
+        users.find(scope: scope, updated: true).count.should eq 1
+      end
+    end
 
-    it "executes a distinct command" do
-      database.should_receive(:command).with(
-        distinct: collection.name,
-        key: "name",
-        query: selector
-      ).and_return("values" => [ "durran", "bernerd" ])
-      query.distinct(:name)
+    describe "#update_all" do
+      it "updates all matching documents" do
+        users.insert(documents)
+        users.find(scope: scope).update_all("$set" => { "updated" => true })
+        users.find(scope: scope, updated: true).count.should eq 2
+      end
     end
-  end
 
-  describe "#count" do
+    describe "#upsert" do
+      context "when a document exists" do
+        before do
+          users.insert(scope: scope, counter: 1)
+        end
 
-    it "executes a count command" do
-      database.should_receive(:command).with(
-        count: collection.name,
-        query: selector
-      ).and_return("n" => 4)
+        it "updates the document" do
+          users.find(scope: scope).upsert("$inc" => { counter: 1 })
+          users.find(scope: scope).one["counter"].should eq 2
+        end
+      end
 
-      query.count
+      context "when no document exists" do
+        it "inserts a document" do
+          users.find(scope: scope).upsert("$inc" => { counter: 1 })
+          users.find(scope: scope).one["counter"].should eq 1
+        end
+      end
     end
 
-    it "returns the count" do
-      database.stub(command: { "n" => 4 })
+    describe "#remove" do
+      it "removes the first matching document" do
+        users.insert(documents)
+        users.find(scope: scope).remove
+        users.find(scope: scope).count.should eq 1
+      end
+    end
 
-      query.count.should eq 4
+    describe "#remove_all" do
+      it "removes all matching documents" do
+        users.insert(documents)
+        users.find(scope: scope).remove_all
+        users.find(scope: scope).count.should eq 0
+      end
     end
   end
 
-  describe "#update" do
-
-    let(:change) do
-      Hash[ a: 1 ]
+  context "with a local connection" do
+    let(:session) do
+      Moped::Session.new %w[127.0.0.1:27017], database: "moped_test"
     end
 
-    it "updates the record matching selector with change" do
-      session.should_receive(:with, :consistency => :strong).
-        and_yield(session)
+    let(:users) { session[:users] }
+    let(:node_for_reads) { :primary }
 
-      session.should_receive(:execute).with do |update|
-        update.flags.should eq []
-        update.selector.should eq query.operation.selector
-        update.update.should eq change
-      end
-
-      query.update change
-    end
-  end
+    include_examples "Query"
 
-  describe "#update_all" do
+    describe "#each" do
+      context "with a limit and large result set" do
+        it "gets more and closes cursors" do
+          11.times do
+            users.insert(scope: scope, large_field: "a"*1_000_000)
+          end
 
-    let(:change) do
-      Hash[ a: 1 ]
-    end
+          stats = Support::Stats.collect do
+            users.find(scope: scope).limit(10).entries
+          end
 
-    it "updates all records matching selector with change" do
-      query.should_receive(:update).with(change, [:multi])
-      query.update_all change
+          stats[:primary].grep(Moped::Protocol::GetMore).count.should eq 1
+          stats[:primary].grep(Moped::Protocol::KillCursors).count.should eq 1
+        end
+      end
     end
   end
 
-  describe "#upsert" do
-
-    let(:change) do
-      Hash[ a: 1 ]
+  context "with a remote connection", mongohq: :auth do
+    before :all do
+      @session = Support::MongoHQ.auth_session
     end
 
-    it "upserts the record matching selector with change" do
-      query.should_receive(:update).with(change, [:upsert])
-      query.upsert change
-    end
+    let(:users) { @session[:users] }
+    let(:node_for_reads) { :primary }
+
+    include_examples "Query"
   end
 
-  describe "#remove" do
+  context "with a remote replica set connection with eventual consistency", mongohq: :replica_set do
+    before :all do
+      @session = Support::MongoHQ.replica_set_session.with(safe: true, consistency: :eventual)
+      @session.command ping: 1
+    end
 
-    it "removes the first matching document" do
-      session.should_receive(:with, :consistency => :strong).
-        and_yield(session)
+    let(:users) { @session[:users] }
+    let(:node_for_reads) { :secondary }
 
-      session.should_receive(:execute).with do |delete|
-        delete.flags.should eq [:remove_first]
-        delete.selector.should eq query.operation.selector
-      end
-
-      query.remove
-    end
+    include_examples "Query"
   end
 
-  describe "#remove_all" do
+  context "with a remote replica set connection with strong consistency", mongohq: :replica_set do
+    before :all do
+      @session = Support::MongoHQ.replica_set_session.with(safe: true, consistency: :strong)
+    end
 
-    it "removes all matching documents" do
-      session.should_receive(:with, :consistency => :strong).
-        and_yield(session)
+    let(:users) { @session[:users] }
+    let(:node_for_reads) { :primary }
 
-      session.should_receive(:execute).with do |delete|
-        delete.flags.should eq []
-        delete.selector.should eq query.operation.selector
-      end
+    include_examples "Query"
+  end
 
-      query.remove_all
+  context "with a local replica set w/ failover", replica_set: true do
+    let(:session) do
+      Moped::Session.new seeds, database: "moped_test"
     end
-  end
 
-  describe "#each" do
+    let(:scope) do
+      object_id
+    end
 
     before do
-      session.should_receive(:with).
-        with(retain_socket: true).and_return(session)
+      # Force connection before recording stats
+      session.command ping: 1
     end
 
-    it "creates a new cursor" do
-      cursor = mock(Moped::Cursor, next: nil)
-      Moped::Cursor.should_receive(:new).
-        with(session, query.operation).and_return(cursor)
+    context "and running with eventual consistency" do
+      it "queries a secondary node" do
+        stats = Support::Stats.collect do
+          session.with(consistency: :eventual)[:users].find(scope: scope).entries
+        end
 
-      query.each
-    end
+        stats[:secondary].grep(Moped::Protocol::Query).count.should eq 1
+        stats[:primary].should be_empty
+      end
+
+      it "sets the slave ok flag" do
+        stats = Support::Stats.collect do
+          session.with(consistency: :eventual)[:users].find(scope: scope).one
+        end
+
+        query = stats[:secondary].grep(Moped::Protocol::Query).first
+        query.flags.should include :slave_ok
+      end
 
-    it "yields all documents in the cursor" do
-      cursor = Moped::Cursor.allocate
-      cursor.stub(:to_enum).and_return([1, 2].to_enum)
+      context "and no secondaries are available" do
+        before do
+          @secondaries.each &:stop
+        end
 
-      Moped::Cursor.stub(new: cursor)
+        it "queries the primary node" do
+          stats = Support::Stats.collect do
+            session.with(consistency: :eventual)[:users].find(scope: scope).entries
+          end
 
-      query.to_a.should eq [1, 2]
+          stats[:primary].grep(Moped::Protocol::Query).count.should eq 1
+        end
+      end
     end
 
-    it "returns an enumerator" do
-      cursor = mock(Moped::Cursor)
-      Moped::Cursor.stub(new: cursor)
+    context "and running with strong consistency" do
+      it "queries the primary node" do
+        stats = Support::Stats.collect do
+          session.with(consistency: :strong)[:users].find(scope: scope).entries
+        end
+
+        stats[:primary].grep(Moped::Protocol::Query).count.should eq 1
+        stats[:secondary].should be_empty
+      end
 
-      query.each.should be_a Enumerator
+      it "does not set the slave ok flag" do
+        stats = Support::Stats.collect do
+          session.with(consistency: :strong)[:users].find(scope: scope).one
+        end
+
+        query = stats[:primary].grep(Moped::Protocol::Query).first
+        query.flags.should_not include :slave_ok
+      end
     end
   end
 end

@@ -1,80 +0,0 @@
-require "spec_helper"
-
-describe Moped::Server do
-
-  describe "#initialize" do
-    let(:server) do
-      described_class.new("localhost:123")
-    end
-
-    it "stores the original address" do
-      server.address.should eq "localhost:123"
-    end
-
-    it "stores the resolved address" do
-      server.resolved_address.should eql "127.0.0.1:123"
-    end
-
-    it "stores the resolved ip" do
-      server.ip_address.should eq "127.0.0.1"
-    end
-
-    it "stores the port" do
-      server.port.should eq 123
-    end
-  end
-
-  describe "==" do
-    context "when ip and port are the same" do
-      it "returns true" do
-        described_class.new("127.0.0.1:999").should eq \
-          described_class.new("localhost:999")
-      end
-    end
-
-    context "when ip and port are different" do
-      it "returns false" do
-        described_class.new("127.0.0.1:1000").should_not eq \
-          described_class.new("localhost:999")
-      end
-    end
-
-    context "when other is not a server" do
-      it "returns false" do
-        described_class.new("127.0.0.1:999").should_not eq 1
-      end
-    end
-  end
-
-  context "when added to a set" do
-    let(:set) { Set.new }
-
-    context "and ip and port are the same" do
-      it "does not add both servers" do
-        set << described_class.new("127.0.0.1:1000")
-        set << described_class.new("127.0.0.1:1000")
-
-        set.length.should eq 1
-      end
-
-      context "and the original address is different" do
-        it "does not add both servers" do
-          set << described_class.new("localhost:1000")
-          set << described_class.new("127.0.0.1:1000")
-
-          set.length.should eq 1
-        end
-      end
-    end
-
-    context "and ip and port are different" do
-      it "adds both servers" do
-        set << described_class.new("127.0.0.1:1000")
-        set << described_class.new("127.0.0.1:2000")
-
-        set.length.should eq 2
-      end
-    end
-  end
-
-end

@@ -1,548 +1,63 @@
 require "spec_helper"
 
 describe Moped::Session do
-
-  let(:seeds) do
-    "127.0.0.1:27017"
-  end
-
-  let(:options) do
-    Hash[database: "test", safe: true, consistency: :eventual]
-  end
-
   let(:session) do
-    described_class.new seeds, options
-  end
-
-  describe "#initialize" do
-
-    it "stores the options provided" do
-      session.options.should eq(options)
-    end
-
-    it "stores the cluster" do
-      session.cluster.should be_a(Moped::Cluster)
-    end
-  end
-
-  describe "#current_database" do
-
-    context "when no database option has been set" do
-
-      let(:session) do
-        described_class.new seeds, {}
-      end
-
-      it "raises an exception" do
-        expect { session.current_database }.to raise_exception
-      end
-    end
-
-    context "when a database option is set" do
-
-      let(:database) do
-        stub
-      end
-
-      before do
-        Moped::Database.should_receive(:new).
-          with(session, options[:database]).and_return(database)
-      end
-
-      it "returns the database from the options" do
-        session.current_database.should eq(database)
-      end
-
-      it "memoizes the database" do
-        database = session.current_database
-        session.current_database.should equal(database)
-      end
-    end
-  end
-
-  describe "#safe?" do
-
-    context "when :safe is not present" do
-
-      before do
-        session.options.delete(:safe)
-      end
-
-      it "returns false" do
-        session.should_not be_safe
-      end
-    end
-
-    context "when :safe is present but false" do
-
-      before do
-        session.options[:safe] = false
-      end
-
-      it "returns false" do
-        session.should_not be_safe
-      end
-    end
-
-    context "when :safe is true" do
-
-      before do
-        session.options[:safe] = true
-      end
-
-      it "returns true" do
-        session.should be_safe
-      end
-    end
-
-    context "when :safe is a hash" do
-
-      before do
-        session.options[:safe] = { fsync: true }
-      end
-
-      it "returns true" do
-        session.should be_safe
-      end
-    end
+    Moped::Session.new %w[127.0.0.1:27017], database: "moped_test"
   end
 
   describe "#use" do
-
-    it "sets the :database option" do
-      session.use :admin
-      session.options[:database].should eq(:admin)
-    end
-
-    context "when there is not already a current database" do
-
-      it "sets the current database" do
-        session.should_receive(:set_current_database).with(:admin)
-        session.use :admin
-      end
+    it "changes the current database" do
+      session.use "moped_test_2"
+      session.command(dbStats: 1)["db"].should eq "moped_test_2"
     end
   end
 
   describe "#with" do
-
-    let(:new_options) do
-      Hash[database: "test-2"]
-    end
-
     context "when called with a block" do
-
-      it "yields a session" do
-        session.with(new_options) do |new_session|
-          new_session.should be_a Moped::Session
-        end
+      it "returns the value from the block" do
+        session.with { :value }.should eq :value
       end
 
-      it "yields a new session" do
-        session.with(new_options) do |new_session|
-          new_session.should_not eql session
+      it "yields a session with the provided options" do
+        session.with(safe: true) do |safe|
+          safe.options[:safe].should eq true
         end
       end
 
-      it "returns the result of the block" do
-        session.with(new_options) { false }.should eq false
-      end
-
-      it "merges the old and new session's options" do
-        session.with(new_options) do |new_session|
-          new_session.options.should eq options.merge(new_options)
-        end
-      end
-
-      it "does not change the original session's options" do
-        original_options = options.dup
-        session.with(new_options) do |new_session|
-          session.options.should eql original_options
-        end
-      end
-
-      it "unmemoizes the current database" do
-        db = session.current_database
-        session.with(new_options) do |new_session|
-          new_session.current_database.should_not eql db
+      it "does not modify the original session" do
+        session.with(database: "other") do |safe|
+          session.options[:database].should eq "moped_test"
         end
       end
     end
 
     context "when called without a block" do
-
-      it "returns a session" do
-        session.with(new_options).should be_a Moped::Session
-      end
-
-      it "returns a new session" do
-        session.with(new_options).should_not eql session
-      end
-
-      it "merges the old and new session's options" do
-        session.with(new_options).options.should eq options.merge(new_options)
-      end
-
-      it "does not change the original session's options" do
-        original_options = options.dup
-        session.with(new_options)
-        session.options.should eql original_options
-      end
-    end
-  end
-
-  describe "#new" do
-
-    let(:new_options) do
-      Hash[database: "test-2"]
-    end
-
-    let(:new_session) do
-      described_class.new seeds, options
-    end
-
-    before do
-      new_session.cluster.stub(:reconnect)
-    end
-
-    it "delegates to #with" do
-      session.should_receive(:with).with(new_options).and_return(new_session)
-      session.new(new_options)
-    end
-
-    it "instructs the cluster to reconnect" do
-      session.stub(with: new_session)
-      new_session.cluster.should_receive(:reconnect)
-      session.new(new_options)
-    end
-
-    context "when called with a block" do
-
-      it "yields the new session" do
-        session.stub(with: new_session)
-        session.new(new_options) do |session|
-          session.should eql new_session
-        end
+      it "returns a session with the provided options" do
+        safe = session.with(safe: true)
+        safe.options[:safe].should eq true
       end
-    end
 
-    context "when called without a block" do
-
-      it "returns the new session" do
-        session.stub(with: new_session)
-        session.new(new_options).should eql new_session
+      it "does not modify the original session" do
+        other = session.with(database: "other")
+        session.options[:database].should eq "moped_test"
       end
     end
   end
 
   describe "#drop" do
-
-    it "delegates to the current database" do
-      database = mock(Moped::Database)
-      session.should_receive(:current_database).and_return(database)
-      database.should_receive(:drop)
-      session.drop
-    end
-  end
-
-  describe "#command" do
-
-    let(:command) do
-      Hash[ismaster: 1]
-    end
-
-    it "delegates to the current database" do
-      database = mock(Moped::Database)
-      session.should_receive(:current_database).and_return(database)
-      database.should_receive(:command).with(command)
-      session.command command
-    end
-  end
-
-  describe "#login" do
-
-    it "delegates to the current database" do
-      database = mock(Moped::Database)
-      session.should_receive(:current_database).and_return(database)
-      database.should_receive(:login).with("username", "password")
-      session.login("username", "password")
-    end
-  end
-
-  describe "#logout" do
-
-    it "delegates to the current database" do
-      database = mock(Moped::Database)
-      session.should_receive(:current_database).and_return(database)
-      database.should_receive(:logout)
-      session.logout
-    end
-  end
-
-  describe "#socket_for" do
-
-    it "delegates to the cluster" do
-      session.cluster.should_receive(:socket_for).with(:read)
-      session.send(:socket_for, :read)
-    end
-
-    context "when retain socket option is set" do
-
-      before do
-        session.options[:retain_socket] = true
-      end
-
-      it "only aquires the socket once" do
-        session.cluster.should_receive(:socket_for).
-          with(:read).once.and_return(mock(Moped::Socket))
-
-        session.send(:socket_for, :read)
-        session.send(:socket_for, :read)
-      end
-    end
-  end
-
-  describe "#simple_query" do
-
-    let(:query) do
-      Moped::Protocol::Query.allocate
-    end
-
-    let(:socket) do
-      mock(Moped::Socket)
-    end
-
-    let(:reply) do
-      Moped::Protocol::Reply.allocate.tap do |reply|
-        reply.documents = [{a: 1}]
+    it "drops the current database" do
+      session.with(database: "moped_test_2") do |session|
+        session.drop.should eq("dropped" => "moped_test_2", "ok" => 1)
       end
     end
-
-    before do
-      session.stub(socket_for: socket)
-      session.stub(query: reply)
-    end
-
-    it "limits the query" do
-      session.should_receive(:query) do |query|
-        query.limit.should eq(-1)
-        reply
-      end
-
-      session.simple_query(query)
-    end
-
-    it "returns the document" do
-      session.simple_query(query).should eq(a: 1)
-    end
   end
 
-  describe "#query" do
-
-    let(:query) do
-      Moped::Protocol::Query.allocate
-    end
-
-    let(:socket) do
-      mock(Moped::Socket)
-    end
-
-    let(:reply) do
-      Moped::Protocol::Reply.allocate.tap do |reply|
-        reply.documents = [{a: 1}]
-      end
-    end
-
-    before do
-      session.stub(socket_for: socket)
-      socket.stub(:execute).and_return(reply)
-    end
-
-    context "when consistency is strong" do
-
-      before do
-        session.options[:consistency] = :strong
-      end
-
-      it "queries the master node" do
-        session.should_receive(:socket_for).with(:write).
-          and_return(socket)
-        session.query(query)
-      end
-    end
-
-    context "when consistency is eventual" do
-
-      before do
-        session.options[:consistency] = :eventual
-      end
-
-      it "queries a slave node" do
-        session.should_receive(:socket_for).with(:read).
-          and_return(socket)
-        session.query(query)
-      end
-
-      context "and query accepts flags" do
-
-        it "sets slave_ok on the query flags" do
-          session.stub(socket_for: socket)
-          socket.should_receive(:execute) do |query|
-            query.flags.should include :slave_ok
-          end
-
-          session.query(query)
-        end
-      end
-
-      context "and query does not accept flags" do
-
-        let(:query) do
-          Moped::Protocol::GetMore.allocate
-        end
-
-        it "doesn't try to set flags" do
-          session.stub(socket_for: socket)
-          expect { session.query(query) }.not_to raise_exception
-        end
-      end
-    end
-
-    context "when reply has :query_failure flag" do
-
-      before do
-        reply.flags = [:query_failure]
-      end
-
-      it "raises a QueryFailure exception" do
-        expect {
-          session.query(query)
-        }.to raise_exception(Moped::Errors::QueryFailure)
+  describe "#command" do
+    it "runs the command on the current database" do
+      session.with(database: "moped_test_2") do |session|
+        session.command(dbStats: 1)["db"].should eq "moped_test_2"
       end
     end
   end
 
-  describe "#execute" do
-
-    let(:operation) do
-      Moped::Protocol::Insert.allocate
-    end
-
-    let(:socket) do
-      mock(Moped::Socket)
-    end
-
-    context "when session is not in safe mode" do
-
-      before do
-        session.options[:safe] = false
-      end
-
-      context "when consistency is strong" do
-
-        before do
-          session.options[:consistency] = :strong
-        end
-
-        it "executes the operation on the master node" do
-          session.should_receive(:socket_for).with(:write).
-            and_return(socket)
-          socket.should_receive(:execute).with(operation)
-
-          session.execute(operation)
-        end
-      end
-
-      context "when consistency is eventual" do
-
-        before do
-          session.options[:consistency] = :eventual
-        end
-
-        it "executes the operation on a slave node" do
-          session.should_receive(:socket_for).with(:read).
-            and_return(socket)
-          socket.should_receive(:execute).with(operation)
-
-          session.execute(operation)
-        end
-      end
-    end
-
-    context "when session is in safe mode" do
-
-      let(:reply) do
-        Moped::Protocol::Reply.allocate.tap do |reply|
-          reply.documents = [{a: 1}]
-        end
-      end
-
-      before do
-        session.options[:safe] = { w: 2 }
-      end
-
-      context "when the operation fails" do
-
-        let(:reply) do
-          Moped::Protocol::Reply.allocate.tap do |reply|
-            reply.documents = [{
-              "err"=>"document to insert can't have $ fields",
-              "code"=>13511,
-              "n"=>0,
-              "connectionId"=>894,
-              "ok"=>1.0
-            }]
-          end
-        end
-
-        it "raises an OperationFailure exception" do
-          session.stub(socket_for: socket)
-          socket.stub(execute: reply)
-
-          expect {
-            session.execute(operation)
-          }.to raise_exception(Moped::Errors::OperationFailure)
-        end
-      end
-
-      context "when consistency is strong" do
-
-        before do
-          session.options[:consistency] = :strong
-        end
-
-        it "executes the operation on the master node" do
-          session.should_receive(:socket_for).with(:write).
-            and_return(socket)
-
-          socket.should_receive(:execute) do |op, query|
-            op.should eq operation
-            query.selector.should eq(getlasterror: 1, w: 2)
-            reply
-          end
-
-          session.execute(operation)
-        end
-      end
-
-      context "when consistency is eventual" do
-
-        before do
-          session.options[:consistency] = :eventual
-        end
-
-        it "executes the operation on a slave node" do
-          session.should_receive(:socket_for).with(:read).
-            and_return(socket)
-
-          socket.should_receive(:execute) do |op, query|
-            op.should eq operation
-            query.selector.should eq(getlasterror: 1, w: 2)
-            reply
-          end
-
-          session.execute(operation)
-        end
-      end
-    end
-  end
 end

@@ -1,462 +0,0 @@
-require "spec_helper"
-
-describe Moped::Socket do
-
-  let!(:server) do
-    TCPServer.new "127.0.0.1", 0
-  end
-
-  let(:socket) do
-    described_class.new "127.0.0.1", server.addr[1]
-  end
-
-  let(:connection) {
-    socket.connection
-  }
-
-  before do
-    socket.connect
-  end
-
-  after do
-    connection.close if connection && !connection.closed?
-    server.close unless server.closed?
-  end
-
-  describe "#initialize" do
-    it "stores the host of the server" do
-      socket.host.should eq "127.0.0.1"
-    end
-
-    it "stores the port of the server" do
-      socket.port.should eq server.addr[1]
-    end
-
-    it "connects to the server" do
-      socket.connection.should_not be_closed
-    end
-  end
-
-  describe "#connect" do
-    context "when node is not running" do
-      let(:bogus_port) do
-        server = TCPServer.new("127.0.0.1", 0)
-        server.addr[1].tap do
-          server.close
-        end
-      end
-
-      let(:socket) do
-        described_class.new "127.0.0.1", bogus_port
-      end
-
-      it "returns false" do
-        socket.connect.should be_false
-      end
-    end
-
-    context "when connection times out" do
-      if RUBY_PLATFORM == "java"
-        let(:timeout_server) do
-          java.net.ServerSocket.new(0, 1)
-        end
-
-        let(:timeout_port) do
-          timeout_server.getLocalPort
-        end
-      else
-        let(:timeout_server) do
-          TCPServer.new("127.0.0.1", 0).tap do |server|
-            server.listen(1)
-          end
-        end
-
-        let(:timeout_port) do
-          timeout_server.addr[1]
-        end
-      end
-
-      let(:timeout_socket) do
-        described_class.new "127.0.0.1", timeout_port
-      end
-
-      before do
-        sockaddr = Socket.pack_sockaddr_in(timeout_port, '127.0.0.1')
-
-        5.times do # flood the server socket
-          ::Socket.new(::Socket::AF_INET, ::Socket::SOCK_STREAM, 0).connect_nonblock(sockaddr) rescue nil
-        end
-      end
-
-      after do
-        timeout_server.close unless timeout_server.closed?
-      end
-
-      it "returns false" do
-        timeout_socket.connect.should be_false
-      end
-    end
-  end
-
-  describe "#alive?" do
-    context "when not connected" do
-      let(:socket) do
-        described_class.new("127.0.0.1", 99999).tap do |socket|
-          socket.stub(:connect)
-        end
-      end
-
-      it "should be false" do
-        socket.should_not be_alive
-      end
-    end
-
-    context "when connected but server goes away" do
-      before do
-        remote = server.accept
-        remote.shutdown
-        # Give the socket time to be notified
-        sleep 0.1
-      end
-
-      it "should be false" do
-        socket.should_not be_alive
-      end
-    end
-
-    context "when connected but server goes away" do
-      before do
-        server.close
-        # Give the socket time to be notified
-        sleep 0.1
-      end
-
-      it "should be false" do
-        socket.should_not be_alive
-      end
-    end
-
-    context "when connect is explicitly closed" do
-      before do
-        socket.close
-      end
-
-      it "should be false" do
-        socket.should_not be_alive
-      end
-    end
-
-    context "when connected and server is open" do
-      it "should be true" do
-        socket.should be_alive
-      end
-    end
-  end
-
-  describe "#execute" do
-    let(:query) { Moped::Protocol::Query.new(:moped, :test, {}) }
-
-    context "when competing threads attempt to query" do
-      let(:messages) do
-        10.times.map do |i|
-          Moped::Protocol::Insert.new(:moped, :test, {}).tap do |query|
-            query.stub(request_id: 123)
-          end
-        end
-      end
-
-      it "never issues a partial write" do
-        socket
-
-        threads = 10.times.map do |i|
-          Thread.new do
-            Thread.current.abort_on_exception = true
-            socket.execute messages[i]
-          end
-        end
-
-        threads.each(&:join)
-
-        sock = server.accept
-        data = sock.read messages.join.length
-
-        messages.each do |message|
-          fail "server received partial write" unless data.include? message.serialize
-        end
-      end
-    end
-  end
-
-  describe "#parse_reply" do
-    let(:raw) do
-      Moped::Protocol::Reply.allocate.tap do |reply|
-        reply.request_id = 1
-        reply.response_to = 1
-        reply.op_code = 1
-        reply.flags = [:await_capable]
-        reply.offset = 4
-        reply.count = 1
-        reply.documents = [{"name" => "John"}]
-      end.serialize
-    end
-
-    let(:reply) do
-      socket.parse_reply(raw.length, raw[4..-1])
-    end
-
-    it "sets the length" do
-      reply.length.should eq raw.length
-    end
-
-    it "sets the response_to" do
-      reply.response_to.should eq 1
-    end
-
-    it "sets the request id" do
-      reply.request_id.should eq 1
-    end
-
-    it "sets the flags" do
-      reply.flags.should eq [:await_capable]
-    end
-
-    it "sets the offset" do
-      reply.offset.should eq 4
-    end
-
-    it "sets the count" do
-      reply.count.should eq 1
-    end
-
-    it "sets the documents" do
-      reply.documents.should eq [{"name" => "John"}]
-    end
-  end
-
-  describe "#close" do
-    let(:exception) { RuntimeError.new }
-    let(:callback) { stub }
-
-    it "closes the connection" do
-      connection.should_receive(:close).at_least(1)
-      socket.close
-    end
-
-    it "marks the socket as dead" do
-      socket.close
-      socket.should_not be_alive
-    end
-  end
-
-  describe "#login" do
-
-    let(:connection) do
-      Support::MockConnection.new
-    end
-
-    before do
-      socket.stub(connection: connection)
-    end
-
-    context "when authentication is successful" do
-      before do
-        # getnonce
-        connection.pending_replies << Hash["nonce" => "123", "ok" => 1]
-        # authenticate
-        connection.pending_replies << Hash["ok" => 1]
-      end
-
-      it "returns true" do
-        socket.login("admin", "username", "password").should be_true
-      end
-
-      it "adds the credentials to the auth cache" do
-        socket.login(:admin, "username", "password")
-        socket.auth.should eq("admin" => ["username", "password"])
-      end
-    end
-
-    context "when a nonce fails to generate" do
-      before do
-        # getnonce
-        connection.pending_replies << Hash["ok" => 0]
-      end
-
-      it "raises an operation failure" do
-        lambda do
-          socket.login(:admin, "username", "password")
-        end.should raise_exception(Moped::Errors::OperationFailure)
-      end
-
-      it "does not add the credentials to the auth cache" do
-        socket.login(:admin, "username", "password") rescue nil
-        socket.auth.should be_empty
-      end
-    end
-
-    context "when authentication fails" do
-      before do
-        # getnonce
-        connection.pending_replies << Hash["nonce" => "123", "ok" => 1]
-        # authenticate
-        connection.pending_replies << Hash["ok" => 0]
-      end
-
-      it "raises an operation failure" do
-        lambda do
-          socket.login(:admin, "username", "password")
-        end.should raise_exception(Moped::Errors::OperationFailure)
-      end
-
-      it "does not add the credentials to the auth cache" do
-        socket.login(:admin, "username", "password") rescue nil
-        socket.auth.should be_empty
-      end
-    end
-
-  end
-
-  describe "#logout" do
-
-    let(:connection) do
-      Support::MockConnection.new
-    end
-
-    before do
-      socket.stub(connection: connection)
-      socket.auth["admin"] = ["username", "password"]
-    end
-
-    context "when logout is successful" do
-      before do
-        connection.pending_replies << Hash["ok" => 1]
-      end
-
-      it "removes the stored credentials" do
-        socket.logout :admin
-        socket.auth.should be_empty
-      end
-    end
-
-    context "when logout is unsuccessful" do
-      before do
-        connection.pending_replies << Hash["ok" => 0]
-      end
-
-      it "does not remove the stored credentials" do
-        socket.logout :admin rescue nil
-        socket.auth.should_not be_empty
-      end
-
-      it "raises an operation failure" do
-        lambda do
-          socket.logout :admin
-        end.should raise_exception(Moped::Errors::OperationFailure)
-      end
-    end
-
-  end
-
-  describe "#apply_auth" do
-    context "when the socket is unauthenticated" do
-      it "logs in with each credential provided" do
-        socket.should_receive(:login).with("admin", "username", "password")
-        socket.should_receive(:login).with("test", "username", "password")
-
-        socket.apply_auth(
-          "admin" => ["username", "password"],
-          "test" => ["username", "password"]
-        )
-      end
-    end
-
-    context "when the socket is authenticated" do
-      before do
-        socket.auth["admin"] = ["username", "password"]
-      end
-
-      context "and a credential is unchanged" do
-        it "does nothing" do
-          socket.should_not_receive(:login)
-          socket.apply_auth("admin" => ["username", "password"])
-        end
-      end
-
-      context "and a credential changes" do
-        it "logs in with the new credentials" do
-          socket.should_receive(:login).with("admin", "newuser", "password")
-          socket.apply_auth("admin" => ["newuser", "password"])
-        end
-      end
-
-      context "and a credential is removed" do
-        it "logs out from the database" do
-          socket.should_receive(:logout).with("admin")
-          socket.apply_auth({})
-        end
-      end
-
-      context "and a credential is added" do
-        it "logs in with the added credentials" do
-          socket.should_receive(:login).with("test", "username", "password")
-          socket.apply_auth(
-            "admin" => ["username", "password"],
-            "test" => ["username", "password"]
-          )
-        end
-      end
-    end
-  end
-
-  describe "instrument" do
-
-    context "when a logger is configured in debug mode" do
-      before do
-        Moped.stub(logger: mock(Logger, debug?: true))
-      end
-
-      it "logs the operations" do
-        socket.should_receive(:log_operations).once
-        socket.instrument([]) {}
-      end
-    end
-
-    context "when a logger is configured but not in debug level" do
-      before do
-        Moped.stub(logger: mock(Logger, debug?: false))
-      end
-
-      it "does not log the operations" do
-        socket.should_receive(:log_operations).never
-        socket.instrument([]) {}
-      end
-    end
-
-    context "when no logger is configured" do
-      before do
-        Moped.stub(logger: nil)
-      end
-
-      it "does not log the operations" do
-        socket.should_receive(:log_operations).never
-        socket.instrument([]) {}
-      end
-    end
-
-    context "when an error occurs" do
-      before do
-        Moped.stub(logger: mock(Logger, debug?: true))
-      end
-
-      it "does not log the operations" do
-        socket.should_receive(:log_operations).never
-
-        lambda do
-          socket.instrument([]) { raise "inner error" }
-        end.should raise_exception("inner error")
-      end
-    end
-
-  end
-
-end

@@ -1 +0,0 @@
-require "spec_helper"

@@ -1,106 +0,0 @@
-require "spec_helper"
-
-describe "testing" do
-  let(:cluster) { Moped::Cluster.new "", false }
-  let(:socket) { Moped::Socket.new "", 99999 }
-  let(:connection) { Support::MockConnection.new }
-
-  before do
-    socket.stub(connection: connection, alive?: true)
-  end
-
-  describe "#sync_socket" do
-  end
-end
-
-__END__
-
-# sequence for single node startup:
-#
-#   connection failure (node not up)
-
-{"ismaster" => true, "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-# sequence for replica set startup pre-initialize:
-#
-#   connection failure
-
-{"ismaster" => false, "secondary" => false, "info" => "can't get local.system.replset config from self or any seed (EMPTYCONFIG)", "isreplicaset" => true, "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-# sequence for replica set startup (master):
-#
-#   connection failure (node not up)
-
-{"ismaster" => false, "secondary" => false, "info" => "can't get local.system.replset config from self or any seed (EMPTYCONFIG)", "isreplicaset" => true, "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-{"setName" => "3ff029114780", "ismaster" => false, "secondary" => true, "hosts" => ["localhost:59246", "localhost:59248", "localhost:59247"], "me" => "localhost:59246", "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-{"ismaster" => false, "secondary" => false, "info" => "Received replSetInitiate - should come online shortly.", "isreplicaset" => true, "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-{"setName" => "3ff029114780", "ismaster" => true, "secondary" => false, "hosts" => ["localhost:59246", "localhost:59248", "localhost:59247"], "primary" => "localhost:59246", "me" => "localhost:59246", "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-# sequence for replica set startup (secondary):
-#
-#   connection failure (node not up)
-
-{"ismaster" => false, "secondary" => false, "info" => "can't get local.system.replset config from self or any seed (EMPTYCONFIG)", "isreplicaset" => true, "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-{"setName" => "3fef4842b608", "ismaster" => false, "secondary" => false, "hosts" => ["localhost:61085", "localhost:61086", "localhost:61084"], "me" => "localhost:61085", "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-{"setName" => "3fef4842b608", "ismaster" => false, "secondary" => false, "hosts" => ["localhost:61085", "localhost:61086", "localhost:61084"], "primary" => "localhost:61084", "me" => "localhost:61085", "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-{"setName" => "3fef4842b608", "ismaster" => false, "secondary" => true, "hosts" => ["localhost:61085", "localhost:61086", "localhost:61084"], "primary" => "localhost:61084", "me" => "localhost:61085", "maxBsonObjectSize" => 16777216, "ok" => 1.0}
-
-__END__
-
-describe Moped::Cluster do
-  context "when connecting to a single seed" do
-    context "and the seed is down"
-
-    context "and that seed is primary" do
-      it "finds the primary node"
-
-      it "adds the secondary to the dynamic seeds"
-      it "adds the arbiter to the dynamic seeds"
-    end
-
-    context "and that seed is secondary" do
-      it "finds the secondary node"
-
-      it "adds the primary to the dynamic seeds"
-      it "adds the arbiter to the dynamic seeds"
-    end
-
-    context "and that seed is an arbiter" do
-      it "adds the primary to the dynamic seeds"
-      it "adds the secondary to the dynamic seeds"
-    end
-  end
-
-  context "when connected to a single seed" do
-    context "and that seed goes down" do
-      it "is able to resync from discovered nodes"
-    end
-  end
-
-  context "when connecting to a replica set" do
-    context "and the replica set is not initiated"
-    context "and the replica set is partially initiated"
-    context "and there is no master node"
-    context "and there is no secondary node"
-  end
-
-  context "when connected to a replica set" do
-    context "and the primary node goes down" do
-      it "issues reads to the secondary"
-    end
-
-    context "and the primary node changes"
-
-    context "and the secondary node goes down" do
-      it "issues inserts to the primary"
-      it "issues reads to the primary"
-    end
-  end
-
-end

@@ -1,9 +1,24 @@
 require "java" if RUBY_PLATFORM == "java"
-require "bundler"
-Bundler.require
+require "rspec"
 
 $:.unshift((Pathname(__FILE__).dirname.parent + "lib").to_s)
 
 require "moped"
+require "support/mongohq"
+require "support/replica_set_simulator"
+require "support/stats"
 
-require "support/mock_connection"
+RSpec.configure do |config|
+  Support::Stats.install!
+
+  config.include Support::ReplicaSetSimulator::Helpers, replica_set: true
+
+  config.filter_run_excluding mongohq: ->(value) do
+    return true if value == :replica_set && !Support::MongoHQ.replica_set_configured?
+    return true if value == :auth && !Support::MongoHQ.auth_node_configured?
+  end
+
+  unless Support::MongoHQ.replica_set_configured? || Support::MongoHQ.auth_node_configured?
+    $stderr.puts Support::MongoHQ.message
+  end
+end

@@ -1,39 +0,0 @@
-module Support
-  class MockConnection
-
-    attr_reader :pending_replies
-
-    def initialize
-      @connected = true
-      @buffer = StringIO.new
-      @pending_replies = []
-    end
-
-    def write(*args)
-
-    end
-
-    def read(*args)
-      while documents = pending_replies.shift
-        reply = Moped::Protocol::Reply.allocate.tap do |reply|
-          documents = [documents] unless documents.is_a? Array
-          reply.documents = documents
-          reply.count = documents.length
-        end
-
-        reply.serialize(@buffer.string)
-      end
-
-      @buffer.read(*args)
-    end
-
-    def closed?
-      !@connected
-    end
-
-    def close
-      @connected = false
-    end
-
-  end
-end

@@ -0,0 +1,61 @@
+module Support
+  module MongoHQ
+    extend self
+
+    def replica_set_configured?
+      ENV["MONGOHQ_REPL_PASS"]
+    end
+
+    def replica_set_seeds
+      [ENV["MONGOHQ_REPL_1_URL"], ENV["MONGOHQ_REPL_2_URL"]]
+    end
+
+    def replica_set_credentials
+      [ENV["MONGOHQ_REPL_USER"], ENV["MONGOHQ_REPL_PASS"]]
+    end
+
+    def replica_set_database
+      ENV["MONGOHQ_REPL_NAME"]
+    end
+
+    def replica_set_session(auth = true)
+      session = Moped::Session.new replica_set_seeds, database: replica_set_database
+      session.login *replica_set_credentials if auth
+      session
+    end
+
+    def auth_seeds
+      [ENV["MONGOHQ_SINGLE_URL"]]
+    end
+
+    def auth_node_configured?
+      ENV["MONGOHQ_SINGLE_PASS"]
+    end
+
+    def auth_credentials
+      [ENV["MONGOHQ_SINGLE_USER"], ENV["MONGOHQ_SINGLE_PASS"]]
+    end
+
+    def auth_database
+      ENV["MONGOHQ_SINGLE_NAME"]
+    end
+
+    def auth_session(auth = true)
+      session = Moped::Session.new auth_seeds, database: auth_database
+      session.login *auth_credentials if auth
+      session
+    end
+
+    def message
+      %Q{
+      ---------------------------------------------------------------------
+      Moped runs specs for authentication and replica sets against MongoHQ.
+
+      If you want to run these specs and need the credentials, contact
+      durran at gmail dot com.
+      ---------------------------------------------------------------------
+      }
+    end
+
+  end
+end

@@ -0,0 +1,306 @@
+module Support
+
+  # This is a helper class for testing replica sets. It works by starting up a
+  # TCP server socket for each desired node. It then proxies all traffic
+  # between a real mongo instance and the client app, with the exception of
+  # ismaster commands, which it returns simulated responses for.
+  class ReplicaSetSimulator
+
+    module Helpers
+      def self.included(context)
+        context.before :all do
+          @replica_set = ReplicaSetSimulator.new
+          @replica_set.start
+
+          @primary, @secondaries = @replica_set.initiate
+        end
+
+        context.after :all do
+          @replica_set.stop
+        end
+
+        context.after :each do
+          @replica_set.nodes.each &:restart
+        end
+
+        context.let :seeds do
+          @replica_set.nodes.map &:address
+        end
+      end
+    end
+
+    attr_reader :nodes
+    attr_reader :manager
+
+    def initialize(nodes = 3)
+      @nodes = nodes.times.map { Node.new(self) }
+      @manager = ConnectionManager.new(@nodes)
+      @mongo = TCPSocket.new "127.0.0.1", 27017
+    end
+
+    # Start the mock replica set.
+    def start
+      @nodes.each &:start
+      @worker = Thread.start do
+        Thread.abort_on_exception = true
+        catch(:shutdown) do
+          loop do
+            Moped.logger.debug "replica_set: waiting for next client"
+            server, client = @manager.next_client
+
+            if server
+              Moped.logger.debug "replica_set: proxying incoming request to mongo"
+              server.proxy(client, @mongo)
+            else
+              Moped.logger.debug "replica_set: no requests; passing"
+              Thread.pass
+            end
+          end
+        end
+      end
+    end
+
+    # Pick a node to be master, and mark the rest as secondary
+    def initiate
+      primary, *secondaries = @nodes.shuffle
+
+      primary.promote
+      secondaries.each &:demote
+
+      return primary, secondaries
+    end
+
+    # Shut down the mock replica set.
+    def stop
+      @manager.shutdown
+      @nodes.each &:stop
+    end
+
+    class Node
+
+      attr_reader :port
+      attr_reader :host
+
+      def initialize(set)
+        @set = set
+        @primary = false
+        @secondary = false
+
+        server = TCPServer.new 0
+        @host = Socket.gethostname
+        @port = server.addr[1]
+        server.close
+      end
+
+      def ==(other)
+        @host == other.host && @port == other.port
+      end
+
+      def address
+        "#{@host}:#{@port}"
+      end
+
+      def primary?
+        @primary
+      end
+
+      def secondary?
+        @secondary
+      end
+
+      def status
+        {
+          "ismaster" => @primary,
+          "secondary" => @secondary,
+          "hosts" => @set.nodes.map(&:address),
+          "me" => address,
+          "maxBsonObjectSize" => 16777216,
+          "ok" => 1.0
+        }
+      end
+
+      def status_reply
+        reply = Moped::Protocol::Reply.new
+        reply.count = 1
+        reply.documents = [status]
+        reply
+      end
+
+      OP_QUERY = 2004
+      OP_GETMORE = 2005
+
+      # Stop and start the node.
+      def restart
+        stop
+        start
+      end
+
+      # Start the node.
+      def start
+        @server = TCPServer.new @port
+      end
+
+      # Stop the node.
+      def stop
+        if @server
+          hiccup
+
+          @server.close
+          @server = nil
+        end
+      end
+      alias close stop
+
+      def accept
+        to_io.accept
+      end
+
+      def closed?
+        !@server || @server.closed?
+      end
+
+      def to_io
+        @server
+      end
+
+      # Mark this node as secondary.
+      def demote
+        @primary = false
+        @secondary = true
+
+        hiccup
+      end
+
+      def hiccup
+        @set.manager.close_clients_for(self)
+      end
+
+      # Mark this node as primary. This also closes any open connections.
+      def promote
+        @primary = true
+        @secondary = false
+
+        hiccup
+      end
+
+      # Proxies a single message from client to the mongo connection.
+      def proxy(client, mongo)
+        incoming_message = client.read(16)
+        length, op_code = incoming_message.unpack("l<x8l<")
+        incoming_message << client.read(length - 16)
+
+        if op_code == OP_QUERY && ismaster_command?(incoming_message)
+          # Intercept the ismaster command and send our own reply.
+          client.write status_reply
+        else
+          # This is a normal command, so proxy it to the real mongo instance.
+          mongo.write incoming_message
+
+          if op_code == OP_QUERY || op_code == OP_GETMORE
+            outgoing_message = mongo.read(4)
+            length, = outgoing_message.unpack('l<')
+            outgoing_message << mongo.read(length - 4)
+
+            client.write outgoing_message
+          end
+        end
+      end
+
+      private
+
+      # Checks a message to see if it's an `ismaster` query.
+      def ismaster_command?(incoming_message)
+        data = StringIO.new(incoming_message)
+        data.read(20) # header and flags
+        data.gets("\x00") # collection name
+        data.read(8) # skip/limit
+
+        selector = Moped::BSON::Document.deserialize(data)
+        selector == { "ismaster" => 1 }
+      end
+    end
+
+    class ConnectionManager
+
+      def initialize(servers)
+        @timeout = 0.1
+        @servers = servers
+        @clients = []
+      end
+
+      def shutdown
+        @servers.each &:close
+        @clients.each &:close
+        @shutdown = true
+      end
+
+      def next_client
+        throw :shutdown if @shutdown
+
+        begin
+          servers = @servers.reject &:closed?
+          clients =  @clients.reject &:closed?
+          Moped.logger.debug "replica_set: selecting on connections"
+          readable, _, errors = Kernel.select(servers + clients, nil, clients, @timeout)
+        rescue IOError, Errno::EBADF, TypeError
+          # Looks like we hit a bad file descriptor or closed connection.
+          Moped.logger.debug "replica_set: io error, retrying"
+          retry
+        end
+
+        return unless readable || errors
+
+        errors.each do |client|
+          client.close
+          @clients.delete client
+        end
+
+        clients, servers = readable.partition { |s| s.class == TCPSocket }
+
+        servers.each do |server|
+          Moped.logger.debug "replica_set: accepting new client for #{server.port}"
+          @clients << server.accept
+        end
+
+        Moped.logger.debug "replica_set: closing dead clients"
+        closed, open = clients.partition &:eof?
+        closed.each { |client| @clients.delete client }
+
+        if client = open.shift
+          Moped.logger.debug "replica_set: finding server for client"
+          server = lookup_server(client)
+
+          Moped.logger.debug "replica_set: sending client #{client.inspect} to #{server.port}"
+          return server, client
+        else
+          nil
+        end
+      end
+
+      def close_clients_for(server)
+        Moped.logger.debug "replica_set: closing open clients on #{server.port}"
+
+        @clients.reject! do |client|
+          port = client.addr(false)[1]
+
+          if port == server.port
+            client.close
+            true
+          else
+            false
+          end
+        end
+      end
+
+      def lookup_server(client)
+        port = client.addr(false)[1]
+
+        @servers.find do |server|
+          server.to_io && server.to_io.addr[1] == port
+        end
+      end
+
+    end
+
+  end
+end

@@ -0,0 +1,51 @@
+module Support
+
+  # Module for recording operations.
+  #
+  #   Support::Stats.install!
+  #
+  #   stats = Support::Stats.collect do
+  #     session.with(safe: true)[:users].insert({})
+  #   end
+  #
+  #   ops = stats["127.0.0.1:27017"]
+  #   ops.size # => 2
+  #   ops[0].class # => Moped::Protocol::Insert
+  #   ops[1].class # => Moped::Protocol::Command
+  #
+  module Stats
+    extend self
+
+    def record(node, operations)
+      key = if node.primary?
+        :primary
+      elsif node.secondary?
+        :secondary
+      else
+        :other
+      end
+
+      @stats[key].concat(operations) if @stats
+    end
+
+    def collect
+      @stats = Hash.new { |hash, key| hash[key] = [] }
+      yield
+      @stats
+    ensure
+      @stats = nil
+    end
+
+    def install!
+      Moped::Node.class_eval <<-EOS
+        alias _logging logging
+
+        def logging(operations, &block)
+          Support::Stats.record(self, operations)
+          _logging(operations, &block)
+        end
+      EOS
+    end
+
+  end
+end