RUBY-236 set op_timeout for socket receive timeouts

README.md

@@ -272,6 +272,14 @@ Notes:
 * Cursors will timeout on the server after 10 minutes. If you need to keep a cursor
   open for more than 10 minutes, specify `:timeout => false` when you create the cursor.
 
+## Socket timeouts
+
+The Ruby driver support timeouts on socket read operations. To enable them, set the
+`:op_timeout` option when you create a `Mongo::Connection` object.
+
+If implementing higher-level timeouts, using tools like `Rack::Timeout`, it's very important
+to call `Mongo::Connection#close` to prevent the subsequent operation from receiving the previous
+request.
 
 # Testing
 

lib/mongo.rb

@@ -76,4 +76,19 @@ if RUBY_PLATFORM =~ /java/
 end
 require 'mongo/gridfs/grid_file_system'
 
-
+# Use SystemTimer on Ruby 1.8
+if !defined?(RUBY_ENGINE) || (RUBY_ENGINE == 'ruby' && RUBY_VERSION < '1.9.0')
+  begin
+    require 'system_timer'
+    Mongo::TimeoutHandler = SystemTimer
+  rescue LoadError
+    warn "Could not load SystemTimer gem. Falling back to timeout.rb." +
+         "SystemTimer is STRONGLY recommended for timeouts in Ruby 1.8.7. " +
+         "See http://ph7spot.com/musings/system-timer for details."
+    require 'timeout'
+    Mongo::TimeoutHandler = Timeout
+  end
+else
+  require 'timeout'
+  Mongo::TimeoutHandler = Timeout
+end

lib/mongo/connection.rb

@@ -68,6 +68,8 @@ module Mongo
     # @option opts [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 (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.
     #
     # @example localhost, 27017
     #   Connection.new
@@ -598,6 +600,9 @@ module Mongo
       @pool_size = opts[:pool_size] || 1
       @timeout   = opts[:timeout]   || 5.0
 
+      # Timeout on socket read operation.
+      @op_timeout = opts[:op_timeout] || nil
+
       # Mutex for synchronizing pool access
       @connection_mutex = Mutex.new
 
@@ -833,19 +838,26 @@ module Mongo
     def receive_message_on_socket(length, socket)
       begin
         message = new_binary_string
-        socket.read(length, message)
+        Mongo::TimeoutHandler.timeout(@op_timeout, OperationTimeout) do
-        raise ConnectionFailure, "connection closed" unless message && message.length > 0
+          socket.read(length, message)
-        if message.length < length
+          raise ConnectionFailure, "connection closed" unless message && message.length > 0
-          chunk = new_binary_string
+          if message.length < length
-          while message.length < length
+            chunk = new_binary_string
-            socket.read(length - message.length, chunk)
+            while message.length < length
-            raise ConnectionFailure, "connection closed" unless chunk.length > 0
+              socket.read(length - message.length, chunk)
-            message << chunk
+              raise ConnectionFailure, "connection closed" unless chunk.length > 0
+              message << chunk
+            end
           end
         end
-        rescue => ex
+      rescue => ex
           close
-          raise ConnectionFailure, "Operation failed with the following exception: #{ex}"
+
+          if ex.class == OperationTimeout
+            raise OperationTimeout, "Timed out waiting in socket read."
+          else
+            raise ConnectionFailure, "Operation failed with the following exception: #{ex}"
+          end
       end
       message
     end

lib/mongo/exceptions.rb

@@ -57,6 +57,9 @@ module Mongo
   # Raised when a database operation fails.
   class OperationFailure < MongoDBError; end
 
+  # Raised when a socket read operation times out.
+  class OperationTimeout < ::Timeout::Error; end
+
   # Raised when a client attempts to perform an invalid operation.
   class InvalidOperation < MongoDBError; end