From c4807a932dcf2aaec95fff9f81794980ada2f51e Mon Sep 17 00:00:00 2001
From: Jens Ehlers <jeh@jehlers.de>
Date: Sat, 19 May 2018 16:22:26 +0200
Subject: [PATCH] [redis] Redis cluster support (#1072)

---
 pom.xml                                       |  2 +-
 redis/README.md                               |  7 +++--
 .../java/com/yahoo/ycsb/db/RedisClient.java   | 30 +++++++++++++++----
 3 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/pom.xml b/pom.xml
index c3d60da0..822d5fa1 100644
--- a/pom.xml
+++ b/pom.xml
@@ -101,7 +101,7 @@ LICENSE file.
     <openjpa.jdbc.version>2.1.1</openjpa.jdbc.version>
     <orientdb.version>2.2.10</orientdb.version>
     <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
-    <redis.version>2.0.0</redis.version>
+    <redis.version>2.9.0</redis.version>
     <riak.version>2.0.5</riak.version>
     <s3.version>1.10.20</s3.version>
     <solr.version>5.5.3</solr.version>
diff --git a/redis/README.md b/redis/README.md
index bc883271..b15a55e7 100644
--- a/redis/README.md
+++ b/redis/README.md
@@ -33,12 +33,15 @@ Git clone YCSB and compile:
 
 ### 4. Provide Redis Connection Parameters
     
-Set the host, port, and password (do not redis auth is not turned on) in the 
-workload you plan to run.
+Set host, port, password, and cluster mode in the workload you plan to run. 
 
 - `redis.host`
 - `redis.port`
 - `redis.password`
+  * Don't set the password if redis auth is disabled.
+- `redis.cluster`
+  * Set the cluster parameter to `true` if redis cluster mode is enabled.
+  * Default is `false`.
 
 Or, you can set configs with the shell command, EG:
 
diff --git a/redis/src/main/java/com/yahoo/ycsb/db/RedisClient.java b/redis/src/main/java/com/yahoo/ycsb/db/RedisClient.java
index a571a2c0..d6cf6f80 100644
--- a/redis/src/main/java/com/yahoo/ycsb/db/RedisClient.java
+++ b/redis/src/main/java/com/yahoo/ycsb/db/RedisClient.java
@@ -29,12 +29,18 @@ import com.yahoo.ycsb.DB;
 import com.yahoo.ycsb.DBException;
 import com.yahoo.ycsb.Status;
 import com.yahoo.ycsb.StringByteIterator;
-
+import redis.clients.jedis.BasicCommands;
+import redis.clients.jedis.HostAndPort;
 import redis.clients.jedis.Jedis;
+import redis.clients.jedis.JedisCluster;
+import redis.clients.jedis.JedisCommands;
 import redis.clients.jedis.Protocol;
 
+import java.io.Closeable;
+import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Properties;
@@ -48,11 +54,12 @@ import java.util.Vector;
  */
 public class RedisClient extends DB {
 
-  private Jedis jedis;
+  private JedisCommands jedis;
 
   public static final String HOST_PROPERTY = "redis.host";
   public static final String PORT_PROPERTY = "redis.port";
   public static final String PASSWORD_PROPERTY = "redis.password";
+  public static final String CLUSTER_PROPERTY = "redis.cluster";
 
   public static final String INDEX_KEY = "_indices";
 
@@ -68,17 +75,28 @@ public class RedisClient extends DB {
     }
     String host = props.getProperty(HOST_PROPERTY);
 
-    jedis = new Jedis(host, port);
-    jedis.connect();
+    boolean clusterEnabled = Boolean.parseBoolean(props.getProperty(CLUSTER_PROPERTY));
+    if (clusterEnabled) {
+      Set<HostAndPort> jedisClusterNodes = new HashSet<>();
+      jedisClusterNodes.add(new HostAndPort(host, port));
+      jedis = new JedisCluster(jedisClusterNodes);
+    } else {
+      jedis = new Jedis(host, port);
+      ((Jedis) jedis).connect();
+    }
 
     String password = props.getProperty(PASSWORD_PROPERTY);
     if (password != null) {
-      jedis.auth(password);
+      ((BasicCommands) jedis).auth(password);
     }
   }
 
   public void cleanup() throws DBException {
-    jedis.disconnect();
+    try {
+      ((Closeable) jedis).close();
+    } catch (IOException e) {
+      throw new DBException("Closing connection failed.");
+    }
   }
 
   /*