From 03303b840947b7d54e501c567009a966fcd26099 Mon Sep 17 00:00:00 2001 From: Kyle Banker Date: Fri, 2 Dec 2011 15:37:05 -0500 Subject: [PATCH] RUBY-374 Close connection on SystemStackError, NoMemoryError, and SystemCallError --- lib/mongo/cursor.rb | 32 +++++++++++----- lib/mongo/networking.rb | 63 ++++++++++++++++++++------------ lib/mongo/repl_set_connection.rb | 26 +++++++++---- test/connection_test.rb | 52 ++++++++++++++++++++++++++ test/replica_sets/basic_test.rb | 48 ++++++++++++++++++++++++ 5 files changed, 180 insertions(+), 41 deletions(-) diff --git a/lib/mongo/cursor.rb b/lib/mongo/cursor.rb index 7001a1d..ceb3eec 100644 --- a/lib/mongo/cursor.rb +++ b/lib/mongo/cursor.rb @@ -518,13 +518,21 @@ module Mongo end def checkout_socket_from_connection - @checkin_connection = true - if @command || @read_preference == :primary - @connection.checkout_writer - else - @read_pool = @connection.read_pool - @connection.checkout_reader + socket = nil + begin + @checkin_connection = true + if @command || @read_preference == :primary + socket = @connection.checkout_writer + 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 = new_pool.checkout - @checkin_read_pool = true + 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 " + diff --git a/lib/mongo/networking.rb b/lib/mongo/networking.rb index db54e6c..eed51f8 100644 --- a/lib/mongo/networking.rb +++ b/lib/mongo/networking.rb @@ -28,19 +28,25 @@ module Mongo add_message_headers(message, operation) packed_message = message.to_s - if connection == :writer - sock = checkout_writer - else - sock = checkout_reader - end - + sock = nil 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 diff --git a/lib/mongo/repl_set_connection.rb b/lib/mongo/repl_set_connection.rb index 79264b7..d2a160c 100644 --- a/lib/mongo/repl_set_connection.rb +++ b/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 - 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) 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) - - if !socket - connect + 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 diff --git a/test/connection_test.rb b/test/connection_test.rb index dd7caae..2997782 100644 --- a/test/connection_test.rb +++ b/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) diff --git a/test/replica_sets/basic_test.rb b/test/replica_sets/basic_test.rb index 1f927e6..82be87c 100644 --- a/test/replica_sets/basic_test.rb +++ b/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