From e49d50acc2ee21f7aa9b3c853c16d7b8bf2f6515 Mon Sep 17 00:00:00 2001
From: Kyle Banker <kylebanker@gmail.com>
Date: Mon, 28 Mar 2011 11:09:27 -0400
Subject: [PATCH] RUBY-236 set op_timeout for socket receive timeouts

---
 README.md               |  8 ++++++++
 lib/mongo.rb            | 17 ++++++++++++++++-
 lib/mongo/connection.rb | 32 ++++++++++++++++++++++----------
 lib/mongo/exceptions.rb |  3 +++
 4 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index 7653278..6ec92bc 100644
--- a/README.md
+++ b/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
 
diff --git a/lib/mongo.rb b/lib/mongo.rb
index ebc715f..cfba033 100644
--- a/lib/mongo.rb
+++ b/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
diff --git a/lib/mongo/connection.rb b/lib/mongo/connection.rb
index 39d483f..48323ce 100644
--- a/lib/mongo/connection.rb
+++ b/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)
-        raise ConnectionFailure, "connection closed" unless message && message.length > 0
-        if message.length < length
-          chunk = new_binary_string
-          while message.length < length
-            socket.read(length - message.length, chunk)
-            raise ConnectionFailure, "connection closed" unless chunk.length > 0
-            message << chunk
+        Mongo::TimeoutHandler.timeout(@op_timeout, OperationTimeout) do
+          socket.read(length, message)
+          raise ConnectionFailure, "connection closed" unless message && message.length > 0
+          if message.length < length
+            chunk = new_binary_string
+            while message.length < length
+              socket.read(length - message.length, 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
diff --git a/lib/mongo/exceptions.rb b/lib/mongo/exceptions.rb
index 21e16d4..0989635 100644
--- a/lib/mongo/exceptions.rb
+++ b/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