Cluster pipelining

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

@@ -0,0 +1,9 @@
+package redis.clients.jedis;
+
+import redis.clients.jedis.commands.ProtocolCommand;
+
+public interface CommandListener {
+
+  void afterCommand(Connection conn, ProtocolCommand cmd, final byte[][] args);
+
+}

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

@@ -37,6 +37,7 @@ public class Connection implements Closeable {
   private SSLSocketFactory sslSocketFactory;
   private SSLParameters sslParameters;
   private HostnameVerifier hostnameVerifier;
+  private CommandListener commandListener;
 
   public Connection() {
   }
@@ -108,6 +109,14 @@ public void rollbackTimeout() {
     }
   }
 
+  public void setCommandListener(CommandListener listener) {
+    this.commandListener = listener;
+  }
+
+  public void clearCommandListener() {
+    this.commandListener = null;
+  }
+
   public Connection sendCommand(final ProtocolCommand cmd, final String... args) {
     final byte[][] bargs = new byte[args.length][];
     for (int i = 0; i < args.length; i++) {
@@ -124,6 +133,12 @@ public Connection sendCommand(final ProtocolCommand cmd, final byte[]... args) {
     try {
       connect();
       Protocol.sendCommand(outputStream, cmd, args);
+
+      // Record commands sent, for retrying in cluster pipeline.
+      if (commandListener != null) {
+        commandListener.afterCommand(this, cmd, args);
+      }
+
       return this;
     } catch (JedisConnectionException ex) {
       /*

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

@@ -105,7 +105,11 @@ public JedisCluster(Set<HostAndPort> jedisClusterNode, int connectionTimeout, in
   public JedisCluster(Set<HostAndPort> jedisClusterNode, int connectionTimeout, int soTimeout,
           int maxAttempts, String password, String clientName, final GenericObjectPoolConfig poolConfig) {
     super(jedisClusterNode, connectionTimeout, soTimeout, maxAttempts, password, clientName, poolConfig);
-}
+  }
+
+  public JedisClusterPipeline pipelined() {
+    return new JedisClusterPipeline(connectionHandler);
+  }
 
   @Override
   public String set(final String key, final String value) {

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

@@ -34,6 +34,14 @@ public Map<String, JedisPool> getNodes() {
     return cache.getNodes();
   }
 
+  public HostAndPort getSlotNode(int slot) {
+    return cache.getSlotNode(slot);
+  }
+
+  public void assignSlotToNode(int slot, HostAndPort targetNode) {
+    cache.assignSlotToNode(slot, targetNode);
+  }
+
   private void initializeSlotsCache(Set<HostAndPort> startNodes, GenericObjectPoolConfig poolConfig, String password, String clientName) {
     for (HostAndPort hostAndPort : startNodes) {
       Jedis jedis = new Jedis(hostAndPort.getHost(), hostAndPort.getPort());

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

@@ -21,6 +21,7 @@
 public class JedisClusterInfoCache {
   private final Map<String, JedisPool> nodes = new HashMap<String, JedisPool>();
   private final Map<Integer, JedisPool> slots = new HashMap<Integer, JedisPool>();
+  private Map<Integer, HostAndPort> mapping = new HashMap<Integer, HostAndPort>();
 
   private final ReentrantReadWriteLock rwl = new ReentrantReadWriteLock();
   private final Lock r = rwl.readLock();
@@ -123,6 +124,7 @@ public void renewClusterSlots(Jedis jedis) {
   private void discoverClusterSlots(Jedis jedis) {
     List<Object> slots = jedis.clusterSlots();
     this.slots.clear();
+    this.mapping.clear();
 
     for (Object slotInfoObj : slots) {
       List<Object> slotInfo = (List<Object>) slotInfoObj;
@@ -205,6 +207,7 @@ public void assignSlotToNode(int slot, HostAndPort targetNode) {
     try {
       JedisPool targetPool = setupNodeIfNotExist(targetNode);
       slots.put(slot, targetPool);
+      mapping.put(slot, targetNode);
     } finally {
       w.unlock();
     }
@@ -216,6 +219,7 @@ public void assignSlotsToNode(List<Integer> targetSlots, HostAndPort targetNode)
       JedisPool targetPool = setupNodeIfNotExist(targetNode);
       for (Integer slot : targetSlots) {
         slots.put(slot, targetPool);
+        mapping.put(slot, targetNode);
       }
     } finally {
       w.unlock();
@@ -277,6 +281,7 @@ public void reset() {
       }
       nodes.clear();
       slots.clear();
+      mapping.clear();
     } finally {
       w.unlock();
     }
@@ -302,4 +307,13 @@ private List<Integer> getAssignedSlotArray(List<Object> slotInfo) {
     }
     return slotNums;
   }
+
+  public HostAndPort getSlotNode(int slot) {
+    r.lock();
+    try {
+      return mapping.get(slot);
+    } finally {
+      r.unlock();
+    }
+  }
 }