You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iotdb.apache.org by ja...@apache.org on 2023/12/12 03:23:41 UTC
(iotdb) 01/01: [IOTDB-6199] Support automatically fetch available DN list in Session & SessionPool
This is an automated email from the ASF dual-hosted git repository.
jackietien pushed a commit to branch IOTDB-6199
in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 4538d517968e3ade4103cd125244068185ed77a8
Author: JackieTien97 <ja...@gmail.com>
AuthorDate: Tue Dec 12 11:23:29 2023 +0800
[IOTDB-6199] Support automatically fetch available DN list in Session & SessionPool
---
.../org/apache/iotdb/isession/INodeSupplier.java | 30 +++
.../java/org/apache/iotdb/isession/ISession.java | 3 +-
.../org/apache/iotdb/session/NodesSupplier.java | 221 +++++++++++++++++++++
.../java/org/apache/iotdb/session/Session.java | 59 +++++-
.../apache/iotdb/session/SessionConnection.java | 22 +-
.../org/apache/iotdb/session/ThriftConnection.java | 190 ++++++++++++++++++
.../org/apache/iotdb/session/pool/SessionPool.java | 87 +++++++-
.../iotdb/session/SessionConnectionTest.java | 13 +-
8 files changed, 612 insertions(+), 13 deletions(-)
diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/INodeSupplier.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/INodeSupplier.java
new file mode 100644
index 00000000000..371630a6b85
--- /dev/null
+++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/INodeSupplier.java
@@ -0,0 +1,30 @@
+/*
+ * 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.isession;
+
+import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+
+import java.util.List;
+import java.util.function.Supplier;
+
+public interface INodeSupplier extends Supplier<List<TEndPoint>> {
+
+ void close();
+}
diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java
index 50c468f7921..9176dbd6827 100644
--- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java
+++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/ISession.java
@@ -60,7 +60,8 @@ public interface ISession extends AutoCloseable {
void open(
boolean enableRPCCompression,
int connectionTimeoutInMs,
- Map<String, TEndPoint> deviceIdToEndpoint)
+ Map<String, TEndPoint> deviceIdToEndpoint,
+ INodeSupplier nodeSupplier)
throws IoTDBConnectionException;
void close() throws IoTDBConnectionException;
diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/NodesSupplier.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/NodesSupplier.java
new file mode 100644
index 00000000000..0ce4af5bc9a
--- /dev/null
+++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/NodesSupplier.java
@@ -0,0 +1,221 @@
+package org.apache.iotdb.session;
+
+/*
+ * 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.
+ */
+
+import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.isession.INodeSupplier;
+import org.apache.iotdb.isession.SessionDataSet;
+import org.apache.iotdb.rpc.IoTDBConnectionException;
+import org.apache.iotdb.rpc.StatementExecutionException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.time.ZoneId;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.CopyOnWriteArraySet;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+
+public class NodesSupplier implements INodeSupplier, Runnable {
+
+ private static final Logger LOGGER = LoggerFactory.getLogger(NodesSupplier.class);
+
+ private static final long UPDATE_PERIOD_IN_S = 10;
+ private static final String SHOW_DATA_NODES_COMMAND = "SHOW DATANODES";
+
+ private static final String STATUS_COLUMN_NAME = "Status";
+
+ private static final String IP_COLUMN_NAME = "RpcAddress";
+
+ private static final String PORT_COLUMN_NAME = "RpcPort";
+
+ private static final String RUNNING_STATUS = "Running";
+
+ // it's ok that TIMEOUT_IN_MS is larger than UPDATE_PERIOD_IN_S, because the next update request
+ // won't be scheduled until last time is done.
+ private static final long TIMEOUT_IN_MS = 60_000;
+
+ private static final int FETCH_SIZE = 10_000;
+
+ // availableNodes won't be updated frequently, so we use CopyOnWriteArraySet which is thread-safe
+ // and is optimized for scenarios of reading more and writing less
+ private volatile Set<TEndPoint> availableNodes = new CopyOnWriteArraySet<>();
+
+ private final boolean useSSL;
+ private final String trustStore;
+ private final String trustStorePwd;
+ private final boolean enableRPCCompression;
+ private final String userName;
+
+ private final String password;
+
+ private final ZoneId zoneId;
+
+ private final int thriftDefaultBufferSize;
+
+ private final int thriftMaxFrameSize;
+
+ private final int connectionTimeoutInMs;
+
+ private final String version;
+
+ private ThriftConnection client;
+
+ public static NodesSupplier createNodeSupplier(
+ List<TEndPoint> endPointList,
+ ScheduledExecutorService executorService,
+ String userName,
+ String password,
+ ZoneId zoneId,
+ int thriftDefaultBufferSize,
+ int thriftMaxFrameSize,
+ int connectionTimeoutInMs,
+ boolean useSSL,
+ String trustStore,
+ String trustStorePwd,
+ boolean enableRPCCompression,
+ String version) {
+
+ NodesSupplier nodesSupplier =
+ new NodesSupplier(
+ endPointList,
+ userName,
+ password,
+ zoneId,
+ thriftDefaultBufferSize,
+ thriftMaxFrameSize,
+ connectionTimeoutInMs,
+ useSSL,
+ trustStore,
+ trustStorePwd,
+ enableRPCCompression,
+ version);
+
+ // call executorService.scheduleAtFixedRate here in a separate line
+ // instead of the constructor to prevent leaking the "this" reference to
+ // another thread, which will cause unsafe publication of this instance.
+ executorService.scheduleAtFixedRate(nodesSupplier, 0, UPDATE_PERIOD_IN_S, TimeUnit.SECONDS);
+
+ return nodesSupplier;
+ }
+
+ private NodesSupplier(
+ List<TEndPoint> endPointList,
+ String userName,
+ String password,
+ ZoneId zoneId,
+ int thriftDefaultBufferSize,
+ int thriftMaxFrameSize,
+ int connectionTimeoutInMs,
+ boolean useSSL,
+ String trustStore,
+ String trustStorePwd,
+ boolean enableRPCCompression,
+ String version) {
+ this.availableNodes.addAll(endPointList);
+ this.userName = userName;
+ this.password = password;
+ this.useSSL = useSSL;
+ this.trustStore = trustStore;
+ this.trustStorePwd = trustStorePwd;
+ this.enableRPCCompression = enableRPCCompression;
+ this.zoneId = zoneId;
+ this.thriftDefaultBufferSize = thriftDefaultBufferSize;
+ this.thriftMaxFrameSize = thriftMaxFrameSize;
+ this.connectionTimeoutInMs = connectionTimeoutInMs;
+ this.version = version;
+ }
+
+ // the method will only be called while reconnecting, so it's ok to copy set each time
+ // 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);
+ }
+
+ @Override
+ public void run() {
+ if (client == null) {
+ for (TEndPoint endPoint : availableNodes) {
+ if (createConnection(endPoint)) {
+ break;
+ }
+ }
+ }
+
+ if (client != null && !updateDataNodeList()) {
+ close();
+ }
+ }
+
+ private boolean createConnection(TEndPoint endPoint) {
+ client =
+ new ThriftConnection(
+ endPoint, thriftDefaultBufferSize, thriftMaxFrameSize, connectionTimeoutInMs);
+ try {
+ client.init(
+ useSSL,
+ trustStore,
+ trustStorePwd,
+ userName,
+ password,
+ enableRPCCompression,
+ zoneId,
+ version);
+ return true;
+ } catch (IoTDBConnectionException e) {
+ LOGGER.warn("Failed to create connection with {}.", endPoint, e);
+ close();
+ return false;
+ }
+ }
+
+ @Override
+ public void close() {
+ client.close();
+ client = null;
+ }
+
+ private boolean updateDataNodeList() {
+ try (SessionDataSet sessionDataSet =
+ client.executeQueryStatement(SHOW_DATA_NODES_COMMAND, TIMEOUT_IN_MS, FETCH_SIZE)) {
+ SessionDataSet.DataIterator iterator = sessionDataSet.iterator();
+ List<TEndPoint> res = new ArrayList<>();
+ while (iterator.next()) {
+ if (RUNNING_STATUS.equals(iterator.getString(STATUS_COLUMN_NAME))) {
+ String ip = iterator.getString(IP_COLUMN_NAME);
+ String port = iterator.getString(PORT_COLUMN_NAME);
+ if (ip != null && port != null) {
+ res.add(new TEndPoint(ip, Integer.parseInt(port)));
+ }
+ }
+ }
+ // replace the older ones.
+ availableNodes = new CopyOnWriteArraySet<>(res);
+ return true;
+ } catch (IoTDBConnectionException | StatementExecutionException e) {
+ LOGGER.warn("Failed to fetch data node list from {}.", client.endPoint, e);
+ return false;
+ }
+ }
+}
diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java
index 80de5ab3acf..142aea1f4fb 100644
--- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java
+++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java
@@ -21,6 +21,7 @@ package org.apache.iotdb.session;
import org.apache.iotdb.common.rpc.thrift.TAggregationType;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.isession.INodeSupplier;
import org.apache.iotdb.isession.ISession;
import org.apache.iotdb.isession.SessionConfig;
import org.apache.iotdb.isession.SessionDataSet;
@@ -90,7 +91,9 @@ import java.util.Map.Entry;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
@@ -146,6 +149,11 @@ public class Session implements ISession {
@SuppressWarnings("squid:S3077") // Non-primitive fields should not be "volatile"
protected volatile Map<TEndPoint, SessionConnection> endPointToSessionConnection;
+ // used to update datanodeList periodically
+ protected volatile ScheduledExecutorService executorService;
+
+ protected INodeSupplier availableNodes;
+
protected boolean enableQueryRedirection = false;
// The version number of the client which used for compatibility in the server
@@ -442,6 +450,31 @@ public class Session implements ISession {
return;
}
+ if (this.executorService != null) {
+ this.executorService.shutdown();
+ this.executorService = null;
+ }
+ if (this.availableNodes != null) {
+ this.availableNodes.close();
+ this.availableNodes = null;
+ }
+
+ this.executorService = Executors.newSingleThreadScheduledExecutor();
+ this.availableNodes =
+ NodesSupplier.createNodeSupplier(
+ getNodeUrls(),
+ executorService,
+ username,
+ password,
+ zoneId,
+ thriftDefaultBufferSize,
+ thriftMaxFrameSize,
+ connectionTimeoutInMs,
+ useSSL,
+ trustStore,
+ trustStorePwd,
+ enableRPCCompression,
+ version.toString());
this.enableRPCCompression = enableRPCCompression;
this.connectionTimeoutInMs = connectionTimeoutInMs;
defaultSessionConnection = constructSessionConnection(this, defaultEndPoint, zoneId);
@@ -454,16 +487,26 @@ public class Session implements ISession {
}
}
+ private List<TEndPoint> getNodeUrls() {
+ if (defaultEndPoint != null) {
+ return Collections.singletonList(defaultEndPoint);
+ } else {
+ return SessionUtils.parseSeedNodeUrls(nodeUrls);
+ }
+ }
+
@Override
public synchronized void open(
boolean enableRPCCompression,
int connectionTimeoutInMs,
- Map<String, TEndPoint> deviceIdToEndpoint)
+ Map<String, TEndPoint> deviceIdToEndpoint,
+ INodeSupplier nodesSupplier)
throws IoTDBConnectionException {
if (!isClosed) {
return;
}
+ this.availableNodes = nodesSupplier;
this.enableRPCCompression = enableRPCCompression;
this.connectionTimeoutInMs = connectionTimeoutInMs;
defaultSessionConnection = constructSessionConnection(this, defaultEndPoint, zoneId);
@@ -490,6 +533,14 @@ public class Session implements ISession {
defaultSessionConnection.close();
}
} finally {
+ // if executorService is null, it means that availableNodes is got from SessionPool and we
+ // shouldn't clean that
+ if (this.executorService != null) {
+ this.executorService.shutdown();
+ this.executorService = null;
+ this.availableNodes.close();
+ this.availableNodes = null;
+ }
isClosed = true;
}
}
@@ -497,9 +548,9 @@ public class Session implements ISession {
public SessionConnection constructSessionConnection(
Session session, TEndPoint endpoint, ZoneId zoneId) throws IoTDBConnectionException {
if (endpoint == null) {
- return new SessionConnection(session, zoneId);
+ return new SessionConnection(session, zoneId, availableNodes);
}
- return new SessionConnection(session, endpoint, zoneId);
+ return new SessionConnection(session, endpoint, zoneId, availableNodes);
}
@Override
@@ -1131,7 +1182,7 @@ public class Session implements ISession {
}
// TODO https://issues.apache.org/jira/browse/IOTDB-1399
- private void removeBrokenSessionConnection(SessionConnection sessionConnection) {
+ public void removeBrokenSessionConnection(SessionConnection sessionConnection) {
// remove the cached broken leader session
if (enableRedirection) {
TEndPoint endPoint = null;
diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java
index 417b60b8650..a4a2bedbacd 100644
--- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java
+++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/SessionConnection.java
@@ -77,8 +77,10 @@ import org.slf4j.LoggerFactory;
import java.security.SecureRandom;
import java.time.ZoneId;
import java.util.ArrayList;
+import java.util.Collections;
import java.util.List;
import java.util.StringJoiner;
+import java.util.function.Supplier;
public class SessionConnection {
@@ -95,15 +97,21 @@ public class SessionConnection {
private List<TEndPoint> endPointList = new ArrayList<>();
private boolean enableRedirect = false;
+ private final Supplier<List<TEndPoint>> availableNodes;
+
// TestOnly
- public SessionConnection() {}
+ public SessionConnection() {
+ availableNodes = Collections::emptyList;
+ }
- public SessionConnection(Session session, TEndPoint endPoint, ZoneId zoneId)
+ public SessionConnection(
+ Session session, TEndPoint endPoint, ZoneId zoneId, Supplier<List<TEndPoint>> availableNodes)
throws IoTDBConnectionException {
this.session = session;
this.endPoint = endPoint;
endPointList.add(endPoint);
this.zoneId = zoneId == null ? ZoneId.systemDefault() : zoneId;
+ this.availableNodes = availableNodes;
try {
init(endPoint, session.useSSL, session.trustStore, session.trustStorePwd);
} catch (IoTDBConnectionException e) {
@@ -111,10 +119,12 @@ public class SessionConnection {
}
}
- public SessionConnection(Session session, ZoneId zoneId) throws IoTDBConnectionException {
+ public SessionConnection(Session session, ZoneId zoneId, Supplier<List<TEndPoint>> availableNodes)
+ throws IoTDBConnectionException {
this.session = session;
this.zoneId = zoneId == null ? ZoneId.systemDefault() : zoneId;
this.endPointList = SessionUtils.parseSeedNodeUrls(session.nodeUrls);
+ this.availableNodes = availableNodes;
initClusterConn();
}
@@ -962,7 +972,6 @@ public class SessionConnection {
if (tryHostNum == endPointList.size()) {
break;
}
- session.defaultEndPoint = endPointList.get(j);
this.endPoint = endPointList.get(j);
if (j == endPointList.size() - 1) {
j = -1;
@@ -979,6 +988,11 @@ public class SessionConnection {
}
}
if (connectedSuccess) {
+ // remove the broken end point
+ session.removeBrokenSessionConnection(this);
+ session.defaultEndPoint = this.endPoint;
+ session.defaultSessionConnection = this;
+ session.endPointToSessionConnection.put(session.defaultEndPoint, this);
break;
}
}
diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/ThriftConnection.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/ThriftConnection.java
new file mode 100644
index 00000000000..e867f961f14
--- /dev/null
+++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/ThriftConnection.java
@@ -0,0 +1,190 @@
+/*
+ * 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 org.apache.iotdb.isession.SessionDataSet;
+import org.apache.iotdb.rpc.IoTDBConnectionException;
+import org.apache.iotdb.rpc.RpcTransportFactory;
+import org.apache.iotdb.rpc.RpcUtils;
+import org.apache.iotdb.rpc.StatementExecutionException;
+import org.apache.iotdb.service.rpc.thrift.IClientRPCService;
+import org.apache.iotdb.service.rpc.thrift.TSCloseSessionReq;
+import org.apache.iotdb.service.rpc.thrift.TSExecuteStatementReq;
+import org.apache.iotdb.service.rpc.thrift.TSExecuteStatementResp;
+import org.apache.iotdb.service.rpc.thrift.TSOpenSessionReq;
+import org.apache.iotdb.service.rpc.thrift.TSOpenSessionResp;
+
+import org.apache.thrift.TException;
+import org.apache.thrift.protocol.TBinaryProtocol;
+import org.apache.thrift.protocol.TCompactProtocol;
+import org.apache.thrift.transport.TTransport;
+import org.apache.thrift.transport.TTransportException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.time.ZoneId;
+
+public class ThriftConnection {
+
+ private static final Logger LOGGER = LoggerFactory.getLogger(ThriftConnection.class);
+
+ protected final TEndPoint endPoint;
+
+ protected final int thriftDefaultBufferSize;
+
+ protected final int thriftMaxFrameSize;
+
+ protected final int connectionTimeoutInMs;
+
+ protected TTransport transport;
+ protected IClientRPCService.Iface client;
+ protected long sessionId;
+ protected long statementId;
+
+ public ThriftConnection(
+ TEndPoint endPoint,
+ int thriftDefaultBufferSize,
+ int thriftMaxFrameSize,
+ int connectionTimeoutInMs) {
+ this.endPoint = endPoint;
+ this.thriftDefaultBufferSize = thriftDefaultBufferSize;
+ this.thriftMaxFrameSize = thriftMaxFrameSize;
+ this.connectionTimeoutInMs = connectionTimeoutInMs;
+ }
+
+ public void init(
+ boolean useSSL,
+ String trustStore,
+ String trustStorePwd,
+ String username,
+ String password,
+ boolean enableRPCCompression,
+ ZoneId zoneId,
+ String version)
+ throws IoTDBConnectionException {
+ RpcTransportFactory.setDefaultBufferCapacity(thriftDefaultBufferSize);
+ RpcTransportFactory.setThriftMaxFrameSize(thriftMaxFrameSize);
+ try {
+ if (useSSL) {
+ transport =
+ RpcTransportFactory.INSTANCE.getTransport(
+ endPoint.getIp(),
+ endPoint.getPort(),
+ connectionTimeoutInMs,
+ trustStore,
+ trustStorePwd);
+ } else {
+ transport =
+ RpcTransportFactory.INSTANCE.getTransport(
+ // as there is a try-catch already, we do not need to use TSocket.wrap
+ endPoint.getIp(), endPoint.getPort(), connectionTimeoutInMs);
+ }
+ if (!transport.isOpen()) {
+ transport.open();
+ }
+ } catch (TTransportException e) {
+ throw new IoTDBConnectionException(e);
+ }
+
+ if (enableRPCCompression) {
+ client = new IClientRPCService.Client(new TCompactProtocol(transport));
+ } else {
+ client = new IClientRPCService.Client(new TBinaryProtocol(transport));
+ }
+ client = RpcUtils.newSynchronizedClient(client);
+
+ TSOpenSessionReq openReq = new TSOpenSessionReq();
+ openReq.setUsername(username);
+ openReq.setPassword(password);
+ openReq.setZoneId(zoneId.toString());
+ openReq.putToConfiguration("version", version);
+
+ try {
+ TSOpenSessionResp openResp = client.openSession(openReq);
+
+ RpcUtils.verifySuccess(openResp.getStatus());
+
+ if (Session.protocolVersion.getValue() != openResp.getServerProtocolVersion().getValue()) {
+ LOGGER.warn(
+ "Protocol differ, Client version is {}}, but Server version is {}",
+ Session.protocolVersion.getValue(),
+ openResp.getServerProtocolVersion().getValue());
+ // less than 0.10
+ if (openResp.getServerProtocolVersion().getValue() == 0) {
+ throw new TException(
+ String.format(
+ "Protocol not supported, Client version is %s, but Server version is %s",
+ Session.protocolVersion.getValue(),
+ openResp.getServerProtocolVersion().getValue()));
+ }
+ }
+
+ sessionId = openResp.getSessionId();
+ statementId = client.requestStatementId(sessionId);
+
+ } catch (Exception e) {
+ transport.close();
+ throw new IoTDBConnectionException(e);
+ }
+ }
+
+ protected SessionDataSet executeQueryStatement(String sql, long timeout, int fetchSize)
+ throws StatementExecutionException, IoTDBConnectionException {
+ TSExecuteStatementReq execReq = new TSExecuteStatementReq(sessionId, sql, statementId);
+ execReq.setFetchSize(fetchSize);
+ execReq.setTimeout(timeout);
+ TSExecuteStatementResp execResp;
+ try {
+ execResp = client.executeQueryStatementV2(execReq);
+ } catch (TException e) {
+ throw new IoTDBConnectionException(e);
+ }
+
+ RpcUtils.verifySuccess(execResp.getStatus());
+ return new SessionDataSet(
+ sql,
+ execResp.getColumns(),
+ execResp.getDataTypeList(),
+ execResp.columnNameIndexMap,
+ execResp.getQueryId(),
+ statementId,
+ client,
+ sessionId,
+ execResp.queryResult,
+ execResp.isIgnoreTimeStamp(),
+ timeout,
+ execResp.moreData,
+ fetchSize);
+ }
+
+ public void close() {
+ if (transport.isOpen()) {
+ try {
+ client.closeSession(new TSCloseSessionReq(sessionId));
+ } catch (TException e) {
+ LOGGER.warn("Closing Session-{} with {} failed", sessionId, endPoint, e);
+ if (transport.isOpen()) {
+ transport.close();
+ }
+ }
+ }
+ }
+}
diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java
index 094da90a2d2..902c8ded417 100644
--- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java
+++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java
@@ -21,6 +21,7 @@ package org.apache.iotdb.session.pool;
import org.apache.iotdb.common.rpc.thrift.TAggregationType;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.isession.INodeSupplier;
import org.apache.iotdb.isession.ISession;
import org.apache.iotdb.isession.SessionConfig;
import org.apache.iotdb.isession.SessionDataSet;
@@ -32,7 +33,9 @@ import org.apache.iotdb.rpc.IoTDBConnectionException;
import org.apache.iotdb.rpc.StatementExecutionException;
import org.apache.iotdb.service.rpc.thrift.TSBackupConfigurationResp;
import org.apache.iotdb.service.rpc.thrift.TSConnectionInfoResp;
+import org.apache.iotdb.session.NodesSupplier;
import org.apache.iotdb.session.Session;
+import org.apache.iotdb.session.util.SessionUtils;
import org.apache.iotdb.tsfile.file.metadata.enums.CompressionType;
import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
import org.apache.iotdb.tsfile.file.metadata.enums.TSEncoding;
@@ -44,11 +47,14 @@ import org.slf4j.LoggerFactory;
import java.io.IOException;
import java.time.ZoneId;
+import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedDeque;
import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
/**
* SessionPool is a wrapper of a Session Set. Using SessionPool, the user do not need to consider
@@ -134,6 +140,11 @@ public class SessionPool implements ISessionPool {
// formatted nodeUrls for logging e.g. "host:port" or "[host:port, host:port, host:port]"
private final String formattedNodeUrls;
+ // used to update datanodeList periodically
+ protected volatile ScheduledExecutorService executorService;
+
+ protected INodeSupplier availableNodes;
+
private static final String INSERT_RECORD_FAIL = "insertRecord failed";
private static final String INSERT_RECORD_ERROR_MSG = "unexpected error in insertRecord";
@@ -356,6 +367,22 @@ public class SessionPool implements ISessionPool {
this.thriftDefaultBufferSize = thriftDefaultBufferSize;
this.thriftMaxFrameSize = thriftMaxFrameSize;
this.formattedNodeUrls = String.format("%s:%s", host, port);
+ this.executorService = Executors.newSingleThreadScheduledExecutor();
+ this.availableNodes =
+ NodesSupplier.createNodeSupplier(
+ Collections.singletonList(new TEndPoint(host, port)),
+ executorService,
+ user,
+ password,
+ zoneId,
+ thriftDefaultBufferSize,
+ thriftMaxFrameSize,
+ connectionTimeoutInMs,
+ useSSL,
+ trustStore,
+ trustStorePwd,
+ enableCompression,
+ version.toString());
}
public SessionPool(
@@ -398,6 +425,22 @@ public class SessionPool implements ISessionPool {
this.useSSL = useSSL;
this.trustStore = trustStore;
this.trustStorePwd = trustStorePwd;
+ this.executorService = Executors.newSingleThreadScheduledExecutor();
+ this.availableNodes =
+ NodesSupplier.createNodeSupplier(
+ Collections.singletonList(new TEndPoint(host, port)),
+ executorService,
+ user,
+ password,
+ zoneId,
+ thriftDefaultBufferSize,
+ thriftMaxFrameSize,
+ connectionTimeoutInMs,
+ useSSL,
+ trustStore,
+ trustStorePwd,
+ enableCompression,
+ version.toString());
}
@SuppressWarnings("squid:S107") // ignore Methods should not have too many parameters
@@ -434,6 +477,22 @@ public class SessionPool implements ISessionPool {
this.thriftDefaultBufferSize = thriftDefaultBufferSize;
this.thriftMaxFrameSize = thriftMaxFrameSize;
this.formattedNodeUrls = nodeUrls.toString();
+ this.executorService = Executors.newSingleThreadScheduledExecutor();
+ this.availableNodes =
+ NodesSupplier.createNodeSupplier(
+ SessionUtils.parseSeedNodeUrls(nodeUrls),
+ executorService,
+ user,
+ password,
+ zoneId,
+ thriftDefaultBufferSize,
+ thriftMaxFrameSize,
+ connectionTimeoutInMs,
+ useSSL,
+ trustStore,
+ trustStorePwd,
+ enableCompression,
+ version.toString());
}
public SessionPool(Builder builder) {
@@ -466,6 +525,22 @@ public class SessionPool implements ISessionPool {
this.useSSL = builder.useSSL;
this.trustStore = builder.trustStore;
this.trustStorePwd = builder.trustStorePwd;
+ this.executorService = Executors.newSingleThreadScheduledExecutor();
+ this.availableNodes =
+ NodesSupplier.createNodeSupplier(
+ SessionUtils.parseSeedNodeUrls(nodeUrls),
+ executorService,
+ user,
+ password,
+ zoneId,
+ thriftDefaultBufferSize,
+ thriftMaxFrameSize,
+ connectionTimeoutInMs,
+ useSSL,
+ trustStore,
+ trustStorePwd,
+ enableCompression,
+ version.toString());
}
private Session constructNewSession() {
@@ -580,7 +655,7 @@ public class SessionPool implements ISessionPool {
session = constructNewSession();
try {
- session.open(enableCompression, connectionTimeoutInMs, deviceIdToEndpoint);
+ session.open(enableCompression, connectionTimeoutInMs, deviceIdToEndpoint, availableNodes);
// avoid someone has called close() the session pool
synchronized (this) {
if (closed) {
@@ -651,6 +726,14 @@ public class SessionPool implements ISessionPool {
LOGGER.warn(CLOSE_THE_SESSION_FAILED, e);
}
}
+ if (this.executorService != null) {
+ this.executorService.shutdown();
+ this.executorService = null;
+ }
+ if (availableNodes != null) {
+ this.availableNodes.close();
+ this.availableNodes = null;
+ }
LOGGER.info("closing the session pool, cleaning queues...");
this.closed = true;
queue.clear();
@@ -677,7 +760,7 @@ public class SessionPool implements ISessionPool {
private void tryConstructNewSession() {
Session session = constructNewSession();
try {
- session.open(enableCompression, connectionTimeoutInMs, deviceIdToEndpoint);
+ session.open(enableCompression, connectionTimeoutInMs, deviceIdToEndpoint, availableNodes);
// avoid someone has called close() the session pool
synchronized (this) {
if (closed) {
diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionConnectionTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionConnectionTest.java
index e5eecb5545e..6cc28c35a80 100644
--- a/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionConnectionTest.java
+++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/SessionConnectionTest.java
@@ -66,6 +66,7 @@ import org.powermock.reflect.Whitebox;
import java.time.ZoneId;
import java.util.Arrays;
+import java.util.Collections;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyLong;
@@ -151,14 +152,22 @@ public class SessionConnectionTest {
@Test(expected = NumberFormatException.class)
public void testBuildSessionConnection() throws IoTDBConnectionException {
session = new Session("local", 12, "root", "root");
- SessionConnection sessionConnection1 = new SessionConnection(session, ZoneId.systemDefault());
+ SessionConnection sessionConnection1 =
+ new SessionConnection(
+ session,
+ ZoneId.systemDefault(),
+ () -> Collections.singletonList(new TEndPoint("local", 12)));
}
@Test(expected = IoTDBConnectionException.class)
public void testBuildSessionConnection2() throws IoTDBConnectionException {
session = new Session("local", 12, "root", "root");
SessionConnection sessionConnection1 =
- new SessionConnection(session, new TEndPoint("localhost", 1234), ZoneId.systemDefault());
+ new SessionConnection(
+ session,
+ new TEndPoint("localhost", 1234),
+ ZoneId.systemDefault(),
+ () -> Collections.singletonList(new TEndPoint("local", 12)));
}
@Test