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
-      @checkin_connection = true
+      socket = nil
-      if @command || @read_preference == :primary
+      begin
-        @connection.checkout_writer
+        @checkin_connection = true
-      else
+        if @command || @read_preference == :primary
-        @read_pool = @connection.read_pool
+          socket = @connection.checkout_writer
-        @connection.checkout_reader
+        else
+          @read_pool = @connection.read_pool
+          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
-        @read_pool = new_pool
+        sock = nil
-        sock = new_pool.checkout
+        begin
-        @checkin_read_pool = true
+          @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,19 +28,25 @@ module Mongo
       add_message_headers(message, operation)
       packed_message = message.to_s
 
-      if connection == :writer
+      sock = nil
-        sock = checkout_writer
-      else
-        sock = checkout_reader
-      end
-
       begin
-        send_message_on_socket(packed_message, sock)
-      ensure
         if connection == :writer
-          checkin_writer(sock)
+          sock = checkout_writer
         else
-          checkin_reader(sock)
+          sock = checkout_reader
+        end
+
+        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
+            checkin_reader(sock)
+          end
         end
       end
     end
@@ -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,24 +113,29 @@ module Mongo
                         read=:primary, exhaust=false)
       request_id = add_message_headers(message, operation)
       packed_message = message.to_s
-      if socket
-        sock = socket
-        should_checkin = false
-      else
-        if command || read == :primary
-          sock = checkout_writer
-        elsif read == :secondary
-          sock = checkout_reader
-        else
-          sock = checkout_tagged(read)
-        end
-        should_checkin = true
-      end
 
       result = ''
+      sock   = nil
       begin
+        if socket
+          sock = socket
+          should_checkin = false
+        else
+          if command || read == :primary
+            sock = checkout_writer
+          elsif read == :secondary
+            sock = checkout_reader
+          else
+            sock = checkout_tagged(read)
+          end
+          should_checkin = true
+        end
+
         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,11 +300,16 @@ module Mongo
     # if no read pool has been defined.
     def checkout_reader
       connect unless connected?
-      socket = get_socket_from_pool(self.read_pool)
+      begin
+        socket = get_socket_from_pool(self.read_pool)
 
-      if !socket
+        if !socket
-        connect
+          connect
-        socket = get_socket_from_pool(self.primary_pool)
+          socket = get_socket_from_pool(self.primary_pool)
+        end
+      rescue => ex
+        checkin(socket) if socket
+        raise ex
       end
 
       if socket
@@ -317,11 +322,16 @@ module Mongo
     # Checkout a socket for writing (i.e., a primary node).
     def checkout_writer
       connect unless connected?
-      socket = get_socket_from_pool(self.primary_pool)
+      begin
-
-      if !socket
-        connect
         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

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