A blocking version of sendCommand for Redis modules (#2321)

Author: sazzad16

src/main/java/redis/clients/jedis/BinaryJedis.java

@@ -42,10 +42,11 @@
 
 public class BinaryJedis implements BasicCommands, BinaryJedisCommands, MultiKeyBinaryCommands,
     AdvancedBinaryJedisCommands, BinaryScriptingCommands, Closeable {
+
   protected Client client = null;
   protected Transaction transaction = null;
   protected Pipeline pipeline = null;
-  private final byte[][] dummyArray = new byte[0][];
+  protected static final byte[][] DUMMY_ARRAY = new byte[0][];
 
   public BinaryJedis() {
     client = new Client();
@@ -4377,12 +4378,6 @@ public   List<byte[]> xclaim(byte[] key, byte[] groupname, byte[] consumername,
     return client.getBinaryMultiBulkReply();  
   }
 
-  public Object sendCommand(ProtocolCommand cmd, byte[]... args) {
-    checkIsInMultiOrPipeline();
-    client.sendCommand(cmd, args);
-    return client.getOne();
-  }
-
   @Override
   public StreamInfo xinfoStream(byte[] key) {
     checkIsInMultiOrPipeline();
@@ -4407,7 +4402,24 @@ public List<StreamConsumersInfo> xinfoConsumers (byte[] key, byte[] group) {
     return BuilderFactory.STREAM_CONSUMERS_INFO_LIST.build(client.getBinaryMultiBulkReply());
   }
 
+  public Object sendCommand(ProtocolCommand cmd, byte[]... args) {
+    checkIsInMultiOrPipeline();
+    client.sendCommand(cmd, args);
+    return client.getOne();
+  }
+
+  public Object sendBlockingCommand(ProtocolCommand cmd, byte[]... args) {
+    checkIsInMultiOrPipeline();
+    client.sendCommand(cmd, args);
+    client.setTimeoutInfinite();
+    try {
+      return client.getOne();
+    } finally {
+      client.rollbackTimeout();
+    }
+  }
+
   public Object sendCommand(ProtocolCommand cmd) {
-    return sendCommand(cmd, dummyArray);
+    return sendCommand(cmd, DUMMY_ARRAY);
   }
 }

src/main/java/redis/clients/jedis/BinaryJedisCluster.java

@@ -2382,4 +2382,13 @@ public Object execute(Jedis connection){
       }
     }.runBinary(sampleKey);
   }
+
+  public Object sendBlockingCommand(final byte[] sampleKey, final ProtocolCommand cmd, final byte[]... args) {
+    return new JedisClusterCommand<Object>(connectionHandler, maxAttempts) {
+      @Override
+      public Object execute(Jedis connection){
+        return connection.sendBlockingCommand(cmd, args);
+      }
+    }.runBinary(sampleKey);
+  }
 }

src/main/java/redis/clients/jedis/BinaryShardedJedis.java

@@ -1134,6 +1134,13 @@ public Object sendCommand(ProtocolCommand cmd, byte[]... args) {
     return j.sendCommand(cmd, args);
   }
 
+  public Object sendBlockingCommand(ProtocolCommand cmd, byte[]... args) {
+    // default since no sample key provided in JedisCommands interface
+    byte[] sampleKey = args.length > 0 ? args[0] : cmd.getRaw();
+    Jedis j = getShard(sampleKey);
+    return j.sendBlockingCommand(cmd, args);
+  }
+
   public Object sendCommand(ProtocolCommand cmd) {
     return sendCommand(cmd, dummyArray);
   }

src/main/java/redis/clients/jedis/Jedis.java

@@ -4104,4 +4104,15 @@ public Object sendCommand(ProtocolCommand cmd, String... args) {
     client.sendCommand(cmd, args);
     return client.getOne();
   }
+
+  public Object sendBlockingCommand(ProtocolCommand cmd, String... args) {
+    checkIsInMultiOrPipeline();
+    client.sendCommand(cmd, args);
+    client.setTimeoutInfinite();
+    try {
+      return client.getOne();
+    } finally {
+      client.rollbackTimeout();
+    }
+  }
 }

