Merge branch 'beedub-write_timeout' into 3.3

Author: djanowski

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,
@@ -433,6 +434,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

@@ -2,6 +2,7 @@
 require "redis/connection/command_helper"
 require "redis/errors"
 require "socket"
+require "timeout"
 
 class Redis
   module Connection
@@ -12,7 +13,7 @@ module SocketMixin
       def initialize(*args)
         super(*args)
 
-        @timeout = nil
+        @timeout = @write_timeout = nil
         @buffer = ""
       end
 
@@ -24,6 +25,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 +68,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 +227,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 +280,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,19 @@ def test_connection_timeout
     assert (Time.now - start_time) <= opts[:timeout]
   end
 
+  driver(:ruby) do
+    def test_write_timeout
+      TCPServer.new("127.0.0.1", 6383)
+
+      assert_raise(Redis::TimeoutError) do
+        Timeout.timeout(1) do
+          redis = Redis.new(:port => 6383, :timeout => 5, :write_timeout => 0.2)
+          redis.set("foo", "1" * 1048576)
+        end
+      end
+    end
+  end
+
   def close_on_ping(seq, options = {})
     $request = 0
 

test/lint/blocking_commands.rb

@@ -35,7 +35,7 @@ def mock(options = {}, &blk)
         :brpoplpush => lambda do |*args|
           sleep options[:delay] if options.has_key?(:delay)
           to_protocol(args.last)
-        end,
+        end
       }
 
       redis_mock(commands, &blk)