add write timeout feature to ruby driver

Author: Brian Hahn

lib/redis/client.rb

@@ -13,6 +13,7 @@ class Client
       :path => nil,
       :timeout => 5.0,
       :connect_timeout => 5.0,
+      :write_timeout => nil,
       :password => nil,
       :db => 0,
       :driver => nil,
@@ -419,6 +420,8 @@ def _parse_options(options)
         options[:timeout]
       end
 
+      options[:write_timeout] = options[:write_timeout] ? options[:write_timeout].to_f : options[:timeout]
+
       options[:db] = options[:db].to_i
       options[:driver] = _parse_driver(options[:driver]) || Connection.drivers.last
 

lib/redis/connection/ruby.rb

@@ -12,7 +12,7 @@ module SocketMixin
       def initialize(*args)
         super(*args)
 
-        @timeout = nil
+        @timeout = @write_timeout = nil
         @buffer = ""
       end
 
@@ -24,6 +24,14 @@ def timeout=(timeout)
         end
       end
 
+      def write_timeout=(timeout)
+        if timeout && timeout > 0
+          @write_timeout = timeout
+        else
+          @write_timeout = nil
+        end
+      end
+
       def read(nbytes)
         result = @buffer.slice!(0, nbytes)
 
@@ -59,6 +67,11 @@ def _read_from_socket(nbytes)
       rescue EOFError
         raise Errno::ECONNRESET
       end
+
+      # UNIXSocket and TCPSocket don't support write timeouts
+      def write(*args)
+        Timeout.timeout(@write_timeout, TimeoutError) { super }
+      end
     end
 
     if defined?(RUBY_ENGINE) && RUBY_ENGINE == "jruby"
@@ -213,6 +226,7 @@ def self.connect(config)
 
         instance = new(sock)
         instance.timeout = config[:timeout]
+        instance.write_timeout = config[:write_timeout]
         instance.set_tcp_keepalive config[:tcp_keepalive]
         instance
       end
@@ -265,6 +279,10 @@ def timeout=(timeout)
         end
       end
 
+      def write_timeout=(timeout)
+        @sock.write_timeout = timeout
+      end
+
       def write(command)
         @sock.write(build_command(command))
       end

test/internals_test.rb

@@ -160,6 +160,30 @@ def test_connection_timeout
     assert (Time.now - start_time) <= opts[:timeout]
   end
 
+  driver(:ruby) do
+    def test_write_timeout
+      redis_mock(:write_timeout => 0.1, :timeout => 1) do |redis|
+        # This is really disgusting, but this essentially mocks the "super" call in write
+        IO.class_eval do
+          alias_method :old_write, :write
+
+          def write(*args)
+            sleep 1
+          end
+        end
+        assert_raise Redis::TimeoutError do
+          redis.rpush("tmp", "test")
+        end
+      end
+    ensure
+      IO.class_eval do
+        undef_method :write
+        alias_method :write, :old_write
+        undef_method :old_write
+      end
+    end
+  end
+
   def close_on_ping(seq, options = {})
     $request = 0