src/main/java/redis/clients/jedis/JedisCluster.java

@@ -2459,5 +2459,13 @@ public Object execute(Jedis connection){
     }.run(sampleKey);
   }
 
+  public Object sendBlockingCommand(final String sampleKey, final ProtocolCommand cmd, final String... args) {
+    return new JedisClusterCommand<Object>(connectionHandler, maxAttempts) {
+      @Override
+      public Object execute(Jedis connection){
+        return connection.sendBlockingCommand(cmd, args);
+      }
+    }.run(sampleKey);
+  }
 
 }

src/main/java/redis/clients/jedis/ShardedJedis.java

@@ -1142,4 +1142,11 @@ public Object sendCommand(ProtocolCommand cmd, String... args) {
     Jedis j = getShard(sampleKey);
     return j.sendCommand(cmd, args);
   }
+
+  public Object sendBlockingCommand(ProtocolCommand cmd, String... args) {
+    // default since no sample key provided in JedisCommands interface
+    String sampleKey = args.length > 0 ? args[0] : cmd.toString();
+    Jedis j = getShard(sampleKey);
+    return j.sendBlockingCommand(cmd, args);
+  }
 }

src/test/java/redis/clients/jedis/tests/commands/AllKindOfValuesCommandsTest.java

@@ -7,6 +7,7 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.fail;
+import static redis.clients.jedis.Protocol.Command.BLPOP;
 import static redis.clients.jedis.Protocol.Command.HGETALL;
 import static redis.clients.jedis.Protocol.Command.GET;
 import static redis.clients.jedis.Protocol.Command.LRANGE;
@@ -886,6 +887,15 @@ public void sendCommandTest(){
     assertEquals("PONG", SafeEncoder.encode((byte[]) jedis.sendCommand(PING)));
   }
 
+  @Test
+  public void sendBlockingCommandTest(){
+    assertNull(jedis.sendBlockingCommand(BLPOP, "foo", Long.toString(1L)));
+
+    jedis.sendCommand(RPUSH, "foo", "bar");
+    assertEquals(Arrays.asList("foo", "bar"), SafeEncoder.encodeObject(jedis.sendBlockingCommand(BLPOP, "foo", Long.toString(1L))));
+
+    assertNull(jedis.sendBlockingCommand(BLPOP, "foo", Long.toString(1L)));
+  }
 
   @Test
   public void encodeCompleteResponse(){

src/test/java/redis/clients/jedis/tests/commands/BinaryValuesCommandsTest.java

@@ -4,6 +4,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
+import static redis.clients.jedis.Protocol.Command.BLPOP;
 import static redis.clients.jedis.Protocol.Command.GET;
 import static redis.clients.jedis.Protocol.Command.LRANGE;
 import static redis.clients.jedis.Protocol.Command.RPUSH;
@@ -17,6 +18,7 @@
 
 import org.junit.Before;
 import org.junit.Test;
+import redis.clients.jedis.Protocol;
 
 import redis.clients.jedis.Protocol.Keyword;
 import redis.clients.jedis.exceptions.JedisDataException;
@@ -323,4 +325,17 @@ public void sendCommandTest(){
     for (int i = 0; i < 3; i++)
       assertArrayEquals(expected.get(i), list.get(i));
   }
+
+  @Test
+  public void sendBlockingCommandTest() {
+    assertNull(jedis.sendBlockingCommand(BLPOP, bfoo, Protocol.toByteArray(1L)));
+
+    jedis.sendCommand(RPUSH, bfoo, bbar);
+    List<byte[]> blpop = (List<byte[]>) jedis.sendBlockingCommand(BLPOP, bfoo, Protocol.toByteArray(1L));
+    assertEquals(2, blpop.size());
+    assertArrayEquals(bfoo, blpop.get(0));
+    assertArrayEquals(bbar, blpop.get(1));
+
+    assertNull(jedis.sendBlockingCommand(BLPOP, bfoo, Protocol.toByteArray(1L)));
+  }
 }