Handle blocking commands inside pipelines and transactions

Author: byroot

lib/redis.rb

@@ -106,7 +106,12 @@ def queue(*command)
   def commit
     synchronize do |client|
       begin
-        client.call_pipelined(@queue[Thread.current.object_id])
+        pipeline = Pipeline.new(client)
+        @queue[Thread.current.object_id].each do |command|
+          pipeline.call(command)
+        end
+
+        client.call_pipelined(pipeline)
       ensure
         @queue.delete(Thread.current.object_id)
       end
@@ -2405,7 +2410,8 @@ def unwatch
   def pipelined
     synchronize do |client|
       begin
-        original, @client = @client, Pipeline.new
+        pipeline = Pipeline.new(@client)
+        original, @client = @client, pipeline
         yield(self)
         original.call_pipeline(@client)
       ensure
@@ -2450,7 +2456,7 @@ def multi
         client.call([:multi])
       else
         begin
-          pipeline = Pipeline::Multi.new
+          pipeline = Pipeline::Multi.new(@client)
           original, @client = @client, pipeline
           yield(self)
           original.call_pipeline(pipeline)

lib/redis/client.rb

@@ -157,12 +157,11 @@ def call_loop(command, timeout = 0)
     end
 
     def call_pipeline(pipeline)
-      commands = pipeline.commands
-      return [] if commands.empty?
+      return [] if pipeline.futures.empty?
 
       with_reconnect pipeline.with_reconnect? do
         begin
-          pipeline.finish(call_pipelined(commands)).tap do
+          pipeline.finish(call_pipelined(pipeline)).tap do
             self.db = pipeline.db if pipeline.db
           end
         rescue ConnectionError => e
@@ -175,8 +174,8 @@ def call_pipeline(pipeline)
       end
     end
 
-    def call_pipelined(commands)
-      return [] if commands.empty?
+    def call_pipelined(pipeline)
+      return [] if pipeline.futures.empty?
 
       # The method #ensure_connected (called from #process) reconnects once on
       # I/O errors. To make an effort in making sure that commands are not
@@ -186,15 +185,21 @@ def call_pipelined(commands)
       # already successfully executed commands. To circumvent this, don't retry
       # after the first reply has been read successfully.
 
+      commands = pipeline.commands
+
       result = Array.new(commands.size)
       reconnect = @reconnect
 
       begin
         exception = nil
 
         process(commands) do
-          commands.size.times do |i|
-            reply = read
+          pipeline.timeouts.each_with_index do |timeout, i|
+            reply = if timeout
+              with_socket_timeout(timeout) { read }
+            else
+              read
+            end
             result[i] = reply
             @reconnect = false
             exception = reply if exception.nil? && reply.is_a?(CommandError)

lib/redis/pipeline.rb

@@ -1,15 +1,21 @@
 class Redis
   class Pipeline
     attr_accessor :db
+    attr_reader :client
 
     attr :futures
 
-    def initialize
+    def initialize(client)
+      @client = client.is_a?(Pipeline) ? client.client : client
       @with_reconnect = true
       @shutdown = false
       @futures = []
     end
 
+    def timeout
+      client.timeout
+    end
+
     def with_reconnect?
       @with_reconnect
     end
@@ -26,15 +32,19 @@ def empty?
       @futures.empty?
     end
 
-    def call(command, &block)
+    def call(command, timeout: nil, &block)
       # A pipeline that contains a shutdown should not raise ECONNRESET when
       # the connection is gone.
       @shutdown = true if command.first == :shutdown
-      future = Future.new(command, block)
+      future = Future.new(command, block, timeout)
       @futures << future
       future
     end
 
+    def call_with_timeout(command, timeout, &block)
+      call(command, timeout: timeout, &block)
+    end
+
     def call_pipeline(pipeline)
       @shutdown = true if pipeline.shutdown?
       @futures.concat(pipeline.futures)
@@ -43,7 +53,11 @@ def call_pipeline(pipeline)
     end
 
     def commands
-      @futures.map { |f| f._command }
+      @futures.map(&:_command)
+    end
+
+    def timeouts
+      @futures.map(&:timeout)
     end
 
     def with_reconnect(val=true)
@@ -89,6 +103,14 @@ def finish(replies)
         end
       end
 
+      def timeouts
+        if empty?
+          []
+        else
+          [nil, *super, nil]
+        end
+      end
+
       def commands
         if empty?
           []
@@ -108,9 +130,12 @@ def initialize
   class Future < BasicObject
     FutureNotReady = ::Redis::FutureNotReady.new
 
-    def initialize(command, transformation)
+    attr_reader :timeout
+
+    def initialize(command, transformation, timeout)
       @command = command
       @transformation = transformation
+      @timeout = timeout
       @object = FutureNotReady
     end
 

test/blocking_commands_test.rb

@@ -36,4 +36,22 @@ def test_brpoplpush_disable_client_timeout
       assert_equal '0', r.brpoplpush('foo', 'bar')
     end
   end
+
+  def test_brpoplpush_in_transaction
+    results = r.multi do
+      r.brpoplpush('foo', 'bar')
+      r.brpoplpush('foo', 'bar', timeout: 2)
+    end
+    assert_equal [nil, nil], results
+  end
+
+  def test_brpoplpush_in_pipeline
+    mock do |r|
+      results = r.pipelined do
+        r.brpoplpush('foo', 'bar')
+        r.brpoplpush('foo', 'bar', timeout: 2)
+      end
+      assert_equal ['0', '2'], results
+    end
+  end
 end

test/lint/blocking_commands.rb

@@ -152,7 +152,7 @@ def test_brpoplpush_timeout_with_old_prototype
     def test_bzpopmin
       target_version('5.0.0') do
         assert_equal ['{szap}foo', 'a', 0.0], r.bzpopmin('{szap}foo', '{szap}bar', 1)
-        assert_equal nil, r.bzpopmin('{szap}aaa', '{szap}bbb', 1)
+        assert_equal nil, r.bzpopmin('{szap}aaa', '{szap}bbb', 2)
       end
     end