added subscribe and psubscribe with timeout

Author: somebody32

lib/redis.rb

@@ -2134,7 +2134,14 @@ def subscribed?
   # Listen for messages published to the given channels.
   def subscribe(*channels, &block)
     synchronize do |client|
-      _subscription(:subscribe, channels, block)
+      _subscription(:subscribe, 0, channels, block)
+    end
+  end
+
+  # Listen for messages published to the given channels. Throw a timeout error if there is no messages for a timeout period.
+  def subscribe_with_timeout(timeout, *channels, &block)
+    synchronize do |client|
+      _subscription(:subscribe_with_timeout, timeout, channels, block)
     end
   end
 
@@ -2149,7 +2156,14 @@ def unsubscribe(*channels)
   # Listen for messages published to channels matching the given patterns.
   def psubscribe(*channels, &block)
     synchronize do |client|
-      _subscription(:psubscribe, channels, block)
+      _subscription(:psubscribe, 0, channels, block)
+    end
+  end
+
+  # Listen for messages published to channels matching the given patterns. Throw a timeout error if there is no messages for a timeout period.
+  def psubscribe_with_timeout(timeout, *channels, &block)
+    synchronize do |client|
+      _subscription(:psubscribe_with_timeout, timeout, channels, block)
     end
   end
 
@@ -2749,12 +2763,16 @@ def _identity
     }
   end
 
-  def _subscription(method, channels, block)
+  def _subscription(method, timeout, channels, block)
     return @client.call([method] + channels) if subscribed?
 
     begin
       original, @client = @client, SubscribedClient.new(@client)
-      @client.send(method, *channels, &block)
+      if timeout > 0
+        @client.send(method, timeout, *channels, &block)
+      else
+        @client.send(method, *channels, &block)
+      end
     ensure
       @client = original
     end

lib/redis/client.rb

@@ -127,10 +127,10 @@ def call(command, &block)
       end
     end
 
-    def call_loop(command)
+    def call_loop(command, timeout = 0)
       error = nil
 
-      result = without_socket_timeout do
+      result = with_socket_timeout(timeout) do
         process([command]) do
           loop do
             reply = read

lib/redis/subscribe.rb

@@ -12,10 +12,18 @@ def subscribe(*channels, &block)
       subscription("subscribe", "unsubscribe", channels, block)
     end
 
+    def subscribe_with_timeout(timeout, *channels, &block)
+      subscription("subscribe", "unsubscribe", channels, block, timeout)
+    end
+
     def psubscribe(*channels, &block)
       subscription("psubscribe", "punsubscribe", channels, block)
     end
 
+    def psubscribe_with_timeout(timeout, *channels, &block)
+      subscription("psubscribe", "punsubscribe", channels, block, timeout)
+    end
+
     def unsubscribe(*channels)
       call([:unsubscribe, *channels])
     end
@@ -26,13 +34,13 @@ def punsubscribe(*channels)
 
   protected
 
-    def subscription(start, stop, channels, block)
+    def subscription(start, stop, channels, block, timeout = 0)
       sub = Subscription.new(&block)
 
       unsubscribed = false
 
       begin
-        @client.call_loop([start, *channels]) do |line|
+        @client.call_loop([start, *channels], timeout) do |line|
           type, *rest = line
           sub.callbacks[type].call(*rest)
           unsubscribed = type == stop && rest.last == 0

test/publish_subscribe_test.rb

@@ -251,4 +251,42 @@ def test_subscribe_past_a_timeout
       assert received
     end
   end
+
+  def test_subscribe_with_timeout
+    assert_raise Redis::TimeoutError do
+      sleep = %{sleep #{OPTIONS[:timeout] * 2}}
+      publish = %{echo "publish foo bar\r\n" | nc localhost #{OPTIONS[:port]}}
+      cmd = [sleep, publish].join("; ")
+
+      IO.popen(cmd, "r+") do |pipe|
+        received = false
+
+        r.subscribe_with_timeout(OPTIONS[:timeout], "foo")  do |on|
+          on.message do |channel, message|
+            received = true
+            r.unsubscribe
+          end
+        end
+      end
+    end
+  end
+
+  def test_psubscribe_with_timeout
+    assert_raise Redis::TimeoutError do
+      sleep = %{sleep #{OPTIONS[:timeout] * 2}}
+      publish = %{echo "publish foo bar\r\n" | nc localhost #{OPTIONS[:port]}}
+      cmd = [sleep, publish].join("; ")
+
+      IO.popen(cmd, "r+") do |pipe|
+        received = false
+
+        r.psubscribe_with_timeout(OPTIONS[:timeout], "f*")  do |on|
+          on.message do |channel, message|
+            received = true
+            r.unsubscribe
+          end
+        end
+      end
+    end
+  end
 end