From b0d5df72fa49546b5f6d9608fe259e1bec8efdc9 Mon Sep 17 00:00:00 2001
From: Kyle Banker <kylebanker@gmail.com>
Date: Wed, 15 Jun 2011 14:20:11 -0400
Subject: [PATCH] RUBY-269 RUBY-275 added connect_timeout option

---
 lib/mongo/connection.rb           | 18 ++++++++++++++++--
 lib/mongo/repl_set_connection.rb  | 16 ++++++++++++++--
 test/connection_test.rb           | 15 +++++++++++++++
 test/replica_sets/connect_test.rb | 16 +++++++++++++++-
 4 files changed, 60 insertions(+), 5 deletions(-)

diff --git a/lib/mongo/connection.rb b/lib/mongo/connection.rb
index 0270654..c76670b 100644
--- a/lib/mongo/connection.rb
+++ b/lib/mongo/connection.rb
@@ -70,6 +70,8 @@ module Mongo
     #   Note: this setting is relevant only for multi-threaded applications (which in Ruby are rare).
     # @option opts [Float] :op_timeout (nil) The number of seconds to wait for a read operation to time out.
     #   Disabled by default.
+    # @option opts [Float] :connect_timeout (nil) The number of seconds to wait before timing out a
+    #   connection attempt.
     #
     # @example localhost, 27017
     #   Connection.new
@@ -628,6 +630,10 @@ module Mongo
       # Timeout on socket read operation.
       @op_timeout = opts[:op_timeout] || nil
 
+      # Timeout on socket connect.
+      @connect_timeout = opts[:connect_timeout] || nil
+
+
       # Mutex for synchronizing pool access
       @connection_mutex = Mutex.new
 
@@ -698,8 +704,16 @@ module Mongo
     def check_is_master(node)
       begin
         host, port = *node
-        socket = TCPSocket.new(host, port)
-        socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+
+        if @connect_timeout
+          Mongo::TimeoutHandler.timeout(@connect_timeout, OperationTimeout) do
+            socket = TCPSocket.new(host, port)
+            socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+          end
+        else
+          socket = TCPSocket.new(host, port)
+          socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+        end
 
         config = self['admin'].command({:ismaster => 1}, :socket => socket)
       rescue OperationFailure, SocketError, SystemCallError, IOError => ex
diff --git a/lib/mongo/repl_set_connection.rb b/lib/mongo/repl_set_connection.rb
index 0e0037f..4f7d7fd 100644
--- a/lib/mongo/repl_set_connection.rb
+++ b/lib/mongo/repl_set_connection.rb
@@ -46,6 +46,10 @@ module Mongo
     # @option options [Float] :timeout (5.0) When all of the connections a pool are checked out,
     #   this is the number of seconds to wait for a new connection to be released before throwing an exception.
     #   Note: this setting is relevant only for multi-threaded applications.
+    # @option opts [Float] :op_timeout (nil) The number of seconds to wait for a read operation to time out.
+    #   Disabled by default.
+    # @option opts [Float] :connect_timeout (nil) The number of seconds to wait before timing out a
+    #   connection attempt.
     #
     # @example Connect to a replica set and provide two seed nodes. Note that the number of seed nodes does
     #   not have to be equal to the number of replica set members. The purpose of seed nodes is to permit
@@ -207,8 +211,16 @@ module Mongo
     def check_is_master(node)
       begin
         host, port = *node
-        socket = TCPSocket.new(host, port)
-        socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+
+        if @connect_timeout
+          Mongo::TimeoutHandler.timeout(@connect_timeout, OperationTimeout) do
+            socket = TCPSocket.new(host, port)
+            socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+          end
+        else
+          socket = TCPSocket.new(host, port)
+          socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+        end
 
         config = self['admin'].command({:ismaster => 1}, :socket => socket)
 
diff --git a/test/connection_test.rb b/test/connection_test.rb
index 7a68c1b..a4f90fe 100644
--- a/test/connection_test.rb
+++ b/test/connection_test.rb
@@ -22,6 +22,21 @@ class TestConnection < Test::Unit::TestCase
     end
   end
 
+  def test_connection_timeout
+    passed = false
+    begin
+      t0 = Time.now
+      Mongo::Connection.new('192.169.169.1', 27017, :connect_timeout => 3)
+    rescue OperationTimeout
+      passed = true
+      t1 = Time.now
+    end
+
+    assert passed
+    assert t1 - t0 < 4
+  end
+
+
   def test_host_port_accessors
     assert_equal @conn.host, TEST_HOST
     assert_equal @conn.port, TEST_PORT
diff --git a/test/replica_sets/connect_test.rb b/test/replica_sets/connect_test.rb
index 272374e..312b3da 100644
--- a/test/replica_sets/connect_test.rb
+++ b/test/replica_sets/connect_test.rb
@@ -27,6 +27,21 @@ class ConnectTest < Test::Unit::TestCase
     end
   end
 
+  def test_connect_timeout
+    passed = false
+    timeout = 3
+    begin
+      t0 = Time.now
+      ReplSetConnection.new(['192.169.169.1', 27017], :connect_timeout => timeout)
+    rescue OperationTimeout
+      passed = true
+      t1 = Time.now
+    end
+
+    assert passed
+    assert t1 - t0 < timeout + 1
+  end
+
   def test_connect
     @conn = ReplSetConnection.new([RS.host, RS.ports[1]], [RS.host, RS.ports[0]],
       [RS.host, RS.ports[2]], :name => RS.name)
@@ -66,7 +81,6 @@ class ConnectTest < Test::Unit::TestCase
       @conn = ReplSetConnection.new([RS.host, RS.ports[0]], [RS.host, RS.ports[1]],
         [RS.host, RS.ports[2]])
     end
-    assert @conn.connected?
   end
 
   def test_connect_with_secondary_node_killed