RUBY-374 Close connection on SystemStackError, NoMemoryError, and SystemCallError

lib/mongo/cursor.rb

@@ -518,13 +518,21 @@ module Mongo
     end
 
     def checkout_socket_from_connection
+      socket = nil
+      begin
         @checkin_connection = true
         if @command || @read_preference == :primary
-        @connection.checkout_writer
+          socket = @connection.checkout_writer
         else
           @read_pool = @connection.read_pool
-        @connection.checkout_reader
+          socket = @connection.checkout_reader
         end
+      rescue SystemStackError, NoMemoryError, SystemCallError => ex
+        @connection.close
+        raise ex
+      end
+
+      socket
     end
 
     def checkout_socket_for_op_get_more
@@ -540,9 +548,15 @@ module Mongo
         pool.host == @read_pool.host && pool.port == @read_pool.port
       end
       if new_pool
+        sock = nil
+        begin
           @read_pool = new_pool
           sock = new_pool.checkout
           @checkin_read_pool = true
+        rescue SystemStackError, NoMemoryError, SystemCallError => ex
+          @connection.close
+          raise ex
+        end
         return sock
       else
         raise Mongo::OperationFailure, "Failure to continue iterating " +

lib/mongo/networking.rb

@@ -28,15 +28,20 @@ module Mongo
       add_message_headers(message, operation)
       packed_message = message.to_s
 
+      sock = nil
+      begin
         if connection == :writer
           sock = checkout_writer
         else
           sock = checkout_reader
         end
 
-      begin
         send_message_on_socket(packed_message, sock)
+      rescue SystemStackError, NoMemoryError, SystemCallError => ex
+        close
+        raise ex
       ensure
+        if sock
           if connection == :writer
             checkin_writer(sock)
           else
@@ -44,6 +49,7 @@ module Mongo
           end
         end
       end
+    end
 
     # Sends a message to the database, waits for a response, and raises
     # an exception if the operation has failed.
@@ -66,14 +72,18 @@ module Mongo
       last_error_id = add_message_headers(last_error_message, Mongo::Constants::OP_QUERY)
 
       packed_message = message.append!(last_error_message).to_s
-      sock = checkout_writer
+      sock = nil
       begin
+        sock = checkout_writer
         send_message_on_socket(packed_message, sock)
         docs, num_received, cursor_id = receive(sock, last_error_id)
         checkin_writer(sock)
       rescue ConnectionFailure, OperationFailure, OperationTimeout => ex
         checkin_writer(sock)
         raise ex
+      rescue SystemStackError, NoMemoryError, SystemCallError => ex
+        close
+        raise ex
       end
 
       if num_received == 1 && (error = docs[0]['err'] || docs[0]['errmsg'])
@@ -103,6 +113,10 @@ module Mongo
                         read=:primary, exhaust=false)
       request_id = add_message_headers(message, operation)
       packed_message = message.to_s
+
+      result = ''
+      sock   = nil
+      begin
         if socket
           sock = socket
           should_checkin = false
@@ -117,10 +131,11 @@ module Mongo
           should_checkin = true
         end
 
-      result = ''
-      begin
         send_message_on_socket(packed_message, sock)
         result = receive(sock, request_id, exhaust)
+      rescue SystemStackError, NoMemoryError, SystemCallError => ex
+        close
+        raise ex
       ensure
         if should_checkin
           if command || read == :primary

lib/mongo/repl_set_connection.rb

@@ -300,12 +300,17 @@ module Mongo
     # if no read pool has been defined.
     def checkout_reader
       connect unless connected?
+      begin
         socket = get_socket_from_pool(self.read_pool)
 
         if !socket
           connect
           socket = get_socket_from_pool(self.primary_pool)
         end
+      rescue => ex
+        checkin(socket) if socket
+        raise ex
+      end
 
       if socket
         socket
@@ -317,12 +322,17 @@ module Mongo
     # Checkout a socket for writing (i.e., a primary node).
     def checkout_writer
       connect unless connected?
