[IOTDB-6327] Random choosing available nodes to send sql requests

iotdb-client/isession/src/main/java/org/apache/iotdb/isession/INodeSupplier.java

@@ -22,9 +22,12 @@
 import org.apache.iotdb.common.rpc.thrift.TEndPoint;
 
 import java.util.List;
+import java.util.Optional;
 import java.util.function.Supplier;
 
 public interface INodeSupplier extends Supplier<List<TEndPoint>> {
 
   void close();
+
+  Optional<TEndPoint> getQueryEndPoint();
 }

iotdb-client/session/src/main/java/org/apache/iotdb/session/DummyNodesSupplier.java

@@ -24,11 +24,14 @@
 
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 
 public class DummyNodesSupplier implements INodeSupplier {
 
   private final List<TEndPoint> availableNodes;
 
+  private final QueryEndPointPolicy policy = new RoundRobinPolicy();
+
   public DummyNodesSupplier(List<TEndPoint> availableNodes) {
     this.availableNodes = Collections.unmodifiableList(availableNodes);
   }
@@ -38,6 +41,15 @@ public void close() {
     // do nothing
   }
 
+  @Override
+  public Optional<TEndPoint> getQueryEndPoint() {
+    if (availableNodes == null || availableNodes.isEmpty()) {
+      return Optional.empty();
+    } else {
+      return Optional.of(policy.chooseOne(availableNodes));
+    }
+  }
+
   @Override
   public List<TEndPoint> get() {
     return availableNodes;

iotdb-client/session/src/main/java/org/apache/iotdb/session/NodesSupplier.java

@@ -28,9 +28,10 @@
 
 import java.time.ZoneId;
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
-import java.util.Set;
-import java.util.concurrent.CopyOnWriteArraySet;
+import java.util.Optional;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 
@@ -55,9 +56,9 @@ public class NodesSupplier implements INodeSupplier, Runnable {
 
   private static final int FETCH_SIZE = 10_000;
 
-  // availableNodes won't be updated frequently, so we use CopyOnWriteArraySet which is thread-safe
+  // availableNodes won't be updated frequently, so we use CopyOnWriteArrayList which is thread-safe
   // and is optimized for scenarios of reading more and writing less
-  private volatile Set<TEndPoint> availableNodes = new CopyOnWriteArraySet<>();
+  private volatile List<TEndPoint> availableNodes = new CopyOnWriteArrayList<>();
 
   private final boolean useSSL;
   private final String trustStore;
@@ -77,6 +78,8 @@ public class NodesSupplier implements INodeSupplier, Runnable {
 
   private final String version;
 
+  private final QueryEndPointPolicy policy = new RoundRobinPolicy();
+
   private ThriftConnection client;
 
   private volatile boolean closed = false;
@@ -133,7 +136,7 @@ private NodesSupplier(
       String trustStorePwd,
       boolean enableRPCCompression,
       String version) {
-    this.availableNodes.addAll(endPointList);
+    this.availableNodes.addAll(new HashSet<>(endPointList));
     this.userName = userName;
     this.password = password;
     this.useSSL = useSSL;
@@ -151,7 +154,7 @@ private NodesSupplier(
   // and the List needn't be thread-safe, because it will only be used in one thread.
   @Override
   public List<TEndPoint> get() {
-    return new ArrayList<>(availableNodes);
+    return availableNodes;
   }
 
   @Override
@@ -213,6 +216,15 @@ public void close() {
     destroyCurrentClient();
   }
 
+  @Override
+  public Optional<TEndPoint> getQueryEndPoint() {
+    if (availableNodes == null || availableNodes.isEmpty()) {
+      return Optional.empty();
+    } else {
+      return Optional.of(policy.chooseOne(get()));
+    }
+  }
+
   private boolean updateDataNodeList() {
     try (SessionDataSet sessionDataSet =
         client.executeQueryStatement(SHOW_DATA_NODES_COMMAND, TIMEOUT_IN_MS, FETCH_SIZE)) {
@@ -231,7 +243,7 @@ private boolean updateDataNodeList() {
       }
       // replace the older ones.
       if (!res.isEmpty()) {
-        availableNodes = new CopyOnWriteArraySet<>(res);
+        availableNodes = res;
       }
       return true;
     } catch (Exception e) {

iotdb-client/session/src/main/java/org/apache/iotdb/session/QueryEndPointPolicy.java

@@ -0,0 +1,29 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.session;
+
+import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+
+import java.util.List;
+
+public interface QueryEndPointPolicy {
+
+  TEndPoint chooseOne(List<TEndPoint> endPointList);
+}

iotdb-client/session/src/main/java/org/apache/iotdb/session/RoundRobinPolicy.java

@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.session;
+
+import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+
+import java.util.List;
+
+public class RoundRobinPolicy implements QueryEndPointPolicy {
+
+  private int index = 0;
+
+  @Override
+  public TEndPoint chooseOne(List<TEndPoint> endPointList) {
+    int tmp = index;
+    if (tmp >= endPointList.size()) {
+      tmp = 0;
+    }
+    index = tmp + 1;
+    return endPointList.get(tmp);
+  }
+}

iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java

@@ -89,6 +89,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
@@ -888,7 +889,7 @@ public SessionDataSet executeQueryStatement(String sql, long timeoutInMs)
   private SessionDataSet executeStatementMayRedirect(String sql, long timeoutInMs)
       throws StatementExecutionException, IoTDBConnectionException {
     try {
-      return defaultSessionConnection.executeQueryStatement(sql, timeoutInMs);
+      return getQuerySessionConnection().executeQueryStatement(sql, timeoutInMs);
     } catch (RedirectException e) {
       handleQueryRedirection(e.getEndPoint());
       if (enableQueryRedirection) {
@@ -905,6 +906,25 @@ private SessionDataSet executeStatementMayRedirect(String sql, long timeoutInMs)
     }
   }
 
+  private SessionConnection getQuerySessionConnection() {
+    Optional<TEndPoint> endPoint =
+        availableNodes == null ? Optional.empty() : availableNodes.getQueryEndPoint();
+    if (!endPoint.isPresent() || endPointToSessionConnection == null) {
+      return defaultSessionConnection;
+    }
+    SessionConnection connection =
+        endPointToSessionConnection.computeIfAbsent(
+            endPoint.get(),
+            k -> {
+              try {
+                return constructSessionConnection(this, endPoint.get(), zoneId);
+              } catch (IoTDBConnectionException ex) {
+                return null;
+              }
+            });
+    return connection == null ? defaultSessionConnection : connection;
+  }
+
   /**
    * execute non query statement
    *