apache · mychaow · Jul 8, 2021 · Jul 7, 2021
diff --git a/example/session/src/main/java/org/apache/iotdb/SessionExample.java b/example/session/src/main/java/org/apache/iotdb/SessionExample.java
@@ -628,4 +628,14 @@ private static void setTimeout() throws StatementExecutionException {
     Session tempSession = new Session(LOCAL_HOST, 6667, "root", "root", 10000, 20000);
     tempSession.setTimeout(60000);
   }
+
+  private static void createClusterSession() throws IoTDBConnectionException {
+    ArrayList<String> nodeList = new ArrayList<>();
+    nodeList.add("127.0.0.1:6669");
+    nodeList.add("127.0.0.1:6667");
+    nodeList.add("127.0.0.1:6668");
+    Session clusterSession = new Session(nodeList, "root", "root");
+    clusterSession.open();
+    clusterSession.close();
+  }
 }
diff --git a/session/src/main/java/org/apache/iotdb/session/Session.java b/session/src/main/java/org/apache/iotdb/session/Session.java
@@ -72,6 +72,7 @@ public class Session {
   public static final String MSG_UNSUPPORTED_DATA_TYPE = "Unsupported data type:";
   public static final String MSG_DONOT_ENABLE_REDIRECT =
       "Query do not enable redirect," + " please confirm the session and server conf.";
+  protected List<String> nodeUrls;
   protected String username;
   protected String password;
   protected int fetchSize;
@@ -237,6 +238,66 @@ public Session(
     this.enableCacheLeader = enableCacheLeader;
   }
 
+  public Session(List<String> nodeUrls, String username, String password) {
+    this(
+        nodeUrls,
+        username,
+        password,
+        Config.DEFAULT_FETCH_SIZE,
+        null,
+        Config.DEFAULT_INITIAL_BUFFER_CAPACITY,
+        Config.DEFAULT_MAX_FRAME_SIZE,
+        Config.DEFAULT_CACHE_LEADER_MODE);
+  }
+
+  /**
+   * Multiple nodeUrl,If one node down, connect to the next one
+   *
+   * @param nodeUrls List<String> Multiple ip:rpcPort eg.127.0.0.1:9001
+   */
+  public Session(List<String> nodeUrls, String username, String password, int fetchSize) {
+    this(
+        nodeUrls,
+        username,
+        password,
+        fetchSize,
+        null,
+        Config.DEFAULT_INITIAL_BUFFER_CAPACITY,
+        Config.DEFAULT_MAX_FRAME_SIZE,
+        Config.DEFAULT_CACHE_LEADER_MODE);
+  }
+
+  public Session(List<String> nodeUrls, String username, String password, ZoneId zoneId) {
+    this(
+        nodeUrls,
+        username,
+        password,
+        Config.DEFAULT_FETCH_SIZE,
+        zoneId,
+        Config.DEFAULT_INITIAL_BUFFER_CAPACITY,
+        Config.DEFAULT_MAX_FRAME_SIZE,
+        Config.DEFAULT_CACHE_LEADER_MODE);
+  }
+
+  public Session(
+      List<String> nodeUrls,
+      String username,
+      String password,
+      int fetchSize,
+      ZoneId zoneId,
+      int thriftDefaultBufferSize,
+      int thriftMaxFrameSize,
+      boolean enableCacheLeader) {
+    this.nodeUrls = nodeUrls;
+    this.username = username;
+    this.password = password;
+    this.fetchSize = fetchSize;
+    this.zoneId = zoneId;
+    this.thriftDefaultBufferSize = thriftDefaultBufferSize;
+    this.thriftMaxFrameSize = thriftMaxFrameSize;
+    this.enableCacheLeader = enableCacheLeader;
+  }
+
   public void setFetchSize(int fetchSize) {
     this.fetchSize = fetchSize;
   }
@@ -291,6 +352,9 @@ public synchronized void close() throws IoTDBConnectionException {
 
   public SessionConnection constructSessionConnection(
       Session session, EndPoint endpoint, ZoneId zoneId) throws IoTDBConnectionException {
+    if (endpoint == null) {
+      return new SessionConnection(session, zoneId);
+    }
     return new SessionConnection(session, endpoint, zoneId);
   }
 

diff --git a/session/src/main/java/org/apache/iotdb/session/SessionConnection.java b/session/src/main/java/org/apache/iotdb/session/SessionConnection.java
@@ -58,7 +58,9 @@
 import org.slf4j.LoggerFactory;
 
 import java.time.ZoneId;
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Random;
 
 public class SessionConnection {
 
@@ -72,6 +74,7 @@ public class SessionConnection {
   private long statementId;
   private ZoneId zoneId;
   private EndPoint endPoint;
+  private List<EndPoint> endPointList = new ArrayList<>();
   private boolean enableRedirect = false;
 
   // TestOnly
@@ -81,10 +84,18 @@ public SessionConnection(Session session, EndPoint endPoint, ZoneId zoneId)
       throws IoTDBConnectionException {
     this.session = session;
     this.endPoint = endPoint;
+    endPointList.add(endPoint);
     this.zoneId = zoneId == null ? ZoneId.systemDefault() : zoneId;
     init(endPoint);
   }
 
+  public SessionConnection(Session session, ZoneId zoneId) throws IoTDBConnectionException {
+    this.session = session;
+    this.zoneId = zoneId == null ? ZoneId.systemDefault() : zoneId;
+    this.endPointList = SessionUtils.parseSeedNodeUrls(session.nodeUrls);
+    initClusterConn();
+  }
+
   private void init(EndPoint endPoint) throws IoTDBConnectionException {
     RpcTransportFactory.setDefaultBufferCapacity(session.thriftDefaultBufferSize);
     RpcTransportFactory.setThriftMaxFrameSize(session.thriftMaxFrameSize);
@@ -143,6 +154,21 @@ private void init(EndPoint endPoint) throws IoTDBConnectionException {
     }
   }
 
+  private void initClusterConn() throws IoTDBConnectionException {
+    for (EndPoint endPoint : endPointList) {
+      try {
+        session.defaultEndPoint = endPoint;
+        init(endPoint);
+      } catch (IoTDBConnectionException e) {
+        if (!reconnect()) {
+          logger.error("Cluster has no nodes to connect");
+          throw new IoTDBConnectionException(e);
+        }
+      }
+      break;
+    }
+  }
+
   public void close() throws IoTDBConnectionException {
     TSCloseSessionReq req = new TSCloseSessionReq(sessionId);
     try {
@@ -699,24 +725,38 @@ protected void testInsertTablets(TSInsertTabletsReq request)
   }
 
   private boolean reconnect() {
-    boolean flag = false;
+    boolean connectedSuccess = false;
+    Random random = new Random();
     for (int i = 1; i <= Config.RETRY_NUM; i++) {
-      try {
-        if (transport != null) {
-          close();
-          init(endPoint);
-          flag = true;
-        }
-      } catch (Exception e) {
-        try {
-          Thread.sleep(Config.RETRY_INTERVAL_MS);
-        } catch (InterruptedException e1) {
-          logger.error("reconnect is interrupted.", e1);
-          Thread.currentThread().interrupt();
+      if (transport != null) {
+        transport.close();
+        int currHostIndex = random.nextInt(endPointList.size());
+        int tryHostNum = 0;
+        for (int j = currHostIndex; j < endPointList.size(); j++) {
+          if (tryHostNum == endPointList.size()) {
+            break;
+          }
+          session.defaultEndPoint = endPointList.get(j);
+          this.endPoint = endPointList.get(j);
+          if (j == endPointList.size() - 1) {
+            j = -1;
+          }
+          tryHostNum++;
+          try {
+            init(endPoint);
+            connectedSuccess = true;
+          } catch (IoTDBConnectionException e) {
+            logger.error("The current node may have been down {},try next node", endPoint);
+            continue;
+          }
+          break;
         }
       }
+      if (connectedSuccess) {
+        break;
+      }
     }
-    return flag;
+    return connectedSuccess;
   }
 
   public boolean isEnableRedirect() {

diff --git a/session/src/main/java/org/apache/iotdb/session/SessionUtils.java b/session/src/main/java/org/apache/iotdb/session/SessionUtils.java
@@ -18,16 +18,24 @@
  */
 package org.apache.iotdb.session;
 
+import org.apache.iotdb.service.rpc.thrift.EndPoint;
 import org.apache.iotdb.tsfile.exception.write.UnSupportedDataTypeException;
 import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
 import org.apache.iotdb.tsfile.utils.Binary;
 import org.apache.iotdb.tsfile.utils.BytesUtils;
 import org.apache.iotdb.tsfile.write.record.Tablet;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
 
 public class SessionUtils {
 
+  private static final Logger logger = LoggerFactory.getLogger(SessionUtils.class);
+
   public static ByteBuffer getTimeBuffer(Tablet tablet) {
     ByteBuffer timeBuffer = ByteBuffer.allocate(tablet.getTimeBytesSize());
     for (int i = 0; i < tablet.rowSize; i++) {
@@ -88,4 +96,31 @@ public static ByteBuffer getValueBuffer(Tablet tablet) {
     valueBuffer.flip();
     return valueBuffer;
   }
+
+  public static List<EndPoint> parseSeedNodeUrls(List<String> nodeUrls) {
+    if (nodeUrls == null) {
+      throw new NumberFormatException("nodeUrls is null");
+    }
+    List<EndPoint> endPointsList = new ArrayList<>();
+    for (String nodeUrl : nodeUrls) {
+      EndPoint endPoint = parseNodeUrl(nodeUrl);
+      endPointsList.add(endPoint);
+    }
+    return endPointsList;
+  }
+
+  private static EndPoint parseNodeUrl(String nodeUrl) {
+    EndPoint endPoint = new EndPoint();
+    String[] split = nodeUrl.split(":");
+    if (split.length != 2) {
+      throw new NumberFormatException("NodeUrl Incorrect format");
+    }
+    String ip = split[0];
+    try {
+      int rpcPort = Integer.parseInt(split[1]);
+      return endPoint.setIp(ip).setPort(rpcPort);
+    } catch (Exception e) {
+      throw new NumberFormatException("NodeUrl Incorrect format");
+    }
+  }
 }
diff --git a/session/src/test/java/org/apache/iotdb/session/IoTDBSessionComplexIT.java b/session/src/test/java/org/apache/iotdb/session/IoTDBSessionComplexIT.java
@@ -800,4 +800,39 @@ private void queryForBatch() throws ClassNotFoundException, SQLException {
       Assert.assertEquals(700, count);
     }
   }
+
+  @Test
+  public void testSessionCluster() throws IoTDBConnectionException, StatementExecutionException {
+    ArrayList<String> nodeList = new ArrayList<>();
+    nodeList.add("127.0.0.1:6669");
+    nodeList.add("127.0.0.1:6667");
+    nodeList.add("127.0.0.1:6668");
+    session = new Session(nodeList, "root", "root");
+    session.open();
+
+    session.setStorageGroup("root.sg1");
+
+    createTimeseries();
+    insertByStr();
+
+    insertViaSQL();
+    queryByDevice("root.sg1.d1");
+
+    session.close();
+  }
+
+  @Test
+  public void testErrorSessionCluster() throws IoTDBConnectionException {
+    ArrayList<String> nodeList = new ArrayList<>();
+    // test Format error
+    nodeList.add("127.0.0.16669");
+    nodeList.add("127.0.0.1:6667");
+    session = new Session(nodeList, "root", "root");
+    try {
+      session.open();
+    } catch (Exception e) {
+      Assert.assertEquals("NodeUrl Incorrect format", e.getMessage());
+    }
+    session.close();
+  }
 }