+      begin
         socket = get_socket_from_pool(self.primary_pool)
 
         if !socket
           connect
           socket = get_socket_from_pool(self.primary_pool)
         end
+      rescue => ex
+        checkin(socket)
+        raise ex
+      end
 
       if socket
         socket

test/connection_test.rb

@@ -278,6 +278,58 @@ class TestConnection < Test::Unit::TestCase
     end
   end
 
+  context "Socket pools" do
+    context "checking out writers" do
+      setup do
+        @con = standard_connection(:pool_size => 10, :timeout => 10)
+        @coll = @con[MONGO_TEST_DB]['test-connection-exceptions']
+      end
+
+      should "close the connection on send_message for major exceptions" do
+        @con.expects(:checkout_writer).raises(SystemStackError)
+        @con.expects(:close)
+        begin
+          @coll.insert({:foo => "bar"})
+        rescue SystemStackError
+        end
+      end
+
+      should "close the connection on send_message_with_safe_check for major exceptions" do
+        @con.expects(:checkout_writer).raises(SystemStackError)
+        @con.expects(:close)
+        begin
+          @coll.insert({:foo => "bar"}, :safe => true)
+        rescue SystemStackError
+        end
+      end
+
+      should "close the connection on receive_message for major exceptions" do
+        @con.expects(:checkout_writer).raises(SystemStackError)
+        @con.expects(:close)
+        begin
+          @coll.find.next
+        rescue SystemStackError
+        end
+      end
+    end
+
+    context "checking out readers" do
+      setup do
+        @con = standard_connection(:pool_size => 10, :timeout => 10, :slave_ok => true)
+        @coll = @con[MONGO_TEST_DB]['test-connection-exceptions']
+      end
+
+      should "close the connection on receive_message for major exceptions" do
+        @con.expects(:checkout_reader).raises(SystemStackError)
+        @con.expects(:close)
+        begin
+          @coll.find.next
+        rescue SystemStackError
+        end
+      end
+    end
+  end
+
   context "Connection exceptions" do
     setup do
       @con = standard_connection(:pool_size => 10, :timeout => 10)

test/replica_sets/basic_test.rb

@@ -43,4 +43,52 @@ class BasicTest < Test::Unit::TestCase
     assert_equal 90, @conn.refresh_interval
     assert_equal @conn.refresh_mode, false
   end
+
+  context "Socket pools" do
+    context "checking out writers" do
+      setup do
+        seeds = [[self.rs.host, self.rs.ports[0]], [self.rs.host, self.rs.ports[1]],
+          [self.rs.host, self.rs.ports[2]]]
+        args = seeds << {:name => self.rs.name}
+        @con = ReplSetConnection.new(*args)
+        @coll = @con[MONGO_TEST_DB]['test-connection-exceptions']
+      end
+
+      should "close the connection on send_message for major exceptions" do
+        @con.expects(:checkout_writer).raises(SystemStackError)
+        @con.expects(:close)
+        begin
+          @coll.insert({:foo => "bar"})
+        rescue SystemStackError
+        end
+      end
+
+      should "close the connection on send_message_with_safe_check for major exceptions" do
+        @con.expects(:checkout_writer).raises(SystemStackError)
+        @con.expects(:close)
+        begin
+          @coll.insert({:foo => "bar"}, :safe => true)
+        rescue SystemStackError
+        end
+      end
+
+      should "close the connection on receive_message for major exceptions" do
+        @con.expects(:checkout_writer).raises(SystemStackError)
+        @con.expects(:close)
+        begin
+          @coll.find.next
+        rescue SystemStackError
+        end
+      end
+
+      should "close the connection on receive_message for major exceptions" do
+        @con.expects(:checkout_reader).raises(SystemStackError)
+        @con.expects(:close)
+        begin
+          @coll.find({}, :read => :secondary).next
+        rescue SystemStackError
+        end
+      end
+    end
+  end
 end