Merge pull request #615 from redis/write-nonblock

Write nonblock

Author: badboy

.travis.yml

@@ -38,27 +38,48 @@ matrix:
     # hiredis
     - rvm: jruby-18mode
       gemfile: .travis/Gemfile
-      env: conn=hiredis REDIS_BRANCH=2.8
+      env: conn=hiredis REDIS_BRANCH=3.0
+    - rvm: jruby-18mode
+      gemfile: .travis/Gemfile
+      env: conn=hiredis REDIS_BRANCH=3.2
+    - rvm: jruby-19mode
+      gemfile: .travis/Gemfile
+      env: conn=hiredis REDIS_BRANCH=3.0
     - rvm: jruby-19mode
       gemfile: .travis/Gemfile
-      env: conn=hiredis REDIS_BRANCH=2.8
+      env: conn=hiredis REDIS_BRANCH=3.2
     - rvm: jruby-9.0.5.0
       gemfile: .travis/Gemfile
-      env: conn=hiredis REDIS_BRANCH=2.8
+      env: conn=hiredis REDIS_BRANCH=3.0
+    - rvm: jruby-9.0.5.0
+      gemfile: .travis/Gemfile
+      env: conn=hiredis REDIS_BRANCH=3.2
 
     # synchrony
     - rvm: 1.8.7
       gemfile: .travis/Gemfile
-      env: conn=synchrony REDIS_BRANCH=2.8
+      env: conn=synchrony REDIS_BRANCH=3.0
+    - rvm: 1.8.7
+      gemfile: .travis/Gemfile
+      env: conn=synchrony REDIS_BRANCH=3.2
+    - rvm: jruby-18mode
+      gemfile: .travis/Gemfile
+      env: conn=synchrony REDIS_BRANCH=3.0
     - rvm: jruby-18mode
       gemfile: .travis/Gemfile
-      env: conn=synchrony REDIS_BRANCH=2.8
+      env: conn=synchrony REDIS_BRANCH=3.2
     - rvm: jruby-19mode
       gemfile: .travis/Gemfile
-      env: conn=synchrony REDIS_BRANCH=2.8
+      env: conn=synchrony REDIS_BRANCH=3.0
+    - rvm: jruby-19mode
+      gemfile: .travis/Gemfile
+      env: conn=synchrony REDIS_BRANCH=3.2
+    - rvm: jruby-9.0.5.0
+      gemfile: .travis/Gemfile
+      env: conn=synchrony REDIS_BRANCH=3.0
     - rvm: jruby-9.0.5.0
       gemfile: .travis/Gemfile
-      env: conn=synchrony REDIS_BRANCH=2.8
+      env: conn=synchrony REDIS_BRANCH=3.2
   allow_failures:
     - rvm: rbx-2
 

lib/redis/connection/ruby.rb

@@ -10,6 +10,16 @@
   # Not all systems have OpenSSL support
 end
 
+if RUBY_VERSION < "1.9.3"
+  class String
+    # Ruby 1.8.7 does not have byteslice, but it handles encodings differently anyway.
+    # We can simply slice the string, which is a byte array there.
+    def byteslice(*args)
+      slice(*args)
+    end
+  end
+end
+
 class Redis
   module Connection
     module SocketMixin
@@ -80,9 +90,34 @@ def _read_from_socket(nbytes)
         raise Errno::ECONNRESET
       end
 
-      # UNIXSocket and TCPSocket don't support write timeouts
-      def write(*args)
-        Timeout.timeout(@write_timeout, TimeoutError) { super }
+      def _write_to_socket(data)
+        begin
+          write_nonblock(data)
+
+        rescue *NBIO_EXCEPTIONS
+          if IO.select(nil, [self], nil, @write_timeout)
+            retry
+          else
+            raise Redis::TimeoutError
+          end
+        end
+
+      rescue EOFError
+        raise Errno::ECONNRESET
+      end
+
+      def write(data)
+        return super(data) unless @write_timeout
+
+        length = data.bytesize
+        total_count = 0
+        loop do
+          count = _write_to_socket(data)
+
+          total_count += count
+          return total_count if total_count >= length
+          data = data.byteslice(count..-1)
+        end
       end
     end
 

test/internals_test.rb

@@ -5,6 +5,7 @@
 class TestInternals < Test::Unit::TestCase
 
   include Helper::Client
+  include Helper::Skipable
 
   def test_logger
     r.ping
@@ -160,6 +161,25 @@ def test_connection_timeout
     assert (Time.now - start_time) <= opts[:timeout]
   end
 
+  driver(:ruby) do
+    def test_write_timeout
+      return skip("Relies on buffer sizes, might be unreliable")
+
+      server = TCPServer.new("127.0.0.1", 0)
+      port   = server.addr[1]
+
+      # Hacky, but we need the buffer size
+      val = TCPSocket.new("127.0.0.1", port).getsockopt(Socket::SOL_SOCKET, Socket::SO_SNDBUF).unpack("i")[0]
+
+      assert_raise(Redis::TimeoutError) do
+        Timeout.timeout(1) do
+          redis = Redis.new(:port => port, :timeout => 5, :write_timeout => 0.1)
+          redis.set("foo", "1" * val*2)
+        end
+      end
+    end
+  end
+
   def close_on_ping(seq, options = {})
     $request = 0