From 5e9353f6bfea79b0be3916919b03da786e20824e Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Wed, 4 Nov 2026 15:41:24 +0800 Subject: [PATCH 01/32] Support internal sessions through inherited client RPC services --- .../isession/IInternalSessionBuilder.java | 32 ++++ .../iotdb/session/i18n/SessionMessages.java | 5 + .../iotdb/session/i18n/SessionMessages.java | 5 + .../apache/iotdb/session/InternalSession.java | 77 +++++++++ .../org/apache/iotdb/session/Session.java | 69 +++++--- .../iotdb/session/TableSessionBuilder.java | 17 +- .../InternalSubscriptionTableSession.java | 83 +++++++++ .../SubscriptionSessionWrapper.java | 2 +- .../SubscriptionTableSession.java | 4 + .../SubscriptionTableSessionBuilder.java | 24 ++- ...oDnInternalServiceAsyncRequestManager.java | 3 +- .../client/sync/SyncDataNodeClientPool.java | 3 +- .../db/protocol/session/SessionManager.java | 161 +++++++++++------- .../handler/InternalServiceThriftHandler.java | 20 ++- .../thrift/impl/ClientRPCServiceImpl.java | 27 +-- .../impl/DataNodeInternalRPCServiceImpl.java | 43 +++-- .../service/DataNodeInternalRPCService.java | 4 +- .../src/main/thrift/datanode.thrift | 11 +- 18 files changed, 459 insertions(+), 131 deletions(-) create mode 100644 iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java create mode 100644 iotdb-client/session/src/main/java/org/apache/iotdb/session/InternalSession.java create mode 100644 iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/InternalSubscriptionTableSession.java diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java new file mode 100644 index 0000000000000..ae67fa4093da0 --- /dev/null +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java @@ -0,0 +1,32 @@ +/* + * 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.rpc.IoTDBConnectionException; + +/** Capability implemented by builders that can construct sessions on internal endpoints. */ +public interface IInternalSessionBuilder { + + /** + * Builds a session using the caller-provided internal node supplier. Whether the returned session + * is already open follows the builder's ordinary build contract. + */ + T buildInternal(INodeSupplier nodeSupplier) throws IoTDBConnectionException; +} diff --git a/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java b/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java index ded0deac31484..382567f378d61 100644 --- a/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java +++ b/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java @@ -20,6 +20,11 @@ package org.apache.iotdb.session.i18n; public final class SessionMessages { + public static final String EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081 = + "No DataNode internal endpoint is available"; + + public static final String EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392 = + "Internal node supplier must not be null"; // Session public static final String NODE_URLS_EMPTY = "nodeUrls shouldn't be empty."; diff --git a/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java b/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java index 838a7e2148329..9975c43a74cd7 100644 --- a/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java +++ b/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java @@ -20,6 +20,11 @@ package org.apache.iotdb.session.i18n; public final class SessionMessages { + public static final String EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081 = + "没有可用的 DataNode 内部端点"; + + public static final String EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392 = + "内部节点 supplier 不能为空"; // Session public static final String NODE_URLS_EMPTY = "nodeUrls 不能为空。"; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/InternalSession.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/InternalSession.java new file mode 100644 index 0000000000000..fdcaed13f3964 --- /dev/null +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/InternalSession.java @@ -0,0 +1,77 @@ +/* + * 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.INodeSupplier; +import org.apache.iotdb.rpc.IoTDBConnectionException; +import org.apache.iotdb.session.i18n.SessionMessages; + +import java.time.ZoneId; +import java.util.List; +import java.util.Objects; + +/** Session using caller-supplied internal endpoints and the caller's username. */ +public class InternalSession extends Session { + + private final INodeSupplier nodeSupplier; + + public InternalSession(AbstractSessionBuilder builder, INodeSupplier nodeSupplier) { + super(builder); + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + password = null; + useEncryptedPassword = false; + enableAutoFetch = false; + enableRedirection = false; + enableQueryRedirection = false; + } + + @Override + public synchronized void open() throws IoTDBConnectionException { + open(enableThriftRpcCompaction, connectionTimeoutInMs); + } + + @Override + protected INodeSupplier createNodeSupplier(boolean enableRPCCompaction, int timeoutInMs) { + return nodeSupplier; + } + + @Override + public SessionConnection constructSessionConnection( + Session session, TEndPoint endpoint, ZoneId zoneId) throws IoTDBConnectionException { + final List candidates = availableNodes.get(); + if (candidates.isEmpty()) { + throw new IoTDBConnectionException( + SessionMessages.EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081); + } + defaultEndPoint = + candidates.contains(endpoint) + ? endpoint.deepCopy() + : nodeSupplier + .getQueryEndPoint() + .filter(candidates::contains) + .orElse(candidates.get(0)) + .deepCopy(); + return super.constructSessionConnection(session, defaultEndPoint, zoneId); + } +} 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 de3ee417afc0d..f415a63246490 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 @@ import org.apache.iotdb.common.rpc.thrift.TAggregationType; import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.isession.IInternalSessionBuilder; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.ISession; import org.apache.iotdb.isession.SessionConfig; @@ -549,33 +550,7 @@ public synchronized void open(boolean enableRPCCompaction, int connectionTimeout this.availableNodes = null; } - if (enableAutoFetch) { - initThreadPool(); - this.availableNodes = - NodesSupplier.createNodeSupplier( - getNodeUrls(), - SessionUtils.parseSeedNodeUrls(initialNodeUrls), - executorService, - username, - password, - zoneId, - thriftDefaultBufferSize, - thriftMaxFrameSize, - connectionTimeoutInMs, - useSSL, - trustStore, - trustStorePwd, - keyStore, - keyStorePwd, - sslProtocol, - enableRPCCompaction, - version.toString()); - } else { - // A static supplier has no separate seed tier. Keep every configured address after reopen. - Set staticNodes = new LinkedHashSet<>(getNodeUrls()); - staticNodes.addAll(SessionUtils.parseSeedNodeUrls(initialNodeUrls)); - this.availableNodes = new DummyNodesSupplier(new ArrayList<>(staticNodes)); - } + this.availableNodes = createNodeSupplier(enableRPCCompaction, connectionTimeoutInMs); this.connectionTimeoutInMs = connectionTimeoutInMs; setDefaultSessionConnection(constructSessionConnection(this, defaultEndPoint, zoneId)); @@ -589,6 +564,37 @@ public synchronized void open(boolean enableRPCCompaction, int connectionTimeout } } + /** Creates the endpoint source used by initial connections and reconnects. */ + protected INodeSupplier createNodeSupplier( + boolean enableRPCCompaction, int connectionTimeoutInMs) { + if (enableAutoFetch) { + initThreadPool(); + return NodesSupplier.createNodeSupplier( + getNodeUrls(), + SessionUtils.parseSeedNodeUrls(initialNodeUrls), + executorService, + username, + password, + zoneId, + thriftDefaultBufferSize, + thriftMaxFrameSize, + connectionTimeoutInMs, + useSSL, + trustStore, + trustStorePwd, + keyStore, + keyStorePwd, + sslProtocol, + enableRPCCompaction, + version.toString()); + } else { + // A static supplier has no separate seed tier. Keep every configured address after reopen. + Set staticNodes = new LinkedHashSet<>(getNodeUrls()); + staticNodes.addAll(SessionUtils.parseSeedNodeUrls(initialNodeUrls)); + return new DummyNodesSupplier(new ArrayList<>(staticNodes)); + } + } + private void initThreadPool() { this.executorService = Executors.newSingleThreadScheduledExecutor( @@ -4445,7 +4451,8 @@ public void setEnableThriftRpcCompaction(boolean enableThriftRpcCompaction) { this.enableThriftRpcCompaction = enableThriftRpcCompaction; } - public static class Builder extends AbstractSessionBuilder { + public static class Builder extends AbstractSessionBuilder + implements IInternalSessionBuilder { public Builder host(String host) { this.host = host; @@ -4595,5 +4602,11 @@ public Session build() { } return new Session(this); } + + /** Builds an unopened internal session using caller-supplied internal endpoints. */ + @Override + public InternalSession buildInternal(INodeSupplier nodeSupplier) { + return new InternalSession(this, nodeSupplier); + } } } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java index a39af4c847a95..856ed7a91bc6b 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java @@ -19,6 +19,8 @@ package org.apache.iotdb.session; +import org.apache.iotdb.isession.IInternalSessionBuilder; +import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.ITableSession; import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.rpc.IoTDBConnectionException; @@ -41,7 +43,8 @@ * *

All configurations have reasonable default values, which can be overridden as needed. */ -public class TableSessionBuilder extends AbstractSessionBuilder { +public class TableSessionBuilder extends AbstractSessionBuilder + implements IInternalSessionBuilder { private boolean enableCompression = false; private int connectionTimeoutInMs = SessionConfig.DEFAULT_CONNECTION_TIMEOUT_MS; @@ -437,7 +440,17 @@ public ITableSession build() throws IoTDBConnectionException { Collections.singletonList(SessionConfig.DEFAULT_HOST + ":" + SessionConfig.DEFAULT_PORT); } this.sqlDialect = TABLE; - Session newSession = new Session(this); + return buildSession(new Session(this)); + } + + /** Builds and opens a table session using caller-supplied internal endpoints. */ + @Override + public ITableSession buildInternal(INodeSupplier nodeSupplier) throws IoTDBConnectionException { + this.sqlDialect = TABLE; + return buildSession(new InternalSession(this, nodeSupplier)); + } + + private ITableSession buildSession(Session newSession) throws IoTDBConnectionException { newSession.enableIoTDBRpcCompression = isIoTDBRpcCompressionEnabled; newSession.tabletCompressionMinRowSize = tabletCompressionMinRowSize; diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/InternalSubscriptionTableSession.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/InternalSubscriptionTableSession.java new file mode 100644 index 0000000000000..13cf8c659ed94 --- /dev/null +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/InternalSubscriptionTableSession.java @@ -0,0 +1,83 @@ +/* + * 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.subscription; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.rpc.IoTDBConnectionException; +import org.apache.iotdb.session.AbstractSessionBuilder; +import org.apache.iotdb.session.Session; +import org.apache.iotdb.session.SessionConnection; +import org.apache.iotdb.session.i18n.SessionMessages; + +import java.time.ZoneId; +import java.util.List; +import java.util.Objects; + +/** Table subscription-management session using the caller's identity and internal endpoints. */ +public class InternalSubscriptionTableSession extends SubscriptionTableSession { + + public InternalSubscriptionTableSession( + AbstractSessionBuilder builder, INodeSupplier nodeSupplier) { + super(new InternalSessionWrapper(builder, nodeSupplier)); + } + + private static class InternalSessionWrapper extends SubscriptionSessionWrapper { + + private final INodeSupplier nodeSupplier; + + private InternalSessionWrapper(AbstractSessionBuilder builder, INodeSupplier nodeSupplier) { + super(builder); + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + password = null; + useEncryptedPassword = false; + enableAutoFetch = false; + enableRedirection = false; + enableQueryRedirection = false; + } + + @Override + protected INodeSupplier createNodeSupplier(boolean enableRPCCompaction, int timeoutInMs) { + return nodeSupplier; + } + + @Override + public SessionConnection constructSessionConnection( + Session session, TEndPoint endpoint, ZoneId zoneId) throws IoTDBConnectionException { + final List candidates = nodeSupplier.get(); + if (candidates.isEmpty()) { + throw new IoTDBConnectionException( + SessionMessages.EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081); + } + defaultEndPoint = + candidates.contains(endpoint) + ? endpoint.deepCopy() + : nodeSupplier + .getQueryEndPoint() + .filter(candidates::contains) + .orElse(candidates.get(0)) + .deepCopy(); + return super.constructSessionConnection(session, defaultEndPoint, zoneId); + } + } +} diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java index d7e715fa62762..018f4ef3a185d 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java @@ -30,7 +30,7 @@ import java.time.ZoneId; import java.util.Objects; -public final class SubscriptionSessionWrapper extends Session { +public class SubscriptionSessionWrapper extends Session { public SubscriptionSessionWrapper(final AbstractSessionBuilder builder) { super(builder); diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSession.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSession.java index 123deea2fa45d..8949c46ff83ad 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSession.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSession.java @@ -36,6 +36,10 @@ public SubscriptionTableSession(final AbstractSessionBuilder builder) { super(new SubscriptionSessionWrapper(builder)); } + protected SubscriptionTableSession(final SubscriptionSessionWrapper session) { + super(session); + } + /////////////////////////////// open & close /////////////////////////////// @Override diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java index e0c8e95968dc9..0f34d27ee4980 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java @@ -19,10 +19,13 @@ package org.apache.iotdb.session.subscription; +import org.apache.iotdb.isession.IInternalSessionBuilder; +import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.rpc.IoTDBConnectionException; import org.apache.iotdb.session.AbstractSessionBuilder; -public class SubscriptionTableSessionBuilder extends AbstractSessionBuilder { +public class SubscriptionTableSessionBuilder extends AbstractSessionBuilder + implements IInternalSessionBuilder { public SubscriptionTableSessionBuilder() { // use table model @@ -73,4 +76,23 @@ public ISubscriptionTableSession build() throws IoTDBConnectionException { session.open(); return session; } + + /** Builds and opens a table subscription-management session on internal endpoints. */ + @Override + public ISubscriptionTableSession buildInternal(INodeSupplier nodeSupplier) + throws IoTDBConnectionException { + final ISubscriptionTableSession session = + new InternalSubscriptionTableSession(this, nodeSupplier); + try { + session.open(); + } catch (IoTDBConnectionException e) { + try { + session.close(); + } catch (Exception closeException) { + e.addSuppressed(closeException); + } + throw e; + } + return session; + } } diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/CnToDnInternalServiceAsyncRequestManager.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/CnToDnInternalServiceAsyncRequestManager.java index 3010f8b524c04..3cb2468c27893 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/CnToDnInternalServiceAsyncRequestManager.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/CnToDnInternalServiceAsyncRequestManager.java @@ -449,7 +449,8 @@ protected void initActionMapBuilder() { actionMapBuilder.put( CnToDnAsyncRequestType.INSERT_RECORD, (req, client, handler) -> - client.insertRecord((TSInsertRecordReq) req, (DataNodeTSStatusRPCHandler) handler)); + client.insertRecordInternal( + (TSInsertRecordReq) req, (DataNodeTSStatusRPCHandler) handler)); actionMapBuilder.put( CnToDnAsyncRequestType.UPDATE_TABLE, (req, client, handler) -> diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/sync/SyncDataNodeClientPool.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/sync/SyncDataNodeClientPool.java index f35b6d906a525..60c58ddea0d2b 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/sync/SyncDataNodeClientPool.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/sync/SyncDataNodeClientPool.java @@ -137,7 +137,8 @@ private void buildActionMap() { CnToDnSyncRequestType.RESET_PEER_LIST, (req, client) -> client.resetPeerList((TResetPeerListReq) req)); actionMapBuilder.put( - CnToDnSyncRequestType.SHOW_CONFIGURATION, (req, client) -> client.showConfiguration()); + CnToDnSyncRequestType.SHOW_CONFIGURATION, + (req, client) -> client.showConfigurationInternal()); actionMapBuilder.put( CnToDnSyncRequestType.SHOW_APPLIED_CONFIGURATIONS, (req, client) -> client.showAppliedConfigurations()); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java index 85d9acae276ec..7017bda16794f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java @@ -157,75 +157,118 @@ public BasicOpenSessionResp login( final TSStatus loginStatus = AuthorityChecker.checkUser(username, password, useEncryptedPassword); if (loginStatus.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) { - // check the version compatibility - if (!tsProtocolVersion.equals(CURRENT_RPC_VERSION)) { - openSessionResp - .sessionId(-1) - .setCode(TSStatusCode.INCOMPATIBLE_VERSION.getStatusCode()) - .setMessage( - DataNodeMiscMessages.VERSION_INCOMPATIBLE_PLEASE_UPGRADE_TO - + IoTDBConstant.VERSION); - } else { - session.setSqlDialect(sqlDialect); - supplySession(session, userId, username, ZoneId.of(zoneId), clientVersion); - String logInMessage = DataNodeMiscMessages.EXCEPTION_LOGIN_SUCCESSFULLY_1C2EF670; - openSessionResp - .sessionId(session.getId()) - .setCode(TSStatusCode.SUCCESS_STATUS.getStatusCode()) - .setMessage(logInMessage); - AUDIT_LOGGER.log( - new AuditLogFields( - userId, - username, - session.getClientAddress(), - AuditEventType.LOGIN, - AuditLogOperation.CONTROL, - true), - () -> - String.format( - DataNodeMiscMessages - .EXCEPTION_ARG_LOGIN_STATUS_ARG_USER_ARG_ID_ARG_OPENS_SESSION_ARG_70586F85, - IoTDBConstant.GLOBAL_DB_NAME, - openSessionResp.getMessage(), - username, - userId, - session)); - LOGGER.info( - DataNodeMiscMessages.LOGIN_STATUS, - IoTDBConstant.GLOBAL_DB_NAME, - openSessionResp.getMessage(), - username, - session); - updateIdleTime(); - if (enableLoginLock) { - loginLockManager.clearFailure(userId, session.getClientAddress()); - } + BasicOpenSessionResp response = + completeLogin( + session, userId, username, zoneId, tsProtocolVersion, clientVersion, sqlDialect); + if (response.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode() && enableLoginLock) { + loginLockManager.clearFailure(userId, session.getClientAddress()); } + return response; } else { openSessionResp.sessionId(-1).setMessage(loginStatus.message).setCode(loginStatus.code); if (enableLoginLock) { loginLockManager.recordFailure(userId, session.getClientAddress()); } - AUDIT_LOGGER.log( - new AuditLogFields( - userId, - username, - session.getClientAddress(), - AuditEventType.LOGIN, - AuditLogOperation.CONTROL, - false), - () -> - String.format( - DataNodeMiscMessages - .EXCEPTION_USER_ARG_ID_ARG_LOGIN_FAILED_WITH_CODE_ARG_ARG_93FAE4FD, - username, - userId, - loginStatus.getCode(), - loginStatus.getMessage())); + logFailedLogin(session, userId, username, loginStatus); } return openSessionResp; } + /** Opens a session for a trusted internal caller without checking a password. */ + public BasicOpenSessionResp loginInternal( + IClientSession session, + String username, + String zoneId, + TSProtocolVersion tsProtocolVersion, + IoTDBConstant.ClientVersion clientVersion, + SqlDialect sqlDialect) { + final long userId = AuthorityChecker.getUserId(username).orElse(-1L); + if (userId == -1) { + BasicOpenSessionResp response = new BasicOpenSessionResp(); + response + .sessionId(-1) + .setCode(TSStatusCode.USER_NOT_EXIST.getStatusCode()) + .setMessage( + String.format( + DataNodeMiscMessages.MISC_EXCEPTION_USER_S_DOES_NOT_EXIST_0CE725D8, username)); + logFailedLogin(session, userId, username, response); + return response; + } + return completeLogin( + session, userId, username, zoneId, tsProtocolVersion, clientVersion, sqlDialect); + } + + private BasicOpenSessionResp completeLogin( + IClientSession session, + long userId, + String username, + String zoneId, + TSProtocolVersion tsProtocolVersion, + IoTDBConstant.ClientVersion clientVersion, + SqlDialect sqlDialect) { + BasicOpenSessionResp response = new BasicOpenSessionResp(); + if (!tsProtocolVersion.equals(CURRENT_RPC_VERSION)) { + response + .sessionId(-1) + .setCode(TSStatusCode.INCOMPATIBLE_VERSION.getStatusCode()) + .setMessage( + DataNodeMiscMessages.VERSION_INCOMPATIBLE_PLEASE_UPGRADE_TO + IoTDBConstant.VERSION); + return response; + } + + session.setSqlDialect(sqlDialect); + supplySession(session, userId, username, ZoneId.of(zoneId), clientVersion); + response + .sessionId(session.getId()) + .setCode(TSStatusCode.SUCCESS_STATUS.getStatusCode()) + .setMessage(DataNodeMiscMessages.EXCEPTION_LOGIN_SUCCESSFULLY_1C2EF670); + AUDIT_LOGGER.log( + new AuditLogFields( + userId, + username, + session.getClientAddress(), + AuditEventType.LOGIN, + AuditLogOperation.CONTROL, + true), + () -> + String.format( + DataNodeMiscMessages + .EXCEPTION_ARG_LOGIN_STATUS_ARG_USER_ARG_ID_ARG_OPENS_SESSION_ARG_70586F85, + IoTDBConstant.GLOBAL_DB_NAME, + response.getMessage(), + username, + userId, + session)); + LOGGER.info( + DataNodeMiscMessages.LOGIN_STATUS, + IoTDBConstant.GLOBAL_DB_NAME, + response.getMessage(), + username, + session); + updateIdleTime(); + return response; + } + + private void logFailedLogin( + IClientSession session, long userId, String username, TSStatus status) { + AUDIT_LOGGER.log( + new AuditLogFields( + userId, + username, + session.getClientAddress(), + AuditEventType.LOGIN, + AuditLogOperation.CONTROL, + false), + () -> + String.format( + DataNodeMiscMessages + .EXCEPTION_USER_ARG_ID_ARG_LOGIN_FAILED_WITH_CODE_ARG_ARG_93FAE4FD, + username, + userId, + status.getCode(), + status.getMessage())); + } + public boolean closeSession(IClientSession session, LongConsumer releaseByQueryId) { return closeSession(session, releaseByQueryId, true); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/handler/InternalServiceThriftHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/handler/InternalServiceThriftHandler.java index 58130eb91cc58..831e12ea17bbd 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/handler/InternalServiceThriftHandler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/handler/InternalServiceThriftHandler.java @@ -21,19 +21,27 @@ import org.apache.iotdb.commons.service.NoopServerContext; import org.apache.iotdb.commons.service.metric.MetricService; +import org.apache.iotdb.db.protocol.session.ClientSession; +import org.apache.iotdb.db.protocol.session.SessionManager; +import org.apache.iotdb.db.protocol.thrift.impl.IClientRPCServiceWithHandler; +import org.apache.iotdb.rpc.TElasticFramedTransport; import org.apache.thrift.protocol.TProtocol; import org.apache.thrift.server.ServerContext; import org.apache.thrift.server.TServerEventHandler; +import org.apache.thrift.transport.TSocket; import org.apache.thrift.transport.TTransport; +import java.net.Socket; import java.util.concurrent.atomic.AtomicLong; public class InternalServiceThriftHandler implements TServerEventHandler { private AtomicLong thriftConnectionNumber = new AtomicLong(0); + private final IClientRPCServiceWithHandler eventHandler; - public InternalServiceThriftHandler() { + public InternalServiceThriftHandler(IClientRPCServiceWithHandler eventHandler) { + this.eventHandler = eventHandler; MetricService.getInstance() .addMetricSet(new InternalServiceThriftHandlerMetrics(thriftConnectionNumber)); } @@ -46,13 +54,21 @@ public void preServe() { @Override public ServerContext createContext(TProtocol tProtocol, TProtocol tProtocol1) { thriftConnectionNumber.incrementAndGet(); + Socket socket = + ((TSocket) ((TElasticFramedTransport) tProtocol1.getTransport()).getSocket()).getSocket(); + SessionManager.getInstance().registerSession(new ClientSession(socket)); return NoopServerContext.INSTANCE; } @Override public void deleteContext( ServerContext serverContext, TProtocol tProtocol, TProtocol tProtocol1) { - thriftConnectionNumber.decrementAndGet(); + try { + eventHandler.handleClientExit(); + } finally { + SessionManager.getInstance().removeCurrSession(); + thriftConnectionNumber.decrementAndGet(); + } } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java index fcb8fc23999ba..3e8927bdf6057 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java @@ -1626,16 +1626,7 @@ public TSOpenSessionResp openSession(TSOpenSessionReq req) throws TException { } Optional database = parseDatabase(req); IClientSession clientSession = SESSION_MANAGER.getCurrSession(); - BasicOpenSessionResp openSessionResp = - SESSION_MANAGER.login( - SESSION_MANAGER.getCurrSession(), - req.username, - req.password, - req.zoneId, - req.client_protocol, - clientVersion, - sqlDialect, - parseUseEncryptedPassword(req)); + BasicOpenSessionResp openSessionResp = login(clientSession, req, clientVersion, sqlDialect); TSStatus tsStatus = RpcUtils.getStatus(openSessionResp.getCode(), openSessionResp.getMessage()); if (tsStatus.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) { @@ -1651,6 +1642,22 @@ public TSOpenSessionResp openSession(TSOpenSessionReq req) throws TException { return resp.setSessionId(openSessionResp.getSessionId()).setConfiguration(configuration); } + protected BasicOpenSessionResp login( + IClientSession session, + TSOpenSessionReq req, + IoTDBConstant.ClientVersion clientVersion, + SqlDialect sqlDialect) { + return SESSION_MANAGER.login( + session, + req.username, + req.password, + req.zoneId, + req.client_protocol, + clientVersion, + sqlDialect, + parseUseEncryptedPassword(req)); + } + private IoTDBConstant.ClientVersion parseClientVersion(TSOpenSessionReq req) { Map configuration = req.configuration; if (configuration != null && configuration.containsKey("version")) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java index 4b1e05f4e857d..2523be59d7096 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java @@ -77,6 +77,7 @@ import org.apache.iotdb.commons.pipe.agent.plugin.meta.PipePluginMeta; import org.apache.iotdb.commons.pipe.agent.task.meta.PipeMeta; import org.apache.iotdb.commons.queryengine.common.SessionInfo; +import org.apache.iotdb.commons.queryengine.common.SqlDialect; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; @@ -119,6 +120,7 @@ import org.apache.iotdb.db.i18n.DataNodeSchemaMessages; import org.apache.iotdb.db.partition.DataPartitionTableGenerator; import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; +import org.apache.iotdb.db.protocol.basic.BasicOpenSessionResp; import org.apache.iotdb.db.protocol.client.ConfigNodeInfo; import org.apache.iotdb.db.protocol.client.cn.DnToCnInternalServiceAsyncRequestManager; import org.apache.iotdb.db.protocol.client.cn.DnToCnRequestType; @@ -350,6 +352,7 @@ import org.apache.iotdb.rpc.TSStatusCode; import org.apache.iotdb.rpc.subscription.exception.SubscriptionException; import org.apache.iotdb.service.rpc.thrift.TSInsertRecordReq; +import org.apache.iotdb.service.rpc.thrift.TSOpenSessionReq; import org.apache.iotdb.trigger.api.enums.FailureStrategy; import org.apache.iotdb.trigger.api.enums.TriggerEvent; @@ -416,7 +419,8 @@ import static org.apache.iotdb.db.utils.ErrorHandlingUtils.onIoTDBException; import static org.apache.iotdb.db.utils.ErrorHandlingUtils.onQueryException; -public class DataNodeInternalRPCServiceImpl implements IDataNodeRPCService.Iface { +public class DataNodeInternalRPCServiceImpl extends ClientRPCServiceImpl + implements IDataNodeRPCService.Iface { private static final Logger LOGGER = LoggerFactory.getLogger(DataNodeInternalRPCServiceImpl.class); @@ -473,6 +477,8 @@ public class DataNodeInternalRPCServiceImpl implements IDataNodeRPCService.Iface private static final String SYSTEM = "system"; + private long consensusWaitTimeoutSeconds = 30; + public DataNodeInternalRPCServiceImpl(DataNodeContext dataNodeContext) { this(dataNodeContext, SystemMetrics.getInstance()); } @@ -485,8 +491,6 @@ public DataNodeInternalRPCServiceImpl(DataNodeContext dataNodeContext) { this.dataNodeContext = dataNodeContext; } - private long consensusWaitTimeoutSeconds = 30; - private TSStatus waitForConsensusStarted() { if (dataNodeContext.isAllConsensusStarted()) { return null; @@ -508,6 +512,16 @@ private TSStatus waitForConsensusStarted() { } } + @Override + protected BasicOpenSessionResp login( + IClientSession session, + TSOpenSessionReq req, + ClientVersion clientVersion, + SqlDialect sqlDialect) { + return SESSION_MANAGER.loginInternal( + session, req.username, req.zoneId, req.client_protocol, clientVersion, sqlDialect); + } + @Override public TSendFragmentInstanceResp sendFragmentInstance(final TSendFragmentInstanceReq req) { LOGGER.debug(DataNodeMiscMessages.RECEIVE_FRAGMENT_INSTANCE, req.getConsensusGroupId()); @@ -1839,9 +1853,10 @@ private TSStatus executeInternalSchemaTask( @Override public TSStatus executeCQ(TExecuteCQ req) { + IClientSession previousSession = SESSION_MANAGER.getCurrSession(); IClientSession session = new InternalClientSession(req.cqId); - SESSION_MANAGER.registerSession(session); + SESSION_MANAGER.setCurrSession(session); SESSION_MANAGER.supplySession( session, @@ -1923,8 +1938,11 @@ public TSStatus executeCQ(TExecuteCQ req) { } catch (Exception e) { return onQueryException(e, "\"" + executedSQL + "\". " + OperationType.EXECUTE_STATEMENT); } finally { - SESSION_MANAGER.closeSession(session, COORDINATOR::cleanupQueryExecution); - SESSION_MANAGER.removeCurrSession(); + try { + SESSION_MANAGER.closeSession(session, COORDINATOR::cleanupQueryExecution); + } finally { + SESSION_MANAGER.restoreSession(previousSession, session); + } } } @@ -2379,11 +2397,6 @@ private List testAllDataNodeExternalServiceConnection( handler, TEST_CONNECTION_RETRY_NUM, TEST_CONNECTION_TIMEOUT_MS, true)); } - @Override - public TSStatus testConnectionEmptyRPC() throws TException { - return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()); - } - @Override public TSStatus updateClusterTopology(TUpdateClusterTopologyReq req) throws TException { clusterTopology.updateTopology(req.getDataNodes(), req.getTopology()); @@ -2810,7 +2823,7 @@ public TSStatus loadConfiguration() throws TException { } @Override - public TShowConfigurationResp showConfiguration() { + public TShowConfigurationResp showConfigurationInternal() { TShowConfigurationResp resp = new TShowConfigurationResp(RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS), ""); try { @@ -3562,7 +3575,7 @@ public TSStatus stopAndClearDataNode() { } @Override - public TSStatus insertRecord(TSInsertRecordReq req) throws TException { + public TSStatus insertRecordInternal(TSInsertRecordReq req) { try { InsertRowStatement statement = StatementGenerator.createStatement(req); SessionInfo sessionInfo = @@ -3606,10 +3619,6 @@ public TSStatus writeAuditLog(TAuditLogReq req) { return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS); } - public void handleClientExit() { - // Do nothing - } - // ==================================================== // Data Partition Table Integrity Check Implementation // ==================================================== diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java index f9b9c6e010110..ef0afbf295651 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java @@ -78,7 +78,7 @@ public void initThriftServiceThread() throws IllegalAccessException { config.getRpcMaxConcurrentClientNum(), config.getThriftServerAwaitTimeForStopService(), new TrustedChannelAuditServerEventHandler( - new InternalServiceThriftHandler(), + new InternalServiceThriftHandler(getImpl()), new TEndPoint(getBindIP(), getBindPort()), DNAuditLogger.getInstance()::recordTrustedChannelFailureAuditLogIfNecessary), config.isRpcThriftCompressionEnable(), @@ -95,7 +95,7 @@ public void initThriftServiceThread() throws IllegalAccessException { getBindPort(), config.getRpcMaxConcurrentClientNum(), config.getThriftServerAwaitTimeForStopService(), - new InternalServiceThriftHandler(), + new InternalServiceThriftHandler(getImpl()), config.isRpcThriftCompressionEnable(), DeepCopyRpcTransportFactory.INSTANCE); } catch (RPCServiceException e) { diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift index 57afa70e72846..7415b96bd652e 100644 --- a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift +++ b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift @@ -931,7 +931,7 @@ struct TFetchDeviceEntrySegmentResp { * END: Used for EXPLAIN ANALYZE **/ -service IDataNodeRPCService { +service IDataNodeRPCService extends client.IClientRPCService { // -----------------------------------For Data Node----------------------------------------------- @@ -1186,7 +1186,7 @@ service IDataNodeRPCService { common.TSStatus clearCache(set cacheClearOptions) - common.TShowConfigurationResp showConfiguration() + common.TShowConfigurationResp showConfigurationInternal() common.TShowAppliedConfigurationsResp showAppliedConfigurations() @@ -1434,17 +1434,14 @@ service IDataNodeRPCService { common.TTestConnectionResp submitInternalTestConnectionTask(common.TNodeLocations nodeLocations) - /** Empty rpc, only for connection test */ - common.TSStatus testConnectionEmptyRPC() - /** * Push cluster topology to this DataNode. * Each DataNode receives only its own reachable set. */ common.TSStatus updateClusterTopology(1:TUpdateClusterTopologyReq req) - /** to write audit log or other events as time series **/ - common.TSStatus insertRecord(1:client.TSInsertRecordReq req); + /** Write internal audit records independently of client session identity. */ + common.TSStatus insertRecordInternal(1:client.TSInsertRecordReq req); /** * Write an audit log entry to the DataNode's AuditEventLogger From a7af2c8629742a1fe04cc1a29d64dcb1de75a34a Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Wed, 4 Nov 2026 18:00:09 +0800 Subject: [PATCH 02/32] Support internal table subscription pull consumers --- .../base/AbstractSubscriptionConsumer.java | 6 ++ .../base/AbstractSubscriptionProviders.java | 83 +++++++++++++++++++ ...InternalSubscriptionTablePullConsumer.java | 76 +++++++++++++++++ .../SubscriptionTablePullConsumerBuilder.java | 11 ++- 4 files changed, 175 insertions(+), 1 deletion(-) create mode 100644 iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/InternalSubscriptionTablePullConsumer.java diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java index 7f60b82a0e225..4a284c44255cc 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java @@ -20,6 +20,7 @@ package org.apache.iotdb.session.subscription.consumer.base; import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.rpc.subscription.config.ConsumerConstant; import org.apache.iotdb.rpc.subscription.config.TopicConfig; @@ -572,6 +573,11 @@ private void seekAfterInternalTopicProgress( /////////////////////////////// subscription provider /////////////////////////////// + /** Returns the complete internal endpoint source, or null to use server discovery. */ + protected INodeSupplier getInternalNodeSupplier() { + return null; + } + protected abstract AbstractSubscriptionProvider constructSubscriptionProvider( final TEndPoint endPoint, final String username, diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java index afb5285b5f577..e332b6c7931fa 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java @@ -34,6 +34,7 @@ import java.util.ArrayList; import java.util.LinkedHashMap; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.Objects; @@ -92,6 +93,11 @@ void openProviders(final AbstractSubscriptionConsumer consumer) throws Subscript // close stale providers closeProviders(false); + if (consumer.getInternalNodeSupplier() != null) { + syncInternalProviders(consumer); + return; + } + final Map connectionFailures = new LinkedHashMap<>(); final List connectionFailureCauses = new ArrayList<>(); for (final TEndPoint endPoint : initialEndpoints) { @@ -457,6 +463,16 @@ private void syncInternal(final AbstractSubscriptionConsumer consumer) { if (consumer.isClosed() || consumer.isFenced()) { return; } + if (consumer.getInternalNodeSupplier() != null) { + try { + syncInternalProviders(consumer); + } catch (final SubscriptionConsumerFencedException e) { + consumer.fence(e); + } catch (final Exception e) { + LOGGER.warn(SubscriptionMessages.FETCH_ENDPOINTS_FAILED, consumer, e, e); + } + return; + } if (hasNoAvailableProviders()) { try { openProviders(consumer); @@ -567,6 +583,73 @@ private void syncInternal(final AbstractSubscriptionConsumer consumer) { } } + /** Reconciles providers against a complete endpoint list without server-side discovery. */ + private void syncInternalProviders(final AbstractSubscriptionConsumer consumer) { + final Set endpoints = + consumer.getInternalNodeSupplier().get().stream() + .map(TEndPoint::deepCopy) + .collect(Collectors.toCollection(LinkedHashSet::new)); + + for (final AbstractSubscriptionProvider provider : getAllProviders()) { + if (!provider.isAvailable() || !endpoints.contains(provider.getEndPoint())) { + try { + closeAndRemoveProvider(provider.getDataNodeId()); + } catch (final Exception e) { + LOGGER.warn(SubscriptionMessages.PROVIDER_CLOSE_FAILED, provider, e, e); + } + } + } + + final Set connectedEndpoints = + getAllAvailableProviders().stream() + .map(AbstractSubscriptionProvider::getEndPoint) + .collect(Collectors.toSet()); + Exception lastFailure = null; + for (final TEndPoint endpoint : endpoints) { + if (isClosing.get() || consumer.isFenced()) { + return; + } + if (connectedEndpoints.contains(endpoint)) { + continue; + } + try { + final AbstractSubscriptionProvider provider = + consumer.constructProviderAndHandshake(endpoint); + if (containsProvider(provider.getDataNodeId())) { + try { + closeAndRemoveProvider(provider.getDataNodeId()); + } catch (final Exception e) { + LOGGER.warn(SubscriptionMessages.PROVIDER_CLOSE_FAILED, provider, e, e); + } + } + addProvider(provider.getDataNodeId(), provider); + } catch (final SubscriptionConsumerFencedException e) { + consumer.fence(e); + throw e; + } catch (final Exception e) { + lastFailure = e; + LOGGER.warn( + SubscriptionMessages.LOG_ARG_FAILED_CREATE_CONNECTION_ARG_BECAUSE_ARG_E536E22A, + consumer, + endpoint, + e, + e); + } + } + + if (hasNoAvailableProviders()) { + throw new SubscriptionConnectionException( + String.format( + SubscriptionMessages + .EXCEPTION_CLUSTER_HAS_NO_AVAILABLE_SUBSCRIPTION_PROVIDERS_CONNECT_INITIAL_ENDPOINTS_ARG_5DB83198, + endpoints), + lastFailure); + } + if (!subscriptionProviders.containsKey(nextDataNodeId)) { + nextDataNodeId = subscriptionProviders.firstKey(); + } + } + @Override public String toString() { final StringBuilder sb = new StringBuilder(); diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/InternalSubscriptionTablePullConsumer.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/InternalSubscriptionTablePullConsumer.java new file mode 100644 index 0000000000000..32987f0b5c836 --- /dev/null +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/InternalSubscriptionTablePullConsumer.java @@ -0,0 +1,76 @@ +/* + * 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.subscription.consumer.table; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.session.i18n.SessionMessages; +import org.apache.iotdb.session.subscription.consumer.base.AbstractSubscriptionProvider; + +import java.util.Objects; + +/** Table pull consumer using the caller's complete internal endpoint source and username. */ +public class InternalSubscriptionTablePullConsumer extends SubscriptionTablePullConsumer { + + private final INodeSupplier nodeSupplier; + + public InternalSubscriptionTablePullConsumer( + final SubscriptionTablePullConsumerBuilder builder, final INodeSupplier nodeSupplier) { + super(builder); + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + } + + @Override + protected INodeSupplier getInternalNodeSupplier() { + return nodeSupplier; + } + + @Override + protected AbstractSubscriptionProvider constructSubscriptionProvider( + final TEndPoint endPoint, + final String username, + final String password, + final String encryptedPassword, + final String consumerId, + final String consumerGroupId, + final String ownerId, + final Long ownerEpoch, + final int thriftMaxFrameSize, + final long heartbeatIntervalMs, + final int connectionTimeoutInMs) { + // The internal RPC service authenticates the username without a password. Each provider's + // existing session wrapper stays pinned to this endpoint, including during reconnects. + return super.constructSubscriptionProvider( + endPoint, + username, + "", + null, + consumerId, + consumerGroupId, + ownerId, + ownerEpoch, + thriftMaxFrameSize, + heartbeatIntervalMs, + connectionTimeoutInMs); + } +} diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java index efc99debb4977..7c5cbcdb8261d 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java @@ -19,12 +19,15 @@ package org.apache.iotdb.session.subscription.consumer.table; +import org.apache.iotdb.isession.IInternalSessionBuilder; +import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.session.subscription.consumer.ISubscriptionTablePullConsumer; import org.apache.iotdb.session.subscription.consumer.base.AbstractSubscriptionPullConsumerBuilder; import java.util.List; -public class SubscriptionTablePullConsumerBuilder extends AbstractSubscriptionPullConsumerBuilder { +public class SubscriptionTablePullConsumerBuilder extends AbstractSubscriptionPullConsumerBuilder + implements IInternalSessionBuilder { @Override public SubscriptionTablePullConsumerBuilder host(final String host) { @@ -152,4 +155,10 @@ public SubscriptionTablePullConsumerBuilder autoCommitIntervalMs( public ISubscriptionTablePullConsumer build() { return new SubscriptionTablePullConsumer(this); } + + /** Builds an unopened consumer using the caller's internal endpoint source and username. */ + @Override + public ISubscriptionTablePullConsumer buildInternal(final INodeSupplier nodeSupplier) { + return new InternalSubscriptionTablePullConsumer(this, nodeSupplier); + } } From 282fc56a2b397b62f054e5c6a701bfdf4f9ab2a4 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Wed, 4 Nov 2026 18:52:47 +0800 Subject: [PATCH 03/32] Support internal table session pools --- .../pool/InternalTableSessionPool.java | 38 ++++++ .../iotdb/session/pool/SessionPool.java | 126 ++++++++---------- .../session/pool/TableSessionPoolBuilder.java | 12 +- 3 files changed, 106 insertions(+), 70 deletions(-) create mode 100644 iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.java diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.java new file mode 100644 index 0000000000000..b4d8ef74a6cd1 --- /dev/null +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.java @@ -0,0 +1,38 @@ +/* + * 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.pool; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.session.i18n.SessionMessages; + +import java.util.Objects; + +/** Table session pool using caller-supplied internal endpoints and the caller's username. */ +public class InternalTableSessionPool extends TableSessionPool { + + InternalTableSessionPool(AbstractSessionPoolBuilder builder, INodeSupplier nodeSupplier) { + super( + new SessionPool( + builder, + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392))); + } +} 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 16b79c6415c4b..088afed4e2352 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 @@ -173,6 +173,8 @@ public class SessionPool implements ISessionPool { private INodeSupplier availableNodes; + private INodeSupplier internalNodeSupplier; + // set to true, means that we will start a background thread to fetch all available (Status is // not Removing) datanodes in cluster, and these available nodes will be used in retrying stage private boolean enableAutoFetch = true; @@ -526,15 +528,20 @@ public SessionPool( } public SessionPool(AbstractSessionPoolBuilder builder) { + this(builder, null); + } + + SessionPool(AbstractSessionPoolBuilder builder, INodeSupplier internalNodeSupplier) { + this.internalNodeSupplier = internalNodeSupplier; this.maxSize = builder.maxSize; this.user = builder.username; - this.password = builder.pw; + this.password = internalNodeSupplier == null ? builder.pw : null; this.fetchSize = builder.fetchSize; this.waitToGetSessionTimeoutInMs = builder.waitToGetSessionTimeoutInMs; this.enableThriftCompression = builder.isThriftRpcCompactionEnabled; this.enableIoTDBRpcCompression = builder.isIoTDBRpcCompressionEnabled; this.zoneId = builder.zoneId; - this.enableRedirection = builder.enableRedirection; + this.enableRedirection = internalNodeSupplier == null && builder.enableRedirection; if (this.enableRedirection) { deviceIdToEndpoint = new ConcurrentHashMap<>(); tableModelDeviceIdToEndpoint = new ConcurrentHashMap<>(); @@ -544,7 +551,7 @@ public SessionPool(AbstractSessionPoolBuilder builder) { this.version = builder.version; this.thriftDefaultBufferSize = builder.thriftDefaultBufferSize; this.thriftMaxFrameSize = builder.thriftMaxFrameSize; - this.enableAutoFetch = builder.enableAutoFetch; + this.enableAutoFetch = internalNodeSupplier == null && builder.enableAutoFetch; this.useSSL = builder.useSSL; this.trustStore = builder.trustStore; this.trustStorePwd = builder.trustStorePwd; @@ -561,7 +568,13 @@ public SessionPool(AbstractSessionPoolBuilder builder) { if (enableAutoFetch) { initThreadPool(); } - if (builder.nodeUrls != null) { + if (internalNodeSupplier != null) { + this.host = null; + this.port = -1; + this.nodeUrls = null; + this.formattedNodeUrls = internalNodeSupplier.getClass().getSimpleName(); + this.availableNodes = internalNodeSupplier; + } else if (builder.nodeUrls != null) { if (builder.nodeUrls.isEmpty()) { throw new IllegalArgumentException(SessionMessages.NODE_URLS_EMPTY); } @@ -590,67 +603,40 @@ public SessionPool(AbstractSessionPoolBuilder builder) { } private Session constructNewSession() { - Session session; + Session.Builder builder = + new Session.Builder() + .username(user) + .password(password) + .fetchSize(fetchSize) + .zoneId(zoneId) + .thriftDefaultBufferSize(thriftDefaultBufferSize) + .thriftMaxFrameSize(thriftMaxFrameSize) + .enableRedirection(enableRedirection) + .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) + .version(version) + .useSSL(useSSL) + .trustStore(trustStore) + .trustStorePwd(trustStorePwd) + .keyStore(keyStore) + .keyStorePwd(keyStorePwd) + .sslProtocol(sslProtocol) + .maxRetryCount(maxRetryCount) + .retryIntervalInMs(retryIntervalInMs) + .sqlDialect(sqlDialect) + .database(database) + .enableAutoCreateSchema(enableAutoCreateSchema) + .timeOut(queryTimeoutInMs) + .enableIoTDBRpcCompression(enableIoTDBRpcCompression) + .enableThriftRpcCompression(enableThriftCompression); + if (internalNodeSupplier != null) { + return builder.buildInternal(internalNodeSupplier); + } if (nodeUrls == null) { - // Construct custom Session - session = - new Session.Builder() - .host(host) - .port(port) - .username(user) - .password(password) - .fetchSize(fetchSize) - .zoneId(zoneId) - .thriftDefaultBufferSize(thriftDefaultBufferSize) - .thriftMaxFrameSize(thriftMaxFrameSize) - .enableRedirection(enableRedirection) - .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) - .version(version) - .useSSL(useSSL) - .trustStore(trustStore) - .trustStorePwd(trustStorePwd) - .keyStore(keyStore) - .keyStorePwd(keyStorePwd) - .sslProtocol(sslProtocol) - .maxRetryCount(maxRetryCount) - .retryIntervalInMs(retryIntervalInMs) - .sqlDialect(sqlDialect) - .database(database) - .enableAutoCreateSchema(enableAutoCreateSchema) - .timeOut(queryTimeoutInMs) - .enableIoTDBRpcCompression(enableIoTDBRpcCompression) - .enableThriftRpcCompression(enableThriftCompression) - .build(); + builder.host(host).port(port); } else { - // Construct redirect-able Session - session = - new Session.Builder() - .nodeUrls(nodeUrls) - .username(user) - .password(password) - .fetchSize(fetchSize) - .zoneId(zoneId) - .thriftDefaultBufferSize(thriftDefaultBufferSize) - .thriftMaxFrameSize(thriftMaxFrameSize) - .enableRedirection(enableRedirection) - .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) - .version(version) - .useSSL(useSSL) - .trustStore(trustStore) - .trustStorePwd(trustStorePwd) - .keyStore(keyStore) - .keyStorePwd(keyStorePwd) - .sslProtocol(sslProtocol) - .maxRetryCount(maxRetryCount) - .retryIntervalInMs(retryIntervalInMs) - .sqlDialect(sqlDialect) - .database(database) - .enableAutoCreateSchema(enableAutoCreateSchema) - .timeOut(queryTimeoutInMs) - .enableIoTDBRpcCompression(enableIoTDBRpcCompression) - .enableThriftRpcCompression(enableThriftCompression) - .build(); + builder.nodeUrls(nodeUrls); } + Session session = builder.build(); session.setEnableQueryRedirection(enableQueryRedirection); return session; } @@ -849,7 +835,9 @@ public synchronized void close() { this.executorService = null; } if (availableNodes != null) { - this.availableNodes.close(); + if (internalNodeSupplier == null) { + this.availableNodes.close(); + } this.availableNodes = null; } LOGGER.info(SessionMessages.CLOSING_SESSION_POOL); @@ -3544,16 +3532,16 @@ public boolean isEnableThriftCompression() { @Override public void setEnableRedirection(boolean enableRedirection) { - this.enableRedirection = enableRedirection; + this.enableRedirection = internalNodeSupplier == null && enableRedirection; if (this.enableRedirection) { deviceIdToEndpoint = new ConcurrentHashMap<>(); tableModelDeviceIdToEndpoint = new ConcurrentHashMap<>(); } for (ISession session : queue) { - session.setEnableRedirection(enableRedirection); + session.setEnableRedirection(this.enableRedirection); } for (ISession session : occupied.keySet()) { - session.setEnableRedirection(enableRedirection); + session.setEnableRedirection(this.enableRedirection); } } @@ -3564,12 +3552,12 @@ public boolean isEnableRedirection() { @Override public void setEnableQueryRedirection(boolean enableQueryRedirection) { - this.enableQueryRedirection = enableQueryRedirection; + this.enableQueryRedirection = internalNodeSupplier == null && enableQueryRedirection; for (ISession session : queue) { - session.setEnableQueryRedirection(enableQueryRedirection); + session.setEnableQueryRedirection(this.enableQueryRedirection); } for (ISession session : occupied.keySet()) { - session.setEnableQueryRedirection(enableQueryRedirection); + session.setEnableQueryRedirection(this.enableQueryRedirection); } } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java index f19c5fc8878a6..294f40b8190c7 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java @@ -19,6 +19,8 @@ package org.apache.iotdb.session.pool; +import org.apache.iotdb.isession.IInternalSessionBuilder; +import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.isession.pool.ITableSessionPool; @@ -34,7 +36,8 @@ * *

All configurations have reasonable default values, which can be overridden as needed. */ -public class TableSessionPoolBuilder extends AbstractSessionPoolBuilder { +public class TableSessionPoolBuilder extends AbstractSessionPoolBuilder + implements IInternalSessionBuilder { /** * Sets the list of node URLs for the IoTDB cluster. @@ -346,4 +349,11 @@ public ITableSessionPool build() { SessionPool sessionPool = new SessionPool(this); return new TableSessionPool(sessionPool); } + + /** Builds a pool using the caller's internal endpoints and username; connections open lazily. */ + @Override + public ITableSessionPool buildInternal(INodeSupplier nodeSupplier) { + this.sqlDialect = "table"; + return new InternalTableSessionPool(this, nodeSupplier); + } } From 2d3bdcdd43e66e94f8b49deb53130b26b1b38697 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 5 Nov 2026 18:08:28 +0800 Subject: [PATCH 04/32] Revert inherited internal client RPC implementation Revert the prior approach that exposed the complete client RPC surface by inheritance and renamed internal RPC methods. The replacement will keep IDataNodeRPCService stable and multiplex the client service on the existing internal endpoint. --- .../isession/IInternalSessionBuilder.java | 32 ---- .../iotdb/session/i18n/SessionMessages.java | 5 - .../iotdb/session/i18n/SessionMessages.java | 5 - .../apache/iotdb/session/InternalSession.java | 77 --------- .../org/apache/iotdb/session/Session.java | 69 +++----- .../iotdb/session/TableSessionBuilder.java | 17 +- .../pool/InternalTableSessionPool.java | 38 ----- .../iotdb/session/pool/SessionPool.java | 126 +++++++------- .../session/pool/TableSessionPoolBuilder.java | 12 +- .../InternalSubscriptionTableSession.java | 83 --------- .../SubscriptionSessionWrapper.java | 2 +- .../SubscriptionTableSession.java | 4 - .../SubscriptionTableSessionBuilder.java | 24 +-- .../base/AbstractSubscriptionConsumer.java | 6 - .../base/AbstractSubscriptionProviders.java | 83 --------- ...InternalSubscriptionTablePullConsumer.java | 76 --------- .../SubscriptionTablePullConsumerBuilder.java | 11 +- ...oDnInternalServiceAsyncRequestManager.java | 3 +- .../client/sync/SyncDataNodeClientPool.java | 3 +- .../db/protocol/session/SessionManager.java | 161 +++++++----------- .../handler/InternalServiceThriftHandler.java | 20 +-- .../thrift/impl/ClientRPCServiceImpl.java | 27 ++- .../impl/DataNodeInternalRPCServiceImpl.java | 43 ++--- .../service/DataNodeInternalRPCService.java | 4 +- .../src/main/thrift/datanode.thrift | 11 +- 25 files changed, 202 insertions(+), 740 deletions(-) delete mode 100644 iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java delete mode 100644 iotdb-client/session/src/main/java/org/apache/iotdb/session/InternalSession.java delete mode 100644 iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.java delete mode 100644 iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/InternalSubscriptionTableSession.java delete mode 100644 iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/InternalSubscriptionTablePullConsumer.java diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java deleted file mode 100644 index ae67fa4093da0..0000000000000 --- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java +++ /dev/null @@ -1,32 +0,0 @@ -/* - * 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.rpc.IoTDBConnectionException; - -/** Capability implemented by builders that can construct sessions on internal endpoints. */ -public interface IInternalSessionBuilder { - - /** - * Builds a session using the caller-provided internal node supplier. Whether the returned session - * is already open follows the builder's ordinary build contract. - */ - T buildInternal(INodeSupplier nodeSupplier) throws IoTDBConnectionException; -} diff --git a/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java b/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java index 382567f378d61..ded0deac31484 100644 --- a/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java +++ b/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java @@ -20,11 +20,6 @@ package org.apache.iotdb.session.i18n; public final class SessionMessages { - public static final String EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081 = - "No DataNode internal endpoint is available"; - - public static final String EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392 = - "Internal node supplier must not be null"; // Session public static final String NODE_URLS_EMPTY = "nodeUrls shouldn't be empty."; diff --git a/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java b/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java index 9975c43a74cd7..838a7e2148329 100644 --- a/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java +++ b/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java @@ -20,11 +20,6 @@ package org.apache.iotdb.session.i18n; public final class SessionMessages { - public static final String EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081 = - "没有可用的 DataNode 内部端点"; - - public static final String EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392 = - "内部节点 supplier 不能为空"; // Session public static final String NODE_URLS_EMPTY = "nodeUrls 不能为空。"; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/InternalSession.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/InternalSession.java deleted file mode 100644 index fdcaed13f3964..0000000000000 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/InternalSession.java +++ /dev/null @@ -1,77 +0,0 @@ -/* - * 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.INodeSupplier; -import org.apache.iotdb.rpc.IoTDBConnectionException; -import org.apache.iotdb.session.i18n.SessionMessages; - -import java.time.ZoneId; -import java.util.List; -import java.util.Objects; - -/** Session using caller-supplied internal endpoints and the caller's username. */ -public class InternalSession extends Session { - - private final INodeSupplier nodeSupplier; - - public InternalSession(AbstractSessionBuilder builder, INodeSupplier nodeSupplier) { - super(builder); - this.nodeSupplier = - Objects.requireNonNull( - nodeSupplier, - SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); - password = null; - useEncryptedPassword = false; - enableAutoFetch = false; - enableRedirection = false; - enableQueryRedirection = false; - } - - @Override - public synchronized void open() throws IoTDBConnectionException { - open(enableThriftRpcCompaction, connectionTimeoutInMs); - } - - @Override - protected INodeSupplier createNodeSupplier(boolean enableRPCCompaction, int timeoutInMs) { - return nodeSupplier; - } - - @Override - public SessionConnection constructSessionConnection( - Session session, TEndPoint endpoint, ZoneId zoneId) throws IoTDBConnectionException { - final List candidates = availableNodes.get(); - if (candidates.isEmpty()) { - throw new IoTDBConnectionException( - SessionMessages.EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081); - } - defaultEndPoint = - candidates.contains(endpoint) - ? endpoint.deepCopy() - : nodeSupplier - .getQueryEndPoint() - .filter(candidates::contains) - .orElse(candidates.get(0)) - .deepCopy(); - return super.constructSessionConnection(session, defaultEndPoint, zoneId); - } -} 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 f415a63246490..de3ee417afc0d 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,7 +21,6 @@ import org.apache.iotdb.common.rpc.thrift.TAggregationType; import org.apache.iotdb.common.rpc.thrift.TEndPoint; -import org.apache.iotdb.isession.IInternalSessionBuilder; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.ISession; import org.apache.iotdb.isession.SessionConfig; @@ -550,7 +549,33 @@ public synchronized void open(boolean enableRPCCompaction, int connectionTimeout this.availableNodes = null; } - this.availableNodes = createNodeSupplier(enableRPCCompaction, connectionTimeoutInMs); + if (enableAutoFetch) { + initThreadPool(); + this.availableNodes = + NodesSupplier.createNodeSupplier( + getNodeUrls(), + SessionUtils.parseSeedNodeUrls(initialNodeUrls), + executorService, + username, + password, + zoneId, + thriftDefaultBufferSize, + thriftMaxFrameSize, + connectionTimeoutInMs, + useSSL, + trustStore, + trustStorePwd, + keyStore, + keyStorePwd, + sslProtocol, + enableRPCCompaction, + version.toString()); + } else { + // A static supplier has no separate seed tier. Keep every configured address after reopen. + Set staticNodes = new LinkedHashSet<>(getNodeUrls()); + staticNodes.addAll(SessionUtils.parseSeedNodeUrls(initialNodeUrls)); + this.availableNodes = new DummyNodesSupplier(new ArrayList<>(staticNodes)); + } this.connectionTimeoutInMs = connectionTimeoutInMs; setDefaultSessionConnection(constructSessionConnection(this, defaultEndPoint, zoneId)); @@ -564,37 +589,6 @@ public synchronized void open(boolean enableRPCCompaction, int connectionTimeout } } - /** Creates the endpoint source used by initial connections and reconnects. */ - protected INodeSupplier createNodeSupplier( - boolean enableRPCCompaction, int connectionTimeoutInMs) { - if (enableAutoFetch) { - initThreadPool(); - return NodesSupplier.createNodeSupplier( - getNodeUrls(), - SessionUtils.parseSeedNodeUrls(initialNodeUrls), - executorService, - username, - password, - zoneId, - thriftDefaultBufferSize, - thriftMaxFrameSize, - connectionTimeoutInMs, - useSSL, - trustStore, - trustStorePwd, - keyStore, - keyStorePwd, - sslProtocol, - enableRPCCompaction, - version.toString()); - } else { - // A static supplier has no separate seed tier. Keep every configured address after reopen. - Set staticNodes = new LinkedHashSet<>(getNodeUrls()); - staticNodes.addAll(SessionUtils.parseSeedNodeUrls(initialNodeUrls)); - return new DummyNodesSupplier(new ArrayList<>(staticNodes)); - } - } - private void initThreadPool() { this.executorService = Executors.newSingleThreadScheduledExecutor( @@ -4451,8 +4445,7 @@ public void setEnableThriftRpcCompaction(boolean enableThriftRpcCompaction) { this.enableThriftRpcCompaction = enableThriftRpcCompaction; } - public static class Builder extends AbstractSessionBuilder - implements IInternalSessionBuilder { + public static class Builder extends AbstractSessionBuilder { public Builder host(String host) { this.host = host; @@ -4602,11 +4595,5 @@ public Session build() { } return new Session(this); } - - /** Builds an unopened internal session using caller-supplied internal endpoints. */ - @Override - public InternalSession buildInternal(INodeSupplier nodeSupplier) { - return new InternalSession(this, nodeSupplier); - } } } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java index 856ed7a91bc6b..a39af4c847a95 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java @@ -19,8 +19,6 @@ package org.apache.iotdb.session; -import org.apache.iotdb.isession.IInternalSessionBuilder; -import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.ITableSession; import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.rpc.IoTDBConnectionException; @@ -43,8 +41,7 @@ * *

All configurations have reasonable default values, which can be overridden as needed. */ -public class TableSessionBuilder extends AbstractSessionBuilder - implements IInternalSessionBuilder { +public class TableSessionBuilder extends AbstractSessionBuilder { private boolean enableCompression = false; private int connectionTimeoutInMs = SessionConfig.DEFAULT_CONNECTION_TIMEOUT_MS; @@ -440,17 +437,7 @@ public ITableSession build() throws IoTDBConnectionException { Collections.singletonList(SessionConfig.DEFAULT_HOST + ":" + SessionConfig.DEFAULT_PORT); } this.sqlDialect = TABLE; - return buildSession(new Session(this)); - } - - /** Builds and opens a table session using caller-supplied internal endpoints. */ - @Override - public ITableSession buildInternal(INodeSupplier nodeSupplier) throws IoTDBConnectionException { - this.sqlDialect = TABLE; - return buildSession(new InternalSession(this, nodeSupplier)); - } - - private ITableSession buildSession(Session newSession) throws IoTDBConnectionException { + Session newSession = new Session(this); newSession.enableIoTDBRpcCompression = isIoTDBRpcCompressionEnabled; newSession.tabletCompressionMinRowSize = tabletCompressionMinRowSize; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.java deleted file mode 100644 index b4d8ef74a6cd1..0000000000000 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.java +++ /dev/null @@ -1,38 +0,0 @@ -/* - * 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.pool; - -import org.apache.iotdb.isession.INodeSupplier; -import org.apache.iotdb.session.i18n.SessionMessages; - -import java.util.Objects; - -/** Table session pool using caller-supplied internal endpoints and the caller's username. */ -public class InternalTableSessionPool extends TableSessionPool { - - InternalTableSessionPool(AbstractSessionPoolBuilder builder, INodeSupplier nodeSupplier) { - super( - new SessionPool( - builder, - Objects.requireNonNull( - nodeSupplier, - SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392))); - } -} 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 088afed4e2352..16b79c6415c4b 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 @@ -173,8 +173,6 @@ public class SessionPool implements ISessionPool { private INodeSupplier availableNodes; - private INodeSupplier internalNodeSupplier; - // set to true, means that we will start a background thread to fetch all available (Status is // not Removing) datanodes in cluster, and these available nodes will be used in retrying stage private boolean enableAutoFetch = true; @@ -528,20 +526,15 @@ public SessionPool( } public SessionPool(AbstractSessionPoolBuilder builder) { - this(builder, null); - } - - SessionPool(AbstractSessionPoolBuilder builder, INodeSupplier internalNodeSupplier) { - this.internalNodeSupplier = internalNodeSupplier; this.maxSize = builder.maxSize; this.user = builder.username; - this.password = internalNodeSupplier == null ? builder.pw : null; + this.password = builder.pw; this.fetchSize = builder.fetchSize; this.waitToGetSessionTimeoutInMs = builder.waitToGetSessionTimeoutInMs; this.enableThriftCompression = builder.isThriftRpcCompactionEnabled; this.enableIoTDBRpcCompression = builder.isIoTDBRpcCompressionEnabled; this.zoneId = builder.zoneId; - this.enableRedirection = internalNodeSupplier == null && builder.enableRedirection; + this.enableRedirection = builder.enableRedirection; if (this.enableRedirection) { deviceIdToEndpoint = new ConcurrentHashMap<>(); tableModelDeviceIdToEndpoint = new ConcurrentHashMap<>(); @@ -551,7 +544,7 @@ public SessionPool(AbstractSessionPoolBuilder builder) { this.version = builder.version; this.thriftDefaultBufferSize = builder.thriftDefaultBufferSize; this.thriftMaxFrameSize = builder.thriftMaxFrameSize; - this.enableAutoFetch = internalNodeSupplier == null && builder.enableAutoFetch; + this.enableAutoFetch = builder.enableAutoFetch; this.useSSL = builder.useSSL; this.trustStore = builder.trustStore; this.trustStorePwd = builder.trustStorePwd; @@ -568,13 +561,7 @@ public SessionPool(AbstractSessionPoolBuilder builder) { if (enableAutoFetch) { initThreadPool(); } - if (internalNodeSupplier != null) { - this.host = null; - this.port = -1; - this.nodeUrls = null; - this.formattedNodeUrls = internalNodeSupplier.getClass().getSimpleName(); - this.availableNodes = internalNodeSupplier; - } else if (builder.nodeUrls != null) { + if (builder.nodeUrls != null) { if (builder.nodeUrls.isEmpty()) { throw new IllegalArgumentException(SessionMessages.NODE_URLS_EMPTY); } @@ -603,40 +590,67 @@ public SessionPool(AbstractSessionPoolBuilder builder) { } private Session constructNewSession() { - Session.Builder builder = - new Session.Builder() - .username(user) - .password(password) - .fetchSize(fetchSize) - .zoneId(zoneId) - .thriftDefaultBufferSize(thriftDefaultBufferSize) - .thriftMaxFrameSize(thriftMaxFrameSize) - .enableRedirection(enableRedirection) - .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) - .version(version) - .useSSL(useSSL) - .trustStore(trustStore) - .trustStorePwd(trustStorePwd) - .keyStore(keyStore) - .keyStorePwd(keyStorePwd) - .sslProtocol(sslProtocol) - .maxRetryCount(maxRetryCount) - .retryIntervalInMs(retryIntervalInMs) - .sqlDialect(sqlDialect) - .database(database) - .enableAutoCreateSchema(enableAutoCreateSchema) - .timeOut(queryTimeoutInMs) - .enableIoTDBRpcCompression(enableIoTDBRpcCompression) - .enableThriftRpcCompression(enableThriftCompression); - if (internalNodeSupplier != null) { - return builder.buildInternal(internalNodeSupplier); - } + Session session; if (nodeUrls == null) { - builder.host(host).port(port); + // Construct custom Session + session = + new Session.Builder() + .host(host) + .port(port) + .username(user) + .password(password) + .fetchSize(fetchSize) + .zoneId(zoneId) + .thriftDefaultBufferSize(thriftDefaultBufferSize) + .thriftMaxFrameSize(thriftMaxFrameSize) + .enableRedirection(enableRedirection) + .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) + .version(version) + .useSSL(useSSL) + .trustStore(trustStore) + .trustStorePwd(trustStorePwd) + .keyStore(keyStore) + .keyStorePwd(keyStorePwd) + .sslProtocol(sslProtocol) + .maxRetryCount(maxRetryCount) + .retryIntervalInMs(retryIntervalInMs) + .sqlDialect(sqlDialect) + .database(database) + .enableAutoCreateSchema(enableAutoCreateSchema) + .timeOut(queryTimeoutInMs) + .enableIoTDBRpcCompression(enableIoTDBRpcCompression) + .enableThriftRpcCompression(enableThriftCompression) + .build(); } else { - builder.nodeUrls(nodeUrls); + // Construct redirect-able Session + session = + new Session.Builder() + .nodeUrls(nodeUrls) + .username(user) + .password(password) + .fetchSize(fetchSize) + .zoneId(zoneId) + .thriftDefaultBufferSize(thriftDefaultBufferSize) + .thriftMaxFrameSize(thriftMaxFrameSize) + .enableRedirection(enableRedirection) + .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) + .version(version) + .useSSL(useSSL) + .trustStore(trustStore) + .trustStorePwd(trustStorePwd) + .keyStore(keyStore) + .keyStorePwd(keyStorePwd) + .sslProtocol(sslProtocol) + .maxRetryCount(maxRetryCount) + .retryIntervalInMs(retryIntervalInMs) + .sqlDialect(sqlDialect) + .database(database) + .enableAutoCreateSchema(enableAutoCreateSchema) + .timeOut(queryTimeoutInMs) + .enableIoTDBRpcCompression(enableIoTDBRpcCompression) + .enableThriftRpcCompression(enableThriftCompression) + .build(); } - Session session = builder.build(); session.setEnableQueryRedirection(enableQueryRedirection); return session; } @@ -835,9 +849,7 @@ public synchronized void close() { this.executorService = null; } if (availableNodes != null) { - if (internalNodeSupplier == null) { - this.availableNodes.close(); - } + this.availableNodes.close(); this.availableNodes = null; } LOGGER.info(SessionMessages.CLOSING_SESSION_POOL); @@ -3532,16 +3544,16 @@ public boolean isEnableThriftCompression() { @Override public void setEnableRedirection(boolean enableRedirection) { - this.enableRedirection = internalNodeSupplier == null && enableRedirection; + this.enableRedirection = enableRedirection; if (this.enableRedirection) { deviceIdToEndpoint = new ConcurrentHashMap<>(); tableModelDeviceIdToEndpoint = new ConcurrentHashMap<>(); } for (ISession session : queue) { - session.setEnableRedirection(this.enableRedirection); + session.setEnableRedirection(enableRedirection); } for (ISession session : occupied.keySet()) { - session.setEnableRedirection(this.enableRedirection); + session.setEnableRedirection(enableRedirection); } } @@ -3552,12 +3564,12 @@ public boolean isEnableRedirection() { @Override public void setEnableQueryRedirection(boolean enableQueryRedirection) { - this.enableQueryRedirection = internalNodeSupplier == null && enableQueryRedirection; + this.enableQueryRedirection = enableQueryRedirection; for (ISession session : queue) { - session.setEnableQueryRedirection(this.enableQueryRedirection); + session.setEnableQueryRedirection(enableQueryRedirection); } for (ISession session : occupied.keySet()) { - session.setEnableQueryRedirection(this.enableQueryRedirection); + session.setEnableQueryRedirection(enableQueryRedirection); } } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java index 294f40b8190c7..f19c5fc8878a6 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java @@ -19,8 +19,6 @@ package org.apache.iotdb.session.pool; -import org.apache.iotdb.isession.IInternalSessionBuilder; -import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.isession.pool.ITableSessionPool; @@ -36,8 +34,7 @@ * *

All configurations have reasonable default values, which can be overridden as needed. */ -public class TableSessionPoolBuilder extends AbstractSessionPoolBuilder - implements IInternalSessionBuilder { +public class TableSessionPoolBuilder extends AbstractSessionPoolBuilder { /** * Sets the list of node URLs for the IoTDB cluster. @@ -349,11 +346,4 @@ public ITableSessionPool build() { SessionPool sessionPool = new SessionPool(this); return new TableSessionPool(sessionPool); } - - /** Builds a pool using the caller's internal endpoints and username; connections open lazily. */ - @Override - public ITableSessionPool buildInternal(INodeSupplier nodeSupplier) { - this.sqlDialect = "table"; - return new InternalTableSessionPool(this, nodeSupplier); - } } diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/InternalSubscriptionTableSession.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/InternalSubscriptionTableSession.java deleted file mode 100644 index 13cf8c659ed94..0000000000000 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/InternalSubscriptionTableSession.java +++ /dev/null @@ -1,83 +0,0 @@ -/* - * 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.subscription; - -import org.apache.iotdb.common.rpc.thrift.TEndPoint; -import org.apache.iotdb.isession.INodeSupplier; -import org.apache.iotdb.rpc.IoTDBConnectionException; -import org.apache.iotdb.session.AbstractSessionBuilder; -import org.apache.iotdb.session.Session; -import org.apache.iotdb.session.SessionConnection; -import org.apache.iotdb.session.i18n.SessionMessages; - -import java.time.ZoneId; -import java.util.List; -import java.util.Objects; - -/** Table subscription-management session using the caller's identity and internal endpoints. */ -public class InternalSubscriptionTableSession extends SubscriptionTableSession { - - public InternalSubscriptionTableSession( - AbstractSessionBuilder builder, INodeSupplier nodeSupplier) { - super(new InternalSessionWrapper(builder, nodeSupplier)); - } - - private static class InternalSessionWrapper extends SubscriptionSessionWrapper { - - private final INodeSupplier nodeSupplier; - - private InternalSessionWrapper(AbstractSessionBuilder builder, INodeSupplier nodeSupplier) { - super(builder); - this.nodeSupplier = - Objects.requireNonNull( - nodeSupplier, - SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); - password = null; - useEncryptedPassword = false; - enableAutoFetch = false; - enableRedirection = false; - enableQueryRedirection = false; - } - - @Override - protected INodeSupplier createNodeSupplier(boolean enableRPCCompaction, int timeoutInMs) { - return nodeSupplier; - } - - @Override - public SessionConnection constructSessionConnection( - Session session, TEndPoint endpoint, ZoneId zoneId) throws IoTDBConnectionException { - final List candidates = nodeSupplier.get(); - if (candidates.isEmpty()) { - throw new IoTDBConnectionException( - SessionMessages.EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081); - } - defaultEndPoint = - candidates.contains(endpoint) - ? endpoint.deepCopy() - : nodeSupplier - .getQueryEndPoint() - .filter(candidates::contains) - .orElse(candidates.get(0)) - .deepCopy(); - return super.constructSessionConnection(session, defaultEndPoint, zoneId); - } - } -} diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java index 018f4ef3a185d..d7e715fa62762 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java @@ -30,7 +30,7 @@ import java.time.ZoneId; import java.util.Objects; -public class SubscriptionSessionWrapper extends Session { +public final class SubscriptionSessionWrapper extends Session { public SubscriptionSessionWrapper(final AbstractSessionBuilder builder) { super(builder); diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSession.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSession.java index 8949c46ff83ad..123deea2fa45d 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSession.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSession.java @@ -36,10 +36,6 @@ public SubscriptionTableSession(final AbstractSessionBuilder builder) { super(new SubscriptionSessionWrapper(builder)); } - protected SubscriptionTableSession(final SubscriptionSessionWrapper session) { - super(session); - } - /////////////////////////////// open & close /////////////////////////////// @Override diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java index 0f34d27ee4980..e0c8e95968dc9 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java @@ -19,13 +19,10 @@ package org.apache.iotdb.session.subscription; -import org.apache.iotdb.isession.IInternalSessionBuilder; -import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.rpc.IoTDBConnectionException; import org.apache.iotdb.session.AbstractSessionBuilder; -public class SubscriptionTableSessionBuilder extends AbstractSessionBuilder - implements IInternalSessionBuilder { +public class SubscriptionTableSessionBuilder extends AbstractSessionBuilder { public SubscriptionTableSessionBuilder() { // use table model @@ -76,23 +73,4 @@ public ISubscriptionTableSession build() throws IoTDBConnectionException { session.open(); return session; } - - /** Builds and opens a table subscription-management session on internal endpoints. */ - @Override - public ISubscriptionTableSession buildInternal(INodeSupplier nodeSupplier) - throws IoTDBConnectionException { - final ISubscriptionTableSession session = - new InternalSubscriptionTableSession(this, nodeSupplier); - try { - session.open(); - } catch (IoTDBConnectionException e) { - try { - session.close(); - } catch (Exception closeException) { - e.addSuppressed(closeException); - } - throw e; - } - return session; - } } diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java index 4a284c44255cc..7f60b82a0e225 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java @@ -20,7 +20,6 @@ package org.apache.iotdb.session.subscription.consumer.base; import org.apache.iotdb.common.rpc.thrift.TEndPoint; -import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.rpc.subscription.config.ConsumerConstant; import org.apache.iotdb.rpc.subscription.config.TopicConfig; @@ -573,11 +572,6 @@ private void seekAfterInternalTopicProgress( /////////////////////////////// subscription provider /////////////////////////////// - /** Returns the complete internal endpoint source, or null to use server discovery. */ - protected INodeSupplier getInternalNodeSupplier() { - return null; - } - protected abstract AbstractSubscriptionProvider constructSubscriptionProvider( final TEndPoint endPoint, final String username, diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java index e332b6c7931fa..afb5285b5f577 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java @@ -34,7 +34,6 @@ import java.util.ArrayList; import java.util.LinkedHashMap; -import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.Objects; @@ -93,11 +92,6 @@ void openProviders(final AbstractSubscriptionConsumer consumer) throws Subscript // close stale providers closeProviders(false); - if (consumer.getInternalNodeSupplier() != null) { - syncInternalProviders(consumer); - return; - } - final Map connectionFailures = new LinkedHashMap<>(); final List connectionFailureCauses = new ArrayList<>(); for (final TEndPoint endPoint : initialEndpoints) { @@ -463,16 +457,6 @@ private void syncInternal(final AbstractSubscriptionConsumer consumer) { if (consumer.isClosed() || consumer.isFenced()) { return; } - if (consumer.getInternalNodeSupplier() != null) { - try { - syncInternalProviders(consumer); - } catch (final SubscriptionConsumerFencedException e) { - consumer.fence(e); - } catch (final Exception e) { - LOGGER.warn(SubscriptionMessages.FETCH_ENDPOINTS_FAILED, consumer, e, e); - } - return; - } if (hasNoAvailableProviders()) { try { openProviders(consumer); @@ -583,73 +567,6 @@ private void syncInternal(final AbstractSubscriptionConsumer consumer) { } } - /** Reconciles providers against a complete endpoint list without server-side discovery. */ - private void syncInternalProviders(final AbstractSubscriptionConsumer consumer) { - final Set endpoints = - consumer.getInternalNodeSupplier().get().stream() - .map(TEndPoint::deepCopy) - .collect(Collectors.toCollection(LinkedHashSet::new)); - - for (final AbstractSubscriptionProvider provider : getAllProviders()) { - if (!provider.isAvailable() || !endpoints.contains(provider.getEndPoint())) { - try { - closeAndRemoveProvider(provider.getDataNodeId()); - } catch (final Exception e) { - LOGGER.warn(SubscriptionMessages.PROVIDER_CLOSE_FAILED, provider, e, e); - } - } - } - - final Set connectedEndpoints = - getAllAvailableProviders().stream() - .map(AbstractSubscriptionProvider::getEndPoint) - .collect(Collectors.toSet()); - Exception lastFailure = null; - for (final TEndPoint endpoint : endpoints) { - if (isClosing.get() || consumer.isFenced()) { - return; - } - if (connectedEndpoints.contains(endpoint)) { - continue; - } - try { - final AbstractSubscriptionProvider provider = - consumer.constructProviderAndHandshake(endpoint); - if (containsProvider(provider.getDataNodeId())) { - try { - closeAndRemoveProvider(provider.getDataNodeId()); - } catch (final Exception e) { - LOGGER.warn(SubscriptionMessages.PROVIDER_CLOSE_FAILED, provider, e, e); - } - } - addProvider(provider.getDataNodeId(), provider); - } catch (final SubscriptionConsumerFencedException e) { - consumer.fence(e); - throw e; - } catch (final Exception e) { - lastFailure = e; - LOGGER.warn( - SubscriptionMessages.LOG_ARG_FAILED_CREATE_CONNECTION_ARG_BECAUSE_ARG_E536E22A, - consumer, - endpoint, - e, - e); - } - } - - if (hasNoAvailableProviders()) { - throw new SubscriptionConnectionException( - String.format( - SubscriptionMessages - .EXCEPTION_CLUSTER_HAS_NO_AVAILABLE_SUBSCRIPTION_PROVIDERS_CONNECT_INITIAL_ENDPOINTS_ARG_5DB83198, - endpoints), - lastFailure); - } - if (!subscriptionProviders.containsKey(nextDataNodeId)) { - nextDataNodeId = subscriptionProviders.firstKey(); - } - } - @Override public String toString() { final StringBuilder sb = new StringBuilder(); diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/InternalSubscriptionTablePullConsumer.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/InternalSubscriptionTablePullConsumer.java deleted file mode 100644 index 32987f0b5c836..0000000000000 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/InternalSubscriptionTablePullConsumer.java +++ /dev/null @@ -1,76 +0,0 @@ -/* - * 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.subscription.consumer.table; - -import org.apache.iotdb.common.rpc.thrift.TEndPoint; -import org.apache.iotdb.isession.INodeSupplier; -import org.apache.iotdb.session.i18n.SessionMessages; -import org.apache.iotdb.session.subscription.consumer.base.AbstractSubscriptionProvider; - -import java.util.Objects; - -/** Table pull consumer using the caller's complete internal endpoint source and username. */ -public class InternalSubscriptionTablePullConsumer extends SubscriptionTablePullConsumer { - - private final INodeSupplier nodeSupplier; - - public InternalSubscriptionTablePullConsumer( - final SubscriptionTablePullConsumerBuilder builder, final INodeSupplier nodeSupplier) { - super(builder); - this.nodeSupplier = - Objects.requireNonNull( - nodeSupplier, - SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); - } - - @Override - protected INodeSupplier getInternalNodeSupplier() { - return nodeSupplier; - } - - @Override - protected AbstractSubscriptionProvider constructSubscriptionProvider( - final TEndPoint endPoint, - final String username, - final String password, - final String encryptedPassword, - final String consumerId, - final String consumerGroupId, - final String ownerId, - final Long ownerEpoch, - final int thriftMaxFrameSize, - final long heartbeatIntervalMs, - final int connectionTimeoutInMs) { - // The internal RPC service authenticates the username without a password. Each provider's - // existing session wrapper stays pinned to this endpoint, including during reconnects. - return super.constructSubscriptionProvider( - endPoint, - username, - "", - null, - consumerId, - consumerGroupId, - ownerId, - ownerEpoch, - thriftMaxFrameSize, - heartbeatIntervalMs, - connectionTimeoutInMs); - } -} diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java index 7c5cbcdb8261d..efc99debb4977 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java @@ -19,15 +19,12 @@ package org.apache.iotdb.session.subscription.consumer.table; -import org.apache.iotdb.isession.IInternalSessionBuilder; -import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.session.subscription.consumer.ISubscriptionTablePullConsumer; import org.apache.iotdb.session.subscription.consumer.base.AbstractSubscriptionPullConsumerBuilder; import java.util.List; -public class SubscriptionTablePullConsumerBuilder extends AbstractSubscriptionPullConsumerBuilder - implements IInternalSessionBuilder { +public class SubscriptionTablePullConsumerBuilder extends AbstractSubscriptionPullConsumerBuilder { @Override public SubscriptionTablePullConsumerBuilder host(final String host) { @@ -155,10 +152,4 @@ public SubscriptionTablePullConsumerBuilder autoCommitIntervalMs( public ISubscriptionTablePullConsumer build() { return new SubscriptionTablePullConsumer(this); } - - /** Builds an unopened consumer using the caller's internal endpoint source and username. */ - @Override - public ISubscriptionTablePullConsumer buildInternal(final INodeSupplier nodeSupplier) { - return new InternalSubscriptionTablePullConsumer(this, nodeSupplier); - } } diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/CnToDnInternalServiceAsyncRequestManager.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/CnToDnInternalServiceAsyncRequestManager.java index 3cb2468c27893..3010f8b524c04 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/CnToDnInternalServiceAsyncRequestManager.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/CnToDnInternalServiceAsyncRequestManager.java @@ -449,8 +449,7 @@ protected void initActionMapBuilder() { actionMapBuilder.put( CnToDnAsyncRequestType.INSERT_RECORD, (req, client, handler) -> - client.insertRecordInternal( - (TSInsertRecordReq) req, (DataNodeTSStatusRPCHandler) handler)); + client.insertRecord((TSInsertRecordReq) req, (DataNodeTSStatusRPCHandler) handler)); actionMapBuilder.put( CnToDnAsyncRequestType.UPDATE_TABLE, (req, client, handler) -> diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/sync/SyncDataNodeClientPool.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/sync/SyncDataNodeClientPool.java index 60c58ddea0d2b..f35b6d906a525 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/sync/SyncDataNodeClientPool.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/sync/SyncDataNodeClientPool.java @@ -137,8 +137,7 @@ private void buildActionMap() { CnToDnSyncRequestType.RESET_PEER_LIST, (req, client) -> client.resetPeerList((TResetPeerListReq) req)); actionMapBuilder.put( - CnToDnSyncRequestType.SHOW_CONFIGURATION, - (req, client) -> client.showConfigurationInternal()); + CnToDnSyncRequestType.SHOW_CONFIGURATION, (req, client) -> client.showConfiguration()); actionMapBuilder.put( CnToDnSyncRequestType.SHOW_APPLIED_CONFIGURATIONS, (req, client) -> client.showAppliedConfigurations()); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java index 7017bda16794f..85d9acae276ec 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java @@ -157,116 +157,73 @@ public BasicOpenSessionResp login( final TSStatus loginStatus = AuthorityChecker.checkUser(username, password, useEncryptedPassword); if (loginStatus.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) { - BasicOpenSessionResp response = - completeLogin( - session, userId, username, zoneId, tsProtocolVersion, clientVersion, sqlDialect); - if (response.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode() && enableLoginLock) { - loginLockManager.clearFailure(userId, session.getClientAddress()); + // check the version compatibility + if (!tsProtocolVersion.equals(CURRENT_RPC_VERSION)) { + openSessionResp + .sessionId(-1) + .setCode(TSStatusCode.INCOMPATIBLE_VERSION.getStatusCode()) + .setMessage( + DataNodeMiscMessages.VERSION_INCOMPATIBLE_PLEASE_UPGRADE_TO + + IoTDBConstant.VERSION); + } else { + session.setSqlDialect(sqlDialect); + supplySession(session, userId, username, ZoneId.of(zoneId), clientVersion); + String logInMessage = DataNodeMiscMessages.EXCEPTION_LOGIN_SUCCESSFULLY_1C2EF670; + openSessionResp + .sessionId(session.getId()) + .setCode(TSStatusCode.SUCCESS_STATUS.getStatusCode()) + .setMessage(logInMessage); + AUDIT_LOGGER.log( + new AuditLogFields( + userId, + username, + session.getClientAddress(), + AuditEventType.LOGIN, + AuditLogOperation.CONTROL, + true), + () -> + String.format( + DataNodeMiscMessages + .EXCEPTION_ARG_LOGIN_STATUS_ARG_USER_ARG_ID_ARG_OPENS_SESSION_ARG_70586F85, + IoTDBConstant.GLOBAL_DB_NAME, + openSessionResp.getMessage(), + username, + userId, + session)); + LOGGER.info( + DataNodeMiscMessages.LOGIN_STATUS, + IoTDBConstant.GLOBAL_DB_NAME, + openSessionResp.getMessage(), + username, + session); + updateIdleTime(); + if (enableLoginLock) { + loginLockManager.clearFailure(userId, session.getClientAddress()); + } } - return response; } else { openSessionResp.sessionId(-1).setMessage(loginStatus.message).setCode(loginStatus.code); if (enableLoginLock) { loginLockManager.recordFailure(userId, session.getClientAddress()); } - logFailedLogin(session, userId, username, loginStatus); - } - return openSessionResp; - } - - /** Opens a session for a trusted internal caller without checking a password. */ - public BasicOpenSessionResp loginInternal( - IClientSession session, - String username, - String zoneId, - TSProtocolVersion tsProtocolVersion, - IoTDBConstant.ClientVersion clientVersion, - SqlDialect sqlDialect) { - final long userId = AuthorityChecker.getUserId(username).orElse(-1L); - if (userId == -1) { - BasicOpenSessionResp response = new BasicOpenSessionResp(); - response - .sessionId(-1) - .setCode(TSStatusCode.USER_NOT_EXIST.getStatusCode()) - .setMessage( + AUDIT_LOGGER.log( + new AuditLogFields( + userId, + username, + session.getClientAddress(), + AuditEventType.LOGIN, + AuditLogOperation.CONTROL, + false), + () -> String.format( - DataNodeMiscMessages.MISC_EXCEPTION_USER_S_DOES_NOT_EXIST_0CE725D8, username)); - logFailedLogin(session, userId, username, response); - return response; - } - return completeLogin( - session, userId, username, zoneId, tsProtocolVersion, clientVersion, sqlDialect); - } - - private BasicOpenSessionResp completeLogin( - IClientSession session, - long userId, - String username, - String zoneId, - TSProtocolVersion tsProtocolVersion, - IoTDBConstant.ClientVersion clientVersion, - SqlDialect sqlDialect) { - BasicOpenSessionResp response = new BasicOpenSessionResp(); - if (!tsProtocolVersion.equals(CURRENT_RPC_VERSION)) { - response - .sessionId(-1) - .setCode(TSStatusCode.INCOMPATIBLE_VERSION.getStatusCode()) - .setMessage( - DataNodeMiscMessages.VERSION_INCOMPATIBLE_PLEASE_UPGRADE_TO + IoTDBConstant.VERSION); - return response; + DataNodeMiscMessages + .EXCEPTION_USER_ARG_ID_ARG_LOGIN_FAILED_WITH_CODE_ARG_ARG_93FAE4FD, + username, + userId, + loginStatus.getCode(), + loginStatus.getMessage())); } - - session.setSqlDialect(sqlDialect); - supplySession(session, userId, username, ZoneId.of(zoneId), clientVersion); - response - .sessionId(session.getId()) - .setCode(TSStatusCode.SUCCESS_STATUS.getStatusCode()) - .setMessage(DataNodeMiscMessages.EXCEPTION_LOGIN_SUCCESSFULLY_1C2EF670); - AUDIT_LOGGER.log( - new AuditLogFields( - userId, - username, - session.getClientAddress(), - AuditEventType.LOGIN, - AuditLogOperation.CONTROL, - true), - () -> - String.format( - DataNodeMiscMessages - .EXCEPTION_ARG_LOGIN_STATUS_ARG_USER_ARG_ID_ARG_OPENS_SESSION_ARG_70586F85, - IoTDBConstant.GLOBAL_DB_NAME, - response.getMessage(), - username, - userId, - session)); - LOGGER.info( - DataNodeMiscMessages.LOGIN_STATUS, - IoTDBConstant.GLOBAL_DB_NAME, - response.getMessage(), - username, - session); - updateIdleTime(); - return response; - } - - private void logFailedLogin( - IClientSession session, long userId, String username, TSStatus status) { - AUDIT_LOGGER.log( - new AuditLogFields( - userId, - username, - session.getClientAddress(), - AuditEventType.LOGIN, - AuditLogOperation.CONTROL, - false), - () -> - String.format( - DataNodeMiscMessages - .EXCEPTION_USER_ARG_ID_ARG_LOGIN_FAILED_WITH_CODE_ARG_ARG_93FAE4FD, - username, - userId, - status.getCode(), - status.getMessage())); + return openSessionResp; } public boolean closeSession(IClientSession session, LongConsumer releaseByQueryId) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/handler/InternalServiceThriftHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/handler/InternalServiceThriftHandler.java index 831e12ea17bbd..58130eb91cc58 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/handler/InternalServiceThriftHandler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/handler/InternalServiceThriftHandler.java @@ -21,27 +21,19 @@ import org.apache.iotdb.commons.service.NoopServerContext; import org.apache.iotdb.commons.service.metric.MetricService; -import org.apache.iotdb.db.protocol.session.ClientSession; -import org.apache.iotdb.db.protocol.session.SessionManager; -import org.apache.iotdb.db.protocol.thrift.impl.IClientRPCServiceWithHandler; -import org.apache.iotdb.rpc.TElasticFramedTransport; import org.apache.thrift.protocol.TProtocol; import org.apache.thrift.server.ServerContext; import org.apache.thrift.server.TServerEventHandler; -import org.apache.thrift.transport.TSocket; import org.apache.thrift.transport.TTransport; -import java.net.Socket; import java.util.concurrent.atomic.AtomicLong; public class InternalServiceThriftHandler implements TServerEventHandler { private AtomicLong thriftConnectionNumber = new AtomicLong(0); - private final IClientRPCServiceWithHandler eventHandler; - public InternalServiceThriftHandler(IClientRPCServiceWithHandler eventHandler) { - this.eventHandler = eventHandler; + public InternalServiceThriftHandler() { MetricService.getInstance() .addMetricSet(new InternalServiceThriftHandlerMetrics(thriftConnectionNumber)); } @@ -54,21 +46,13 @@ public void preServe() { @Override public ServerContext createContext(TProtocol tProtocol, TProtocol tProtocol1) { thriftConnectionNumber.incrementAndGet(); - Socket socket = - ((TSocket) ((TElasticFramedTransport) tProtocol1.getTransport()).getSocket()).getSocket(); - SessionManager.getInstance().registerSession(new ClientSession(socket)); return NoopServerContext.INSTANCE; } @Override public void deleteContext( ServerContext serverContext, TProtocol tProtocol, TProtocol tProtocol1) { - try { - eventHandler.handleClientExit(); - } finally { - SessionManager.getInstance().removeCurrSession(); - thriftConnectionNumber.decrementAndGet(); - } + thriftConnectionNumber.decrementAndGet(); } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java index 3e8927bdf6057..fcb8fc23999ba 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java @@ -1626,7 +1626,16 @@ public TSOpenSessionResp openSession(TSOpenSessionReq req) throws TException { } Optional database = parseDatabase(req); IClientSession clientSession = SESSION_MANAGER.getCurrSession(); - BasicOpenSessionResp openSessionResp = login(clientSession, req, clientVersion, sqlDialect); + BasicOpenSessionResp openSessionResp = + SESSION_MANAGER.login( + SESSION_MANAGER.getCurrSession(), + req.username, + req.password, + req.zoneId, + req.client_protocol, + clientVersion, + sqlDialect, + parseUseEncryptedPassword(req)); TSStatus tsStatus = RpcUtils.getStatus(openSessionResp.getCode(), openSessionResp.getMessage()); if (tsStatus.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) { @@ -1642,22 +1651,6 @@ public TSOpenSessionResp openSession(TSOpenSessionReq req) throws TException { return resp.setSessionId(openSessionResp.getSessionId()).setConfiguration(configuration); } - protected BasicOpenSessionResp login( - IClientSession session, - TSOpenSessionReq req, - IoTDBConstant.ClientVersion clientVersion, - SqlDialect sqlDialect) { - return SESSION_MANAGER.login( - session, - req.username, - req.password, - req.zoneId, - req.client_protocol, - clientVersion, - sqlDialect, - parseUseEncryptedPassword(req)); - } - private IoTDBConstant.ClientVersion parseClientVersion(TSOpenSessionReq req) { Map configuration = req.configuration; if (configuration != null && configuration.containsKey("version")) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java index 2523be59d7096..4b1e05f4e857d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java @@ -77,7 +77,6 @@ import org.apache.iotdb.commons.pipe.agent.plugin.meta.PipePluginMeta; import org.apache.iotdb.commons.pipe.agent.task.meta.PipeMeta; import org.apache.iotdb.commons.queryengine.common.SessionInfo; -import org.apache.iotdb.commons.queryengine.common.SqlDialect; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; @@ -120,7 +119,6 @@ import org.apache.iotdb.db.i18n.DataNodeSchemaMessages; import org.apache.iotdb.db.partition.DataPartitionTableGenerator; import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; -import org.apache.iotdb.db.protocol.basic.BasicOpenSessionResp; import org.apache.iotdb.db.protocol.client.ConfigNodeInfo; import org.apache.iotdb.db.protocol.client.cn.DnToCnInternalServiceAsyncRequestManager; import org.apache.iotdb.db.protocol.client.cn.DnToCnRequestType; @@ -352,7 +350,6 @@ import org.apache.iotdb.rpc.TSStatusCode; import org.apache.iotdb.rpc.subscription.exception.SubscriptionException; import org.apache.iotdb.service.rpc.thrift.TSInsertRecordReq; -import org.apache.iotdb.service.rpc.thrift.TSOpenSessionReq; import org.apache.iotdb.trigger.api.enums.FailureStrategy; import org.apache.iotdb.trigger.api.enums.TriggerEvent; @@ -419,8 +416,7 @@ import static org.apache.iotdb.db.utils.ErrorHandlingUtils.onIoTDBException; import static org.apache.iotdb.db.utils.ErrorHandlingUtils.onQueryException; -public class DataNodeInternalRPCServiceImpl extends ClientRPCServiceImpl - implements IDataNodeRPCService.Iface { +public class DataNodeInternalRPCServiceImpl implements IDataNodeRPCService.Iface { private static final Logger LOGGER = LoggerFactory.getLogger(DataNodeInternalRPCServiceImpl.class); @@ -477,8 +473,6 @@ public class DataNodeInternalRPCServiceImpl extends ClientRPCServiceImpl private static final String SYSTEM = "system"; - private long consensusWaitTimeoutSeconds = 30; - public DataNodeInternalRPCServiceImpl(DataNodeContext dataNodeContext) { this(dataNodeContext, SystemMetrics.getInstance()); } @@ -491,6 +485,8 @@ public DataNodeInternalRPCServiceImpl(DataNodeContext dataNodeContext) { this.dataNodeContext = dataNodeContext; } + private long consensusWaitTimeoutSeconds = 30; + private TSStatus waitForConsensusStarted() { if (dataNodeContext.isAllConsensusStarted()) { return null; @@ -512,16 +508,6 @@ private TSStatus waitForConsensusStarted() { } } - @Override - protected BasicOpenSessionResp login( - IClientSession session, - TSOpenSessionReq req, - ClientVersion clientVersion, - SqlDialect sqlDialect) { - return SESSION_MANAGER.loginInternal( - session, req.username, req.zoneId, req.client_protocol, clientVersion, sqlDialect); - } - @Override public TSendFragmentInstanceResp sendFragmentInstance(final TSendFragmentInstanceReq req) { LOGGER.debug(DataNodeMiscMessages.RECEIVE_FRAGMENT_INSTANCE, req.getConsensusGroupId()); @@ -1853,10 +1839,9 @@ private TSStatus executeInternalSchemaTask( @Override public TSStatus executeCQ(TExecuteCQ req) { - IClientSession previousSession = SESSION_MANAGER.getCurrSession(); IClientSession session = new InternalClientSession(req.cqId); - SESSION_MANAGER.setCurrSession(session); + SESSION_MANAGER.registerSession(session); SESSION_MANAGER.supplySession( session, @@ -1938,11 +1923,8 @@ public TSStatus executeCQ(TExecuteCQ req) { } catch (Exception e) { return onQueryException(e, "\"" + executedSQL + "\". " + OperationType.EXECUTE_STATEMENT); } finally { - try { - SESSION_MANAGER.closeSession(session, COORDINATOR::cleanupQueryExecution); - } finally { - SESSION_MANAGER.restoreSession(previousSession, session); - } + SESSION_MANAGER.closeSession(session, COORDINATOR::cleanupQueryExecution); + SESSION_MANAGER.removeCurrSession(); } } @@ -2397,6 +2379,11 @@ private List testAllDataNodeExternalServiceConnection( handler, TEST_CONNECTION_RETRY_NUM, TEST_CONNECTION_TIMEOUT_MS, true)); } + @Override + public TSStatus testConnectionEmptyRPC() throws TException { + return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()); + } + @Override public TSStatus updateClusterTopology(TUpdateClusterTopologyReq req) throws TException { clusterTopology.updateTopology(req.getDataNodes(), req.getTopology()); @@ -2823,7 +2810,7 @@ public TSStatus loadConfiguration() throws TException { } @Override - public TShowConfigurationResp showConfigurationInternal() { + public TShowConfigurationResp showConfiguration() { TShowConfigurationResp resp = new TShowConfigurationResp(RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS), ""); try { @@ -3575,7 +3562,7 @@ public TSStatus stopAndClearDataNode() { } @Override - public TSStatus insertRecordInternal(TSInsertRecordReq req) { + public TSStatus insertRecord(TSInsertRecordReq req) throws TException { try { InsertRowStatement statement = StatementGenerator.createStatement(req); SessionInfo sessionInfo = @@ -3619,6 +3606,10 @@ public TSStatus writeAuditLog(TAuditLogReq req) { return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS); } + public void handleClientExit() { + // Do nothing + } + // ==================================================== // Data Partition Table Integrity Check Implementation // ==================================================== diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java index ef0afbf295651..f9b9c6e010110 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java @@ -78,7 +78,7 @@ public void initThriftServiceThread() throws IllegalAccessException { config.getRpcMaxConcurrentClientNum(), config.getThriftServerAwaitTimeForStopService(), new TrustedChannelAuditServerEventHandler( - new InternalServiceThriftHandler(getImpl()), + new InternalServiceThriftHandler(), new TEndPoint(getBindIP(), getBindPort()), DNAuditLogger.getInstance()::recordTrustedChannelFailureAuditLogIfNecessary), config.isRpcThriftCompressionEnable(), @@ -95,7 +95,7 @@ public void initThriftServiceThread() throws IllegalAccessException { getBindPort(), config.getRpcMaxConcurrentClientNum(), config.getThriftServerAwaitTimeForStopService(), - new InternalServiceThriftHandler(getImpl()), + new InternalServiceThriftHandler(), config.isRpcThriftCompressionEnable(), DeepCopyRpcTransportFactory.INSTANCE); } catch (RPCServiceException e) { diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift index 7415b96bd652e..57afa70e72846 100644 --- a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift +++ b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift @@ -931,7 +931,7 @@ struct TFetchDeviceEntrySegmentResp { * END: Used for EXPLAIN ANALYZE **/ -service IDataNodeRPCService extends client.IClientRPCService { +service IDataNodeRPCService { // -----------------------------------For Data Node----------------------------------------------- @@ -1186,7 +1186,7 @@ service IDataNodeRPCService extends client.IClientRPCService { common.TSStatus clearCache(set cacheClearOptions) - common.TShowConfigurationResp showConfigurationInternal() + common.TShowConfigurationResp showConfiguration() common.TShowAppliedConfigurationsResp showAppliedConfigurations() @@ -1434,14 +1434,17 @@ service IDataNodeRPCService extends client.IClientRPCService { common.TTestConnectionResp submitInternalTestConnectionTask(common.TNodeLocations nodeLocations) + /** Empty rpc, only for connection test */ + common.TSStatus testConnectionEmptyRPC() + /** * Push cluster topology to this DataNode. * Each DataNode receives only its own reachable set. */ common.TSStatus updateClusterTopology(1:TUpdateClusterTopologyReq req) - /** Write internal audit records independently of client session identity. */ - common.TSStatus insertRecordInternal(1:client.TSInsertRecordReq req); + /** to write audit log or other events as time series **/ + common.TSStatus insertRecord(1:client.TSInsertRecordReq req); /** * Write an audit log entry to the DataNode's AuditEventLogger From e90897169d9512e48abe30479a68cbb315f970ad Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 5 Nov 2026 18:12:19 +0800 Subject: [PATCH 05/32] Use multiplexed client RPC service on internal endpoint Preserve IDataNodeRPCService as the default internal processor and expose IClientRPCService through a service-name prefix on the same endpoint. Add internal session, table pool, subscription, reconnect, and SSL support with passwordless internal login constrained by account policy. --- .../isession/IInternalSessionBuilder.java | 32 +++++ .../apache/iotdb/rpc/RpcServiceConstants.java | 29 ++++ .../iotdb/session/i18n/SessionMessages.java | 6 + .../iotdb/session/i18n/SessionMessages.java | 6 + .../apache/iotdb/session/InternalSession.java | 75 ++++++++++ .../org/apache/iotdb/session/Session.java | 21 ++- .../iotdb/session/SessionConnection.java | 13 +- .../iotdb/session/TableSessionBuilder.java | 28 +++- .../pool/InternalTableSessionPool.java | 39 +++++ .../iotdb/session/pool/SessionPool.java | 136 +++++++++--------- .../session/pool/TableSessionPoolBuilder.java | 21 ++- .../iotdb/session/InternalSessionTest.java | 63 ++++++++ .../subscription/config/ConsumerConstant.java | 7 + .../InternalSubscriptionTableSession.java | 86 +++++++++++ .../SubscriptionSessionWrapper.java | 24 +++- .../SubscriptionTableSession.java | 4 + .../SubscriptionTableSessionBuilder.java | 54 ++++++- .../SubscriptionTreeSessionBuilder.java | 30 ++++ .../base/AbstractSubscriptionConsumer.java | 39 +++++ .../AbstractSubscriptionConsumerBuilder.java | 37 +++++ .../base/AbstractSubscriptionProvider.java | 14 ++ .../base/AbstractSubscriptionProviders.java | 106 ++++++++++++++ ...InternalSubscriptionTablePullConsumer.java | 80 +++++++++++ .../SubscriptionTablePullConsumerBuilder.java | 47 +++++- .../SubscriptionTablePushConsumerBuilder.java | 36 +++++ .../SubscriptionTreePullConsumerBuilder.java | 36 +++++ .../SubscriptionTreePushConsumerBuilder.java | 36 +++++ .../db/protocol/session/SessionManager.java | 73 ++++++++++ .../InternalClientRPCServiceProcessor.java | 56 ++++++++ .../handler/InternalServiceThriftHandler.java | 19 ++- .../thrift/impl/ClientRPCServiceImpl.java | 36 +++-- .../impl/InternalClientRPCServiceImpl.java | 29 ++++ .../service/DataNodeInternalRPCService.java | 27 +++- .../InternalServiceThriftHandlerTest.java | 57 ++++++++ 34 files changed, 1302 insertions(+), 100 deletions(-) create mode 100644 iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java create mode 100644 iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/RpcServiceConstants.java create mode 100644 iotdb-client/session/src/main/java/org/apache/iotdb/session/InternalSession.java create mode 100644 iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.java create mode 100644 iotdb-client/session/src/test/java/org/apache/iotdb/session/InternalSessionTest.java create mode 100644 iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/InternalSubscriptionTableSession.java create mode 100644 iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/InternalSubscriptionTablePullConsumer.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/InternalClientRPCServiceProcessor.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/InternalClientRPCServiceImpl.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/thrift/handler/InternalServiceThriftHandlerTest.java diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java new file mode 100644 index 0000000000000..ae67fa4093da0 --- /dev/null +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java @@ -0,0 +1,32 @@ +/* + * 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.rpc.IoTDBConnectionException; + +/** Capability implemented by builders that can construct sessions on internal endpoints. */ +public interface IInternalSessionBuilder { + + /** + * Builds a session using the caller-provided internal node supplier. Whether the returned session + * is already open follows the builder's ordinary build contract. + */ + T buildInternal(INodeSupplier nodeSupplier) throws IoTDBConnectionException; +} diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/RpcServiceConstants.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/RpcServiceConstants.java new file mode 100644 index 0000000000000..69ffda9602fa0 --- /dev/null +++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/RpcServiceConstants.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.rpc; + +/** Shared constants for Thrift service names. */ +public final class RpcServiceConstants { + + /** Multiplexed service name used to reach the client RPC service on the internal endpoint. */ + public static final String INTERNAL_CLIENT_RPC_SERVICE = "IClientRPCService"; + + private RpcServiceConstants() {} +} diff --git a/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java b/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java index ded0deac31484..e256a062265a2 100644 --- a/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java +++ b/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java @@ -21,6 +21,12 @@ public final class SessionMessages { + public static final String EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081 = + "No DataNode internal endpoint is available"; + + public static final String EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392 = + "Internal node supplier must not be null"; + // Session public static final String NODE_URLS_EMPTY = "nodeUrls shouldn't be empty."; public static final String REDIRECT_TWICE_LOG = "redirect twice"; diff --git a/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java b/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java index 838a7e2148329..32d4eda5083ef 100644 --- a/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java +++ b/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java @@ -21,6 +21,12 @@ public final class SessionMessages { + public static final String EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081 = + "没有可用的 DataNode 内部端点"; + + public static final String EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392 = + "内部节点 supplier 不能为空"; + // Session public static final String NODE_URLS_EMPTY = "nodeUrls 不能为空。"; public static final String REDIRECT_TWICE_LOG = "重定向了两次"; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/InternalSession.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/InternalSession.java new file mode 100644 index 0000000000000..504ab6005eacf --- /dev/null +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/InternalSession.java @@ -0,0 +1,75 @@ +/* + * 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.INodeSupplier; +import org.apache.iotdb.rpc.IoTDBConnectionException; +import org.apache.iotdb.rpc.RpcServiceConstants; +import org.apache.iotdb.session.i18n.SessionMessages; + +import java.time.ZoneId; +import java.util.List; +import java.util.Objects; + +/** Session using caller-supplied internal endpoints and the caller's username. */ +public class InternalSession extends Session { + + private final INodeSupplier nodeSupplier; + + public InternalSession(AbstractSessionBuilder builder, INodeSupplier nodeSupplier) { + super(builder); + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + password = null; + useEncryptedPassword = false; + enableAutoFetch = false; + enableRedirection = false; + enableQueryRedirection = false; + rpcServiceName = RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE; + } + + @Override + public synchronized void open(boolean enableRPCCompaction, int connectionTimeoutInMs) + throws IoTDBConnectionException { + super.open(enableRPCCompaction, connectionTimeoutInMs, null, null, nodeSupplier); + } + + @Override + public SessionConnection constructSessionConnection( + Session session, TEndPoint endpoint, ZoneId zoneId) throws IoTDBConnectionException { + final List candidates = nodeSupplier.get(); + if (candidates.isEmpty()) { + throw new IoTDBConnectionException( + SessionMessages.EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081); + } + defaultEndPoint = + candidates.contains(endpoint) + ? endpoint.deepCopy() + : nodeSupplier + .getQueryEndPoint() + .filter(candidates::contains) + .orElse(candidates.get(0)) + .deepCopy(); + return super.constructSessionConnection(session, defaultEndPoint, zoneId); + } +} 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 de3ee417afc0d..5263f0af4bd47 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 @@ import org.apache.iotdb.common.rpc.thrift.TAggregationType; import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.isession.IInternalSessionBuilder; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.ISession; import org.apache.iotdb.isession.SessionConfig; @@ -187,6 +188,9 @@ public class Session implements ISession { protected INodeSupplier availableNodes; + /** Multiplexed Thrift service name used by internal sessions. */ + protected String rpcServiceName; + protected boolean enableQueryRedirection = false; // The version number of the client which used for compatibility in the server @@ -4437,6 +4441,10 @@ protected void setDefaultSessionConnection(SessionConnection defaultSessionConne this.defaultSessionConnection = defaultSessionConnection; } + protected void setRpcServiceName(String rpcServiceName) { + this.rpcServiceName = rpcServiceName; + } + public boolean isEnableThriftRpcCompaction() { return enableThriftRpcCompaction; } @@ -4445,7 +4453,8 @@ public void setEnableThriftRpcCompaction(boolean enableThriftRpcCompaction) { this.enableThriftRpcCompaction = enableThriftRpcCompaction; } - public static class Builder extends AbstractSessionBuilder { + public static class Builder extends AbstractSessionBuilder + implements IInternalSessionBuilder { public Builder host(String host) { this.host = host; @@ -4595,5 +4604,15 @@ public Session build() { } return new Session(this); } + + /** Builds an unopened internal session using caller-supplied internal endpoints. */ + @Override + public InternalSession buildInternal(INodeSupplier nodeSupplier) { + return createInternalSession(nodeSupplier); + } + + protected InternalSession createInternalSession(INodeSupplier nodeSupplier) { + return new InternalSession(this, nodeSupplier); + } } } 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 49ecfef949440..3548924a29dd3 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 @@ -74,6 +74,8 @@ import org.apache.thrift.TException; import org.apache.thrift.protocol.TBinaryProtocol; import org.apache.thrift.protocol.TCompactProtocol; +import org.apache.thrift.protocol.TMultiplexedProtocol; +import org.apache.thrift.protocol.TProtocol; import org.apache.thrift.transport.TTransport; import org.apache.thrift.transport.TTransportException; import org.apache.tsfile.utils.Pair; @@ -231,11 +233,14 @@ private void init( throw new IoTDBConnectionException(e); } - if (session.enableThriftRpcCompaction) { - client = new IClientRPCService.Client(new TCompactProtocol(transport)); - } else { - client = new IClientRPCService.Client(new TBinaryProtocol(transport)); + TProtocol protocol = + session.enableThriftRpcCompaction + ? new TCompactProtocol(transport) + : new TBinaryProtocol(transport); + if (session.rpcServiceName != null) { + protocol = new TMultiplexedProtocol(protocol, session.rpcServiceName); } + client = new IClientRPCService.Client(protocol); client = RpcUtils.newSynchronizedClient(client); TSOpenSessionReq openReq = new TSOpenSessionReq(); diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java index a39af4c847a95..e5a90efda7169 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java @@ -19,6 +19,8 @@ package org.apache.iotdb.session; +import org.apache.iotdb.isession.IInternalSessionBuilder; +import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.ITableSession; import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.rpc.IoTDBConnectionException; @@ -41,7 +43,8 @@ * *

All configurations have reasonable default values, which can be overridden as needed. */ -public class TableSessionBuilder extends AbstractSessionBuilder { +public class TableSessionBuilder extends AbstractSessionBuilder + implements IInternalSessionBuilder { private boolean enableCompression = false; private int connectionTimeoutInMs = SessionConfig.DEFAULT_CONNECTION_TIMEOUT_MS; @@ -432,21 +435,40 @@ public TableSessionBuilder tabletCompressionMinRowSize(int tabletCompressionMinR * @throws IoTDBConnectionException if an error occurs while establishing the connection. */ public ITableSession build() throws IoTDBConnectionException { + return wrapTableSession(constructNewSession()); + } + + /** Builds and opens a table session using caller-supplied internal endpoints. */ + @Override + public ITableSession buildInternal(INodeSupplier nodeSupplier) throws IoTDBConnectionException { + this.sqlDialect = TABLE; + return wrapTableSession(openSession(new InternalSession(this, nodeSupplier))); + } + + protected Session constructNewSession() throws IoTDBConnectionException { if (nodeUrls == null) { this.nodeUrls = Collections.singletonList(SessionConfig.DEFAULT_HOST + ":" + SessionConfig.DEFAULT_PORT); } this.sqlDialect = TABLE; - Session newSession = new Session(this); + return openSession(new Session(this)); + } + + protected ITableSession wrapTableSession(Session session) { + return new TableSession(session); + } + + private Session openSession(Session newSession) throws IoTDBConnectionException { newSession.enableIoTDBRpcCompression = isIoTDBRpcCompressionEnabled; newSession.tabletCompressionMinRowSize = tabletCompressionMinRowSize; try { newSession.open(isThriftRpcCompactionEnabled, connectionTimeoutInMs); + return newSession; } catch (IoTDBConnectionException e) { newSession.close(); throw e; } - return new TableSession(newSession); } + } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.java new file mode 100644 index 0000000000000..be84296f18f8e --- /dev/null +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.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.pool; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.session.i18n.SessionMessages; + +import java.util.Objects; + +/** Table session pool using caller-supplied internal endpoints and the caller's username. */ +public class InternalTableSessionPool extends TableSessionPool { + + protected InternalTableSessionPool( + AbstractSessionPoolBuilder builder, INodeSupplier nodeSupplier) { + super( + new SessionPool( + builder, + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392))); + } +} 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 16b79c6415c4b..5fef4471221f6 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 @@ -173,6 +173,8 @@ public class SessionPool implements ISessionPool { private INodeSupplier availableNodes; + private INodeSupplier internalNodeSupplier; + // set to true, means that we will start a background thread to fetch all available (Status is // not Removing) datanodes in cluster, and these available nodes will be used in retrying stage private boolean enableAutoFetch = true; @@ -526,15 +528,20 @@ public SessionPool( } public SessionPool(AbstractSessionPoolBuilder builder) { + this(builder, null); + } + + protected SessionPool(AbstractSessionPoolBuilder builder, INodeSupplier internalNodeSupplier) { + this.internalNodeSupplier = internalNodeSupplier; this.maxSize = builder.maxSize; this.user = builder.username; - this.password = builder.pw; + this.password = internalNodeSupplier == null ? builder.pw : null; this.fetchSize = builder.fetchSize; this.waitToGetSessionTimeoutInMs = builder.waitToGetSessionTimeoutInMs; this.enableThriftCompression = builder.isThriftRpcCompactionEnabled; this.enableIoTDBRpcCompression = builder.isIoTDBRpcCompressionEnabled; this.zoneId = builder.zoneId; - this.enableRedirection = builder.enableRedirection; + this.enableRedirection = internalNodeSupplier == null && builder.enableRedirection; if (this.enableRedirection) { deviceIdToEndpoint = new ConcurrentHashMap<>(); tableModelDeviceIdToEndpoint = new ConcurrentHashMap<>(); @@ -544,7 +551,7 @@ public SessionPool(AbstractSessionPoolBuilder builder) { this.version = builder.version; this.thriftDefaultBufferSize = builder.thriftDefaultBufferSize; this.thriftMaxFrameSize = builder.thriftMaxFrameSize; - this.enableAutoFetch = builder.enableAutoFetch; + this.enableAutoFetch = internalNodeSupplier == null && builder.enableAutoFetch; this.useSSL = builder.useSSL; this.trustStore = builder.trustStore; this.trustStorePwd = builder.trustStorePwd; @@ -561,7 +568,13 @@ public SessionPool(AbstractSessionPoolBuilder builder) { if (enableAutoFetch) { initThreadPool(); } - if (builder.nodeUrls != null) { + if (internalNodeSupplier != null) { + this.host = null; + this.port = -1; + this.nodeUrls = null; + this.formattedNodeUrls = internalNodeSupplier.getClass().getSimpleName(); + this.availableNodes = internalNodeSupplier; + } else if (builder.nodeUrls != null) { if (builder.nodeUrls.isEmpty()) { throw new IllegalArgumentException(SessionMessages.NODE_URLS_EMPTY); } @@ -589,72 +602,53 @@ public SessionPool(AbstractSessionPoolBuilder builder) { } } - private Session constructNewSession() { - Session session; + protected ISession constructNewSession() { + Session.Builder builder = + new Session.Builder() + .username(user) + .password(password) + .fetchSize(fetchSize) + .zoneId(zoneId) + .thriftDefaultBufferSize(thriftDefaultBufferSize) + .thriftMaxFrameSize(thriftMaxFrameSize) + .enableRedirection(enableRedirection) + .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) + .version(version) + .useSSL(useSSL) + .trustStore(trustStore) + .trustStorePwd(trustStorePwd) + .keyStore(keyStore) + .keyStorePwd(keyStorePwd) + .sslProtocol(sslProtocol) + .maxRetryCount(maxRetryCount) + .retryIntervalInMs(retryIntervalInMs) + .sqlDialect(sqlDialect) + .database(database) + .enableAutoCreateSchema(enableAutoCreateSchema) + .timeOut(queryTimeoutInMs) + .enableIoTDBRpcCompression(enableIoTDBRpcCompression) + .enableThriftRpcCompression(enableThriftCompression); + if (isInternalSessionPool()) { + return builder.buildInternal(internalNodeSupplier); + } if (nodeUrls == null) { - // Construct custom Session - session = - new Session.Builder() - .host(host) - .port(port) - .username(user) - .password(password) - .fetchSize(fetchSize) - .zoneId(zoneId) - .thriftDefaultBufferSize(thriftDefaultBufferSize) - .thriftMaxFrameSize(thriftMaxFrameSize) - .enableRedirection(enableRedirection) - .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) - .version(version) - .useSSL(useSSL) - .trustStore(trustStore) - .trustStorePwd(trustStorePwd) - .keyStore(keyStore) - .keyStorePwd(keyStorePwd) - .sslProtocol(sslProtocol) - .maxRetryCount(maxRetryCount) - .retryIntervalInMs(retryIntervalInMs) - .sqlDialect(sqlDialect) - .database(database) - .enableAutoCreateSchema(enableAutoCreateSchema) - .timeOut(queryTimeoutInMs) - .enableIoTDBRpcCompression(enableIoTDBRpcCompression) - .enableThriftRpcCompression(enableThriftCompression) - .build(); + builder.host(host).port(port); } else { - // Construct redirect-able Session - session = - new Session.Builder() - .nodeUrls(nodeUrls) - .username(user) - .password(password) - .fetchSize(fetchSize) - .zoneId(zoneId) - .thriftDefaultBufferSize(thriftDefaultBufferSize) - .thriftMaxFrameSize(thriftMaxFrameSize) - .enableRedirection(enableRedirection) - .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) - .version(version) - .useSSL(useSSL) - .trustStore(trustStore) - .trustStorePwd(trustStorePwd) - .keyStore(keyStore) - .keyStorePwd(keyStorePwd) - .sslProtocol(sslProtocol) - .maxRetryCount(maxRetryCount) - .retryIntervalInMs(retryIntervalInMs) - .sqlDialect(sqlDialect) - .database(database) - .enableAutoCreateSchema(enableAutoCreateSchema) - .timeOut(queryTimeoutInMs) - .enableIoTDBRpcCompression(enableIoTDBRpcCompression) - .enableThriftRpcCompression(enableThriftCompression) - .build(); + builder.nodeUrls(nodeUrls); } + Session session = builder.build(); session.setEnableQueryRedirection(enableQueryRedirection); return session; } + protected boolean isInternalSessionPool() { + return internalNodeSupplier != null; + } + + protected INodeSupplier getInternalNodeSupplier() { + return internalNodeSupplier; + } + private void initThreadPool() { this.executorService = Executors.newSingleThreadScheduledExecutor( @@ -849,7 +843,9 @@ public synchronized void close() { this.executorService = null; } if (availableNodes != null) { - this.availableNodes.close(); + if (!isInternalSessionPool()) { + this.availableNodes.close(); + } this.availableNodes = null; } LOGGER.info(SessionMessages.CLOSING_SESSION_POOL); @@ -3544,16 +3540,16 @@ public boolean isEnableThriftCompression() { @Override public void setEnableRedirection(boolean enableRedirection) { - this.enableRedirection = enableRedirection; + this.enableRedirection = !isInternalSessionPool() && enableRedirection; if (this.enableRedirection) { deviceIdToEndpoint = new ConcurrentHashMap<>(); tableModelDeviceIdToEndpoint = new ConcurrentHashMap<>(); } for (ISession session : queue) { - session.setEnableRedirection(enableRedirection); + session.setEnableRedirection(this.enableRedirection); } for (ISession session : occupied.keySet()) { - session.setEnableRedirection(enableRedirection); + session.setEnableRedirection(this.enableRedirection); } } @@ -3564,12 +3560,12 @@ public boolean isEnableRedirection() { @Override public void setEnableQueryRedirection(boolean enableQueryRedirection) { - this.enableQueryRedirection = enableQueryRedirection; + this.enableQueryRedirection = !isInternalSessionPool() && enableQueryRedirection; for (ISession session : queue) { - session.setEnableQueryRedirection(enableQueryRedirection); + session.setEnableQueryRedirection(this.enableQueryRedirection); } for (ISession session : occupied.keySet()) { - session.setEnableQueryRedirection(enableQueryRedirection); + session.setEnableQueryRedirection(this.enableQueryRedirection); } } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java index f19c5fc8878a6..ff2502e3220ce 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java @@ -19,12 +19,16 @@ package org.apache.iotdb.session.pool; +import org.apache.iotdb.isession.IInternalSessionBuilder; +import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.isession.pool.ITableSessionPool; +import org.apache.iotdb.session.i18n.SessionMessages; import java.time.ZoneId; import java.util.Collections; import java.util.List; +import java.util.Objects; /** * A builder class for constructing instances of {@link ITableSessionPool}. @@ -34,7 +38,8 @@ * *

All configurations have reasonable default values, which can be overridden as needed. */ -public class TableSessionPoolBuilder extends AbstractSessionPoolBuilder { +public class TableSessionPoolBuilder extends AbstractSessionPoolBuilder + implements IInternalSessionBuilder { /** * Sets the list of node URLs for the IoTDB cluster. @@ -346,4 +351,18 @@ public ITableSessionPool build() { SessionPool sessionPool = new SessionPool(this); return new TableSessionPool(sessionPool); } + + /** Builds a pool using caller-supplied internal endpoints; connections open lazily. */ + @Override + public ITableSessionPool buildInternal(INodeSupplier nodeSupplier) { + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + this.sqlDialect = "table"; + return createInternalTableSessionPool(nodeSupplier); + } + + protected ITableSessionPool createInternalTableSessionPool(INodeSupplier nodeSupplier) { + return new InternalTableSessionPool(this, nodeSupplier); + } } diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/InternalSessionTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/InternalSessionTest.java new file mode 100644 index 0000000000000..2c79bc764c0e8 --- /dev/null +++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/InternalSessionTest.java @@ -0,0 +1,63 @@ +/* + * 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.isession.INodeSupplier; +import org.apache.iotdb.isession.pool.ITableSessionPool; +import org.apache.iotdb.rpc.RpcServiceConstants; +import org.apache.iotdb.session.pool.TableSessionPoolBuilder; + +import org.junit.Test; +import org.powermock.reflect.Whitebox; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; + +public class InternalSessionTest { + + @Test + public void testInternalTableSessionPoolDoesNotCloseCallerSupplier() { + final INodeSupplier nodeSupplier = mock(INodeSupplier.class); + final ITableSessionPool pool = new TableSessionPoolBuilder().buildInternal(nodeSupplier); + + pool.close(); + + verify(nodeSupplier, never()).close(); + } + + @Test + public void testBuildInternalConfiguresMultiplexedTransport() { + final INodeSupplier nodeSupplier = mock(INodeSupplier.class); + final InternalSession session = + new Session.Builder().host("127.0.0.1").port(10730).buildInternal(nodeSupplier); + + assertEquals( + RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE, + Whitebox.getInternalState(session, "rpcServiceName")); + assertNull(Whitebox.getInternalState(session, "password")); + assertFalse((Boolean) Whitebox.getInternalState(session, "enableAutoFetch")); + assertFalse((Boolean) Whitebox.getInternalState(session, "enableRedirection")); + assertFalse((Boolean) Whitebox.getInternalState(session, "enableQueryRedirection")); + } +} diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/config/ConsumerConstant.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/config/ConsumerConstant.java index 6ecb56ba63c69..2e38e7b088f44 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/config/ConsumerConstant.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/rpc/subscription/config/ConsumerConstant.java @@ -38,6 +38,13 @@ public class ConsumerConstant { public static final String PASSWORD_KEY = "password"; public static final String ENCRYPTED_PASSWORD_KEY = "encrypted-password"; + public static final String USE_SSL_KEY = "use-ssl"; + public static final String TRUST_STORE_KEY = "trust-store"; + public static final String TRUST_STORE_PWD_KEY = "trust-store-pwd"; + public static final String KEY_STORE_KEY = "key-store"; + public static final String KEY_STORE_PWD_KEY = "key-store-pwd"; + public static final String SSL_PROTOCOL_KEY = "ssl-protocol"; + public static final String CONSUMER_ID_KEY = "consumer-id"; public static final String CONSUMER_GROUP_ID_KEY = "group-id"; public static final String CONSUMER_INSTANCE_ID_KEY = "consumer-instance-id"; diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/InternalSubscriptionTableSession.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/InternalSubscriptionTableSession.java new file mode 100644 index 0000000000000..16e2bd6f20cec --- /dev/null +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/InternalSubscriptionTableSession.java @@ -0,0 +1,86 @@ +/* + * 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.subscription; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.rpc.IoTDBConnectionException; +import org.apache.iotdb.rpc.RpcServiceConstants; +import org.apache.iotdb.session.AbstractSessionBuilder; +import org.apache.iotdb.session.Session; +import org.apache.iotdb.session.SessionConnection; +import org.apache.iotdb.session.i18n.SessionMessages; + +import java.time.ZoneId; +import java.util.List; +import java.util.Objects; + +/** Table subscription-management session using the caller's identity and internal endpoints. */ +public class InternalSubscriptionTableSession extends SubscriptionTableSession { + + public InternalSubscriptionTableSession( + AbstractSessionBuilder builder, INodeSupplier nodeSupplier) { + super(new InternalSessionWrapper(builder, nodeSupplier)); + } + + private static class InternalSessionWrapper extends SubscriptionSessionWrapper { + + private final INodeSupplier nodeSupplier; + + private InternalSessionWrapper(AbstractSessionBuilder builder, INodeSupplier nodeSupplier) { + super(builder); + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + password = null; + useEncryptedPassword = false; + enableAutoFetch = false; + enableRedirection = false; + enableQueryRedirection = false; + rpcServiceName = RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE; + } + + @Override + public synchronized void open(boolean enableRPCCompaction, int connectionTimeoutInMs) + throws IoTDBConnectionException { + super.open(enableRPCCompaction, connectionTimeoutInMs, null, null, nodeSupplier); + } + + @Override + public SessionConnection constructSessionConnection( + Session session, TEndPoint endpoint, ZoneId zoneId) throws IoTDBConnectionException { + final List candidates = nodeSupplier.get(); + if (candidates.isEmpty()) { + throw new IoTDBConnectionException( + SessionMessages.EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081); + } + defaultEndPoint = + candidates.contains(endpoint) + ? endpoint.deepCopy() + : nodeSupplier + .getQueryEndPoint() + .filter(candidates::contains) + .orElse(candidates.get(0)) + .deepCopy(); + return super.constructSessionConnection(session, defaultEndPoint, zoneId); + } + } +} diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java index d7e715fa62762..6538eea370686 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java @@ -30,14 +30,14 @@ import java.time.ZoneId; import java.util.Objects; -public final class SubscriptionSessionWrapper extends Session { +public class SubscriptionSessionWrapper extends Session { public SubscriptionSessionWrapper(final AbstractSessionBuilder builder) { super(builder); } public void open() throws IoTDBConnectionException { - super.open(enableThriftRpcCompaction, connectionTimeoutInMs); + open(enableThriftRpcCompaction, connectionTimeoutInMs); } public void close() throws IoTDBConnectionException { @@ -70,4 +70,24 @@ public SubscriptionSessionConnection getSessionConnection() throws IoTDBConnecti public int getThriftMaxFrameSize() { return thriftMaxFrameSize; } + + @Override + public void setRpcServiceName(final String rpcServiceName) { + super.setRpcServiceName(rpcServiceName); + } + + public void setSslConfig( + final boolean useSSL, + final String trustStore, + final String trustStorePwd, + final String keyStore, + final String keyStorePwd, + final String sslProtocol) { + this.useSSL = useSSL; + this.trustStore = trustStore; + this.trustStorePwd = trustStorePwd; + this.keyStore = keyStore; + this.keyStorePwd = keyStorePwd; + this.sslProtocol = sslProtocol; + } } diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSession.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSession.java index 123deea2fa45d..8949c46ff83ad 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSession.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSession.java @@ -36,6 +36,10 @@ public SubscriptionTableSession(final AbstractSessionBuilder builder) { super(new SubscriptionSessionWrapper(builder)); } + protected SubscriptionTableSession(final SubscriptionSessionWrapper session) { + super(session); + } + /////////////////////////////// open & close /////////////////////////////// @Override diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java index e0c8e95968dc9..f54a5aae25ad5 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java @@ -19,10 +19,13 @@ package org.apache.iotdb.session.subscription; +import org.apache.iotdb.isession.IInternalSessionBuilder; +import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.rpc.IoTDBConnectionException; import org.apache.iotdb.session.AbstractSessionBuilder; -public class SubscriptionTableSessionBuilder extends AbstractSessionBuilder { +public class SubscriptionTableSessionBuilder extends AbstractSessionBuilder + implements IInternalSessionBuilder { public SubscriptionTableSessionBuilder() { // use table model @@ -63,6 +66,36 @@ public SubscriptionTableSessionBuilder thriftMaxFrameSize(final int thriftMaxFra return this; } + public SubscriptionTableSessionBuilder useSSL(final boolean useSSL) { + super.useSSL = useSSL; + return this; + } + + public SubscriptionTableSessionBuilder trustStore(final String trustStore) { + super.trustStore = trustStore; + return this; + } + + public SubscriptionTableSessionBuilder trustStorePwd(final String trustStorePwd) { + super.trustStorePwd = trustStorePwd; + return this; + } + + public SubscriptionTableSessionBuilder keyStore(final String keyStore) { + super.keyStore = keyStore; + return this; + } + + public SubscriptionTableSessionBuilder keyStorePwd(final String keyStorePwd) { + super.keyStorePwd = keyStorePwd; + return this; + } + + public SubscriptionTableSessionBuilder sslProtocol(final String sslProtocol) { + super.sslProtocol = sslProtocol; + return this; + } + public SubscriptionTableSessionBuilder connectionTimeoutInMs(final int connectionTimeoutInMs) { super.connectionTimeoutInMs = connectionTimeoutInMs; return this; @@ -73,4 +106,23 @@ public ISubscriptionTableSession build() throws IoTDBConnectionException { session.open(); return session; } + + /** Builds and opens a table subscription-management session on internal endpoints. */ + @Override + public ISubscriptionTableSession buildInternal(INodeSupplier nodeSupplier) + throws IoTDBConnectionException { + final ISubscriptionTableSession session = + new InternalSubscriptionTableSession(this, nodeSupplier); + try { + session.open(); + } catch (IoTDBConnectionException e) { + try { + session.close(); + } catch (Exception closeException) { + e.addSuppressed(closeException); + } + throw e; + } + return session; + } } diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTreeSessionBuilder.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTreeSessionBuilder.java index 098b718bc6932..bf90be85871c0 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTreeSessionBuilder.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTreeSessionBuilder.java @@ -62,6 +62,36 @@ public SubscriptionTreeSessionBuilder thriftMaxFrameSize(final int thriftMaxFram return this; } + public SubscriptionTreeSessionBuilder useSSL(final boolean useSSL) { + super.useSSL = useSSL; + return this; + } + + public SubscriptionTreeSessionBuilder trustStore(final String trustStore) { + super.trustStore = trustStore; + return this; + } + + public SubscriptionTreeSessionBuilder trustStorePwd(final String trustStorePwd) { + super.trustStorePwd = trustStorePwd; + return this; + } + + public SubscriptionTreeSessionBuilder keyStore(final String keyStore) { + super.keyStore = keyStore; + return this; + } + + public SubscriptionTreeSessionBuilder keyStorePwd(final String keyStorePwd) { + super.keyStorePwd = keyStorePwd; + return this; + } + + public SubscriptionTreeSessionBuilder sslProtocol(final String sslProtocol) { + super.sslProtocol = sslProtocol; + return this; + } + public SubscriptionTreeSessionBuilder connectionTimeoutInMs(final int connectionTimeoutInMs) { super.connectionTimeoutInMs = connectionTimeoutInMs; return this; diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java index 7f60b82a0e225..438c38ffc6d4d 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java @@ -20,6 +20,7 @@ package org.apache.iotdb.session.subscription.consumer.base; import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.rpc.subscription.config.ConsumerConstant; import org.apache.iotdb.rpc.subscription.config.TopicConfig; @@ -147,6 +148,13 @@ abstract class AbstractSubscriptionConsumer implements AutoCloseable { private final int connectionTimeoutInMs; private final int maxPollParallelism; + private final boolean useSSL; + private final String trustStore; + private final String trustStorePwd; + private final String keyStore; + private final String keyStorePwd; + private final String sslProtocol; + /** * The latest watermark timestamp received from the server. Updated when WATERMARK events are * processed and stripped. Consumer users can query this to check timestamp progress. @@ -251,6 +259,13 @@ protected AbstractSubscriptionConsumer(final AbstractSubscriptionConsumerBuilder this.thriftMaxFrameSize = builder.thriftMaxFrameSize; this.connectionTimeoutInMs = builder.connectionTimeoutInMs; this.maxPollParallelism = builder.maxPollParallelism; + + this.useSSL = builder.useSSL; + this.trustStore = builder.trustStore; + this.trustStorePwd = builder.trustStorePwd; + this.keyStore = builder.keyStore; + this.keyStorePwd = builder.keyStorePwd; + this.sslProtocol = builder.sslProtocol; } protected AbstractSubscriptionConsumer( @@ -269,6 +284,18 @@ protected AbstractSubscriptionConsumer( properties.getOrDefault( ConsumerConstant.PASSWORD_KEY, SessionConfig.DEFAULT_PASSWORD)) .encryptedPassword((String) properties.get(ConsumerConstant.ENCRYPTED_PASSWORD_KEY)) + .useSSL( + Boolean.parseBoolean( + String.valueOf( + properties.getOrDefault(ConsumerConstant.USE_SSL_KEY, Boolean.FALSE)))) + .trustStore((String) properties.get(ConsumerConstant.TRUST_STORE_KEY)) + .trustStorePwd((String) properties.get(ConsumerConstant.TRUST_STORE_PWD_KEY)) + .keyStore((String) properties.get(ConsumerConstant.KEY_STORE_KEY)) + .keyStorePwd((String) properties.get(ConsumerConstant.KEY_STORE_PWD_KEY)) + .sslProtocol( + (String) + properties.getOrDefault( + ConsumerConstant.SSL_PROTOCOL_KEY, SessionConfig.DEFAULT_SSL_PROTOCOL)) .consumerId((String) properties.get(ConsumerConstant.CONSUMER_ID_KEY)) .consumerGroupId((String) properties.get(ConsumerConstant.CONSUMER_GROUP_ID_KEY)) .ownerId((String) properties.get(ConsumerConstant.OWNER_ID_KEY)) @@ -572,6 +599,16 @@ private void seekAfterInternalTopicProgress( /////////////////////////////// subscription provider /////////////////////////////// + /** Returns the complete internal endpoint source, or null to use server discovery. */ + protected INodeSupplier getInternalNodeSupplier() { + return null; + } + + /** Returns the multiplexed client RPC service name for internal providers. */ + protected String getRpcServiceName() { + return null; + } + protected abstract AbstractSubscriptionProvider constructSubscriptionProvider( final TEndPoint endPoint, final String username, @@ -601,6 +638,8 @@ AbstractSubscriptionProvider constructProviderAndHandshake(final TEndPoint endPo this.heartbeatIntervalMs, this.connectionTimeoutInMs); provider.setConsumerInstanceId(consumerInstanceId); + provider.setRpcServiceName(getRpcServiceName()); + provider.setSslConfig(useSSL, trustStore, trustStorePwd, keyStore, keyStorePwd, sslProtocol); try { provider.handshake(); } catch (final Exception e) { diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumerBuilder.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumerBuilder.java index a9e9d38f334c9..357173eb3604f 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumerBuilder.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumerBuilder.java @@ -55,6 +55,13 @@ public class AbstractSubscriptionConsumerBuilder { protected int connectionTimeoutInMs = ConsumerConstant.CONNECTION_TIMEOUT_MS_DEFAULT_VALUE; protected int maxPollParallelism = ConsumerConstant.MAX_POLL_PARALLELISM_DEFAULT_VALUE; + protected boolean useSSL = false; + protected String trustStore; + protected String trustStorePwd; + protected String keyStore; + protected String keyStorePwd; + protected String sslProtocol = SessionConfig.DEFAULT_SSL_PROTOCOL; + public AbstractSubscriptionConsumerBuilder host(final String host) { this.host = host; return this; @@ -171,6 +178,36 @@ public AbstractSubscriptionConsumerBuilder connectionTimeoutInMs( return this; } + public AbstractSubscriptionConsumerBuilder useSSL(final boolean useSSL) { + this.useSSL = useSSL; + return this; + } + + public AbstractSubscriptionConsumerBuilder trustStore(final String trustStore) { + this.trustStore = trustStore; + return this; + } + + public AbstractSubscriptionConsumerBuilder trustStorePwd(final String trustStorePwd) { + this.trustStorePwd = trustStorePwd; + return this; + } + + public AbstractSubscriptionConsumerBuilder keyStore(final String keyStore) { + this.keyStore = keyStore; + return this; + } + + public AbstractSubscriptionConsumerBuilder keyStorePwd(final String keyStorePwd) { + this.keyStorePwd = keyStorePwd; + return this; + } + + public AbstractSubscriptionConsumerBuilder sslProtocol(final String sslProtocol) { + this.sslProtocol = sslProtocol; + return this; + } + public AbstractSubscriptionConsumerBuilder maxPollParallelism(final int maxPollParallelism) { // Here the minimum value of max poll parallelism is set to 1 instead of 0, in order to use a // single thread to execute poll whenever there are idle resources available, thereby diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java index a70acb0761691..46ff7448be8ed 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProvider.java @@ -179,6 +179,20 @@ SubscriptionSessionConnection getSessionConnection() throws IoTDBConnectionExcep return session.getSessionConnection(); } + void setRpcServiceName(final String rpcServiceName) { + session.setRpcServiceName(rpcServiceName); + } + + void setSslConfig( + final boolean useSSL, + final String trustStore, + final String trustStorePwd, + final String keyStore, + final String keyStorePwd, + final String sslProtocol) { + session.setSslConfig(useSSL, trustStore, trustStorePwd, keyStore, keyStorePwd, sslProtocol); + } + private TPipeSubscribeResp pipeSubscribe(final TPipeSubscribeReq req) throws TException, IoTDBConnectionException { rpcLock.lock(); diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java index afb5285b5f577..43f05dcf100b0 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionProviders.java @@ -34,6 +34,7 @@ import java.util.ArrayList; import java.util.LinkedHashMap; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.Objects; @@ -92,6 +93,11 @@ void openProviders(final AbstractSubscriptionConsumer consumer) throws Subscript // close stale providers closeProviders(false); + if (consumer.getInternalNodeSupplier() != null) { + syncInternalProviders(consumer); + return; + } + final Map connectionFailures = new LinkedHashMap<>(); final List connectionFailureCauses = new ArrayList<>(); for (final TEndPoint endPoint : initialEndpoints) { @@ -457,6 +463,16 @@ private void syncInternal(final AbstractSubscriptionConsumer consumer) { if (consumer.isClosed() || consumer.isFenced()) { return; } + if (consumer.getInternalNodeSupplier() != null) { + try { + syncInternalProviders(consumer); + } catch (final SubscriptionConsumerFencedException e) { + consumer.fence(e); + } catch (final Exception e) { + LOGGER.warn(SubscriptionMessages.FETCH_ENDPOINTS_FAILED, consumer, e, e); + } + return; + } if (hasNoAvailableProviders()) { try { openProviders(consumer); @@ -567,6 +583,96 @@ private void syncInternal(final AbstractSubscriptionConsumer consumer) { } } + /** Reconciles providers against a complete endpoint list without server-side discovery. */ + private void syncInternalProviders(final AbstractSubscriptionConsumer consumer) { + final Set endpoints = + consumer.getInternalNodeSupplier().get().stream() + .map(TEndPoint::deepCopy) + .collect(Collectors.toCollection(LinkedHashSet::new)); + + for (final AbstractSubscriptionProvider provider : getAllProviders()) { + if (!provider.isAvailable() || !endpoints.contains(provider.getEndPoint())) { + try { + closeAndRemoveProvider(provider.getDataNodeId()); + } catch (final Exception e) { + LOGGER.warn(SubscriptionMessages.PROVIDER_CLOSE_FAILED, provider, e, e); + } + continue; + } + try { + consumer.subscribedTopics = provider.heartbeat().getTopics(); + provider.setAvailable(); + } catch (final SubscriptionConsumerFencedException e) { + consumer.fence(e); + throw e; + } catch (final Exception e) { + LOGGER.warn( + SubscriptionMessages + .LOG_ARG_FAILED_SENDING_HEARTBEAT_SUBSCRIPTION_PROVIDER_ARG_BECAUSE_ARG_SET_0B38FB1F, + consumer, + provider, + e, + e); + provider.setUnavailable(); + try { + closeAndRemoveProvider(provider.getDataNodeId()); + } catch (final Exception closeException) { + LOGGER.warn( + SubscriptionMessages.PROVIDER_CLOSE_FAILED, provider, closeException, closeException); + } + } + } + + final Set connectedEndpoints = + getAllAvailableProviders().stream() + .map(AbstractSubscriptionProvider::getEndPoint) + .collect(Collectors.toSet()); + Exception lastFailure = null; + for (final TEndPoint endpoint : endpoints) { + if (isClosing.get() || consumer.isFenced()) { + return; + } + if (connectedEndpoints.contains(endpoint)) { + continue; + } + try { + final AbstractSubscriptionProvider provider = + consumer.constructProviderAndHandshake(endpoint); + if (containsProvider(provider.getDataNodeId())) { + try { + closeAndRemoveProvider(provider.getDataNodeId()); + } catch (final Exception e) { + LOGGER.warn(SubscriptionMessages.PROVIDER_CLOSE_FAILED, provider, e, e); + } + } + addProvider(provider.getDataNodeId(), provider); + } catch (final SubscriptionConsumerFencedException e) { + consumer.fence(e); + throw e; + } catch (final Exception e) { + lastFailure = e; + LOGGER.warn( + SubscriptionMessages.LOG_ARG_FAILED_CREATE_CONNECTION_ARG_BECAUSE_ARG_E536E22A, + consumer, + endpoint, + e, + e); + } + } + + if (hasNoAvailableProviders()) { + throw new SubscriptionConnectionException( + String.format( + SubscriptionMessages + .EXCEPTION_CLUSTER_HAS_NO_AVAILABLE_SUBSCRIPTION_PROVIDERS_CONNECT_INITIAL_ENDPOINTS_ARG_5DB83198, + endpoints), + lastFailure); + } + if (!subscriptionProviders.containsKey(nextDataNodeId)) { + nextDataNodeId = subscriptionProviders.firstKey(); + } + } + @Override public String toString() { final StringBuilder sb = new StringBuilder(); diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/InternalSubscriptionTablePullConsumer.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/InternalSubscriptionTablePullConsumer.java new file mode 100644 index 0000000000000..7b2b22e2f5b08 --- /dev/null +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/InternalSubscriptionTablePullConsumer.java @@ -0,0 +1,80 @@ +/* + * 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.subscription.consumer.table; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.rpc.RpcServiceConstants; +import org.apache.iotdb.session.i18n.SessionMessages; +import org.apache.iotdb.session.subscription.consumer.base.AbstractSubscriptionProvider; + +import java.util.Objects; + +/** Table pull consumer using the caller's complete internal endpoint source and username. */ +public class InternalSubscriptionTablePullConsumer extends SubscriptionTablePullConsumer { + + private final INodeSupplier nodeSupplier; + + public InternalSubscriptionTablePullConsumer( + final SubscriptionTablePullConsumerBuilder builder, final INodeSupplier nodeSupplier) { + super(builder); + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + } + + @Override + protected INodeSupplier getInternalNodeSupplier() { + return nodeSupplier; + } + + @Override + protected String getRpcServiceName() { + return RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE; + } + + @Override + protected AbstractSubscriptionProvider constructSubscriptionProvider( + final TEndPoint endPoint, + final String username, + final String password, + final String encryptedPassword, + final String consumerId, + final String consumerGroupId, + final String ownerId, + final Long ownerEpoch, + final int thriftMaxFrameSize, + final long heartbeatIntervalMs, + final int connectionTimeoutInMs) { + return super.constructSubscriptionProvider( + endPoint, + username, + "", + null, + consumerId, + consumerGroupId, + ownerId, + ownerEpoch, + thriftMaxFrameSize, + heartbeatIntervalMs, + connectionTimeoutInMs); + } +} diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java index efc99debb4977..a710551bae33c 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java @@ -19,12 +19,15 @@ package org.apache.iotdb.session.subscription.consumer.table; +import org.apache.iotdb.isession.IInternalSessionBuilder; +import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.session.subscription.consumer.ISubscriptionTablePullConsumer; import org.apache.iotdb.session.subscription.consumer.base.AbstractSubscriptionPullConsumerBuilder; import java.util.List; -public class SubscriptionTablePullConsumerBuilder extends AbstractSubscriptionPullConsumerBuilder { +public class SubscriptionTablePullConsumerBuilder extends AbstractSubscriptionPullConsumerBuilder + implements IInternalSessionBuilder { @Override public SubscriptionTablePullConsumerBuilder host(final String host) { @@ -130,6 +133,42 @@ public SubscriptionTablePullConsumerBuilder connectionTimeoutInMs( return this; } + @Override + public SubscriptionTablePullConsumerBuilder useSSL(final boolean useSSL) { + super.useSSL(useSSL); + return this; + } + + @Override + public SubscriptionTablePullConsumerBuilder trustStore(final String trustStore) { + super.trustStore(trustStore); + return this; + } + + @Override + public SubscriptionTablePullConsumerBuilder trustStorePwd(final String trustStorePwd) { + super.trustStorePwd(trustStorePwd); + return this; + } + + @Override + public SubscriptionTablePullConsumerBuilder keyStore(final String keyStore) { + super.keyStore(keyStore); + return this; + } + + @Override + public SubscriptionTablePullConsumerBuilder keyStorePwd(final String keyStorePwd) { + super.keyStorePwd(keyStorePwd); + return this; + } + + @Override + public SubscriptionTablePullConsumerBuilder sslProtocol(final String sslProtocol) { + super.sslProtocol(sslProtocol); + return this; + } + @Override public SubscriptionTablePullConsumerBuilder maxPollParallelism(final int maxPollParallelism) { super.maxPollParallelism(maxPollParallelism); @@ -152,4 +191,10 @@ public SubscriptionTablePullConsumerBuilder autoCommitIntervalMs( public ISubscriptionTablePullConsumer build() { return new SubscriptionTablePullConsumer(this); } + + /** Builds an unopened consumer using the caller's internal endpoint source and username. */ + @Override + public ISubscriptionTablePullConsumer buildInternal(final INodeSupplier nodeSupplier) { + return new InternalSubscriptionTablePullConsumer(this, nodeSupplier); + } } diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePushConsumerBuilder.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePushConsumerBuilder.java index 143b34056af9e..55e9633656c80 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePushConsumerBuilder.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePushConsumerBuilder.java @@ -132,6 +132,42 @@ public SubscriptionTablePushConsumerBuilder connectionTimeoutInMs( return this; } + @Override + public SubscriptionTablePushConsumerBuilder useSSL(final boolean useSSL) { + super.useSSL(useSSL); + return this; + } + + @Override + public SubscriptionTablePushConsumerBuilder trustStore(final String trustStore) { + super.trustStore(trustStore); + return this; + } + + @Override + public SubscriptionTablePushConsumerBuilder trustStorePwd(final String trustStorePwd) { + super.trustStorePwd(trustStorePwd); + return this; + } + + @Override + public SubscriptionTablePushConsumerBuilder keyStore(final String keyStore) { + super.keyStore(keyStore); + return this; + } + + @Override + public SubscriptionTablePushConsumerBuilder keyStorePwd(final String keyStorePwd) { + super.keyStorePwd(keyStorePwd); + return this; + } + + @Override + public SubscriptionTablePushConsumerBuilder sslProtocol(final String sslProtocol) { + super.sslProtocol(sslProtocol); + return this; + } + @Override public SubscriptionTablePushConsumerBuilder maxPollParallelism(final int maxPollParallelism) { super.maxPollParallelism(maxPollParallelism); diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/tree/SubscriptionTreePullConsumerBuilder.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/tree/SubscriptionTreePullConsumerBuilder.java index 2d057b7bfbd80..f6761e6db8366 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/tree/SubscriptionTreePullConsumerBuilder.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/tree/SubscriptionTreePullConsumerBuilder.java @@ -130,6 +130,42 @@ public SubscriptionTreePullConsumerBuilder connectionTimeoutInMs( return this; } + @Override + public SubscriptionTreePullConsumerBuilder useSSL(final boolean useSSL) { + super.useSSL(useSSL); + return this; + } + + @Override + public SubscriptionTreePullConsumerBuilder trustStore(final String trustStore) { + super.trustStore(trustStore); + return this; + } + + @Override + public SubscriptionTreePullConsumerBuilder trustStorePwd(final String trustStorePwd) { + super.trustStorePwd(trustStorePwd); + return this; + } + + @Override + public SubscriptionTreePullConsumerBuilder keyStore(final String keyStore) { + super.keyStore(keyStore); + return this; + } + + @Override + public SubscriptionTreePullConsumerBuilder keyStorePwd(final String keyStorePwd) { + super.keyStorePwd(keyStorePwd); + return this; + } + + @Override + public SubscriptionTreePullConsumerBuilder sslProtocol(final String sslProtocol) { + super.sslProtocol(sslProtocol); + return this; + } + @Override public SubscriptionTreePullConsumerBuilder maxPollParallelism(final int maxPollParallelism) { super.maxPollParallelism(maxPollParallelism); diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/tree/SubscriptionTreePushConsumerBuilder.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/tree/SubscriptionTreePushConsumerBuilder.java index 9067f3e7ae564..50bd5b82a6c92 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/tree/SubscriptionTreePushConsumerBuilder.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/tree/SubscriptionTreePushConsumerBuilder.java @@ -132,6 +132,42 @@ public SubscriptionTreePushConsumerBuilder connectionTimeoutInMs( return this; } + @Override + public SubscriptionTreePushConsumerBuilder useSSL(final boolean useSSL) { + super.useSSL(useSSL); + return this; + } + + @Override + public SubscriptionTreePushConsumerBuilder trustStore(final String trustStore) { + super.trustStore(trustStore); + return this; + } + + @Override + public SubscriptionTreePushConsumerBuilder trustStorePwd(final String trustStorePwd) { + super.trustStorePwd(trustStorePwd); + return this; + } + + @Override + public SubscriptionTreePushConsumerBuilder keyStore(final String keyStore) { + super.keyStore(keyStore); + return this; + } + + @Override + public SubscriptionTreePushConsumerBuilder keyStorePwd(final String keyStorePwd) { + super.keyStorePwd(keyStorePwd); + return this; + } + + @Override + public SubscriptionTreePushConsumerBuilder sslProtocol(final String sslProtocol) { + super.sslProtocol(sslProtocol); + return this; + } + @Override public SubscriptionTreePushConsumerBuilder maxPollParallelism(final int maxPollParallelism) { super.maxPollParallelism(maxPollParallelism); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java index 85d9acae276ec..4495c0d296153 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java @@ -226,6 +226,79 @@ public BasicOpenSessionResp login( return openSessionResp; } + /** Opens a session without password verification for trusted internal client RPC callers. */ + public BasicOpenSessionResp loginInternal( + final IClientSession session, + final String username, + final String zoneId, + final TSProtocolVersion tsProtocolVersion, + final IoTDBConstant.ClientVersion clientVersion, + final SqlDialect sqlDialect) { + final BasicOpenSessionResp openSessionResp = new BasicOpenSessionResp(); + final long userId = AuthorityChecker.getUserId(username).orElse(-1L); + if (userId == -1) { + final TSStatus status = + RpcUtils.getStatus( + TSStatusCode.USER_NOT_EXIST, + String.format( + DataNodeMiscMessages.MISC_EXCEPTION_USER_S_DOES_NOT_EXIST_0CE725D8, username)); + openSessionResp.sessionId(-1).setCode(status.getCode()).setMessage(status.getMessage()); + return openSessionResp; + } + + final TSStatus loginPolicyStatus = + AuthorityChecker.getAccessControl().allowUserToLogin(username); + if (loginPolicyStatus.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + openSessionResp + .sessionId(-1) + .setCode(loginPolicyStatus.getCode()) + .setMessage(loginPolicyStatus.getMessage()); + return openSessionResp; + } + + if (!tsProtocolVersion.equals(CURRENT_RPC_VERSION)) { + openSessionResp + .sessionId(-1) + .setCode(TSStatusCode.INCOMPATIBLE_VERSION.getStatusCode()) + .setMessage( + DataNodeMiscMessages.VERSION_INCOMPATIBLE_PLEASE_UPGRADE_TO + IoTDBConstant.VERSION); + return openSessionResp; + } + + session.setSqlDialect(sqlDialect); + supplySession(session, userId, username, ZoneId.of(zoneId), clientVersion); + final String logInMessage = DataNodeMiscMessages.EXCEPTION_LOGIN_SUCCESSFULLY_1C2EF670; + openSessionResp + .sessionId(session.getId()) + .setCode(TSStatusCode.SUCCESS_STATUS.getStatusCode()) + .setMessage(logInMessage); + AUDIT_LOGGER.log( + new AuditLogFields( + userId, + username, + session.getClientAddress(), + AuditEventType.LOGIN, + AuditLogOperation.CONTROL, + true), + () -> + String.format( + DataNodeMiscMessages + .EXCEPTION_ARG_LOGIN_STATUS_ARG_USER_ARG_ID_ARG_OPENS_SESSION_ARG_70586F85, + IoTDBConstant.GLOBAL_DB_NAME, + openSessionResp.getMessage(), + username, + userId, + session)); + LOGGER.info( + DataNodeMiscMessages.LOGIN_STATUS, + IoTDBConstant.GLOBAL_DB_NAME, + openSessionResp.getMessage(), + username, + session); + updateIdleTime(); + return openSessionResp; + } + public boolean closeSession(IClientSession session, LongConsumer releaseByQueryId) { return closeSession(session, releaseByQueryId, true); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/InternalClientRPCServiceProcessor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/InternalClientRPCServiceProcessor.java new file mode 100644 index 0000000000000..6417c41fe5144 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/InternalClientRPCServiceProcessor.java @@ -0,0 +1,56 @@ +/* + * 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.db.protocol.thrift; + +import org.apache.iotdb.db.protocol.session.ClientSession; +import org.apache.iotdb.db.protocol.session.SessionManager; +import org.apache.iotdb.db.protocol.thrift.impl.IClientRPCServiceWithHandler; +import org.apache.iotdb.rpc.TElasticFramedTransport; + +import org.apache.thrift.TException; +import org.apache.thrift.protocol.TProtocol; +import org.apache.thrift.transport.TSocket; +import org.apache.thrift.transport.TTransport; + +import java.net.Socket; + +/** Registers a client session only for connections using the multiplexed client RPC service. */ +public class InternalClientRPCServiceProcessor extends ProcessorWithMetrics { + + public InternalClientRPCServiceProcessor(IClientRPCServiceWithHandler iface) { + super(iface); + } + + @Override + public void process(TProtocol in, TProtocol out) throws TException { + registerSessionIfNecessary(out); + super.process(in, out); + } + + private void registerSessionIfNecessary(TProtocol protocol) { + if (SessionManager.getInstance().getCurrSession() != null) { + return; + } + + TTransport socketTransport = ((TElasticFramedTransport) protocol.getTransport()).getSocket(); + Socket socket = ((TSocket) socketTransport).getSocket(); + SessionManager.getInstance().registerSession(new ClientSession(socket)); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/handler/InternalServiceThriftHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/handler/InternalServiceThriftHandler.java index 58130eb91cc58..f215bdbd16a6d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/handler/InternalServiceThriftHandler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/handler/InternalServiceThriftHandler.java @@ -21,6 +21,9 @@ import org.apache.iotdb.commons.service.NoopServerContext; import org.apache.iotdb.commons.service.metric.MetricService; +import org.apache.iotdb.db.protocol.session.IClientSession; +import org.apache.iotdb.db.protocol.session.SessionManager; +import org.apache.iotdb.db.protocol.thrift.impl.IClientRPCServiceWithHandler; import org.apache.thrift.protocol.TProtocol; import org.apache.thrift.server.ServerContext; @@ -31,9 +34,11 @@ public class InternalServiceThriftHandler implements TServerEventHandler { - private AtomicLong thriftConnectionNumber = new AtomicLong(0); + private final AtomicLong thriftConnectionNumber = new AtomicLong(0); + private final IClientRPCServiceWithHandler eventHandler; - public InternalServiceThriftHandler() { + public InternalServiceThriftHandler(IClientRPCServiceWithHandler eventHandler) { + this.eventHandler = eventHandler; MetricService.getInstance() .addMetricSet(new InternalServiceThriftHandlerMetrics(thriftConnectionNumber)); } @@ -52,7 +57,15 @@ public ServerContext createContext(TProtocol tProtocol, TProtocol tProtocol1) { @Override public void deleteContext( ServerContext serverContext, TProtocol tProtocol, TProtocol tProtocol1) { - thriftConnectionNumber.decrementAndGet(); + try { + IClientSession session = SessionManager.getInstance().getCurrSession(); + if (session != null) { + eventHandler.handleClientExit(); + } + } finally { + SessionManager.getInstance().removeCurrSession(); + thriftConnectionNumber.decrementAndGet(); + } } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java index fcb8fc23999ba..52a39c4d86f5f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java @@ -1626,16 +1626,7 @@ public TSOpenSessionResp openSession(TSOpenSessionReq req) throws TException { } Optional database = parseDatabase(req); IClientSession clientSession = SESSION_MANAGER.getCurrSession(); - BasicOpenSessionResp openSessionResp = - SESSION_MANAGER.login( - SESSION_MANAGER.getCurrSession(), - req.username, - req.password, - req.zoneId, - req.client_protocol, - clientVersion, - sqlDialect, - parseUseEncryptedPassword(req)); + BasicOpenSessionResp openSessionResp = login(clientSession, req, clientVersion, sqlDialect); TSStatus tsStatus = RpcUtils.getStatus(openSessionResp.getCode(), openSessionResp.getMessage()); if (tsStatus.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) { @@ -1651,6 +1642,31 @@ public TSOpenSessionResp openSession(TSOpenSessionReq req) throws TException { return resp.setSessionId(openSessionResp.getSessionId()).setConfiguration(configuration); } + protected BasicOpenSessionResp login( + IClientSession session, + TSOpenSessionReq req, + IoTDBConstant.ClientVersion clientVersion, + SqlDialect sqlDialect) { + if (useInternalSessionLogin()) { + return SESSION_MANAGER.loginInternal( + session, req.username, req.zoneId, req.client_protocol, clientVersion, sqlDialect); + } + return SESSION_MANAGER.login( + session, + req.username, + req.password, + null, + req.zoneId, + req.client_protocol, + clientVersion, + sqlDialect, + parseUseEncryptedPassword(req)); + } + + protected boolean useInternalSessionLogin() { + return false; + } + private IoTDBConstant.ClientVersion parseClientVersion(TSOpenSessionReq req) { Map configuration = req.configuration; if (configuration != null && configuration.containsKey("version")) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/InternalClientRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/InternalClientRPCServiceImpl.java new file mode 100644 index 0000000000000..788eac9a376b5 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/InternalClientRPCServiceImpl.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.db.protocol.thrift.impl; + +/** Client RPC handler exposed through the multiplexed internal endpoint. */ +public class InternalClientRPCServiceImpl extends ClientRPCServiceImpl { + + @Override + protected boolean useInternalSessionLogin() { + return true; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java index f9b9c6e010110..666cdc47a50da 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java @@ -32,12 +32,18 @@ import org.apache.iotdb.db.audit.DNAuditLogger; import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.protocol.thrift.InternalClientRPCServiceProcessor; import org.apache.iotdb.db.protocol.thrift.handler.InternalServiceThriftHandler; import org.apache.iotdb.db.protocol.thrift.impl.DataNodeInternalRPCServiceImpl; +import org.apache.iotdb.db.protocol.thrift.impl.IClientRPCServiceWithHandler; +import org.apache.iotdb.db.protocol.thrift.impl.InternalClientRPCServiceImpl; import org.apache.iotdb.db.service.DataNode.DataNodeContext; import org.apache.iotdb.db.service.metrics.DataNodeInternalRPCServiceMetrics; import org.apache.iotdb.mpp.rpc.thrift.IDataNodeRPCService.Processor; import org.apache.iotdb.rpc.DeepCopyRpcTransportFactory; +import org.apache.iotdb.rpc.RpcServiceConstants; + +import org.apache.thrift.TMultiplexedProcessor; import java.util.concurrent.atomic.AtomicReference; @@ -48,9 +54,11 @@ public class DataNodeInternalRPCService extends ThriftService private static final CommonConfig commonConfig = CommonDescriptor.getInstance().getConfig(); protected final AtomicReference impl = new AtomicReference<>(); + protected final AtomicReference clientImpl = + new AtomicReference<>(); protected DataNodeContext dataNodeContext; - private DataNodeInternalRPCService() {} + protected DataNodeInternalRPCService() {} @Override public ServiceType getID() { @@ -60,8 +68,14 @@ public ServiceType getID() { @Override public void initTProcessor() { DataNodeInternalRPCServiceImpl service = getImpl(); + IClientRPCServiceWithHandler clientService = getClientImpl(); initSyncedServiceImpl(null); - processor = new Processor<>(service); + TMultiplexedProcessor multiplexedProcessor = new TMultiplexedProcessor(); + multiplexedProcessor.registerProcessor( + RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE, + new InternalClientRPCServiceProcessor(clientService)); + multiplexedProcessor.registerDefault(new Processor<>(service)); + processor = multiplexedProcessor; } @Override @@ -78,7 +92,7 @@ public void initThriftServiceThread() throws IllegalAccessException { config.getRpcMaxConcurrentClientNum(), config.getThriftServerAwaitTimeForStopService(), new TrustedChannelAuditServerEventHandler( - new InternalServiceThriftHandler(), + new InternalServiceThriftHandler(getClientImpl()), new TEndPoint(getBindIP(), getBindPort()), DNAuditLogger.getInstance()::recordTrustedChannelFailureAuditLogIfNecessary), config.isRpcThriftCompressionEnable(), @@ -95,7 +109,7 @@ public void initThriftServiceThread() throws IllegalAccessException { getBindPort(), config.getRpcMaxConcurrentClientNum(), config.getThriftServerAwaitTimeForStopService(), - new InternalServiceThriftHandler(), + new InternalServiceThriftHandler(getClientImpl()), config.isRpcThriftCompressionEnable(), DeepCopyRpcTransportFactory.INSTANCE); } catch (RPCServiceException e) { @@ -121,6 +135,11 @@ public DataNodeInternalRPCServiceImpl getImpl() { return impl.get(); } + protected IClientRPCServiceWithHandler getClientImpl() { + clientImpl.compareAndSet(null, new InternalClientRPCServiceImpl()); + return clientImpl.get(); + } + private static class DataNodeInternalRPCServiceHolder { private static final DataNodeInternalRPCService INSTANCE = new DataNodeInternalRPCService(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/thrift/handler/InternalServiceThriftHandlerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/thrift/handler/InternalServiceThriftHandlerTest.java new file mode 100644 index 0000000000000..7ee06d084667d --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/thrift/handler/InternalServiceThriftHandlerTest.java @@ -0,0 +1,57 @@ +/* + * 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.db.protocol.thrift.handler; + +import org.apache.iotdb.db.protocol.session.IClientSession; +import org.apache.iotdb.db.protocol.session.SessionManager; +import org.apache.iotdb.db.protocol.thrift.impl.IClientRPCServiceWithHandler; + +import org.junit.Test; + +import static org.junit.Assert.assertNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; + +public class InternalServiceThriftHandlerTest { + + @Test + public void testDeleteContextCleansUpInvalidatedSession() { + final IClientRPCServiceWithHandler clientHandler = mock(IClientRPCServiceWithHandler.class); + final InternalServiceThriftHandler handler = new InternalServiceThriftHandler(clientHandler); + final SessionManager sessionManager = SessionManager.getInstance(); + final IClientSession previousSession = sessionManager.getCurrSession(); + final IClientSession session = mock(IClientSession.class); + + handler.createContext(null, null); + sessionManager.setCurrSession(session); + try { + handler.deleteContext(null, null, null); + + verify(clientHandler).handleClientExit(); + assertNull(sessionManager.getCurrSession()); + } finally { + if (previousSession != null) { + sessionManager.setCurrSession(previousSession); + } else { + sessionManager.removeCurrSession(); + } + } + } +} From d012cc015f2cdbc4f0bbabc7abd93f70a79e426e Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 5 Nov 2026 18:16:57 +0800 Subject: [PATCH 06/32] Fix Spotless formatting for internal session builders --- .../java/org/apache/iotdb/session/TableSessionBuilder.java | 1 - .../org/apache/iotdb/session/pool/TableSessionPoolBuilder.java | 3 +-- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java index e5a90efda7169..66c52b011a3a0 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java @@ -470,5 +470,4 @@ private Session openSession(Session newSession) throws IoTDBConnectionException throw e; } } - } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java index ff2502e3220ce..d0104675d0a8d 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java @@ -356,8 +356,7 @@ public ITableSessionPool build() { @Override public ITableSessionPool buildInternal(INodeSupplier nodeSupplier) { Objects.requireNonNull( - nodeSupplier, - SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + nodeSupplier, SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); this.sqlDialect = "table"; return createInternalTableSessionPool(nodeSupplier); } From 1f1971cab556f25fd285daf9e818e7377d458c58 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 5 Nov 2026 18:21:18 +0800 Subject: [PATCH 07/32] Fix SessionPool internal session type --- .../main/java/org/apache/iotdb/session/pool/SessionPool.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 5fef4471221f6..b86804c8c142e 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 @@ -872,7 +872,7 @@ public void closeResultSet(SessionDataSetWrapper wrapper) { @SuppressWarnings({"squid:S2446"}) private void tryConstructNewSession() { - Session session = constructNewSession(); + ISession session = constructNewSession(); try { session.open( enableThriftCompression, From f9e559a1091fbbd1d5e5a1dc23eee436524a1692 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 5 Nov 2026 18:26:45 +0800 Subject: [PATCH 08/32] Fix internal login compatibility with baseline --- .../iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java | 1 - 1 file changed, 1 deletion(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java index 52a39c4d86f5f..db832dcbcebe4 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java @@ -1655,7 +1655,6 @@ protected BasicOpenSessionResp login( session, req.username, req.password, - null, req.zoneId, req.client_protocol, clientVersion, From 62d38247db88d817138d25556a17b76cfaf5ecd3 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 5 Nov 2026 18:43:20 +0800 Subject: [PATCH 09/32] Test internal client RPC routing and SSL propagation Verify multiplexed and default processor routing, session reconnects, subscription SSL forwarding, and an internal-session integration path. --- .../iotdb/session/it/InternalSessionIT.java | 97 ++++++++++++ .../session/InternalSessionProtocolTest.java | 149 ++++++++++++++++++ .../base/AbstractSubscriptionConsumer.java | 10 +- .../base/SubscriptionSslConfigTest.java | 113 +++++++++++++ .../service/DataNodeInternalRPCService.java | 19 ++- .../DataNodeInternalRPCServiceTest.java | 100 ++++++++++++ 6 files changed, 479 insertions(+), 9 deletions(-) create mode 100644 integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionIT.java create mode 100644 iotdb-client/session/src/test/java/org/apache/iotdb/session/InternalSessionProtocolTest.java create mode 100644 iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionSslConfigTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/service/DataNodeInternalRPCServiceTest.java diff --git a/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionIT.java b/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionIT.java new file mode 100644 index 0000000000000..b0f34cd524336 --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionIT.java @@ -0,0 +1,97 @@ +/* + * 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.it; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.itbase.category.ClusterIT; +import org.apache.iotdb.itbase.category.LocalStandaloneIT; +import org.apache.iotdb.session.DummyNodesSupplier; +import org.apache.iotdb.session.Session; + +import org.apache.tsfile.enums.TSDataType; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.util.Collections; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +@RunWith(IoTDBTestRunner.class) +@Category({LocalStandaloneIT.class, ClusterIT.class}) +public class InternalSessionIT { + + private static final String DATABASE = "root.internal_session_it"; + + @Before + public void setUp() throws Exception { + EnvFactory.getEnv().initClusterEnvironment(); + } + + @After + public void tearDown() throws Exception { + EnvFactory.getEnv().cleanClusterEnvironment(); + } + + @Test + public void testInternalSessionUsesInternalPort() throws Exception { + final DataNodeWrapper dataNode = EnvFactory.getEnv().getDataNodeWrapper(0); + final INodeSupplier nodeSupplier = + new DummyNodesSupplier( + Collections.singletonList( + new TEndPoint(dataNode.getInternalAddress(), dataNode.getInternalPort()))); + + try (Session session = + new Session.Builder() + .host(dataNode.getInternalAddress()) + .port(dataNode.getInternalPort()) + .username("root") + .buildInternal(nodeSupplier)) { + session.open(); + try { + session.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + session.executeNonQueryStatement( + "CREATE TIMESERIES " + DATABASE + ".d1.s1 WITH DATATYPE=INT64, ENCODING=RLE"); + session.insertRecord( + DATABASE + ".d1", + 1L, + Collections.singletonList("s1"), + Collections.singletonList(TSDataType.INT64), + Collections.singletonList(1L)); + + try (SessionDataSet dataSet = + session.executeQueryStatement("SELECT s1 FROM " + DATABASE + ".d1")) { + assertTrue(dataSet.hasNext()); + assertEquals(1L, dataSet.next().getFields().get(0).getLongV()); + } + } finally { + session.executeNonQueryStatement("DELETE DATABASE " + DATABASE); + } + } + } +} diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/InternalSessionProtocolTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/InternalSessionProtocolTest.java new file mode 100644 index 0000000000000..b1331c30f32e9 --- /dev/null +++ b/iotdb-client/session/src/test/java/org/apache/iotdb/session/InternalSessionProtocolTest.java @@ -0,0 +1,149 @@ +/* + * 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.common.rpc.thrift.TSStatus; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.rpc.DeepCopyRpcTransportFactory; +import org.apache.iotdb.rpc.RpcServiceConstants; +import org.apache.iotdb.rpc.TSStatusCode; +import org.apache.iotdb.service.rpc.thrift.IClientRPCService; +import org.apache.iotdb.service.rpc.thrift.TSCloseSessionReq; +import org.apache.iotdb.service.rpc.thrift.TSOpenSessionReq; +import org.apache.iotdb.service.rpc.thrift.TSOpenSessionResp; +import org.apache.iotdb.service.rpc.thrift.TSProtocolVersion; + +import org.apache.thrift.TMultiplexedProcessor; +import org.apache.thrift.protocol.TBinaryProtocol; +import org.apache.thrift.server.TThreadPoolServer; +import org.apache.thrift.transport.TServerSocket; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.powermock.reflect.Whitebox; + +import java.net.ServerSocket; +import java.util.Collections; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class InternalSessionProtocolTest { + + private final AtomicInteger openSessionCalls = new AtomicInteger(); + private TestServer server; + + @Before + public void setUp() throws Exception { + server = new TestServer(openSessionCalls); + } + + @After + public void tearDown() throws Exception { + server.close(); + } + + @Test + public void testReconnectRetainsMultiplexedClientServiceName() throws Exception { + final INodeSupplier nodeSupplier = + new DummyNodesSupplier(Collections.singletonList(server.endpoint)); + try (InternalSession session = + new Session.Builder() + .host(server.endpoint.getIp()) + .port(server.endpoint.getPort()) + .username("root") + .buildInternal(nodeSupplier)) { + session.open(); + + final SessionConnection connection = + Whitebox.getInternalState(session, "defaultSessionConnection"); + assertTrue(Whitebox.invokeMethod(connection, "reconnect")); + assertEquals(2, openSessionCalls.get()); + } + } + + private static final class TestServer implements AutoCloseable { + private final TEndPoint endpoint; + private final TServerSocket socket; + private final TThreadPoolServer server; + private final ExecutorService workers = Executors.newCachedThreadPool(); + private final Thread thread; + + private TestServer(final AtomicInteger openSessionCalls) throws Exception { + final ServerSocket listener = new ServerSocket(0); + endpoint = new TEndPoint("127.0.0.1", listener.getLocalPort()); + socket = new TServerSocket(listener); + + final IClientRPCService.Iface handler = mock(IClientRPCService.Iface.class); + when(handler.openSession(any(TSOpenSessionReq.class))) + .thenAnswer( + invocation -> { + openSessionCalls.incrementAndGet(); + return new TSOpenSessionResp( + new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()), + TSProtocolVersion.IOTDB_SERVICE_PROTOCOL_V3) + .setSessionId(1); + }); + when(handler.closeSession(any(TSCloseSessionReq.class))) + .thenReturn(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode())); + + final TMultiplexedProcessor processor = new TMultiplexedProcessor(); + processor.registerProcessor( + RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE, + new IClientRPCService.Processor<>(handler)); + + server = + new TThreadPoolServer( + new TThreadPoolServer.Args(socket) + .processor(processor) + .protocolFactory(new TBinaryProtocol.Factory()) + .transportFactory(DeepCopyRpcTransportFactory.INSTANCE) + .executorService(workers)); + final CountDownLatch ready = new CountDownLatch(1); + thread = + new Thread( + () -> { + ready.countDown(); + server.serve(); + }, + "internal-session-protocol-test-server"); + thread.setDaemon(true); + thread.start(); + assertTrue(ready.await(5, TimeUnit.SECONDS)); + } + + @Override + public void close() throws Exception { + server.stop(); + socket.close(); + workers.shutdownNow(); + thread.join(5_000); + } + } +} diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java index 438c38ffc6d4d..3c4eec7fedfea 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/base/AbstractSubscriptionConsumer.java @@ -622,6 +622,12 @@ protected abstract AbstractSubscriptionProvider constructSubscriptionProvider( final long heartbeatIntervalMs, final int connectionTimeoutInMs); + void configureProvider(final AbstractSubscriptionProvider provider) { + provider.setConsumerInstanceId(consumerInstanceId); + provider.setRpcServiceName(getRpcServiceName()); + provider.setSslConfig(useSSL, trustStore, trustStorePwd, keyStore, keyStorePwd, sslProtocol); + } + AbstractSubscriptionProvider constructProviderAndHandshake(final TEndPoint endPoint) throws SubscriptionException { final AbstractSubscriptionProvider provider = @@ -637,9 +643,7 @@ AbstractSubscriptionProvider constructProviderAndHandshake(final TEndPoint endPo this.thriftMaxFrameSize, this.heartbeatIntervalMs, this.connectionTimeoutInMs); - provider.setConsumerInstanceId(consumerInstanceId); - provider.setRpcServiceName(getRpcServiceName()); - provider.setSslConfig(useSSL, trustStore, trustStorePwd, keyStore, keyStorePwd, sslProtocol); + configureProvider(provider); try { provider.handshake(); } catch (final Exception e) { diff --git a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionSslConfigTest.java b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionSslConfigTest.java new file mode 100644 index 0000000000000..bbd5418a218c1 --- /dev/null +++ b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionSslConfigTest.java @@ -0,0 +1,113 @@ +/* + * 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.subscription.consumer.base; + +import org.apache.iotdb.rpc.subscription.config.ConsumerConstant; +import org.apache.iotdb.session.subscription.SubscriptionSessionWrapper; +import org.apache.iotdb.session.subscription.SubscriptionTableSessionBuilder; +import org.apache.iotdb.session.subscription.consumer.table.SubscriptionTablePullConsumerBuilder; +import org.apache.iotdb.session.subscription.consumer.tree.SubscriptionTreePullConsumer; + +import org.junit.Test; +import org.powermock.reflect.Whitebox; + +import java.util.Properties; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; + +public class SubscriptionSslConfigTest { + + @Test + public void testConsumerBuilderPropagatesSslConfigToProvider() { + final SubscriptionTablePullConsumerBuilder builder = + new SubscriptionTablePullConsumerBuilder() + .host("127.0.0.1") + .port(6667) + .username("root") + .password("root") + .consumerId("consumer") + .consumerGroupId("group") + .useSSL(true) + .trustStore("trust-store") + .trustStorePwd("trust-password") + .keyStore("key-store") + .keyStorePwd("key-password") + .sslProtocol("TLSv1.2"); + final AbstractSubscriptionConsumer consumer = (AbstractSubscriptionConsumer) builder.build(); + final AbstractSubscriptionProvider provider = mock(AbstractSubscriptionProvider.class); + + consumer.configureProvider(provider); + + verify(provider) + .setSslConfig( + true, "trust-store", "trust-password", "key-store", "key-password", "TLSv1.2"); + } + + @Test + public void testConsumerPropertiesPopulateSslConfig() { + final Properties properties = new Properties(); + properties.put(ConsumerConstant.HOST_KEY, "127.0.0.1"); + properties.put(ConsumerConstant.PORT_KEY, 6667); + properties.put(ConsumerConstant.USERNAME_KEY, "root"); + properties.put(ConsumerConstant.PASSWORD_KEY, "root"); + properties.put(ConsumerConstant.CONSUMER_ID_KEY, "consumer"); + properties.put(ConsumerConstant.CONSUMER_GROUP_ID_KEY, "group"); + properties.put(ConsumerConstant.USE_SSL_KEY, "true"); + properties.put(ConsumerConstant.TRUST_STORE_KEY, "trust-store"); + properties.put(ConsumerConstant.TRUST_STORE_PWD_KEY, "trust-password"); + properties.put(ConsumerConstant.KEY_STORE_KEY, "key-store"); + properties.put(ConsumerConstant.KEY_STORE_PWD_KEY, "key-password"); + properties.put(ConsumerConstant.SSL_PROTOCOL_KEY, "TLSv1.2"); + + final AbstractSubscriptionConsumer consumer = + (AbstractSubscriptionConsumer) new SubscriptionTreePullConsumer(properties); + + assertTrue((Boolean) Whitebox.getInternalState(consumer, "useSSL")); + assertEquals("trust-store", Whitebox.getInternalState(consumer, "trustStore")); + assertEquals("trust-password", Whitebox.getInternalState(consumer, "trustStorePwd")); + assertEquals("key-store", Whitebox.getInternalState(consumer, "keyStore")); + assertEquals("key-password", Whitebox.getInternalState(consumer, "keyStorePwd")); + assertEquals("TLSv1.2", Whitebox.getInternalState(consumer, "sslProtocol")); + } + + @Test + public void testSubscriptionSessionWrapperRetainsSslConfigBeforeOpen() { + final SubscriptionTableSessionBuilder builder = + new SubscriptionTableSessionBuilder() + .host("127.0.0.1") + .port(6667) + .username("root") + .password("root"); + final SubscriptionSessionWrapper session = new SubscriptionSessionWrapper(builder); + + session.setSslConfig( + true, "trust-store", "trust-password", "key-store", "key-password", "TLSv1.2"); + + assertTrue((Boolean) Whitebox.getInternalState(session, "useSSL")); + assertEquals("trust-store", Whitebox.getInternalState(session, "trustStore")); + assertEquals("trust-password", Whitebox.getInternalState(session, "trustStorePwd")); + assertEquals("key-store", Whitebox.getInternalState(session, "keyStore")); + assertEquals("key-password", Whitebox.getInternalState(session, "keyStorePwd")); + assertEquals("TLSv1.2", Whitebox.getInternalState(session, "sslProtocol")); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java index 666cdc47a50da..c8329c016c937 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java @@ -44,6 +44,7 @@ import org.apache.iotdb.rpc.RpcServiceConstants; import org.apache.thrift.TMultiplexedProcessor; +import org.apache.thrift.TProcessor; import java.util.concurrent.atomic.AtomicReference; @@ -70,12 +71,9 @@ public void initTProcessor() { DataNodeInternalRPCServiceImpl service = getImpl(); IClientRPCServiceWithHandler clientService = getClientImpl(); initSyncedServiceImpl(null); - TMultiplexedProcessor multiplexedProcessor = new TMultiplexedProcessor(); - multiplexedProcessor.registerProcessor( - RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE, - new InternalClientRPCServiceProcessor(clientService)); - multiplexedProcessor.registerDefault(new Processor<>(service)); - processor = multiplexedProcessor; + processor = + createMultiplexedProcessor( + new Processor<>(service), new InternalClientRPCServiceProcessor(clientService)); } @Override @@ -135,6 +133,15 @@ public DataNodeInternalRPCServiceImpl getImpl() { return impl.get(); } + static TProcessor createMultiplexedProcessor( + final TProcessor controlProcessor, final TProcessor clientProcessor) { + final TMultiplexedProcessor multiplexedProcessor = new TMultiplexedProcessor(); + multiplexedProcessor.registerProcessor( + RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE, clientProcessor); + multiplexedProcessor.registerDefault(controlProcessor); + return multiplexedProcessor; + } + protected IClientRPCServiceWithHandler getClientImpl() { clientImpl.compareAndSet(null, new InternalClientRPCServiceImpl()); return clientImpl.get(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/service/DataNodeInternalRPCServiceTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/service/DataNodeInternalRPCServiceTest.java new file mode 100644 index 0000000000000..1855bd6752c1f --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/service/DataNodeInternalRPCServiceTest.java @@ -0,0 +1,100 @@ +/* + * 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.db.service; + +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.mpp.rpc.thrift.IDataNodeRPCService; +import org.apache.iotdb.rpc.RpcServiceConstants; +import org.apache.iotdb.rpc.TSStatusCode; +import org.apache.iotdb.service.rpc.thrift.IClientRPCService; +import org.apache.iotdb.service.rpc.thrift.TSOpenSessionReq; +import org.apache.iotdb.service.rpc.thrift.TSOpenSessionResp; +import org.apache.iotdb.service.rpc.thrift.TSProtocolVersion; + +import org.apache.thrift.TProcessor; +import org.apache.thrift.protocol.TBinaryProtocol; +import org.apache.thrift.protocol.TMultiplexedProtocol; +import org.apache.thrift.transport.TMemoryBuffer; +import org.junit.Test; + +import static org.junit.Assert.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class DataNodeInternalRPCServiceTest { + + @Test + public void testMultiplexedClientAndDefaultInternalRoutes() throws Exception { + final IClientRPCService.Iface clientHandler = mock(IClientRPCService.Iface.class); + final IDataNodeRPCService.Iface controlHandler = mock(IDataNodeRPCService.Iface.class); + final TSOpenSessionResp openSessionResp = + new TSOpenSessionResp( + new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()), + TSProtocolVersion.IOTDB_SERVICE_PROTOCOL_V3) + .setSessionId(10); + when(clientHandler.openSession(any())).thenReturn(openSessionResp); + when(controlHandler.testConnectionEmptyRPC()) + .thenReturn(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode())); + + final TProcessor processor = + DataNodeInternalRPCService.createMultiplexedProcessor( + new IDataNodeRPCService.Processor<>(controlHandler), + new IClientRPCService.Processor<>(clientHandler)); + + final TMemoryBuffer clientRequest = new TMemoryBuffer(1024); + final TMemoryBuffer clientResponse = new TMemoryBuffer(1024); + final IClientRPCService.Client client = + new IClientRPCService.Client( + new TMultiplexedProtocol( + new TBinaryProtocol(clientResponse), + RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE), + new TMultiplexedProtocol( + new TBinaryProtocol(clientRequest), + RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE)); + client.send_openSession( + new TSOpenSessionReq() + .setClient_protocol(TSProtocolVersion.IOTDB_SERVICE_PROTOCOL_V3) + .setUsername("root") + .setZoneId("UTC")); + + processor.process(new TBinaryProtocol(clientRequest), new TBinaryProtocol(clientResponse)); + + assertEquals(10, client.recv_openSession().getSessionId()); + verify(clientHandler).openSession(any()); + verify(controlHandler, never()).testConnectionEmptyRPC(); + + final TMemoryBuffer controlRequest = new TMemoryBuffer(1024); + final TMemoryBuffer controlResponse = new TMemoryBuffer(1024); + final IDataNodeRPCService.Client controlClient = + new IDataNodeRPCService.Client( + new TBinaryProtocol(controlResponse), new TBinaryProtocol(controlRequest)); + controlClient.send_testConnectionEmptyRPC(); + + processor.process(new TBinaryProtocol(controlRequest), new TBinaryProtocol(controlResponse)); + + assertEquals( + TSStatusCode.SUCCESS_STATUS.getStatusCode(), + controlClient.recv_testConnectionEmptyRPC().getCode()); + verify(controlHandler).testConnectionEmptyRPC(); + } +} From ea7d8d1ddfdeafc50617a47a10e109700a71ad5e Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 5 Nov 2026 18:53:23 +0800 Subject: [PATCH 10/32] Remove unavailable test dependencies from subscription SSL test --- .../base/SubscriptionSslConfigTest.java | 124 +++++++++++++++--- 1 file changed, 103 insertions(+), 21 deletions(-) diff --git a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionSslConfigTest.java b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionSslConfigTest.java index bbd5418a218c1..744867252c064 100644 --- a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionSslConfigTest.java +++ b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionSslConfigTest.java @@ -19,21 +19,22 @@ package org.apache.iotdb.session.subscription.consumer.base; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.rpc.subscription.config.ConsumerConstant; +import org.apache.iotdb.session.AbstractSessionBuilder; import org.apache.iotdb.session.subscription.SubscriptionSessionWrapper; import org.apache.iotdb.session.subscription.SubscriptionTableSessionBuilder; +import org.apache.iotdb.session.subscription.SubscriptionTreeSessionBuilder; import org.apache.iotdb.session.subscription.consumer.table.SubscriptionTablePullConsumerBuilder; import org.apache.iotdb.session.subscription.consumer.tree.SubscriptionTreePullConsumer; import org.junit.Test; -import org.powermock.reflect.Whitebox; +import java.lang.reflect.Field; import java.util.Properties; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.verify; public class SubscriptionSslConfigTest { @@ -54,17 +55,20 @@ public void testConsumerBuilderPropagatesSslConfigToProvider() { .keyStorePwd("key-password") .sslProtocol("TLSv1.2"); final AbstractSubscriptionConsumer consumer = (AbstractSubscriptionConsumer) builder.build(); - final AbstractSubscriptionProvider provider = mock(AbstractSubscriptionProvider.class); + final RecordingProvider provider = new RecordingProvider(); consumer.configureProvider(provider); - verify(provider) - .setSslConfig( - true, "trust-store", "trust-password", "key-store", "key-password", "TLSv1.2"); + assertTrue(provider.useSSL); + assertEquals("trust-store", provider.trustStore); + assertEquals("trust-password", provider.trustStorePwd); + assertEquals("key-store", provider.keyStore); + assertEquals("key-password", provider.keyStorePwd); + assertEquals("TLSv1.2", provider.sslProtocol); } @Test - public void testConsumerPropertiesPopulateSslConfig() { + public void testConsumerPropertiesPopulateSslConfig() throws Exception { final Properties properties = new Properties(); properties.put(ConsumerConstant.HOST_KEY, "127.0.0.1"); properties.put(ConsumerConstant.PORT_KEY, 6667); @@ -82,16 +86,16 @@ public void testConsumerPropertiesPopulateSslConfig() { final AbstractSubscriptionConsumer consumer = (AbstractSubscriptionConsumer) new SubscriptionTreePullConsumer(properties); - assertTrue((Boolean) Whitebox.getInternalState(consumer, "useSSL")); - assertEquals("trust-store", Whitebox.getInternalState(consumer, "trustStore")); - assertEquals("trust-password", Whitebox.getInternalState(consumer, "trustStorePwd")); - assertEquals("key-store", Whitebox.getInternalState(consumer, "keyStore")); - assertEquals("key-password", Whitebox.getInternalState(consumer, "keyStorePwd")); - assertEquals("TLSv1.2", Whitebox.getInternalState(consumer, "sslProtocol")); + assertTrue((Boolean) getField(consumer, "useSSL")); + assertEquals("trust-store", getField(consumer, "trustStore")); + assertEquals("trust-password", getField(consumer, "trustStorePwd")); + assertEquals("key-store", getField(consumer, "keyStore")); + assertEquals("key-password", getField(consumer, "keyStorePwd")); + assertEquals("TLSv1.2", getField(consumer, "sslProtocol")); } @Test - public void testSubscriptionSessionWrapperRetainsSslConfigBeforeOpen() { + public void testSubscriptionSessionWrapperRetainsSslConfigBeforeOpen() throws Exception { final SubscriptionTableSessionBuilder builder = new SubscriptionTableSessionBuilder() .host("127.0.0.1") @@ -103,11 +107,89 @@ public void testSubscriptionSessionWrapperRetainsSslConfigBeforeOpen() { session.setSslConfig( true, "trust-store", "trust-password", "key-store", "key-password", "TLSv1.2"); - assertTrue((Boolean) Whitebox.getInternalState(session, "useSSL")); - assertEquals("trust-store", Whitebox.getInternalState(session, "trustStore")); - assertEquals("trust-password", Whitebox.getInternalState(session, "trustStorePwd")); - assertEquals("key-store", Whitebox.getInternalState(session, "keyStore")); - assertEquals("key-password", Whitebox.getInternalState(session, "keyStorePwd")); - assertEquals("TLSv1.2", Whitebox.getInternalState(session, "sslProtocol")); + assertTrue((Boolean) getField(session, "useSSL")); + assertEquals("trust-store", getField(session, "trustStore")); + assertEquals("trust-password", getField(session, "trustStorePwd")); + assertEquals("key-store", getField(session, "keyStore")); + assertEquals("key-password", getField(session, "keyStorePwd")); + assertEquals("TLSv1.2", getField(session, "sslProtocol")); + } + + private static Object getField(final Object target, final String fieldName) throws Exception { + for (Class type = target.getClass(); type != null; type = type.getSuperclass()) { + try { + final Field field = type.getDeclaredField(fieldName); + field.setAccessible(true); + return field.get(target); + } catch (NoSuchFieldException ignored) { + // Search the superclass. + } + } + throw new NoSuchFieldException(fieldName); + } + + private static final class RecordingProvider extends AbstractSubscriptionProvider { + + private boolean useSSL; + private String trustStore; + private String trustStorePwd; + private String keyStore; + private String keyStorePwd; + private String sslProtocol; + + private RecordingProvider() { + super( + new TEndPoint("127.0.0.1", 6667), + "root", + "root", + null, + "consumer", + "group", + null, + null, + 1024 * 1024, + 1000, + 1000); + } + + @Override + protected AbstractSessionBuilder constructSubscriptionSessionBuilder( + final String host, + final int port, + final String username, + final String password, + final String encryptedPassword, + final int thriftMaxFrameSize, + final int connectionTimeoutInMs) { + return new SubscriptionTreeSessionBuilder() + .host(host) + .port(port) + .username(username) + .password(password) + .thriftMaxFrameSize(thriftMaxFrameSize) + .connectionTimeoutInMs(connectionTimeoutInMs); + } + + @Override + void setRpcServiceName(final String rpcServiceName) {} + + @Override + void setSslConfig( + final boolean useSSL, + final String trustStore, + final String trustStorePwd, + final String keyStore, + final String keyStorePwd, + final String sslProtocol) { + this.useSSL = useSSL; + this.trustStore = trustStore; + this.trustStorePwd = trustStorePwd; + this.keyStore = keyStore; + this.keyStorePwd = keyStorePwd; + this.sslProtocol = sslProtocol; + } + + @Override + void setConsumerInstanceId(final String consumerInstanceId) {} } } From fabd1e8dc1a2c0351e273da5f1d6760d948740c6 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 11:01:12 +0800 Subject: [PATCH 11/32] Move internal session builders to node commons --- .../iotdb/session/it/InternalSessionIT.java | 5 +- .../org/apache/iotdb/session/Session.java | 110 ++++++++++------ .../iotdb/session/TableSessionBuilder.java | 14 +- .../iotdb/session/pool/SessionPool.java | 122 +++++++++--------- .../session/pool/TableSessionPoolBuilder.java | 23 +--- .../iotdb/session/InternalSessionTest.java | 63 --------- .../InternalSubscriptionTableSession.java | 86 ------------ .../SubscriptionTableSessionBuilder.java | 26 +--- .../SubscriptionTablePullConsumerBuilder.java | 13 +- .../base/SubscriptionSslConfigTest.java | 62 ++------- iotdb-core/node-commons/pom.xml | 5 + .../internal}/session/InternalSession.java | 45 +++---- .../session/InternalSessionBuilder.java | 33 +++-- .../session/InternalTableSessionBuilder.java | 51 ++++++++ .../InternalTableSessionPoolBuilder.java | 66 ++++++++++ .../InternalSubscriptionSessionWrapper.java | 69 ++++++++++ ...InternalSubscriptionTablePullConsumer.java | 9 +- ...lSubscriptionTablePullConsumerBuilder.java | 50 +++++++ .../InternalSubscriptionTableSession.java | 18 +-- ...ternalSubscriptionTableSessionBuilder.java | 49 +++++++ .../session/InternalSessionProtocolTest.java | 23 +++- 21 files changed, 521 insertions(+), 421 deletions(-) delete mode 100644 iotdb-client/session/src/test/java/org/apache/iotdb/session/InternalSessionTest.java delete mode 100644 iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/InternalSubscriptionTableSession.java rename {iotdb-client/session/src/main/java/org/apache/iotdb => iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal}/session/InternalSession.java (58%) rename iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.java => iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java (57%) create mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java create mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java create mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java rename {iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table => iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription}/InternalSubscriptionTablePullConsumer.java (86%) create mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java rename iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java => iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSession.java (61%) create mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java rename {iotdb-client/session/src/test/java/org/apache/iotdb => iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal}/session/InternalSessionProtocolTest.java (88%) diff --git a/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionIT.java b/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionIT.java index b0f34cd524336..5f8c9bfdcef4a 100644 --- a/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionIT.java @@ -20,6 +20,7 @@ package org.apache.iotdb.session.it; import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.client.internal.session.InternalSessionBuilder; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.SessionDataSet; import org.apache.iotdb.it.env.EnvFactory; @@ -67,11 +68,11 @@ public void testInternalSessionUsesInternalPort() throws Exception { new TEndPoint(dataNode.getInternalAddress(), dataNode.getInternalPort()))); try (Session session = - new Session.Builder() + new InternalSessionBuilder(nodeSupplier) .host(dataNode.getInternalAddress()) .port(dataNode.getInternalPort()) .username("root") - .buildInternal(nodeSupplier)) { + .build()) { session.open(); try { session.executeNonQueryStatement("CREATE DATABASE " + DATABASE); 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 5263f0af4bd47..d9b5687775647 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,7 +21,6 @@ import org.apache.iotdb.common.rpc.thrift.TAggregationType; import org.apache.iotdb.common.rpc.thrift.TEndPoint; -import org.apache.iotdb.isession.IInternalSessionBuilder; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.ISession; import org.apache.iotdb.isession.SessionConfig; @@ -32,6 +31,7 @@ import org.apache.iotdb.rpc.IoTDBConnectionException; import org.apache.iotdb.rpc.NoValidValueException; import org.apache.iotdb.rpc.RedirectException; +import org.apache.iotdb.rpc.RpcServiceConstants; import org.apache.iotdb.rpc.StatementExecutionException; import org.apache.iotdb.rpc.UrlUtils; import org.apache.iotdb.service.rpc.thrift.TCreateTimeseriesUsingSchemaTemplateReq; @@ -544,42 +544,23 @@ public synchronized void open(boolean enableRPCCompaction, int connectionTimeout return; } + final boolean closeAvailableNodes = executorService != null; if (this.executorService != null) { this.executorService.shutdown(); this.executorService = null; } if (this.availableNodes != null) { - this.availableNodes.close(); + if (closeAvailableNodes) { + this.availableNodes.close(); + } this.availableNodes = null; } if (enableAutoFetch) { initThreadPool(); - this.availableNodes = - NodesSupplier.createNodeSupplier( - getNodeUrls(), - SessionUtils.parseSeedNodeUrls(initialNodeUrls), - executorService, - username, - password, - zoneId, - thriftDefaultBufferSize, - thriftMaxFrameSize, - connectionTimeoutInMs, - useSSL, - trustStore, - trustStorePwd, - keyStore, - keyStorePwd, - sslProtocol, - enableRPCCompaction, - version.toString()); - } else { - // A static supplier has no separate seed tier. Keep every configured address after reopen. - Set staticNodes = new LinkedHashSet<>(getNodeUrls()); - staticNodes.addAll(SessionUtils.parseSeedNodeUrls(initialNodeUrls)); - this.availableNodes = new DummyNodesSupplier(new ArrayList<>(staticNodes)); } + this.availableNodes = getNodeSupplier(enableRPCCompaction, connectionTimeoutInMs); + this.defaultEndPoint = resolveDefaultEndPoint(this.defaultEndPoint, this.availableNodes); this.connectionTimeoutInMs = connectionTimeoutInMs; setDefaultSessionConnection(constructSessionConnection(this, defaultEndPoint, zoneId)); @@ -610,6 +591,48 @@ private void initThreadPool() { }); } + /** + * Creates the endpoint source used by this session. Subclasses may override this hook to provide + * a server-managed endpoint source without exposing it through the client builder API. + */ + protected INodeSupplier getNodeSupplier( + final boolean enableRPCCompaction, final int connectionTimeoutInMs) { + if (enableAutoFetch) { + return NodesSupplier.createNodeSupplier( + getNodeUrls(), + SessionUtils.parseSeedNodeUrls(initialNodeUrls), + executorService, + username, + password, + zoneId, + thriftDefaultBufferSize, + thriftMaxFrameSize, + connectionTimeoutInMs, + useSSL, + trustStore, + trustStorePwd, + keyStore, + keyStorePwd, + sslProtocol, + enableRPCCompaction, + version.toString()); + } + // A static supplier has no separate seed tier. Keep every configured address after reopen. + Set staticNodes = new LinkedHashSet<>(getNodeUrls()); + staticNodes.addAll(SessionUtils.parseSeedNodeUrls(initialNodeUrls)); + return new DummyNodesSupplier(new ArrayList<>(staticNodes)); + } + + /** + * Selects the endpoint used for the initial connection. Internal sessions that support endpoint + * discovery override this hook so a stale builder endpoint cannot be used. + */ + protected TEndPoint resolveDefaultEndPoint( + final TEndPoint requestedEndPoint, final INodeSupplier nodeSupplier) + throws IoTDBConnectionException { + return requestedEndPoint; + } + private static List shuffleNodeUrls(List endPoints) { try { Collections.shuffle(endPoints); @@ -692,12 +715,16 @@ public synchronized void close() throws IoTDBConnectionException { getDefaultSessionConnection().close(); setDefaultSessionConnection(null); } finally { - // if executorService is null, it means that availableNodes is got from SessionPool and we - // shouldn't clean that + // A node supplier created outside this session, such as one from SessionPool or a server + // internal session, is not owned by this transport. if (this.executorService != null) { this.executorService.shutdown(); this.executorService = null; - this.availableNodes.close(); + if (this.availableNodes != null) { + this.availableNodes.close(); + this.availableNodes = null; + } + } else { this.availableNodes = null; } isClosed = true; @@ -4445,6 +4472,16 @@ protected void setRpcServiceName(String rpcServiceName) { this.rpcServiceName = rpcServiceName; } + /** Applies the common transport and authentication settings for server internal sessions. */ + protected final void configureInternalSession() { + password = null; + useEncryptedPassword = false; + enableAutoFetch = false; + enableRedirection = false; + enableQueryRedirection = false; + rpcServiceName = RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE; + } + public boolean isEnableThriftRpcCompaction() { return enableThriftRpcCompaction; } @@ -4453,8 +4490,7 @@ public void setEnableThriftRpcCompaction(boolean enableThriftRpcCompaction) { this.enableThriftRpcCompaction = enableThriftRpcCompaction; } - public static class Builder extends AbstractSessionBuilder - implements IInternalSessionBuilder { + public static class Builder extends AbstractSessionBuilder { public Builder host(String host) { this.host = host; @@ -4602,17 +4638,11 @@ public Session build() { SessionMessages .EXCEPTION_YOU_SHOULD_SPECIFY_EITHER_NODEURLS_HOST_RPCPORT_BUT_NOT_BOTH_77E7B084); } - return new Session(this); + return createSession(); } - /** Builds an unopened internal session using caller-supplied internal endpoints. */ - @Override - public InternalSession buildInternal(INodeSupplier nodeSupplier) { - return createInternalSession(nodeSupplier); - } - - protected InternalSession createInternalSession(INodeSupplier nodeSupplier) { - return new InternalSession(this, nodeSupplier); + protected Session createSession() { + return new Session(this); } } } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java index 66c52b011a3a0..b41233ebe3933 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java @@ -19,8 +19,6 @@ package org.apache.iotdb.session; -import org.apache.iotdb.isession.IInternalSessionBuilder; -import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.ITableSession; import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.rpc.IoTDBConnectionException; @@ -43,8 +41,7 @@ * *

All configurations have reasonable default values, which can be overridden as needed. */ -public class TableSessionBuilder extends AbstractSessionBuilder - implements IInternalSessionBuilder { +public class TableSessionBuilder extends AbstractSessionBuilder { private boolean enableCompression = false; private int connectionTimeoutInMs = SessionConfig.DEFAULT_CONNECTION_TIMEOUT_MS; @@ -438,13 +435,6 @@ public ITableSession build() throws IoTDBConnectionException { return wrapTableSession(constructNewSession()); } - /** Builds and opens a table session using caller-supplied internal endpoints. */ - @Override - public ITableSession buildInternal(INodeSupplier nodeSupplier) throws IoTDBConnectionException { - this.sqlDialect = TABLE; - return wrapTableSession(openSession(new InternalSession(this, nodeSupplier))); - } - protected Session constructNewSession() throws IoTDBConnectionException { if (nodeUrls == null) { this.nodeUrls = @@ -458,7 +448,7 @@ protected ITableSession wrapTableSession(Session session) { return new TableSession(session); } - private Session openSession(Session newSession) throws IoTDBConnectionException { + protected Session openSession(Session newSession) throws IoTDBConnectionException { newSession.enableIoTDBRpcCompression = isIoTDBRpcCompressionEnabled; newSession.tabletCompressionMinRowSize = tabletCompressionMinRowSize; 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 b86804c8c142e..3b5dd914b2d21 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 @@ -157,6 +157,7 @@ public class SessionPool implements ISessionPool { private final int connectionTimeoutInMs; private final boolean enableIoTDBRpcCompression; private final boolean enableThriftCompression; + private final boolean internalSessionPool; // whether the queue is closed. private boolean closed; @@ -173,8 +174,6 @@ public class SessionPool implements ISessionPool { private INodeSupplier availableNodes; - private INodeSupplier internalNodeSupplier; - // set to true, means that we will start a background thread to fetch all available (Status is // not Removing) datanodes in cluster, and these available nodes will be used in retrying stage private boolean enableAutoFetch = true; @@ -528,20 +527,21 @@ public SessionPool( } public SessionPool(AbstractSessionPoolBuilder builder) { - this(builder, null); + this(builder, false); } - protected SessionPool(AbstractSessionPoolBuilder builder, INodeSupplier internalNodeSupplier) { - this.internalNodeSupplier = internalNodeSupplier; + protected SessionPool( + final AbstractSessionPoolBuilder builder, final boolean internalSessionPool) { + this.internalSessionPool = internalSessionPool; this.maxSize = builder.maxSize; this.user = builder.username; - this.password = internalNodeSupplier == null ? builder.pw : null; + this.password = internalSessionPool ? null : builder.pw; this.fetchSize = builder.fetchSize; this.waitToGetSessionTimeoutInMs = builder.waitToGetSessionTimeoutInMs; this.enableThriftCompression = builder.isThriftRpcCompactionEnabled; this.enableIoTDBRpcCompression = builder.isIoTDBRpcCompressionEnabled; this.zoneId = builder.zoneId; - this.enableRedirection = internalNodeSupplier == null && builder.enableRedirection; + this.enableRedirection = !internalSessionPool && builder.enableRedirection; if (this.enableRedirection) { deviceIdToEndpoint = new ConcurrentHashMap<>(); tableModelDeviceIdToEndpoint = new ConcurrentHashMap<>(); @@ -551,7 +551,7 @@ protected SessionPool(AbstractSessionPoolBuilder builder, INodeSupplier internal this.version = builder.version; this.thriftDefaultBufferSize = builder.thriftDefaultBufferSize; this.thriftMaxFrameSize = builder.thriftMaxFrameSize; - this.enableAutoFetch = internalNodeSupplier == null && builder.enableAutoFetch; + this.enableAutoFetch = !internalSessionPool && builder.enableAutoFetch; this.useSSL = builder.useSSL; this.trustStore = builder.trustStore; this.trustStorePwd = builder.trustStorePwd; @@ -568,12 +568,11 @@ protected SessionPool(AbstractSessionPoolBuilder builder, INodeSupplier internal if (enableAutoFetch) { initThreadPool(); } - if (internalNodeSupplier != null) { + if (internalSessionPool) { this.host = null; this.port = -1; this.nodeUrls = null; - this.formattedNodeUrls = internalNodeSupplier.getClass().getSimpleName(); - this.availableNodes = internalNodeSupplier; + this.formattedNodeUrls = builder.getClass().getSimpleName(); } else if (builder.nodeUrls != null) { if (builder.nodeUrls.isEmpty()) { throw new IllegalArgumentException(SessionMessages.NODE_URLS_EMPTY); @@ -602,51 +601,60 @@ protected SessionPool(AbstractSessionPoolBuilder builder, INodeSupplier internal } } - protected ISession constructNewSession() { - Session.Builder builder = - new Session.Builder() - .username(user) - .password(password) - .fetchSize(fetchSize) - .zoneId(zoneId) - .thriftDefaultBufferSize(thriftDefaultBufferSize) - .thriftMaxFrameSize(thriftMaxFrameSize) - .enableRedirection(enableRedirection) - .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) - .version(version) - .useSSL(useSSL) - .trustStore(trustStore) - .trustStorePwd(trustStorePwd) - .keyStore(keyStore) - .keyStorePwd(keyStorePwd) - .sslProtocol(sslProtocol) - .maxRetryCount(maxRetryCount) - .retryIntervalInMs(retryIntervalInMs) - .sqlDialect(sqlDialect) - .database(database) - .enableAutoCreateSchema(enableAutoCreateSchema) - .timeOut(queryTimeoutInMs) - .enableIoTDBRpcCompression(enableIoTDBRpcCompression) - .enableThriftRpcCompression(enableThriftCompression); - if (isInternalSessionPool()) { - return builder.buildInternal(internalNodeSupplier); - } + protected Session.Builder createSessionBuilder() { + return new Session.Builder() + .username(user) + .password(password) + .fetchSize(fetchSize) + .zoneId(zoneId) + .thriftDefaultBufferSize(thriftDefaultBufferSize) + .thriftMaxFrameSize(thriftMaxFrameSize) + .enableRedirection(enableRedirection) + .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) + .version(version) + .useSSL(useSSL) + .trustStore(trustStore) + .trustStorePwd(trustStorePwd) + .keyStore(keyStore) + .keyStorePwd(keyStorePwd) + .sslProtocol(sslProtocol) + .maxRetryCount(maxRetryCount) + .retryIntervalInMs(retryIntervalInMs) + .sqlDialect(sqlDialect) + .database(database) + .enableAutoCreateSchema(enableAutoCreateSchema) + .timeOut(queryTimeoutInMs) + .enableIoTDBRpcCompression(enableIoTDBRpcCompression) + .enableThriftRpcCompression(enableThriftCompression); + } + + protected Session constructNewSession() { + final Session.Builder builder = createSessionBuilder(); if (nodeUrls == null) { builder.host(host).port(port); } else { builder.nodeUrls(nodeUrls); } - Session session = builder.build(); + final Session session = builder.build(); session.setEnableQueryRedirection(enableQueryRedirection); return session; } - protected boolean isInternalSessionPool() { - return internalNodeSupplier != null; + protected void openSession(final ISession session) throws IoTDBConnectionException { + if (internalSessionPool) { + session.open(enableThriftCompression, connectionTimeoutInMs); + } else { + session.open( + enableThriftCompression, + connectionTimeoutInMs, + deviceIdToEndpoint, + tableModelDeviceIdToEndpoint, + availableNodes); + } } - protected INodeSupplier getInternalNodeSupplier() { - return internalNodeSupplier; + protected final boolean isInternalSessionPool() { + return internalSessionPool; } private void initThreadPool() { @@ -758,12 +766,7 @@ private ISession getSession() throws IoTDBConnectionException { session = constructNewSession(); try { - session.open( - enableThriftCompression, - connectionTimeoutInMs, - deviceIdToEndpoint, - tableModelDeviceIdToEndpoint, - availableNodes); + openSession(session); // avoid someone has called close() the session pool synchronized (this) { if (closed) { @@ -843,9 +846,7 @@ public synchronized void close() { this.executorService = null; } if (availableNodes != null) { - if (!isInternalSessionPool()) { - this.availableNodes.close(); - } + this.availableNodes.close(); this.availableNodes = null; } LOGGER.info(SessionMessages.CLOSING_SESSION_POOL); @@ -872,14 +873,9 @@ public void closeResultSet(SessionDataSetWrapper wrapper) { @SuppressWarnings({"squid:S2446"}) private void tryConstructNewSession() { - ISession session = constructNewSession(); + Session session = constructNewSession(); try { - session.open( - enableThriftCompression, - connectionTimeoutInMs, - deviceIdToEndpoint, - tableModelDeviceIdToEndpoint, - availableNodes); + openSession(session); // avoid someone has called close() the session pool synchronized (this) { if (closed) { @@ -3540,7 +3536,7 @@ public boolean isEnableThriftCompression() { @Override public void setEnableRedirection(boolean enableRedirection) { - this.enableRedirection = !isInternalSessionPool() && enableRedirection; + this.enableRedirection = !internalSessionPool && enableRedirection; if (this.enableRedirection) { deviceIdToEndpoint = new ConcurrentHashMap<>(); tableModelDeviceIdToEndpoint = new ConcurrentHashMap<>(); @@ -3560,7 +3556,7 @@ public boolean isEnableRedirection() { @Override public void setEnableQueryRedirection(boolean enableQueryRedirection) { - this.enableQueryRedirection = !isInternalSessionPool() && enableQueryRedirection; + this.enableQueryRedirection = !internalSessionPool && enableQueryRedirection; for (ISession session : queue) { session.setEnableQueryRedirection(this.enableQueryRedirection); } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java index d0104675d0a8d..9f4b1364d02ec 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java @@ -19,16 +19,12 @@ package org.apache.iotdb.session.pool; -import org.apache.iotdb.isession.IInternalSessionBuilder; -import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.isession.pool.ITableSessionPool; -import org.apache.iotdb.session.i18n.SessionMessages; import java.time.ZoneId; import java.util.Collections; import java.util.List; -import java.util.Objects; /** * A builder class for constructing instances of {@link ITableSessionPool}. @@ -38,8 +34,7 @@ * *

All configurations have reasonable default values, which can be overridden as needed. */ -public class TableSessionPoolBuilder extends AbstractSessionPoolBuilder - implements IInternalSessionBuilder { +public class TableSessionPoolBuilder extends AbstractSessionPoolBuilder { /** * Sets the list of node URLs for the IoTDB cluster. @@ -348,20 +343,14 @@ public ITableSessionPool build() { Collections.singletonList(SessionConfig.DEFAULT_HOST + ":" + SessionConfig.DEFAULT_PORT); } this.sqlDialect = "table"; - SessionPool sessionPool = new SessionPool(this); - return new TableSessionPool(sessionPool); + return createTableSessionPool(createSessionPool()); } - /** Builds a pool using caller-supplied internal endpoints; connections open lazily. */ - @Override - public ITableSessionPool buildInternal(INodeSupplier nodeSupplier) { - Objects.requireNonNull( - nodeSupplier, SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); - this.sqlDialect = "table"; - return createInternalTableSessionPool(nodeSupplier); + protected SessionPool createSessionPool() { + return new SessionPool(this); } - protected ITableSessionPool createInternalTableSessionPool(INodeSupplier nodeSupplier) { - return new InternalTableSessionPool(this, nodeSupplier); + protected ITableSessionPool createTableSessionPool(final SessionPool sessionPool) { + return new TableSessionPool(sessionPool); } } diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/InternalSessionTest.java b/iotdb-client/session/src/test/java/org/apache/iotdb/session/InternalSessionTest.java deleted file mode 100644 index 2c79bc764c0e8..0000000000000 --- a/iotdb-client/session/src/test/java/org/apache/iotdb/session/InternalSessionTest.java +++ /dev/null @@ -1,63 +0,0 @@ -/* - * 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.isession.INodeSupplier; -import org.apache.iotdb.isession.pool.ITableSessionPool; -import org.apache.iotdb.rpc.RpcServiceConstants; -import org.apache.iotdb.session.pool.TableSessionPoolBuilder; - -import org.junit.Test; -import org.powermock.reflect.Whitebox; - -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertNull; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.verify; - -public class InternalSessionTest { - - @Test - public void testInternalTableSessionPoolDoesNotCloseCallerSupplier() { - final INodeSupplier nodeSupplier = mock(INodeSupplier.class); - final ITableSessionPool pool = new TableSessionPoolBuilder().buildInternal(nodeSupplier); - - pool.close(); - - verify(nodeSupplier, never()).close(); - } - - @Test - public void testBuildInternalConfiguresMultiplexedTransport() { - final INodeSupplier nodeSupplier = mock(INodeSupplier.class); - final InternalSession session = - new Session.Builder().host("127.0.0.1").port(10730).buildInternal(nodeSupplier); - - assertEquals( - RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE, - Whitebox.getInternalState(session, "rpcServiceName")); - assertNull(Whitebox.getInternalState(session, "password")); - assertFalse((Boolean) Whitebox.getInternalState(session, "enableAutoFetch")); - assertFalse((Boolean) Whitebox.getInternalState(session, "enableRedirection")); - assertFalse((Boolean) Whitebox.getInternalState(session, "enableQueryRedirection")); - } -} diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/InternalSubscriptionTableSession.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/InternalSubscriptionTableSession.java deleted file mode 100644 index 16e2bd6f20cec..0000000000000 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/InternalSubscriptionTableSession.java +++ /dev/null @@ -1,86 +0,0 @@ -/* - * 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.subscription; - -import org.apache.iotdb.common.rpc.thrift.TEndPoint; -import org.apache.iotdb.isession.INodeSupplier; -import org.apache.iotdb.rpc.IoTDBConnectionException; -import org.apache.iotdb.rpc.RpcServiceConstants; -import org.apache.iotdb.session.AbstractSessionBuilder; -import org.apache.iotdb.session.Session; -import org.apache.iotdb.session.SessionConnection; -import org.apache.iotdb.session.i18n.SessionMessages; - -import java.time.ZoneId; -import java.util.List; -import java.util.Objects; - -/** Table subscription-management session using the caller's identity and internal endpoints. */ -public class InternalSubscriptionTableSession extends SubscriptionTableSession { - - public InternalSubscriptionTableSession( - AbstractSessionBuilder builder, INodeSupplier nodeSupplier) { - super(new InternalSessionWrapper(builder, nodeSupplier)); - } - - private static class InternalSessionWrapper extends SubscriptionSessionWrapper { - - private final INodeSupplier nodeSupplier; - - private InternalSessionWrapper(AbstractSessionBuilder builder, INodeSupplier nodeSupplier) { - super(builder); - this.nodeSupplier = - Objects.requireNonNull( - nodeSupplier, - SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); - password = null; - useEncryptedPassword = false; - enableAutoFetch = false; - enableRedirection = false; - enableQueryRedirection = false; - rpcServiceName = RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE; - } - - @Override - public synchronized void open(boolean enableRPCCompaction, int connectionTimeoutInMs) - throws IoTDBConnectionException { - super.open(enableRPCCompaction, connectionTimeoutInMs, null, null, nodeSupplier); - } - - @Override - public SessionConnection constructSessionConnection( - Session session, TEndPoint endpoint, ZoneId zoneId) throws IoTDBConnectionException { - final List candidates = nodeSupplier.get(); - if (candidates.isEmpty()) { - throw new IoTDBConnectionException( - SessionMessages.EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081); - } - defaultEndPoint = - candidates.contains(endpoint) - ? endpoint.deepCopy() - : nodeSupplier - .getQueryEndPoint() - .filter(candidates::contains) - .orElse(candidates.get(0)) - .deepCopy(); - return super.constructSessionConnection(session, defaultEndPoint, zoneId); - } - } -} diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java index f54a5aae25ad5..5a74a564d02a1 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java @@ -19,13 +19,10 @@ package org.apache.iotdb.session.subscription; -import org.apache.iotdb.isession.IInternalSessionBuilder; -import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.rpc.IoTDBConnectionException; import org.apache.iotdb.session.AbstractSessionBuilder; -public class SubscriptionTableSessionBuilder extends AbstractSessionBuilder - implements IInternalSessionBuilder { +public class SubscriptionTableSessionBuilder extends AbstractSessionBuilder { public SubscriptionTableSessionBuilder() { // use table model @@ -102,27 +99,12 @@ public SubscriptionTableSessionBuilder connectionTimeoutInMs(final int connectio } public ISubscriptionTableSession build() throws IoTDBConnectionException { - final ISubscriptionTableSession session = new SubscriptionTableSession(this); + final ISubscriptionTableSession session = createSubscriptionSession(); session.open(); return session; } - /** Builds and opens a table subscription-management session on internal endpoints. */ - @Override - public ISubscriptionTableSession buildInternal(INodeSupplier nodeSupplier) - throws IoTDBConnectionException { - final ISubscriptionTableSession session = - new InternalSubscriptionTableSession(this, nodeSupplier); - try { - session.open(); - } catch (IoTDBConnectionException e) { - try { - session.close(); - } catch (Exception closeException) { - e.addSuppressed(closeException); - } - throw e; - } - return session; + protected ISubscriptionTableSession createSubscriptionSession() { + return new SubscriptionTableSession(this); } } diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java index a710551bae33c..dbb1e020ab431 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java @@ -19,15 +19,12 @@ package org.apache.iotdb.session.subscription.consumer.table; -import org.apache.iotdb.isession.IInternalSessionBuilder; -import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.session.subscription.consumer.ISubscriptionTablePullConsumer; import org.apache.iotdb.session.subscription.consumer.base.AbstractSubscriptionPullConsumerBuilder; import java.util.List; -public class SubscriptionTablePullConsumerBuilder extends AbstractSubscriptionPullConsumerBuilder - implements IInternalSessionBuilder { +public class SubscriptionTablePullConsumerBuilder extends AbstractSubscriptionPullConsumerBuilder { @Override public SubscriptionTablePullConsumerBuilder host(final String host) { @@ -189,12 +186,10 @@ public SubscriptionTablePullConsumerBuilder autoCommitIntervalMs( } public ISubscriptionTablePullConsumer build() { - return new SubscriptionTablePullConsumer(this); + return createConsumer(); } - /** Builds an unopened consumer using the caller's internal endpoint source and username. */ - @Override - public ISubscriptionTablePullConsumer buildInternal(final INodeSupplier nodeSupplier) { - return new InternalSubscriptionTablePullConsumer(this, nodeSupplier); + protected ISubscriptionTablePullConsumer createConsumer() { + return new SubscriptionTablePullConsumer(this); } } diff --git a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionSslConfigTest.java b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionSslConfigTest.java index 744867252c064..a29657ff0c344 100644 --- a/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionSslConfigTest.java +++ b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionSslConfigTest.java @@ -22,15 +22,12 @@ import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.rpc.subscription.config.ConsumerConstant; import org.apache.iotdb.session.AbstractSessionBuilder; -import org.apache.iotdb.session.subscription.SubscriptionSessionWrapper; -import org.apache.iotdb.session.subscription.SubscriptionTableSessionBuilder; import org.apache.iotdb.session.subscription.SubscriptionTreeSessionBuilder; import org.apache.iotdb.session.subscription.consumer.table.SubscriptionTablePullConsumerBuilder; import org.apache.iotdb.session.subscription.consumer.tree.SubscriptionTreePullConsumer; import org.junit.Test; -import java.lang.reflect.Field; import java.util.Properties; import static org.junit.Assert.assertEquals; @@ -59,16 +56,11 @@ public void testConsumerBuilderPropagatesSslConfigToProvider() { consumer.configureProvider(provider); - assertTrue(provider.useSSL); - assertEquals("trust-store", provider.trustStore); - assertEquals("trust-password", provider.trustStorePwd); - assertEquals("key-store", provider.keyStore); - assertEquals("key-password", provider.keyStorePwd); - assertEquals("TLSv1.2", provider.sslProtocol); + assertSslConfig(provider); } @Test - public void testConsumerPropertiesPopulateSslConfig() throws Exception { + public void testConsumerPropertiesPropagateSslConfigToProvider() { final Properties properties = new Properties(); properties.put(ConsumerConstant.HOST_KEY, "127.0.0.1"); properties.put(ConsumerConstant.PORT_KEY, 6667); @@ -85,47 +77,20 @@ public void testConsumerPropertiesPopulateSslConfig() throws Exception { final AbstractSubscriptionConsumer consumer = (AbstractSubscriptionConsumer) new SubscriptionTreePullConsumer(properties); + final RecordingProvider provider = new RecordingProvider(); - assertTrue((Boolean) getField(consumer, "useSSL")); - assertEquals("trust-store", getField(consumer, "trustStore")); - assertEquals("trust-password", getField(consumer, "trustStorePwd")); - assertEquals("key-store", getField(consumer, "keyStore")); - assertEquals("key-password", getField(consumer, "keyStorePwd")); - assertEquals("TLSv1.2", getField(consumer, "sslProtocol")); - } + consumer.configureProvider(provider); - @Test - public void testSubscriptionSessionWrapperRetainsSslConfigBeforeOpen() throws Exception { - final SubscriptionTableSessionBuilder builder = - new SubscriptionTableSessionBuilder() - .host("127.0.0.1") - .port(6667) - .username("root") - .password("root"); - final SubscriptionSessionWrapper session = new SubscriptionSessionWrapper(builder); - - session.setSslConfig( - true, "trust-store", "trust-password", "key-store", "key-password", "TLSv1.2"); - - assertTrue((Boolean) getField(session, "useSSL")); - assertEquals("trust-store", getField(session, "trustStore")); - assertEquals("trust-password", getField(session, "trustStorePwd")); - assertEquals("key-store", getField(session, "keyStore")); - assertEquals("key-password", getField(session, "keyStorePwd")); - assertEquals("TLSv1.2", getField(session, "sslProtocol")); + assertSslConfig(provider); } - private static Object getField(final Object target, final String fieldName) throws Exception { - for (Class type = target.getClass(); type != null; type = type.getSuperclass()) { - try { - final Field field = type.getDeclaredField(fieldName); - field.setAccessible(true); - return field.get(target); - } catch (NoSuchFieldException ignored) { - // Search the superclass. - } - } - throw new NoSuchFieldException(fieldName); + private static void assertSslConfig(final RecordingProvider provider) { + assertTrue(provider.useSSL); + assertEquals("trust-store", provider.trustStore); + assertEquals("trust-password", provider.trustStorePwd); + assertEquals("key-store", provider.keyStore); + assertEquals("key-password", provider.keyStorePwd); + assertEquals("TLSv1.2", provider.sslProtocol); } private static final class RecordingProvider extends AbstractSubscriptionProvider { @@ -170,9 +135,6 @@ protected AbstractSessionBuilder constructSubscriptionSessionBuilder( .connectionTimeoutInMs(connectionTimeoutInMs); } - @Override - void setRpcServiceName(final String rpcServiceName) {} - @Override void setSslConfig( final boolean useSSL, diff --git a/iotdb-core/node-commons/pom.xml b/iotdb-core/node-commons/pom.xml index 3bb68a29562c3..3e8a5b87b2d3c 100644 --- a/iotdb-core/node-commons/pom.xml +++ b/iotdb-core/node-commons/pom.xml @@ -45,6 +45,11 @@ iotdb-subscription 2.0.11-SNAPSHOT + + org.apache.iotdb + iotdb-session + 2.0.11-SNAPSHOT + org.apache.tsfile common diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/InternalSession.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java similarity index 58% rename from iotdb-client/session/src/main/java/org/apache/iotdb/session/InternalSession.java rename to iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java index 504ab6005eacf..403a977ddc59c 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/InternalSession.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java @@ -17,59 +17,54 @@ * under the License. */ -package org.apache.iotdb.session; +package org.apache.iotdb.commons.client.internal.session; import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.rpc.IoTDBConnectionException; -import org.apache.iotdb.rpc.RpcServiceConstants; +import org.apache.iotdb.session.AbstractSessionBuilder; +import org.apache.iotdb.session.Session; import org.apache.iotdb.session.i18n.SessionMessages; -import java.time.ZoneId; import java.util.List; import java.util.Objects; -/** Session using caller-supplied internal endpoints and the caller's username. */ +/** Base session for trusted server-to-server client RPC connections. */ public class InternalSession extends Session { private final INodeSupplier nodeSupplier; - public InternalSession(AbstractSessionBuilder builder, INodeSupplier nodeSupplier) { + protected InternalSession( + final AbstractSessionBuilder builder, final INodeSupplier nodeSupplier) { super(builder); + configureInternalSession(); this.nodeSupplier = Objects.requireNonNull( nodeSupplier, SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); - password = null; - useEncryptedPassword = false; - enableAutoFetch = false; - enableRedirection = false; - enableQueryRedirection = false; - rpcServiceName = RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE; } @Override - public synchronized void open(boolean enableRPCCompaction, int connectionTimeoutInMs) - throws IoTDBConnectionException { - super.open(enableRPCCompaction, connectionTimeoutInMs, null, null, nodeSupplier); + protected final INodeSupplier getNodeSupplier( + final boolean enableRPCCompaction, final int connectionTimeoutInMs) { + return nodeSupplier; } @Override - public SessionConnection constructSessionConnection( - Session session, TEndPoint endpoint, ZoneId zoneId) throws IoTDBConnectionException { + protected final TEndPoint resolveDefaultEndPoint( + final TEndPoint requestedEndPoint, final INodeSupplier nodeSupplier) + throws IoTDBConnectionException { final List candidates = nodeSupplier.get(); if (candidates.isEmpty()) { throw new IoTDBConnectionException( SessionMessages.EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081); } - defaultEndPoint = - candidates.contains(endpoint) - ? endpoint.deepCopy() - : nodeSupplier - .getQueryEndPoint() - .filter(candidates::contains) - .orElse(candidates.get(0)) - .deepCopy(); - return super.constructSessionConnection(session, defaultEndPoint, zoneId); + return candidates.contains(requestedEndPoint) + ? requestedEndPoint.deepCopy() + : nodeSupplier + .getQueryEndPoint() + .filter(candidates::contains) + .orElse(candidates.get(0)) + .deepCopy(); } } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java similarity index 57% rename from iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.java rename to iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java index be84296f18f8e..fd4309e611dec 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/InternalTableSessionPool.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java @@ -17,23 +17,32 @@ * under the License. */ -package org.apache.iotdb.session.pool; +package org.apache.iotdb.commons.client.internal.session; import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.session.Session; import org.apache.iotdb.session.i18n.SessionMessages; import java.util.Objects; -/** Table session pool using caller-supplied internal endpoints and the caller's username. */ -public class InternalTableSessionPool extends TableSessionPool { - - protected InternalTableSessionPool( - AbstractSessionPoolBuilder builder, INodeSupplier nodeSupplier) { - super( - new SessionPool( - builder, - Objects.requireNonNull( - nodeSupplier, - SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392))); +/** Server-side builder for tree-model internal sessions. */ +public class InternalSessionBuilder extends Session.Builder { + + private final INodeSupplier nodeSupplier; + + public InternalSessionBuilder(final INodeSupplier nodeSupplier) { + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + } + + protected final INodeSupplier getNodeSupplier() { + return nodeSupplier; + } + + @Override + protected Session createSession() { + return new InternalSession(this, getNodeSupplier()); } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java new file mode 100644 index 0000000000000..561c7c3538ad4 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java @@ -0,0 +1,51 @@ +/* + * 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.commons.client.internal.session; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.rpc.IoTDBConnectionException; +import org.apache.iotdb.session.Session; +import org.apache.iotdb.session.TableSessionBuilder; +import org.apache.iotdb.session.i18n.SessionMessages; + +import java.util.Objects; + +/** Server-side builder for table-model internal sessions. */ +public class InternalTableSessionBuilder extends TableSessionBuilder { + + private final INodeSupplier nodeSupplier; + + public InternalTableSessionBuilder(final INodeSupplier nodeSupplier) { + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + } + + protected final INodeSupplier getNodeSupplier() { + return nodeSupplier; + } + + @Override + protected Session constructNewSession() throws IoTDBConnectionException { + this.sqlDialect = "table"; + return openSession(new InternalSession(this, getNodeSupplier())); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java new file mode 100644 index 0000000000000..cf9793d46e8ed --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java @@ -0,0 +1,66 @@ +/* + * 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.commons.client.internal.session; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.session.Session; +import org.apache.iotdb.session.i18n.SessionMessages; +import org.apache.iotdb.session.pool.SessionPool; +import org.apache.iotdb.session.pool.TableSessionPoolBuilder; + +import java.util.Objects; + +/** Server-side builder for lazily-created table-model internal session pools. */ +public class InternalTableSessionPoolBuilder extends TableSessionPoolBuilder { + + private final INodeSupplier nodeSupplier; + + public InternalTableSessionPoolBuilder(final INodeSupplier nodeSupplier) { + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + } + + protected final INodeSupplier getNodeSupplier() { + return nodeSupplier; + } + + @Override + protected SessionPool createSessionPool() { + return new InternalSessionPool(this, getNodeSupplier()); + } + + private static final class InternalSessionPool extends SessionPool { + + private final INodeSupplier nodeSupplier; + + private InternalSessionPool( + final InternalTableSessionPoolBuilder builder, final INodeSupplier nodeSupplier) { + super(builder, true); + this.nodeSupplier = nodeSupplier; + } + + @Override + protected Session constructNewSession() { + return new InternalSession(createSessionBuilder(), nodeSupplier); + } + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java new file mode 100644 index 0000000000000..0e00bd9c0e313 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java @@ -0,0 +1,69 @@ +/* + * 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.commons.client.internal.subscription; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.rpc.IoTDBConnectionException; +import org.apache.iotdb.session.AbstractSessionBuilder; +import org.apache.iotdb.session.i18n.SessionMessages; +import org.apache.iotdb.session.subscription.SubscriptionSessionWrapper; + +import java.util.List; +import java.util.Objects; + +final class InternalSubscriptionSessionWrapper extends SubscriptionSessionWrapper { + + private final INodeSupplier nodeSupplier; + + InternalSubscriptionSessionWrapper( + final AbstractSessionBuilder builder, final INodeSupplier nodeSupplier) { + super(builder); + configureInternalSession(); + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + } + + @Override + protected INodeSupplier getNodeSupplier( + final boolean enableRPCCompaction, final int connectionTimeoutInMs) { + return nodeSupplier; + } + + @Override + protected TEndPoint resolveDefaultEndPoint( + final TEndPoint requestedEndPoint, final INodeSupplier availableNodes) + throws IoTDBConnectionException { + final List candidates = availableNodes.get(); + if (candidates.isEmpty()) { + throw new IoTDBConnectionException( + SessionMessages.EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081); + } + return candidates.contains(requestedEndPoint) + ? requestedEndPoint.deepCopy() + : availableNodes + .getQueryEndPoint() + .filter(candidates::contains) + .orElse(candidates.get(0)) + .deepCopy(); + } +} diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/InternalSubscriptionTablePullConsumer.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumer.java similarity index 86% rename from iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/InternalSubscriptionTablePullConsumer.java rename to iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumer.java index 7b2b22e2f5b08..1b688fab18b60 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/InternalSubscriptionTablePullConsumer.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumer.java @@ -17,22 +17,23 @@ * under the License. */ -package org.apache.iotdb.session.subscription.consumer.table; +package org.apache.iotdb.commons.client.internal.subscription; import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.rpc.RpcServiceConstants; import org.apache.iotdb.session.i18n.SessionMessages; import org.apache.iotdb.session.subscription.consumer.base.AbstractSubscriptionProvider; +import org.apache.iotdb.session.subscription.consumer.table.SubscriptionTablePullConsumer; +import org.apache.iotdb.session.subscription.consumer.table.SubscriptionTablePullConsumerBuilder; import java.util.Objects; -/** Table pull consumer using the caller's complete internal endpoint source and username. */ -public class InternalSubscriptionTablePullConsumer extends SubscriptionTablePullConsumer { +final class InternalSubscriptionTablePullConsumer extends SubscriptionTablePullConsumer { private final INodeSupplier nodeSupplier; - public InternalSubscriptionTablePullConsumer( + InternalSubscriptionTablePullConsumer( final SubscriptionTablePullConsumerBuilder builder, final INodeSupplier nodeSupplier) { super(builder); this.nodeSupplier = diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java new file mode 100644 index 0000000000000..0d136ea78a684 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java @@ -0,0 +1,50 @@ +/* + * 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.commons.client.internal.subscription; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.session.i18n.SessionMessages; +import org.apache.iotdb.session.subscription.consumer.ISubscriptionTablePullConsumer; +import org.apache.iotdb.session.subscription.consumer.table.SubscriptionTablePullConsumerBuilder; + +import java.util.Objects; + +/** Server-side builder for table subscription pull consumers. */ +public class InternalSubscriptionTablePullConsumerBuilder + extends SubscriptionTablePullConsumerBuilder { + + private final INodeSupplier nodeSupplier; + + public InternalSubscriptionTablePullConsumerBuilder(final INodeSupplier nodeSupplier) { + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + } + + protected final INodeSupplier getNodeSupplier() { + return nodeSupplier; + } + + @Override + protected ISubscriptionTablePullConsumer createConsumer() { + return new InternalSubscriptionTablePullConsumer(this, getNodeSupplier()); + } +} diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSession.java similarity index 61% rename from iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java rename to iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSession.java index ae67fa4093da0..f4aa5bbd30ce6 100644 --- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/IInternalSessionBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSession.java @@ -17,16 +17,16 @@ * under the License. */ -package org.apache.iotdb.isession; +package org.apache.iotdb.commons.client.internal.subscription; -import org.apache.iotdb.rpc.IoTDBConnectionException; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.session.AbstractSessionBuilder; +import org.apache.iotdb.session.subscription.SubscriptionTableSession; -/** Capability implemented by builders that can construct sessions on internal endpoints. */ -public interface IInternalSessionBuilder { +final class InternalSubscriptionTableSession extends SubscriptionTableSession { - /** - * Builds a session using the caller-provided internal node supplier. Whether the returned session - * is already open follows the builder's ordinary build contract. - */ - T buildInternal(INodeSupplier nodeSupplier) throws IoTDBConnectionException; + InternalSubscriptionTableSession( + final AbstractSessionBuilder builder, final INodeSupplier nodeSupplier) { + super(new InternalSubscriptionSessionWrapper(builder, nodeSupplier)); + } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java new file mode 100644 index 0000000000000..b34e07b621a4e --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java @@ -0,0 +1,49 @@ +/* + * 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.commons.client.internal.subscription; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.session.i18n.SessionMessages; +import org.apache.iotdb.session.subscription.ISubscriptionTableSession; +import org.apache.iotdb.session.subscription.SubscriptionTableSessionBuilder; + +import java.util.Objects; + +/** Server-side builder for table subscription-management sessions. */ +public class InternalSubscriptionTableSessionBuilder extends SubscriptionTableSessionBuilder { + + private final INodeSupplier nodeSupplier; + + public InternalSubscriptionTableSessionBuilder(final INodeSupplier nodeSupplier) { + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + } + + protected final INodeSupplier getNodeSupplier() { + return nodeSupplier; + } + + @Override + protected ISubscriptionTableSession createSubscriptionSession() { + return new InternalSubscriptionTableSession(this, getNodeSupplier()); + } +} diff --git a/iotdb-client/session/src/test/java/org/apache/iotdb/session/InternalSessionProtocolTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalSessionProtocolTest.java similarity index 88% rename from iotdb-client/session/src/test/java/org/apache/iotdb/session/InternalSessionProtocolTest.java rename to iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalSessionProtocolTest.java index b1331c30f32e9..996a47d902d28 100644 --- a/iotdb-client/session/src/test/java/org/apache/iotdb/session/InternalSessionProtocolTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalSessionProtocolTest.java @@ -17,7 +17,7 @@ * under the License. */ -package org.apache.iotdb.session; +package org.apache.iotdb.commons.client.internal.session; import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.common.rpc.thrift.TSStatus; @@ -30,6 +30,9 @@ import org.apache.iotdb.service.rpc.thrift.TSOpenSessionReq; import org.apache.iotdb.service.rpc.thrift.TSOpenSessionResp; import org.apache.iotdb.service.rpc.thrift.TSProtocolVersion; +import org.apache.iotdb.session.DummyNodesSupplier; +import org.apache.iotdb.session.Session; +import org.apache.iotdb.session.SessionConnection; import org.apache.thrift.TMultiplexedProcessor; import org.apache.thrift.protocol.TBinaryProtocol; @@ -38,8 +41,8 @@ import org.junit.After; import org.junit.Before; import org.junit.Test; -import org.powermock.reflect.Whitebox; +import java.lang.reflect.Field; import java.net.ServerSocket; import java.util.Collections; import java.util.concurrent.CountDownLatch; @@ -73,21 +76,27 @@ public void tearDown() throws Exception { public void testReconnectRetainsMultiplexedClientServiceName() throws Exception { final INodeSupplier nodeSupplier = new DummyNodesSupplier(Collections.singletonList(server.endpoint)); - try (InternalSession session = - new Session.Builder() + try (Session session = + new InternalSessionBuilder(nodeSupplier) .host(server.endpoint.getIp()) .port(server.endpoint.getPort()) .username("root") - .buildInternal(nodeSupplier)) { + .build()) { session.open(); final SessionConnection connection = - Whitebox.getInternalState(session, "defaultSessionConnection"); - assertTrue(Whitebox.invokeMethod(connection, "reconnect")); + (SessionConnection) getField(session, "defaultSessionConnection"); + assertTrue(connection.reconnect()); assertEquals(2, openSessionCalls.get()); } } + private static Object getField(final Object target, final String fieldName) throws Exception { + final Field field = Session.class.getDeclaredField(fieldName); + field.setAccessible(true); + return field.get(target); + } + private static final class TestServer implements AutoCloseable { private final TEndPoint endpoint; private final TServerSocket socket; From f1a7762bbd6f2a68f29aef08a59ca2a6d126e982 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 11:20:34 +0800 Subject: [PATCH 12/32] Refine internal session extension points --- .../java/org/apache/iotdb/session/Session.java | 4 ---- .../apache/iotdb/session/TableSessionBuilder.java | 6 +----- .../org/apache/iotdb/session/pool/SessionPool.java | 6 +----- .../iotdb/session/pool/TableSessionPool.java | 2 +- .../session/pool/TableSessionPoolBuilder.java | 10 +++------- .../subscription/SubscriptionSessionWrapper.java | 2 +- .../session/InternalTableSessionPoolBuilder.java | 14 ++++++++++++-- 7 files changed, 19 insertions(+), 25 deletions(-) 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 d9b5687775647..df71f23e47caa 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 @@ -4468,10 +4468,6 @@ protected void setDefaultSessionConnection(SessionConnection defaultSessionConne this.defaultSessionConnection = defaultSessionConnection; } - protected void setRpcServiceName(String rpcServiceName) { - this.rpcServiceName = rpcServiceName; - } - /** Applies the common transport and authentication settings for server internal sessions. */ protected final void configureInternalSession() { password = null; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java index b41233ebe3933..cc6b6bf440a21 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java @@ -432,7 +432,7 @@ public TableSessionBuilder tabletCompressionMinRowSize(int tabletCompressionMinR * @throws IoTDBConnectionException if an error occurs while establishing the connection. */ public ITableSession build() throws IoTDBConnectionException { - return wrapTableSession(constructNewSession()); + return new TableSession(constructNewSession()); } protected Session constructNewSession() throws IoTDBConnectionException { @@ -444,10 +444,6 @@ protected Session constructNewSession() throws IoTDBConnectionException { return openSession(new Session(this)); } - protected ITableSession wrapTableSession(Session session) { - return new TableSession(session); - } - protected Session openSession(Session newSession) throws IoTDBConnectionException { newSession.enableIoTDBRpcCompression = isIoTDBRpcCompressionEnabled; newSession.tabletCompressionMinRowSize = tabletCompressionMinRowSize; 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 3b5dd914b2d21..91cdd9b324bd0 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 @@ -640,7 +640,7 @@ protected Session constructNewSession() { return session; } - protected void openSession(final ISession session) throws IoTDBConnectionException { + private void openSession(final ISession session) throws IoTDBConnectionException { if (internalSessionPool) { session.open(enableThriftCompression, connectionTimeoutInMs); } else { @@ -653,10 +653,6 @@ protected void openSession(final ISession session) throws IoTDBConnectionExcepti } } - protected final boolean isInternalSessionPool() { - return internalSessionPool; - } - private void initThreadPool() { this.executorService = Executors.newSingleThreadScheduledExecutor( diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPool.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPool.java index 4e08f202a14ad..5e3a53c4462ff 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPool.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPool.java @@ -27,7 +27,7 @@ public class TableSessionPool implements ITableSessionPool { private final SessionPool sessionPool; - TableSessionPool(SessionPool sessionPool) { + protected TableSessionPool(SessionPool sessionPool) { this.sessionPool = sessionPool; } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java index 9f4b1364d02ec..5b7764f52ad1e 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java @@ -343,14 +343,10 @@ public ITableSessionPool build() { Collections.singletonList(SessionConfig.DEFAULT_HOST + ":" + SessionConfig.DEFAULT_PORT); } this.sqlDialect = "table"; - return createTableSessionPool(createSessionPool()); + return createSessionPool(); } - protected SessionPool createSessionPool() { - return new SessionPool(this); - } - - protected ITableSessionPool createTableSessionPool(final SessionPool sessionPool) { - return new TableSessionPool(sessionPool); + protected ITableSessionPool createSessionPool() { + return new TableSessionPool(new SessionPool(this)); } } diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java index 6538eea370686..272671cabd946 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java @@ -73,7 +73,7 @@ public int getThriftMaxFrameSize() { @Override public void setRpcServiceName(final String rpcServiceName) { - super.setRpcServiceName(rpcServiceName); + this.rpcServiceName = rpcServiceName; } public void setSslConfig( diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java index cf9793d46e8ed..8bc9c1d2f17b5 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java @@ -20,9 +20,11 @@ package org.apache.iotdb.commons.client.internal.session; import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.isession.pool.ITableSessionPool; import org.apache.iotdb.session.Session; import org.apache.iotdb.session.i18n.SessionMessages; import org.apache.iotdb.session.pool.SessionPool; +import org.apache.iotdb.session.pool.TableSessionPool; import org.apache.iotdb.session.pool.TableSessionPoolBuilder; import java.util.Objects; @@ -44,8 +46,16 @@ protected final INodeSupplier getNodeSupplier() { } @Override - protected SessionPool createSessionPool() { - return new InternalSessionPool(this, getNodeSupplier()); + protected ITableSessionPool createSessionPool() { + return new InternalTableSessionPool(this, getNodeSupplier()); + } + + private static final class InternalTableSessionPool extends TableSessionPool { + + private InternalTableSessionPool( + final InternalTableSessionPoolBuilder builder, final INodeSupplier nodeSupplier) { + super(new InternalSessionPool(builder, nodeSupplier)); + } } private static final class InternalSessionPool extends SessionPool { From 5760a366e86bd97959c8ac995698e432919f115d Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 11:54:22 +0800 Subject: [PATCH 13/32] Use session hooks for internal pool behavior --- .../org/apache/iotdb/session/Session.java | 7 +- .../iotdb/session/pool/SessionPool.java | 129 ++++++++---------- .../SubscriptionSessionWrapper.java | 1 - .../InternalTableSessionPoolBuilder.java | 9 +- 4 files changed, 65 insertions(+), 81 deletions(-) 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 df71f23e47caa..bcc1b38630c7c 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 @@ -556,9 +556,6 @@ public synchronized void open(boolean enableRPCCompaction, int connectionTimeout this.availableNodes = null; } - if (enableAutoFetch) { - initThreadPool(); - } this.availableNodes = getNodeSupplier(enableRPCCompaction, connectionTimeoutInMs); this.defaultEndPoint = resolveDefaultEndPoint(this.defaultEndPoint, this.availableNodes); @@ -598,6 +595,7 @@ private void initThreadPool() { protected INodeSupplier getNodeSupplier( final boolean enableRPCCompaction, final int connectionTimeoutInMs) { if (enableAutoFetch) { + initThreadPool(); return NodesSupplier.createNodeSupplier( getNodeUrls(), SessionUtils.parseSeedNodeUrls(initialNodeUrls), @@ -4472,9 +4470,6 @@ protected void setDefaultSessionConnection(SessionConnection defaultSessionConne protected final void configureInternalSession() { password = null; useEncryptedPassword = false; - enableAutoFetch = false; - enableRedirection = false; - enableQueryRedirection = false; rpcServiceName = RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE; } 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 91cdd9b324bd0..590e6658e5c28 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 @@ -157,7 +157,6 @@ public class SessionPool implements ISessionPool { private final int connectionTimeoutInMs; private final boolean enableIoTDBRpcCompression; private final boolean enableThriftCompression; - private final boolean internalSessionPool; // whether the queue is closed. private boolean closed; @@ -173,6 +172,7 @@ public class SessionPool implements ISessionPool { private volatile ScheduledExecutorService executorService; private INodeSupplier availableNodes; + private boolean nodeSupplierInitialized; // set to true, means that we will start a background thread to fetch all available (Status is // not Removing) datanodes in cluster, and these available nodes will be used in retrying stage @@ -432,8 +432,7 @@ public SessionPool( this.thriftDefaultBufferSize = thriftDefaultBufferSize; this.thriftMaxFrameSize = thriftMaxFrameSize; this.formattedNodeUrls = UrlUtils.formatTEndPointIpv4AndIpv6Url(host, port); - initThreadPool(); - initAvailableNodes(Collections.singletonList(new TEndPoint(host, port))); + initializeNodeSupplier(); } public SessionPool( @@ -479,8 +478,7 @@ public SessionPool( this.trustStore = trustStore; this.trustStorePwd = trustStorePwd; this.sslProtocol = SessionConfig.DEFAULT_SSL_PROTOCOL; - initThreadPool(); - initAvailableNodes(Collections.singletonList(new TEndPoint(host, port))); + initializeNodeSupplier(); } @SuppressWarnings("squid:S107") // ignore Methods should not have too many parameters @@ -522,26 +520,19 @@ public SessionPool( this.thriftDefaultBufferSize = thriftDefaultBufferSize; this.thriftMaxFrameSize = thriftMaxFrameSize; this.formattedNodeUrls = this.nodeUrls.toString(); - initThreadPool(); - initAvailableNodes(SessionUtils.parseSeedNodeUrls(this.nodeUrls)); + initializeNodeSupplier(); } public SessionPool(AbstractSessionPoolBuilder builder) { - this(builder, false); - } - - protected SessionPool( - final AbstractSessionPoolBuilder builder, final boolean internalSessionPool) { - this.internalSessionPool = internalSessionPool; this.maxSize = builder.maxSize; this.user = builder.username; - this.password = internalSessionPool ? null : builder.pw; + this.password = builder.pw; this.fetchSize = builder.fetchSize; this.waitToGetSessionTimeoutInMs = builder.waitToGetSessionTimeoutInMs; this.enableThriftCompression = builder.isThriftRpcCompactionEnabled; this.enableIoTDBRpcCompression = builder.isIoTDBRpcCompressionEnabled; this.zoneId = builder.zoneId; - this.enableRedirection = !internalSessionPool && builder.enableRedirection; + this.enableRedirection = builder.enableRedirection; if (this.enableRedirection) { deviceIdToEndpoint = new ConcurrentHashMap<>(); tableModelDeviceIdToEndpoint = new ConcurrentHashMap<>(); @@ -551,7 +542,7 @@ protected SessionPool( this.version = builder.version; this.thriftDefaultBufferSize = builder.thriftDefaultBufferSize; this.thriftMaxFrameSize = builder.thriftMaxFrameSize; - this.enableAutoFetch = !internalSessionPool && builder.enableAutoFetch; + this.enableAutoFetch = builder.enableAutoFetch; this.useSSL = builder.useSSL; this.trustStore = builder.trustStore; this.trustStorePwd = builder.trustStorePwd; @@ -565,15 +556,7 @@ protected SessionPool( this.queryTimeoutInMs = builder.timeOut; this.enableAutoCreateSchema = builder.enableAutoCreateSchema; - if (enableAutoFetch) { - initThreadPool(); - } - if (internalSessionPool) { - this.host = null; - this.port = -1; - this.nodeUrls = null; - this.formattedNodeUrls = builder.getClass().getSimpleName(); - } else if (builder.nodeUrls != null) { + if (builder.nodeUrls != null) { if (builder.nodeUrls.isEmpty()) { throw new IllegalArgumentException(SessionMessages.NODE_URLS_EMPTY); } @@ -581,23 +564,11 @@ protected SessionPool( this.host = null; this.port = -1; this.formattedNodeUrls = builder.nodeUrls.toString(); - if (enableAutoFetch) { - initAvailableNodes(SessionUtils.parseSeedNodeUrls(nodeUrls)); - } else { - this.availableNodes = new DummyNodesSupplier(SessionUtils.parseSeedNodeUrls(nodeUrls)); - } - } else { this.host = builder.host; this.port = builder.rpcPort; this.nodeUrls = null; this.formattedNodeUrls = UrlUtils.formatTEndPointIpv4AndIpv6Url(host, port); - if (enableAutoFetch) { - initAvailableNodes(Collections.singletonList(new TEndPoint(host, port))); - } else { - this.availableNodes = - new DummyNodesSupplier(Collections.singletonList(new TEndPoint(host, port))); - } } } @@ -640,17 +611,14 @@ protected Session constructNewSession() { return session; } - private void openSession(final ISession session) throws IoTDBConnectionException { - if (internalSessionPool) { - session.open(enableThriftCompression, connectionTimeoutInMs); - } else { - session.open( - enableThriftCompression, - connectionTimeoutInMs, - deviceIdToEndpoint, - tableModelDeviceIdToEndpoint, - availableNodes); - } + protected void openSession(final ISession session) throws IoTDBConnectionException { + initializeNodeSupplier(); + session.open( + enableThriftCompression, + connectionTimeoutInMs, + deviceIdToEndpoint, + tableModelDeviceIdToEndpoint, + availableNodes); } private void initThreadPool() { @@ -670,25 +638,40 @@ private void initThreadPool() { }); } - private void initAvailableNodes(List endPointList) { - this.availableNodes = - NodesSupplier.createNodeSupplier( - endPointList, - executorService, - user, - password, - zoneId, - thriftDefaultBufferSize, - thriftMaxFrameSize, - connectionTimeoutInMs, - useSSL, - trustStore, - trustStorePwd, - keyStore, - keyStorePwd, - sslProtocol, - enableThriftCompression, - version.toString()); + private synchronized void initializeNodeSupplier() { + if (nodeSupplierInitialized) { + return; + } + final List endPointList = + nodeUrls == null + ? Collections.singletonList(new TEndPoint(host, port)) + : SessionUtils.parseSeedNodeUrls(nodeUrls); + this.availableNodes = createNodeSupplier(endPointList); + this.nodeSupplierInitialized = true; + } + + protected INodeSupplier createNodeSupplier(final List endPointList) { + if (enableAutoFetch) { + initThreadPool(); + return NodesSupplier.createNodeSupplier( + endPointList, + executorService, + user, + password, + zoneId, + thriftDefaultBufferSize, + thriftMaxFrameSize, + connectionTimeoutInMs, + useSSL, + trustStore, + trustStorePwd, + keyStore, + keyStorePwd, + sslProtocol, + enableThriftCompression, + version.toString()); + } + return new DummyNodesSupplier(endPointList); } // if this method throws an exception, either the server is broken, or the ip/port/user/password @@ -3532,16 +3515,16 @@ public boolean isEnableThriftCompression() { @Override public void setEnableRedirection(boolean enableRedirection) { - this.enableRedirection = !internalSessionPool && enableRedirection; + this.enableRedirection = enableRedirection; if (this.enableRedirection) { deviceIdToEndpoint = new ConcurrentHashMap<>(); tableModelDeviceIdToEndpoint = new ConcurrentHashMap<>(); } for (ISession session : queue) { - session.setEnableRedirection(this.enableRedirection); + session.setEnableRedirection(enableRedirection); } for (ISession session : occupied.keySet()) { - session.setEnableRedirection(this.enableRedirection); + session.setEnableRedirection(enableRedirection); } } @@ -3552,12 +3535,12 @@ public boolean isEnableRedirection() { @Override public void setEnableQueryRedirection(boolean enableQueryRedirection) { - this.enableQueryRedirection = !internalSessionPool && enableQueryRedirection; + this.enableQueryRedirection = enableQueryRedirection; for (ISession session : queue) { - session.setEnableQueryRedirection(this.enableQueryRedirection); + session.setEnableQueryRedirection(enableQueryRedirection); } for (ISession session : occupied.keySet()) { - session.setEnableQueryRedirection(this.enableQueryRedirection); + session.setEnableQueryRedirection(enableQueryRedirection); } } diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java index 272671cabd946..e7a51d237a0c1 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java @@ -71,7 +71,6 @@ public int getThriftMaxFrameSize() { return thriftMaxFrameSize; } - @Override public void setRpcServiceName(final String rpcServiceName) { this.rpcServiceName = rpcServiceName; } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java index 8bc9c1d2f17b5..23e400b320de3 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java @@ -20,7 +20,9 @@ package org.apache.iotdb.commons.client.internal.session; import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.isession.ISession; import org.apache.iotdb.isession.pool.ITableSessionPool; +import org.apache.iotdb.rpc.IoTDBConnectionException; import org.apache.iotdb.session.Session; import org.apache.iotdb.session.i18n.SessionMessages; import org.apache.iotdb.session.pool.SessionPool; @@ -64,7 +66,7 @@ private static final class InternalSessionPool extends SessionPool { private InternalSessionPool( final InternalTableSessionPoolBuilder builder, final INodeSupplier nodeSupplier) { - super(builder, true); + super(builder); this.nodeSupplier = nodeSupplier; } @@ -72,5 +74,10 @@ private InternalSessionPool( protected Session constructNewSession() { return new InternalSession(createSessionBuilder(), nodeSupplier); } + + @Override + protected void openSession(final ISession session) throws IoTDBConnectionException { + session.open(isEnableThriftCompression(), getConnectionTimeoutInMs()); + } } } From 9af3ad0867c4469aaa6316765e63b88e94047ba6 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 12:06:28 +0800 Subject: [PATCH 14/32] Simplify internal node supplier ownership --- .../java/org/apache/iotdb/session/Session.java | 17 ++++++----------- .../InternalTableSessionPoolBuilder.java | 8 ++++---- 2 files changed, 10 insertions(+), 15 deletions(-) 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 bcc1b38630c7c..3c1e14ba19309 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 @@ -544,15 +544,12 @@ public synchronized void open(boolean enableRPCCompaction, int connectionTimeout return; } - final boolean closeAvailableNodes = executorService != null; if (this.executorService != null) { this.executorService.shutdown(); this.executorService = null; } if (this.availableNodes != null) { - if (closeAvailableNodes) { - this.availableNodes.close(); - } + this.availableNodes.close(); this.availableNodes = null; } @@ -659,6 +656,7 @@ public synchronized void open( } this.availableNodes = nodesSupplier; + this.defaultEndPoint = resolveDefaultEndPoint(this.defaultEndPoint, this.availableNodes); this.enableThriftRpcCompaction = enableThriftRpcCompaction; this.connectionTimeoutInMs = connectionTimeoutInMs; setDefaultSessionConnection(constructSessionConnection(this, defaultEndPoint, zoneId)); @@ -685,6 +683,7 @@ public synchronized void open( } this.availableNodes = nodesSupplier; + this.defaultEndPoint = resolveDefaultEndPoint(this.defaultEndPoint, this.availableNodes); this.enableThriftRpcCompaction = enableThriftRpcCompaction; this.connectionTimeoutInMs = connectionTimeoutInMs; setDefaultSessionConnection(constructSessionConnection(this, defaultEndPoint, zoneId)); @@ -713,16 +712,12 @@ public synchronized void close() throws IoTDBConnectionException { getDefaultSessionConnection().close(); setDefaultSessionConnection(null); } finally { - // A node supplier created outside this session, such as one from SessionPool or a server - // internal session, is not owned by this transport. if (this.executorService != null) { this.executorService.shutdown(); this.executorService = null; - if (this.availableNodes != null) { - this.availableNodes.close(); - this.availableNodes = null; - } - } else { + } + if (this.availableNodes != null) { + this.availableNodes.close(); this.availableNodes = null; } isClosed = true; diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java index 23e400b320de3..70b107e6cdb9e 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java @@ -19,16 +19,16 @@ package org.apache.iotdb.commons.client.internal.session; +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.pool.ITableSessionPool; -import org.apache.iotdb.rpc.IoTDBConnectionException; import org.apache.iotdb.session.Session; import org.apache.iotdb.session.i18n.SessionMessages; import org.apache.iotdb.session.pool.SessionPool; import org.apache.iotdb.session.pool.TableSessionPool; import org.apache.iotdb.session.pool.TableSessionPoolBuilder; +import java.util.List; import java.util.Objects; /** Server-side builder for lazily-created table-model internal session pools. */ @@ -76,8 +76,8 @@ protected Session constructNewSession() { } @Override - protected void openSession(final ISession session) throws IoTDBConnectionException { - session.open(isEnableThriftCompression(), getConnectionTimeoutInMs()); + protected INodeSupplier createNodeSupplier(final List ignored) { + return nodeSupplier; } } } From 7be93fc482f07312760939eb06ffacc374ca3d3c Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 12:14:14 +0800 Subject: [PATCH 15/32] Support configurable internal session endpoints --- .../iotdb/session/i18n/SessionMessages.java | 3 ++ .../iotdb/session/i18n/SessionMessages.java | 2 ++ .../org/apache/iotdb/session/Session.java | 21 ++--------- .../SubscriptionSessionWrapper.java | 3 +- .../internal/session/InternalSession.java | 36 ++++--------------- .../session/InternalSessionBuilder.java | 15 ++++++++ .../session/InternalTableSessionBuilder.java | 15 ++++++++ .../InternalTableSessionPoolBuilder.java | 18 ++++++++-- .../InternalSubscriptionSessionWrapper.java | 36 ++++--------------- ...InternalSubscriptionTablePullConsumer.java | 8 +---- ...lSubscriptionTablePullConsumerBuilder.java | 15 ++++++++ ...ternalSubscriptionTableSessionBuilder.java | 21 +++++++++++ 12 files changed, 104 insertions(+), 89 deletions(-) diff --git a/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java b/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java index e256a062265a2..341aae9e3de62 100644 --- a/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java +++ b/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java @@ -309,6 +309,9 @@ private SessionMessages() {} public static final String EXCEPTION_SESSION_CAN_NOT_CONNECT_TO_ARG_226FBA66 = "Session can not connect to {}"; + public static final String EXCEPTION_END_POINT_MUST_NOT_BE_NULL_908A4B27 = + "Endpoint must not be null."; + public static final String EXCEPTION_ALL_VALUES_ARE_NULL_AND_THIS_SUBMISSION_IS_IGNORED_DEVICEID_IS_ARG_TIME_IS_ARG_MEASUREMENTS_IS_ARG_5524221C = "All values are null and this submission is ignored,deviceId is [{}],time is" diff --git a/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java b/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java index 32d4eda5083ef..ed04bf7126f02 100644 --- a/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java +++ b/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java @@ -306,6 +306,8 @@ private SessionMessages() {} public static final String EXCEPTION_SESSION_CAN_NOT_CONNECT_TO_ARG_226FBA66 = "Session 无法连接到 {}"; + public static final String EXCEPTION_END_POINT_MUST_NOT_BE_NULL_908A4B27 = "Endpoint 不能为空。"; + public static final String EXCEPTION_ALL_VALUES_ARE_NULL_AND_THIS_SUBMISSION_IS_IGNORED_DEVICEID_IS_ARG_TIME_IS_ARG_MEASUREMENTS_IS_ARG_5524221C = "所有值均为空，本次提交被忽略，deviceId 为 [{}]，time 为 [{}]，measurements 为 [{}]"; 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 3c1e14ba19309..a4f523fb44ef7 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 @@ -31,7 +31,6 @@ import org.apache.iotdb.rpc.IoTDBConnectionException; import org.apache.iotdb.rpc.NoValidValueException; import org.apache.iotdb.rpc.RedirectException; -import org.apache.iotdb.rpc.RpcServiceConstants; import org.apache.iotdb.rpc.StatementExecutionException; import org.apache.iotdb.rpc.UrlUtils; import org.apache.iotdb.service.rpc.thrift.TCreateTimeseriesUsingSchemaTemplateReq; @@ -554,7 +553,6 @@ public synchronized void open(boolean enableRPCCompaction, int connectionTimeout } this.availableNodes = getNodeSupplier(enableRPCCompaction, connectionTimeoutInMs); - this.defaultEndPoint = resolveDefaultEndPoint(this.defaultEndPoint, this.availableNodes); this.connectionTimeoutInMs = connectionTimeoutInMs; setDefaultSessionConnection(constructSessionConnection(this, defaultEndPoint, zoneId)); @@ -618,16 +616,6 @@ protected INodeSupplier getNodeSupplier( return new DummyNodesSupplier(new ArrayList<>(staticNodes)); } - /** - * Selects the endpoint used for the initial connection. Internal sessions that support endpoint - * discovery override this hook so a stale builder endpoint cannot be used. - */ - protected TEndPoint resolveDefaultEndPoint( - final TEndPoint requestedEndPoint, final INodeSupplier nodeSupplier) - throws IoTDBConnectionException { - return requestedEndPoint; - } - private static List shuffleNodeUrls(List endPoints) { try { Collections.shuffle(endPoints); @@ -656,7 +644,6 @@ public synchronized void open( } this.availableNodes = nodesSupplier; - this.defaultEndPoint = resolveDefaultEndPoint(this.defaultEndPoint, this.availableNodes); this.enableThriftRpcCompaction = enableThriftRpcCompaction; this.connectionTimeoutInMs = connectionTimeoutInMs; setDefaultSessionConnection(constructSessionConnection(this, defaultEndPoint, zoneId)); @@ -683,7 +670,6 @@ public synchronized void open( } this.availableNodes = nodesSupplier; - this.defaultEndPoint = resolveDefaultEndPoint(this.defaultEndPoint, this.availableNodes); this.enableThriftRpcCompaction = enableThriftRpcCompaction; this.connectionTimeoutInMs = connectionTimeoutInMs; setDefaultSessionConnection(constructSessionConnection(this, defaultEndPoint, zoneId)); @@ -4461,11 +4447,8 @@ protected void setDefaultSessionConnection(SessionConnection defaultSessionConne this.defaultSessionConnection = defaultSessionConnection; } - /** Applies the common transport and authentication settings for server internal sessions. */ - protected final void configureInternalSession() { - password = null; - useEncryptedPassword = false; - rpcServiceName = RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE; + protected void setRpcServiceName(final String rpcServiceName) { + this.rpcServiceName = rpcServiceName; } public boolean isEnableThriftRpcCompaction() { diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java index e7a51d237a0c1..6538eea370686 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java @@ -71,8 +71,9 @@ public int getThriftMaxFrameSize() { return thriftMaxFrameSize; } + @Override public void setRpcServiceName(final String rpcServiceName) { - this.rpcServiceName = rpcServiceName; + super.setRpcServiceName(rpcServiceName); } public void setSslConfig( diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java index 403a977ddc59c..f80accedadc56 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java @@ -19,15 +19,10 @@ package org.apache.iotdb.commons.client.internal.session; -import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; -import org.apache.iotdb.rpc.IoTDBConnectionException; +import org.apache.iotdb.rpc.RpcServiceConstants; import org.apache.iotdb.session.AbstractSessionBuilder; import org.apache.iotdb.session.Session; -import org.apache.iotdb.session.i18n.SessionMessages; - -import java.util.List; -import java.util.Objects; /** Base session for trusted server-to-server client RPC connections. */ public class InternalSession extends Session { @@ -37,34 +32,15 @@ public class InternalSession extends Session { protected InternalSession( final AbstractSessionBuilder builder, final INodeSupplier nodeSupplier) { super(builder); - configureInternalSession(); - this.nodeSupplier = - Objects.requireNonNull( - nodeSupplier, - SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + setRpcServiceName(RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE); + this.nodeSupplier = nodeSupplier; } @Override protected final INodeSupplier getNodeSupplier( final boolean enableRPCCompaction, final int connectionTimeoutInMs) { - return nodeSupplier; - } - - @Override - protected final TEndPoint resolveDefaultEndPoint( - final TEndPoint requestedEndPoint, final INodeSupplier nodeSupplier) - throws IoTDBConnectionException { - final List candidates = nodeSupplier.get(); - if (candidates.isEmpty()) { - throw new IoTDBConnectionException( - SessionMessages.EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081); - } - return candidates.contains(requestedEndPoint) - ? requestedEndPoint.deepCopy() - : nodeSupplier - .getQueryEndPoint() - .filter(candidates::contains) - .orElse(candidates.get(0)) - .deepCopy(); + return nodeSupplier != null + ? nodeSupplier + : super.getNodeSupplier(enableRPCCompaction, connectionTimeoutInMs); } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java index fd4309e611dec..90c9c2bbf4a5a 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java @@ -19,6 +19,7 @@ package org.apache.iotdb.commons.client.internal.session; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.session.Session; import org.apache.iotdb.session.i18n.SessionMessages; @@ -30,6 +31,10 @@ public class InternalSessionBuilder extends Session.Builder { private final INodeSupplier nodeSupplier; + public InternalSessionBuilder() { + this.nodeSupplier = null; + } + public InternalSessionBuilder(final INodeSupplier nodeSupplier) { this.nodeSupplier = Objects.requireNonNull( @@ -37,6 +42,16 @@ public InternalSessionBuilder(final INodeSupplier nodeSupplier) { SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); } + public InternalSessionBuilder endPoint(final TEndPoint endPoint) { + final TEndPoint endpoint = + Objects.requireNonNull( + endPoint, SessionMessages.EXCEPTION_END_POINT_MUST_NOT_BE_NULL_908A4B27); + this.host = endpoint.getIp(); + this.rpcPort = endpoint.getPort(); + this.nodeUrls = null; + return this; + } + protected final INodeSupplier getNodeSupplier() { return nodeSupplier; } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java index 561c7c3538ad4..0175c37fa37a0 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java @@ -19,6 +19,7 @@ package org.apache.iotdb.commons.client.internal.session; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.rpc.IoTDBConnectionException; import org.apache.iotdb.session.Session; @@ -32,6 +33,10 @@ public class InternalTableSessionBuilder extends TableSessionBuilder { private final INodeSupplier nodeSupplier; + public InternalTableSessionBuilder() { + this.nodeSupplier = null; + } + public InternalTableSessionBuilder(final INodeSupplier nodeSupplier) { this.nodeSupplier = Objects.requireNonNull( @@ -39,6 +44,16 @@ public InternalTableSessionBuilder(final INodeSupplier nodeSupplier) { SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); } + public InternalTableSessionBuilder endPoint(final TEndPoint endPoint) { + final TEndPoint endpoint = + Objects.requireNonNull( + endPoint, SessionMessages.EXCEPTION_END_POINT_MUST_NOT_BE_NULL_908A4B27); + this.host = endpoint.getIp(); + this.rpcPort = endpoint.getPort(); + this.nodeUrls = null; + return this; + } + protected final INodeSupplier getNodeSupplier() { return nodeSupplier; } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java index 70b107e6cdb9e..3879dc5f500d4 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java @@ -36,6 +36,10 @@ public class InternalTableSessionPoolBuilder extends TableSessionPoolBuilder { private final INodeSupplier nodeSupplier; + public InternalTableSessionPoolBuilder() { + this.nodeSupplier = null; + } + public InternalTableSessionPoolBuilder(final INodeSupplier nodeSupplier) { this.nodeSupplier = Objects.requireNonNull( @@ -43,6 +47,16 @@ public InternalTableSessionPoolBuilder(final INodeSupplier nodeSupplier) { SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); } + public InternalTableSessionPoolBuilder endPoint(final TEndPoint endPoint) { + final TEndPoint endpoint = + Objects.requireNonNull( + endPoint, SessionMessages.EXCEPTION_END_POINT_MUST_NOT_BE_NULL_908A4B27); + this.host = endpoint.getIp(); + this.rpcPort = endpoint.getPort(); + this.nodeUrls = null; + return this; + } + protected final INodeSupplier getNodeSupplier() { return nodeSupplier; } @@ -76,8 +90,8 @@ protected Session constructNewSession() { } @Override - protected INodeSupplier createNodeSupplier(final List ignored) { - return nodeSupplier; + protected INodeSupplier createNodeSupplier(final List endPointList) { + return nodeSupplier != null ? nodeSupplier : super.createNodeSupplier(endPointList); } } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java index 0e00bd9c0e313..1876c4a4c2d60 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java @@ -19,16 +19,11 @@ package org.apache.iotdb.commons.client.internal.subscription; -import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; -import org.apache.iotdb.rpc.IoTDBConnectionException; +import org.apache.iotdb.rpc.RpcServiceConstants; import org.apache.iotdb.session.AbstractSessionBuilder; -import org.apache.iotdb.session.i18n.SessionMessages; import org.apache.iotdb.session.subscription.SubscriptionSessionWrapper; -import java.util.List; -import java.util.Objects; - final class InternalSubscriptionSessionWrapper extends SubscriptionSessionWrapper { private final INodeSupplier nodeSupplier; @@ -36,34 +31,15 @@ final class InternalSubscriptionSessionWrapper extends SubscriptionSessionWrappe InternalSubscriptionSessionWrapper( final AbstractSessionBuilder builder, final INodeSupplier nodeSupplier) { super(builder); - configureInternalSession(); - this.nodeSupplier = - Objects.requireNonNull( - nodeSupplier, - SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + setRpcServiceName(RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE); + this.nodeSupplier = nodeSupplier; } @Override protected INodeSupplier getNodeSupplier( final boolean enableRPCCompaction, final int connectionTimeoutInMs) { - return nodeSupplier; - } - - @Override - protected TEndPoint resolveDefaultEndPoint( - final TEndPoint requestedEndPoint, final INodeSupplier availableNodes) - throws IoTDBConnectionException { - final List candidates = availableNodes.get(); - if (candidates.isEmpty()) { - throw new IoTDBConnectionException( - SessionMessages.EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081); - } - return candidates.contains(requestedEndPoint) - ? requestedEndPoint.deepCopy() - : availableNodes - .getQueryEndPoint() - .filter(candidates::contains) - .orElse(candidates.get(0)) - .deepCopy(); + return nodeSupplier != null + ? nodeSupplier + : super.getNodeSupplier(enableRPCCompaction, connectionTimeoutInMs); } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumer.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumer.java index 1b688fab18b60..9269c9232ece8 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumer.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumer.java @@ -22,13 +22,10 @@ import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.rpc.RpcServiceConstants; -import org.apache.iotdb.session.i18n.SessionMessages; import org.apache.iotdb.session.subscription.consumer.base.AbstractSubscriptionProvider; import org.apache.iotdb.session.subscription.consumer.table.SubscriptionTablePullConsumer; import org.apache.iotdb.session.subscription.consumer.table.SubscriptionTablePullConsumerBuilder; -import java.util.Objects; - final class InternalSubscriptionTablePullConsumer extends SubscriptionTablePullConsumer { private final INodeSupplier nodeSupplier; @@ -36,10 +33,7 @@ final class InternalSubscriptionTablePullConsumer extends SubscriptionTablePullC InternalSubscriptionTablePullConsumer( final SubscriptionTablePullConsumerBuilder builder, final INodeSupplier nodeSupplier) { super(builder); - this.nodeSupplier = - Objects.requireNonNull( - nodeSupplier, - SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + this.nodeSupplier = nodeSupplier; } @Override diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java index 0d136ea78a684..f6fcfac2b6f22 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java @@ -19,6 +19,7 @@ package org.apache.iotdb.commons.client.internal.subscription; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.session.i18n.SessionMessages; import org.apache.iotdb.session.subscription.consumer.ISubscriptionTablePullConsumer; @@ -32,6 +33,10 @@ public class InternalSubscriptionTablePullConsumerBuilder private final INodeSupplier nodeSupplier; + public InternalSubscriptionTablePullConsumerBuilder() { + this.nodeSupplier = null; + } + public InternalSubscriptionTablePullConsumerBuilder(final INodeSupplier nodeSupplier) { this.nodeSupplier = Objects.requireNonNull( @@ -39,6 +44,16 @@ public InternalSubscriptionTablePullConsumerBuilder(final INodeSupplier nodeSupp SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); } + public InternalSubscriptionTablePullConsumerBuilder endPoint(final TEndPoint endPoint) { + final TEndPoint endpoint = + Objects.requireNonNull( + endPoint, SessionMessages.EXCEPTION_END_POINT_MUST_NOT_BE_NULL_908A4B27); + this.host = endpoint.getIp(); + this.port = endpoint.getPort(); + this.nodeUrls = null; + return this; + } + protected final INodeSupplier getNodeSupplier() { return nodeSupplier; } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java index b34e07b621a4e..9e443dd3f5d87 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java @@ -19,11 +19,13 @@ package org.apache.iotdb.commons.client.internal.subscription; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.session.i18n.SessionMessages; import org.apache.iotdb.session.subscription.ISubscriptionTableSession; import org.apache.iotdb.session.subscription.SubscriptionTableSessionBuilder; +import java.util.List; import java.util.Objects; /** Server-side builder for table subscription-management sessions. */ @@ -31,6 +33,10 @@ public class InternalSubscriptionTableSessionBuilder extends SubscriptionTableSe private final INodeSupplier nodeSupplier; + public InternalSubscriptionTableSessionBuilder() { + this.nodeSupplier = null; + } + public InternalSubscriptionTableSessionBuilder(final INodeSupplier nodeSupplier) { this.nodeSupplier = Objects.requireNonNull( @@ -38,6 +44,21 @@ public InternalSubscriptionTableSessionBuilder(final INodeSupplier nodeSupplier) SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); } + public InternalSubscriptionTableSessionBuilder endPoint(final TEndPoint endPoint) { + final TEndPoint endpoint = + Objects.requireNonNull( + endPoint, SessionMessages.EXCEPTION_END_POINT_MUST_NOT_BE_NULL_908A4B27); + this.host = endpoint.getIp(); + this.rpcPort = endpoint.getPort(); + this.nodeUrls = null; + return this; + } + + public InternalSubscriptionTableSessionBuilder nodeUrls(final List nodeUrls) { + this.nodeUrls = nodeUrls; + return this; + } + protected final INodeSupplier getNodeSupplier() { return nodeSupplier; } From 4a49e615a4cdf7fd09c682bb55b7e79a65542ff9 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 14:15:04 +0800 Subject: [PATCH 16/32] Use getter hook for RPC service name --- .../org/apache/iotdb/session/Session.java | 64 +++++++++---------- .../iotdb/session/SessionConnection.java | 5 +- .../SubscriptionSessionWrapper.java | 10 ++- .../internal/session/InternalSession.java | 6 +- .../InternalSubscriptionSessionWrapper.java | 6 +- 5 files changed, 52 insertions(+), 39 deletions(-) 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 a4f523fb44ef7..28b09b102ba79 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 @@ -187,9 +187,6 @@ public class Session implements ISession { protected INodeSupplier availableNodes; - /** Multiplexed Thrift service name used by internal sessions. */ - protected String rpcServiceName; - protected boolean enableQueryRedirection = false; // The version number of the client which used for compatibility in the server @@ -552,7 +549,14 @@ public synchronized void open(boolean enableRPCCompaction, int connectionTimeout this.availableNodes = null; } - this.availableNodes = getNodeSupplier(enableRPCCompaction, connectionTimeoutInMs); + if (enableAutoFetch) { + this.availableNodes = getNodeSupplier(enableRPCCompaction, connectionTimeoutInMs); + } else { + // A static supplier has no separate seed tier. Keep every configured address after reopen. + Set staticNodes = new LinkedHashSet<>(getNodeUrls()); + staticNodes.addAll(SessionUtils.parseSeedNodeUrls(initialNodeUrls)); + this.availableNodes = new DummyNodesSupplier(new ArrayList<>(staticNodes)); + } this.connectionTimeoutInMs = connectionTimeoutInMs; setDefaultSessionConnection(constructSessionConnection(this, defaultEndPoint, zoneId)); @@ -584,36 +588,30 @@ private void initThreadPool() { } /** - * Creates the endpoint source used by this session. Subclasses may override this hook to provide - * a server-managed endpoint source without exposing it through the client builder API. + * Creates the auto-fetch endpoint source used by this session. Subclasses may override this hook + * to provide a server-managed endpoint source without exposing it through the client API. */ protected INodeSupplier getNodeSupplier( final boolean enableRPCCompaction, final int connectionTimeoutInMs) { - if (enableAutoFetch) { - initThreadPool(); - return NodesSupplier.createNodeSupplier( - getNodeUrls(), - SessionUtils.parseSeedNodeUrls(initialNodeUrls), - executorService, - username, - password, - zoneId, - thriftDefaultBufferSize, - thriftMaxFrameSize, - connectionTimeoutInMs, - useSSL, - trustStore, - trustStorePwd, - keyStore, - keyStorePwd, - sslProtocol, - enableRPCCompaction, - version.toString()); - } - // A static supplier has no separate seed tier. Keep every configured address after reopen. - Set staticNodes = new LinkedHashSet<>(getNodeUrls()); - staticNodes.addAll(SessionUtils.parseSeedNodeUrls(initialNodeUrls)); - return new DummyNodesSupplier(new ArrayList<>(staticNodes)); + initThreadPool(); + return NodesSupplier.createNodeSupplier( + getNodeUrls(), + SessionUtils.parseSeedNodeUrls(initialNodeUrls), + executorService, + username, + password, + zoneId, + thriftDefaultBufferSize, + thriftMaxFrameSize, + connectionTimeoutInMs, + useSSL, + trustStore, + trustStorePwd, + keyStore, + keyStorePwd, + sslProtocol, + enableRPCCompaction, + version.toString()); } private static List shuffleNodeUrls(List endPoints) { @@ -4447,8 +4445,8 @@ protected void setDefaultSessionConnection(SessionConnection defaultSessionConne this.defaultSessionConnection = defaultSessionConnection; } - protected void setRpcServiceName(final String rpcServiceName) { - this.rpcServiceName = rpcServiceName; + protected String getRpcServiceName() { + return null; } public boolean isEnableThriftRpcCompaction() { 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 3548924a29dd3..a79ec4338db03 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 @@ -237,8 +237,9 @@ private void init( session.enableThriftRpcCompaction ? new TCompactProtocol(transport) : new TBinaryProtocol(transport); - if (session.rpcServiceName != null) { - protocol = new TMultiplexedProtocol(protocol, session.rpcServiceName); + final String rpcServiceName = session.getRpcServiceName(); + if (rpcServiceName != null) { + protocol = new TMultiplexedProtocol(protocol, rpcServiceName); } client = new IClientRPCService.Client(protocol); client = RpcUtils.newSynchronizedClient(client); diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java index 6538eea370686..2d1ad9ab23b25 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java @@ -32,6 +32,8 @@ public class SubscriptionSessionWrapper extends Session { + private String rpcServiceName; + public SubscriptionSessionWrapper(final AbstractSessionBuilder builder) { super(builder); } @@ -71,9 +73,13 @@ public int getThriftMaxFrameSize() { return thriftMaxFrameSize; } - @Override public void setRpcServiceName(final String rpcServiceName) { - super.setRpcServiceName(rpcServiceName); + this.rpcServiceName = rpcServiceName; + } + + @Override + protected String getRpcServiceName() { + return rpcServiceName; } public void setSslConfig( diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java index f80accedadc56..fbb2ae7b14659 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java @@ -32,10 +32,14 @@ public class InternalSession extends Session { protected InternalSession( final AbstractSessionBuilder builder, final INodeSupplier nodeSupplier) { super(builder); - setRpcServiceName(RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE); this.nodeSupplier = nodeSupplier; } + @Override + protected String getRpcServiceName() { + return RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE; + } + @Override protected final INodeSupplier getNodeSupplier( final boolean enableRPCCompaction, final int connectionTimeoutInMs) { diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java index 1876c4a4c2d60..1078b951bb50e 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java @@ -31,10 +31,14 @@ final class InternalSubscriptionSessionWrapper extends SubscriptionSessionWrappe InternalSubscriptionSessionWrapper( final AbstractSessionBuilder builder, final INodeSupplier nodeSupplier) { super(builder); - setRpcServiceName(RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE); this.nodeSupplier = nodeSupplier; } + @Override + protected String getRpcServiceName() { + return RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE; + } + @Override protected INodeSupplier getNodeSupplier( final boolean enableRPCCompaction, final int connectionTimeoutInMs) { From 8ac650d78fabd2187414ced785c66b9397d8cd48 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 14:34:16 +0800 Subject: [PATCH 17/32] Use node initialization hooks for internal sessions --- .../org/apache/iotdb/session/Session.java | 51 ++++---- .../iotdb/session/pool/SessionPool.java | 121 ++++++++++-------- .../internal/session/InternalSession.java | 17 ++- .../InternalTableSessionPoolBuilder.java | 10 +- .../InternalSubscriptionSessionWrapper.java | 17 ++- 5 files changed, 122 insertions(+), 94 deletions(-) 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 28b09b102ba79..d4bd438e9291b 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 @@ -550,7 +550,8 @@ public synchronized void open(boolean enableRPCCompaction, int connectionTimeout } if (enableAutoFetch) { - this.availableNodes = getNodeSupplier(enableRPCCompaction, connectionTimeoutInMs); + initThreadPool(); + initAvailableNodes(enableRPCCompaction, connectionTimeoutInMs); } else { // A static supplier has no separate seed tier. Keep every configured address after reopen. Set staticNodes = new LinkedHashSet<>(getNodeUrls()); @@ -570,7 +571,8 @@ public synchronized void open(boolean enableRPCCompaction, int connectionTimeout } } - private void initThreadPool() { + /** Creates the thread pool used to refresh the available node list. */ + protected void initThreadPool() { this.executorService = Executors.newSingleThreadScheduledExecutor( r -> { @@ -587,31 +589,28 @@ private void initThreadPool() { }); } - /** - * Creates the auto-fetch endpoint source used by this session. Subclasses may override this hook - * to provide a server-managed endpoint source without exposing it through the client API. - */ - protected INodeSupplier getNodeSupplier( + /** Initializes the auto-fetch endpoint source used by this session. */ + protected void initAvailableNodes( final boolean enableRPCCompaction, final int connectionTimeoutInMs) { - initThreadPool(); - return NodesSupplier.createNodeSupplier( - getNodeUrls(), - SessionUtils.parseSeedNodeUrls(initialNodeUrls), - executorService, - username, - password, - zoneId, - thriftDefaultBufferSize, - thriftMaxFrameSize, - connectionTimeoutInMs, - useSSL, - trustStore, - trustStorePwd, - keyStore, - keyStorePwd, - sslProtocol, - enableRPCCompaction, - version.toString()); + this.availableNodes = + NodesSupplier.createNodeSupplier( + getNodeUrls(), + SessionUtils.parseSeedNodeUrls(initialNodeUrls), + executorService, + username, + password, + zoneId, + thriftDefaultBufferSize, + thriftMaxFrameSize, + connectionTimeoutInMs, + useSSL, + trustStore, + trustStorePwd, + keyStore, + keyStorePwd, + sslProtocol, + enableRPCCompaction, + version.toString()); } private static List shuffleNodeUrls(List endPoints) { 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 590e6658e5c28..ebb8664d3c75e 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 @@ -171,8 +171,7 @@ public class SessionPool implements ISessionPool { @SuppressWarnings("java:S3077") private volatile ScheduledExecutorService executorService; - private INodeSupplier availableNodes; - private boolean nodeSupplierInitialized; + protected INodeSupplier availableNodes; // set to true, means that we will start a background thread to fetch all available (Status is // not Removing) datanodes in cluster, and these available nodes will be used in retrying stage @@ -432,7 +431,8 @@ public SessionPool( this.thriftDefaultBufferSize = thriftDefaultBufferSize; this.thriftMaxFrameSize = thriftMaxFrameSize; this.formattedNodeUrls = UrlUtils.formatTEndPointIpv4AndIpv6Url(host, port); - initializeNodeSupplier(); + initThreadPool(); + initAvailableNodes(Collections.singletonList(new TEndPoint(host, port))); } public SessionPool( @@ -478,7 +478,8 @@ public SessionPool( this.trustStore = trustStore; this.trustStorePwd = trustStorePwd; this.sslProtocol = SessionConfig.DEFAULT_SSL_PROTOCOL; - initializeNodeSupplier(); + initThreadPool(); + initAvailableNodes(Collections.singletonList(new TEndPoint(host, port))); } @SuppressWarnings("squid:S107") // ignore Methods should not have too many parameters @@ -520,10 +521,16 @@ public SessionPool( this.thriftDefaultBufferSize = thriftDefaultBufferSize; this.thriftMaxFrameSize = thriftMaxFrameSize; this.formattedNodeUrls = this.nodeUrls.toString(); - initializeNodeSupplier(); + initThreadPool(); + initAvailableNodes(SessionUtils.parseSeedNodeUrls(this.nodeUrls)); } public SessionPool(AbstractSessionPoolBuilder builder) { + this(builder, null); + } + + protected SessionPool( + final AbstractSessionPoolBuilder builder, final INodeSupplier suppliedNodes) { this.maxSize = builder.maxSize; this.user = builder.username; this.password = builder.pw; @@ -556,6 +563,9 @@ public SessionPool(AbstractSessionPoolBuilder builder) { this.queryTimeoutInMs = builder.timeOut; this.enableAutoCreateSchema = builder.enableAutoCreateSchema; + if (enableAutoFetch && suppliedNodes == null) { + initThreadPool(); + } if (builder.nodeUrls != null) { if (builder.nodeUrls.isEmpty()) { throw new IllegalArgumentException(SessionMessages.NODE_URLS_EMPTY); @@ -564,11 +574,31 @@ public SessionPool(AbstractSessionPoolBuilder builder) { this.host = null; this.port = -1; this.formattedNodeUrls = builder.nodeUrls.toString(); + if (enableAutoFetch) { + if (suppliedNodes == null) { + initAvailableNodes(SessionUtils.parseSeedNodeUrls(nodeUrls)); + } else { + this.availableNodes = suppliedNodes; + } + } else { + this.availableNodes = new DummyNodesSupplier(SessionUtils.parseSeedNodeUrls(nodeUrls)); + } + } else { this.host = builder.host; this.port = builder.rpcPort; this.nodeUrls = null; this.formattedNodeUrls = UrlUtils.formatTEndPointIpv4AndIpv6Url(host, port); + if (enableAutoFetch) { + if (suppliedNodes == null) { + initAvailableNodes(Collections.singletonList(new TEndPoint(host, port))); + } else { + this.availableNodes = suppliedNodes; + } + } else { + this.availableNodes = + new DummyNodesSupplier(Collections.singletonList(new TEndPoint(host, port))); + } } } @@ -611,17 +641,8 @@ protected Session constructNewSession() { return session; } - protected void openSession(final ISession session) throws IoTDBConnectionException { - initializeNodeSupplier(); - session.open( - enableThriftCompression, - connectionTimeoutInMs, - deviceIdToEndpoint, - tableModelDeviceIdToEndpoint, - availableNodes); - } - - private void initThreadPool() { + /** Creates the thread pool used to refresh the available node list. */ + protected void initThreadPool() { this.executorService = Executors.newSingleThreadScheduledExecutor( r -> { @@ -638,40 +659,26 @@ private void initThreadPool() { }); } - private synchronized void initializeNodeSupplier() { - if (nodeSupplierInitialized) { - return; - } - final List endPointList = - nodeUrls == null - ? Collections.singletonList(new TEndPoint(host, port)) - : SessionUtils.parseSeedNodeUrls(nodeUrls); - this.availableNodes = createNodeSupplier(endPointList); - this.nodeSupplierInitialized = true; - } - - protected INodeSupplier createNodeSupplier(final List endPointList) { - if (enableAutoFetch) { - initThreadPool(); - return NodesSupplier.createNodeSupplier( - endPointList, - executorService, - user, - password, - zoneId, - thriftDefaultBufferSize, - thriftMaxFrameSize, - connectionTimeoutInMs, - useSSL, - trustStore, - trustStorePwd, - keyStore, - keyStorePwd, - sslProtocol, - enableThriftCompression, - version.toString()); - } - return new DummyNodesSupplier(endPointList); + /** Initializes the auto-fetch endpoint source used by this pool. */ + protected void initAvailableNodes(final List endPointList) { + this.availableNodes = + NodesSupplier.createNodeSupplier( + endPointList, + executorService, + user, + password, + zoneId, + thriftDefaultBufferSize, + thriftMaxFrameSize, + connectionTimeoutInMs, + useSSL, + trustStore, + trustStorePwd, + keyStore, + keyStorePwd, + sslProtocol, + enableThriftCompression, + version.toString()); } // if this method throws an exception, either the server is broken, or the ip/port/user/password @@ -745,7 +752,12 @@ private ISession getSession() throws IoTDBConnectionException { session = constructNewSession(); try { - openSession(session); + session.open( + enableThriftCompression, + connectionTimeoutInMs, + deviceIdToEndpoint, + tableModelDeviceIdToEndpoint, + availableNodes); // avoid someone has called close() the session pool synchronized (this) { if (closed) { @@ -854,7 +866,12 @@ public void closeResultSet(SessionDataSetWrapper wrapper) { private void tryConstructNewSession() { Session session = constructNewSession(); try { - openSession(session); + session.open( + enableThriftCompression, + connectionTimeoutInMs, + deviceIdToEndpoint, + tableModelDeviceIdToEndpoint, + availableNodes); // avoid someone has called close() the session pool synchronized (this) { if (closed) { diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java index fbb2ae7b14659..f08455bc00df2 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java @@ -41,10 +41,19 @@ protected String getRpcServiceName() { } @Override - protected final INodeSupplier getNodeSupplier( + protected void initThreadPool() { + if (nodeSupplier == null) { + super.initThreadPool(); + } + } + + @Override + protected void initAvailableNodes( final boolean enableRPCCompaction, final int connectionTimeoutInMs) { - return nodeSupplier != null - ? nodeSupplier - : super.getNodeSupplier(enableRPCCompaction, connectionTimeoutInMs); + if (nodeSupplier == null) { + super.initAvailableNodes(enableRPCCompaction, connectionTimeoutInMs); + } else { + this.availableNodes = nodeSupplier; + } } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java index 3879dc5f500d4..92c3ab35c59ec 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java @@ -28,10 +28,9 @@ import org.apache.iotdb.session.pool.TableSessionPool; import org.apache.iotdb.session.pool.TableSessionPoolBuilder; -import java.util.List; import java.util.Objects; -/** Server-side builder for lazily-created table-model internal session pools. */ +/** Server-side builder for table-model internal session pools. */ public class InternalTableSessionPoolBuilder extends TableSessionPoolBuilder { private final INodeSupplier nodeSupplier; @@ -80,7 +79,7 @@ private static final class InternalSessionPool extends SessionPool { private InternalSessionPool( final InternalTableSessionPoolBuilder builder, final INodeSupplier nodeSupplier) { - super(builder); + super(builder, nodeSupplier); this.nodeSupplier = nodeSupplier; } @@ -88,10 +87,5 @@ private InternalSessionPool( protected Session constructNewSession() { return new InternalSession(createSessionBuilder(), nodeSupplier); } - - @Override - protected INodeSupplier createNodeSupplier(final List endPointList) { - return nodeSupplier != null ? nodeSupplier : super.createNodeSupplier(endPointList); - } } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java index 1078b951bb50e..d88e449630863 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java @@ -40,10 +40,19 @@ protected String getRpcServiceName() { } @Override - protected INodeSupplier getNodeSupplier( + protected void initThreadPool() { + if (nodeSupplier == null) { + super.initThreadPool(); + } + } + + @Override + protected void initAvailableNodes( final boolean enableRPCCompaction, final int connectionTimeoutInMs) { - return nodeSupplier != null - ? nodeSupplier - : super.getNodeSupplier(enableRPCCompaction, connectionTimeoutInMs); + if (nodeSupplier == null) { + super.initAvailableNodes(enableRPCCompaction, connectionTimeoutInMs); + } else { + this.availableNodes = nodeSupplier; + } } } From 78e7151ae93420762b70af5b4c8f711992aaa195 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 14:47:22 +0800 Subject: [PATCH 18/32] Refine internal session pool construction --- .../org/apache/iotdb/session/Session.java | 6 +- .../iotdb/session/TableSessionBuilder.java | 15 +-- .../iotdb/session/pool/SessionPool.java | 103 +++++++++++------- .../SubscriptionSessionWrapper.java | 2 +- .../session/InternalSessionBuilder.java | 6 +- .../internal/session/InternalSessionPool.java | 41 +++++++ .../session/InternalSessionPoolBuilder.java | 45 ++++++++ .../session/InternalTableSessionBuilder.java | 10 +- .../session/InternalTableSessionPool.java | 29 +++++ .../InternalTableSessionPoolBuilder.java | 33 +----- ...lSubscriptionTablePullConsumerBuilder.java | 6 +- ...ternalSubscriptionTableSessionBuilder.java | 6 +- 12 files changed, 196 insertions(+), 106 deletions(-) create mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java create mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java create mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java 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 d4bd438e9291b..b478e3311e24a 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 @@ -571,7 +571,6 @@ public synchronized void open(boolean enableRPCCompaction, int connectionTimeout } } - /** Creates the thread pool used to refresh the available node list. */ protected void initThreadPool() { this.executorService = Executors.newSingleThreadScheduledExecutor( @@ -589,7 +588,6 @@ protected void initThreadPool() { }); } - /** Initializes the auto-fetch endpoint source used by this session. */ protected void initAvailableNodes( final boolean enableRPCCompaction, final int connectionTimeoutInMs) { this.availableNodes = @@ -695,11 +693,11 @@ public synchronized void close() throws IoTDBConnectionException { getDefaultSessionConnection().close(); setDefaultSessionConnection(null); } 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; - } - if (this.availableNodes != null) { this.availableNodes.close(); this.availableNodes = null; } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java index cc6b6bf440a21..a52cb227070a0 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/TableSessionBuilder.java @@ -432,28 +432,25 @@ public TableSessionBuilder tabletCompressionMinRowSize(int tabletCompressionMinR * @throws IoTDBConnectionException if an error occurs while establishing the connection. */ public ITableSession build() throws IoTDBConnectionException { - return new TableSession(constructNewSession()); - } - - protected Session constructNewSession() throws IoTDBConnectionException { if (nodeUrls == null) { this.nodeUrls = Collections.singletonList(SessionConfig.DEFAULT_HOST + ":" + SessionConfig.DEFAULT_PORT); } this.sqlDialect = TABLE; - return openSession(new Session(this)); - } - - protected Session openSession(Session newSession) throws IoTDBConnectionException { + Session newSession = createSession(); newSession.enableIoTDBRpcCompression = isIoTDBRpcCompressionEnabled; newSession.tabletCompressionMinRowSize = tabletCompressionMinRowSize; try { newSession.open(isThriftRpcCompactionEnabled, connectionTimeoutInMs); - return newSession; } catch (IoTDBConnectionException e) { newSession.close(); throw e; } + return new TableSession(newSession); + } + + protected Session createSession() { + return new Session(this); } } 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 ebb8664d3c75e..c98dc17194ddf 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 @@ -171,7 +171,7 @@ public class SessionPool implements ISessionPool { @SuppressWarnings("java:S3077") private volatile ScheduledExecutorService executorService; - protected INodeSupplier availableNodes; + private INodeSupplier availableNodes; // set to true, means that we will start a background thread to fetch all available (Status is // not Removing) datanodes in cluster, and these available nodes will be used in retrying stage @@ -602,47 +602,77 @@ protected SessionPool( } } - protected Session.Builder createSessionBuilder() { - return new Session.Builder() - .username(user) - .password(password) - .fetchSize(fetchSize) - .zoneId(zoneId) - .thriftDefaultBufferSize(thriftDefaultBufferSize) - .thriftMaxFrameSize(thriftMaxFrameSize) - .enableRedirection(enableRedirection) - .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) - .version(version) - .useSSL(useSSL) - .trustStore(trustStore) - .trustStorePwd(trustStorePwd) - .keyStore(keyStore) - .keyStorePwd(keyStorePwd) - .sslProtocol(sslProtocol) - .maxRetryCount(maxRetryCount) - .retryIntervalInMs(retryIntervalInMs) - .sqlDialect(sqlDialect) - .database(database) - .enableAutoCreateSchema(enableAutoCreateSchema) - .timeOut(queryTimeoutInMs) - .enableIoTDBRpcCompression(enableIoTDBRpcCompression) - .enableThriftRpcCompression(enableThriftCompression); - } - - protected Session constructNewSession() { - final Session.Builder builder = createSessionBuilder(); + private Session constructNewSession() { + Session session; if (nodeUrls == null) { - builder.host(host).port(port); + // Construct custom Session + session = + createSessionBuilder() + .host(host) + .port(port) + .username(user) + .password(password) + .fetchSize(fetchSize) + .zoneId(zoneId) + .thriftDefaultBufferSize(thriftDefaultBufferSize) + .thriftMaxFrameSize(thriftMaxFrameSize) + .enableRedirection(enableRedirection) + .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) + .version(version) + .useSSL(useSSL) + .trustStore(trustStore) + .trustStorePwd(trustStorePwd) + .keyStore(keyStore) + .keyStorePwd(keyStorePwd) + .sslProtocol(sslProtocol) + .maxRetryCount(maxRetryCount) + .retryIntervalInMs(retryIntervalInMs) + .sqlDialect(sqlDialect) + .database(database) + .enableAutoCreateSchema(enableAutoCreateSchema) + .timeOut(queryTimeoutInMs) + .enableIoTDBRpcCompression(enableIoTDBRpcCompression) + .enableThriftRpcCompression(enableThriftCompression) + .build(); } else { - builder.nodeUrls(nodeUrls); + // Construct redirect-able Session + session = + createSessionBuilder() + .nodeUrls(nodeUrls) + .username(user) + .password(password) + .fetchSize(fetchSize) + .zoneId(zoneId) + .thriftDefaultBufferSize(thriftDefaultBufferSize) + .thriftMaxFrameSize(thriftMaxFrameSize) + .enableRedirection(enableRedirection) + .enableRecordsAutoConvertTablet(enableRecordsAutoConvertTablet) + .version(version) + .useSSL(useSSL) + .trustStore(trustStore) + .trustStorePwd(trustStorePwd) + .keyStore(keyStore) + .keyStorePwd(keyStorePwd) + .sslProtocol(sslProtocol) + .maxRetryCount(maxRetryCount) + .retryIntervalInMs(retryIntervalInMs) + .sqlDialect(sqlDialect) + .database(database) + .enableAutoCreateSchema(enableAutoCreateSchema) + .timeOut(queryTimeoutInMs) + .enableIoTDBRpcCompression(enableIoTDBRpcCompression) + .enableThriftRpcCompression(enableThriftCompression) + .build(); } - final Session session = builder.build(); session.setEnableQueryRedirection(enableQueryRedirection); return session; } - /** Creates the thread pool used to refresh the available node list. */ - protected void initThreadPool() { + protected Session.Builder createSessionBuilder() { + return new Session.Builder(); + } + + private void initThreadPool() { this.executorService = Executors.newSingleThreadScheduledExecutor( r -> { @@ -659,8 +689,7 @@ protected void initThreadPool() { }); } - /** Initializes the auto-fetch endpoint source used by this pool. */ - protected void initAvailableNodes(final List endPointList) { + private void initAvailableNodes(List endPointList) { this.availableNodes = NodesSupplier.createNodeSupplier( endPointList, diff --git a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java index 2d1ad9ab23b25..a14b8eb5aace1 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java @@ -39,7 +39,7 @@ public SubscriptionSessionWrapper(final AbstractSessionBuilder builder) { } public void open() throws IoTDBConnectionException { - open(enableThriftRpcCompaction, connectionTimeoutInMs); + super.open(enableThriftRpcCompaction, connectionTimeoutInMs); } public void close() throws IoTDBConnectionException { diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java index 90c9c2bbf4a5a..a5b5da83edebf 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java @@ -52,12 +52,8 @@ public InternalSessionBuilder endPoint(final TEndPoint endPoint) { return this; } - protected final INodeSupplier getNodeSupplier() { - return nodeSupplier; - } - @Override protected Session createSession() { - return new InternalSession(this, getNodeSupplier()); + return new InternalSession(this, nodeSupplier); } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java new file mode 100644 index 0000000000000..e5deab9b5520c --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java @@ -0,0 +1,41 @@ +/* + * 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.commons.client.internal.session; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.session.Session; +import org.apache.iotdb.session.pool.SessionPool; + +final class InternalSessionPool extends SessionPool { + + private final INodeSupplier nodeSupplier; + + InternalSessionPool(final InternalSessionPoolBuilder builder) { + super(builder.getSessionPoolBuilder(), builder.getNodeSupplier()); + this.nodeSupplier = builder.getNodeSupplier(); + } + + @Override + protected Session.Builder createSessionBuilder() { + return nodeSupplier == null + ? new InternalSessionBuilder() + : new InternalSessionBuilder(nodeSupplier); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java new file mode 100644 index 0000000000000..c79b77489d12c --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java @@ -0,0 +1,45 @@ +/* + * 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.commons.client.internal.session; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.session.pool.AbstractSessionPoolBuilder; + +import java.util.Objects; + +final class InternalSessionPoolBuilder { + + private final AbstractSessionPoolBuilder sessionPoolBuilder; + private final INodeSupplier nodeSupplier; + + InternalSessionPoolBuilder( + final AbstractSessionPoolBuilder sessionPoolBuilder, final INodeSupplier nodeSupplier) { + this.sessionPoolBuilder = Objects.requireNonNull(sessionPoolBuilder); + this.nodeSupplier = nodeSupplier; + } + + AbstractSessionPoolBuilder getSessionPoolBuilder() { + return sessionPoolBuilder; + } + + INodeSupplier getNodeSupplier() { + return nodeSupplier; + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java index 0175c37fa37a0..79ce2af635a50 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java @@ -21,7 +21,6 @@ import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; -import org.apache.iotdb.rpc.IoTDBConnectionException; import org.apache.iotdb.session.Session; import org.apache.iotdb.session.TableSessionBuilder; import org.apache.iotdb.session.i18n.SessionMessages; @@ -54,13 +53,8 @@ public InternalTableSessionBuilder endPoint(final TEndPoint endPoint) { return this; } - protected final INodeSupplier getNodeSupplier() { - return nodeSupplier; - } - @Override - protected Session constructNewSession() throws IoTDBConnectionException { - this.sqlDialect = "table"; - return openSession(new InternalSession(this, getNodeSupplier())); + protected Session createSession() { + return new InternalSession(this, nodeSupplier); } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java new file mode 100644 index 0000000000000..21ee0bbcde74a --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.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.commons.client.internal.session; + +import org.apache.iotdb.session.pool.TableSessionPool; + +final class InternalTableSessionPool extends TableSessionPool { + + InternalTableSessionPool(final InternalSessionPoolBuilder builder) { + super(new InternalSessionPool(builder)); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java index 92c3ab35c59ec..706677a334217 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java @@ -22,10 +22,7 @@ import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.pool.ITableSessionPool; -import org.apache.iotdb.session.Session; import org.apache.iotdb.session.i18n.SessionMessages; -import org.apache.iotdb.session.pool.SessionPool; -import org.apache.iotdb.session.pool.TableSessionPool; import org.apache.iotdb.session.pool.TableSessionPoolBuilder; import java.util.Objects; @@ -56,36 +53,8 @@ public InternalTableSessionPoolBuilder endPoint(final TEndPoint endPoint) { return this; } - protected final INodeSupplier getNodeSupplier() { - return nodeSupplier; - } - @Override protected ITableSessionPool createSessionPool() { - return new InternalTableSessionPool(this, getNodeSupplier()); - } - - private static final class InternalTableSessionPool extends TableSessionPool { - - private InternalTableSessionPool( - final InternalTableSessionPoolBuilder builder, final INodeSupplier nodeSupplier) { - super(new InternalSessionPool(builder, nodeSupplier)); - } - } - - private static final class InternalSessionPool extends SessionPool { - - private final INodeSupplier nodeSupplier; - - private InternalSessionPool( - final InternalTableSessionPoolBuilder builder, final INodeSupplier nodeSupplier) { - super(builder, nodeSupplier); - this.nodeSupplier = nodeSupplier; - } - - @Override - protected Session constructNewSession() { - return new InternalSession(createSessionBuilder(), nodeSupplier); - } + return new InternalTableSessionPool(new InternalSessionPoolBuilder(this, nodeSupplier)); } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java index f6fcfac2b6f22..24798528ddca2 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java @@ -54,12 +54,8 @@ public InternalSubscriptionTablePullConsumerBuilder endPoint(final TEndPoint end return this; } - protected final INodeSupplier getNodeSupplier() { - return nodeSupplier; - } - @Override protected ISubscriptionTablePullConsumer createConsumer() { - return new InternalSubscriptionTablePullConsumer(this, getNodeSupplier()); + return new InternalSubscriptionTablePullConsumer(this, nodeSupplier); } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java index 9e443dd3f5d87..9f5e6f20aff48 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java @@ -59,12 +59,8 @@ public InternalSubscriptionTableSessionBuilder nodeUrls(final List nodeU return this; } - protected final INodeSupplier getNodeSupplier() { - return nodeSupplier; - } - @Override protected ISubscriptionTableSession createSubscriptionSession() { - return new InternalSubscriptionTableSession(this, getNodeSupplier()); + return new InternalSubscriptionTableSession(this, nodeSupplier); } } From 7f264573fe4dae2824d7cd5b576f51b2f735fb56 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 15:01:42 +0800 Subject: [PATCH 19/32] Use overrides for internal session pool initialization --- .../iotdb/session/pool/SessionPool.java | 28 +++------ .../internal/session/InternalSessionPool.java | 30 ++++++++-- .../session/InternalSessionPoolBuilder.java | 36 ++++++++--- .../session/InternalTableSessionPool.java | 29 --------- .../InternalTableSessionPoolBuilder.java | 60 ------------------- 5 files changed, 62 insertions(+), 121 deletions(-) delete mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java delete mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java 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 c98dc17194ddf..5ef0eac2b3aef 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 @@ -171,7 +171,9 @@ public class SessionPool implements ISessionPool { @SuppressWarnings("java:S3077") private volatile ScheduledExecutorService executorService; - private INodeSupplier availableNodes; + protected INodeSupplier availableNodes; + + protected final AbstractSessionPoolBuilder sessionPoolBuilder; // set to true, means that we will start a background thread to fetch all available (Status is // not Removing) datanodes in cluster, and these available nodes will be used in retrying stage @@ -526,11 +528,7 @@ public SessionPool( } public SessionPool(AbstractSessionPoolBuilder builder) { - this(builder, null); - } - - protected SessionPool( - final AbstractSessionPoolBuilder builder, final INodeSupplier suppliedNodes) { + this.sessionPoolBuilder = builder; this.maxSize = builder.maxSize; this.user = builder.username; this.password = builder.pw; @@ -563,7 +561,7 @@ protected SessionPool( this.queryTimeoutInMs = builder.timeOut; this.enableAutoCreateSchema = builder.enableAutoCreateSchema; - if (enableAutoFetch && suppliedNodes == null) { + if (enableAutoFetch) { initThreadPool(); } if (builder.nodeUrls != null) { @@ -575,11 +573,7 @@ protected SessionPool( this.port = -1; this.formattedNodeUrls = builder.nodeUrls.toString(); if (enableAutoFetch) { - if (suppliedNodes == null) { - initAvailableNodes(SessionUtils.parseSeedNodeUrls(nodeUrls)); - } else { - this.availableNodes = suppliedNodes; - } + initAvailableNodes(SessionUtils.parseSeedNodeUrls(nodeUrls)); } else { this.availableNodes = new DummyNodesSupplier(SessionUtils.parseSeedNodeUrls(nodeUrls)); } @@ -590,11 +584,7 @@ protected SessionPool( this.nodeUrls = null; this.formattedNodeUrls = UrlUtils.formatTEndPointIpv4AndIpv6Url(host, port); if (enableAutoFetch) { - if (suppliedNodes == null) { - initAvailableNodes(Collections.singletonList(new TEndPoint(host, port))); - } else { - this.availableNodes = suppliedNodes; - } + initAvailableNodes(Collections.singletonList(new TEndPoint(host, port))); } else { this.availableNodes = new DummyNodesSupplier(Collections.singletonList(new TEndPoint(host, port))); @@ -672,7 +662,7 @@ protected Session.Builder createSessionBuilder() { return new Session.Builder(); } - private void initThreadPool() { + protected void initThreadPool() { this.executorService = Executors.newSingleThreadScheduledExecutor( r -> { @@ -689,7 +679,7 @@ private void initThreadPool() { }); } - private void initAvailableNodes(List endPointList) { + protected void initAvailableNodes(List endPointList) { this.availableNodes = NodesSupplier.createNodeSupplier( endPointList, diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java index e5deab9b5520c..28d644a049a9b 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java @@ -19,23 +19,45 @@ package org.apache.iotdb.commons.client.internal.session; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.session.Session; import org.apache.iotdb.session.pool.SessionPool; -final class InternalSessionPool extends SessionPool { +import java.util.List; - private final INodeSupplier nodeSupplier; +public final class InternalSessionPool extends SessionPool { InternalSessionPool(final InternalSessionPoolBuilder builder) { - super(builder.getSessionPoolBuilder(), builder.getNodeSupplier()); - this.nodeSupplier = builder.getNodeSupplier(); + super(builder); + } + + @Override + protected void initThreadPool() { + if (getInternalBuilder().getNodeSupplier() == null) { + super.initThreadPool(); + } + } + + @Override + protected void initAvailableNodes(final List endPointList) { + final INodeSupplier nodeSupplier = getInternalBuilder().getNodeSupplier(); + if (nodeSupplier == null) { + super.initAvailableNodes(endPointList); + } else { + this.availableNodes = nodeSupplier; + } } @Override protected Session.Builder createSessionBuilder() { + final INodeSupplier nodeSupplier = getInternalBuilder().getNodeSupplier(); return nodeSupplier == null ? new InternalSessionBuilder() : new InternalSessionBuilder(nodeSupplier); } + + private InternalSessionPoolBuilder getInternalBuilder() { + return (InternalSessionPoolBuilder) sessionPoolBuilder; + } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java index c79b77489d12c..a45b4a3d7ebfe 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java @@ -19,27 +19,45 @@ package org.apache.iotdb.commons.client.internal.session; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; -import org.apache.iotdb.session.pool.AbstractSessionPoolBuilder; +import org.apache.iotdb.session.i18n.SessionMessages; +import org.apache.iotdb.session.pool.SessionPool; import java.util.Objects; -final class InternalSessionPoolBuilder { +/** Server-side builder for tree-model internal session pools. */ +public class InternalSessionPoolBuilder extends SessionPool.Builder { - private final AbstractSessionPoolBuilder sessionPoolBuilder; private final INodeSupplier nodeSupplier; - InternalSessionPoolBuilder( - final AbstractSessionPoolBuilder sessionPoolBuilder, final INodeSupplier nodeSupplier) { - this.sessionPoolBuilder = Objects.requireNonNull(sessionPoolBuilder); - this.nodeSupplier = nodeSupplier; + public InternalSessionPoolBuilder() { + this.nodeSupplier = null; } - AbstractSessionPoolBuilder getSessionPoolBuilder() { - return sessionPoolBuilder; + public InternalSessionPoolBuilder(final INodeSupplier nodeSupplier) { + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + } + + public InternalSessionPoolBuilder endPoint(final TEndPoint endPoint) { + final TEndPoint endpoint = + Objects.requireNonNull( + endPoint, SessionMessages.EXCEPTION_END_POINT_MUST_NOT_BE_NULL_908A4B27); + this.host = endpoint.getIp(); + this.rpcPort = endpoint.getPort(); + this.nodeUrls = null; + return this; } INodeSupplier getNodeSupplier() { return nodeSupplier; } + + @Override + public InternalSessionPool build() { + return new InternalSessionPool(this); + } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java deleted file mode 100644 index 21ee0bbcde74a..0000000000000 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java +++ /dev/null @@ -1,29 +0,0 @@ -/* - * 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.commons.client.internal.session; - -import org.apache.iotdb.session.pool.TableSessionPool; - -final class InternalTableSessionPool extends TableSessionPool { - - InternalTableSessionPool(final InternalSessionPoolBuilder builder) { - super(new InternalSessionPool(builder)); - } -} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java deleted file mode 100644 index 706677a334217..0000000000000 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java +++ /dev/null @@ -1,60 +0,0 @@ -/* - * 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.commons.client.internal.session; - -import org.apache.iotdb.common.rpc.thrift.TEndPoint; -import org.apache.iotdb.isession.INodeSupplier; -import org.apache.iotdb.isession.pool.ITableSessionPool; -import org.apache.iotdb.session.i18n.SessionMessages; -import org.apache.iotdb.session.pool.TableSessionPoolBuilder; - -import java.util.Objects; - -/** Server-side builder for table-model internal session pools. */ -public class InternalTableSessionPoolBuilder extends TableSessionPoolBuilder { - - private final INodeSupplier nodeSupplier; - - public InternalTableSessionPoolBuilder() { - this.nodeSupplier = null; - } - - public InternalTableSessionPoolBuilder(final INodeSupplier nodeSupplier) { - this.nodeSupplier = - Objects.requireNonNull( - nodeSupplier, - SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); - } - - public InternalTableSessionPoolBuilder endPoint(final TEndPoint endPoint) { - final TEndPoint endpoint = - Objects.requireNonNull( - endPoint, SessionMessages.EXCEPTION_END_POINT_MUST_NOT_BE_NULL_908A4B27); - this.host = endpoint.getIp(); - this.rpcPort = endpoint.getPort(); - this.nodeUrls = null; - return this; - } - - @Override - protected ITableSessionPool createSessionPool() { - return new InternalTableSessionPool(new InternalSessionPoolBuilder(this, nodeSupplier)); - } -} From 20252f832f34348533f8359c426b707e51940324 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 15:33:56 +0800 Subject: [PATCH 20/32] Support custom node suppliers in internal session pools --- .../pool/AbstractSessionPoolBuilder.java | 2 + .../iotdb/session/pool/SessionPool.java | 35 +++++++-------- .../session/InternalSessionBuilder.java | 11 ----- .../internal/session/InternalSessionPool.java | 30 +------------ .../session/InternalSessionPoolBuilder.java | 21 +-------- .../session/InternalTableSessionBuilder.java | 11 ----- .../session/InternalTableSessionPool.java | 43 ++++++++++++++++++ .../InternalTableSessionPoolBuilder.java | 45 +++++++++++++++++++ ...lSubscriptionTablePullConsumerBuilder.java | 11 ----- ...ternalSubscriptionTableSessionBuilder.java | 11 ----- 10 files changed, 108 insertions(+), 112 deletions(-) create mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java create mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/AbstractSessionPoolBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/AbstractSessionPoolBuilder.java index 91c93d8dd63b8..f8ce0c3d54c5b 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/AbstractSessionPoolBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/AbstractSessionPoolBuilder.java @@ -19,6 +19,7 @@ package org.apache.iotdb.session.pool; +import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.session.AbstractSessionBuilder; @@ -26,4 +27,5 @@ public class AbstractSessionPoolBuilder extends AbstractSessionBuilder { int maxSize = SessionConfig.DEFAULT_SESSION_POOL_MAX_SIZE; long waitToGetSessionTimeoutInMs = 60_000; int connectionTimeoutInMs = SessionConfig.DEFAULT_CONNECTION_TIMEOUT_MS; + protected INodeSupplier nodeSupplier; } 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 5ef0eac2b3aef..35dd03886d932 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 @@ -171,9 +171,7 @@ public class SessionPool implements ISessionPool { @SuppressWarnings("java:S3077") private volatile ScheduledExecutorService executorService; - protected INodeSupplier availableNodes; - - protected final AbstractSessionPoolBuilder sessionPoolBuilder; + private INodeSupplier availableNodes; // set to true, means that we will start a background thread to fetch all available (Status is // not Removing) datanodes in cluster, and these available nodes will be used in retrying stage @@ -528,7 +526,6 @@ public SessionPool( } public SessionPool(AbstractSessionPoolBuilder builder) { - this.sessionPoolBuilder = builder; this.maxSize = builder.maxSize; this.user = builder.username; this.password = builder.pw; @@ -561,9 +558,11 @@ public SessionPool(AbstractSessionPoolBuilder builder) { this.queryTimeoutInMs = builder.timeOut; this.enableAutoCreateSchema = builder.enableAutoCreateSchema; - if (enableAutoFetch) { + final INodeSupplier suppliedNodes = builder.nodeSupplier; + if (enableAutoFetch && suppliedNodes == null) { initThreadPool(); } + final List endPointList; if (builder.nodeUrls != null) { if (builder.nodeUrls.isEmpty()) { throw new IllegalArgumentException(SessionMessages.NODE_URLS_EMPTY); @@ -572,23 +571,21 @@ public SessionPool(AbstractSessionPoolBuilder builder) { this.host = null; this.port = -1; this.formattedNodeUrls = builder.nodeUrls.toString(); - if (enableAutoFetch) { - initAvailableNodes(SessionUtils.parseSeedNodeUrls(nodeUrls)); - } else { - this.availableNodes = new DummyNodesSupplier(SessionUtils.parseSeedNodeUrls(nodeUrls)); - } - + endPointList = SessionUtils.parseSeedNodeUrls(nodeUrls); } else { this.host = builder.host; this.port = builder.rpcPort; this.nodeUrls = null; this.formattedNodeUrls = UrlUtils.formatTEndPointIpv4AndIpv6Url(host, port); - if (enableAutoFetch) { - initAvailableNodes(Collections.singletonList(new TEndPoint(host, port))); - } else { - this.availableNodes = - new DummyNodesSupplier(Collections.singletonList(new TEndPoint(host, port))); - } + endPointList = Collections.singletonList(new TEndPoint(host, port)); + } + + if (suppliedNodes != null) { + this.availableNodes = suppliedNodes; + } else if (enableAutoFetch) { + initAvailableNodes(endPointList); + } else { + this.availableNodes = new DummyNodesSupplier(endPointList); } } @@ -662,7 +659,7 @@ protected Session.Builder createSessionBuilder() { return new Session.Builder(); } - protected void initThreadPool() { + private void initThreadPool() { this.executorService = Executors.newSingleThreadScheduledExecutor( r -> { @@ -679,7 +676,7 @@ protected void initThreadPool() { }); } - protected void initAvailableNodes(List endPointList) { + private void initAvailableNodes(List endPointList) { this.availableNodes = NodesSupplier.createNodeSupplier( endPointList, diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java index a5b5da83edebf..0c81de3594172 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java @@ -19,7 +19,6 @@ package org.apache.iotdb.commons.client.internal.session; -import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.session.Session; import org.apache.iotdb.session.i18n.SessionMessages; @@ -42,16 +41,6 @@ public InternalSessionBuilder(final INodeSupplier nodeSupplier) { SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); } - public InternalSessionBuilder endPoint(final TEndPoint endPoint) { - final TEndPoint endpoint = - Objects.requireNonNull( - endPoint, SessionMessages.EXCEPTION_END_POINT_MUST_NOT_BE_NULL_908A4B27); - this.host = endpoint.getIp(); - this.rpcPort = endpoint.getPort(); - this.nodeUrls = null; - return this; - } - @Override protected Session createSession() { return new InternalSession(this, nodeSupplier); diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java index 28d644a049a9b..ea4a27a9ab6f6 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java @@ -19,45 +19,17 @@ package org.apache.iotdb.commons.client.internal.session; -import org.apache.iotdb.common.rpc.thrift.TEndPoint; -import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.session.Session; import org.apache.iotdb.session.pool.SessionPool; -import java.util.List; - public final class InternalSessionPool extends SessionPool { InternalSessionPool(final InternalSessionPoolBuilder builder) { super(builder); } - @Override - protected void initThreadPool() { - if (getInternalBuilder().getNodeSupplier() == null) { - super.initThreadPool(); - } - } - - @Override - protected void initAvailableNodes(final List endPointList) { - final INodeSupplier nodeSupplier = getInternalBuilder().getNodeSupplier(); - if (nodeSupplier == null) { - super.initAvailableNodes(endPointList); - } else { - this.availableNodes = nodeSupplier; - } - } - @Override protected Session.Builder createSessionBuilder() { - final INodeSupplier nodeSupplier = getInternalBuilder().getNodeSupplier(); - return nodeSupplier == null - ? new InternalSessionBuilder() - : new InternalSessionBuilder(nodeSupplier); - } - - private InternalSessionPoolBuilder getInternalBuilder() { - return (InternalSessionPoolBuilder) sessionPoolBuilder; + return new InternalSessionBuilder(); } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java index a45b4a3d7ebfe..6f032cdaad1da 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java @@ -19,7 +19,6 @@ package org.apache.iotdb.commons.client.internal.session; -import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.session.i18n.SessionMessages; import org.apache.iotdb.session.pool.SessionPool; @@ -29,11 +28,7 @@ /** Server-side builder for tree-model internal session pools. */ public class InternalSessionPoolBuilder extends SessionPool.Builder { - private final INodeSupplier nodeSupplier; - - public InternalSessionPoolBuilder() { - this.nodeSupplier = null; - } + public InternalSessionPoolBuilder() {} public InternalSessionPoolBuilder(final INodeSupplier nodeSupplier) { this.nodeSupplier = @@ -42,20 +37,6 @@ public InternalSessionPoolBuilder(final INodeSupplier nodeSupplier) { SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); } - public InternalSessionPoolBuilder endPoint(final TEndPoint endPoint) { - final TEndPoint endpoint = - Objects.requireNonNull( - endPoint, SessionMessages.EXCEPTION_END_POINT_MUST_NOT_BE_NULL_908A4B27); - this.host = endpoint.getIp(); - this.rpcPort = endpoint.getPort(); - this.nodeUrls = null; - return this; - } - - INodeSupplier getNodeSupplier() { - return nodeSupplier; - } - @Override public InternalSessionPool build() { return new InternalSessionPool(this); diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java index 79ce2af635a50..161f16eba8e60 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java @@ -19,7 +19,6 @@ package org.apache.iotdb.commons.client.internal.session; -import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.session.Session; import org.apache.iotdb.session.TableSessionBuilder; @@ -43,16 +42,6 @@ public InternalTableSessionBuilder(final INodeSupplier nodeSupplier) { SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); } - public InternalTableSessionBuilder endPoint(final TEndPoint endPoint) { - final TEndPoint endpoint = - Objects.requireNonNull( - endPoint, SessionMessages.EXCEPTION_END_POINT_MUST_NOT_BE_NULL_908A4B27); - this.host = endpoint.getIp(); - this.rpcPort = endpoint.getPort(); - this.nodeUrls = null; - return this; - } - @Override protected Session createSession() { return new InternalSession(this, nodeSupplier); diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java new file mode 100644 index 0000000000000..38c6309ed167b --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java @@ -0,0 +1,43 @@ +/* + * 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.commons.client.internal.session; + +import org.apache.iotdb.isession.ITableSession; +import org.apache.iotdb.isession.pool.ITableSessionPool; +import org.apache.iotdb.rpc.IoTDBConnectionException; +import org.apache.iotdb.session.Session; +import org.apache.iotdb.session.pool.SessionPool; + +final class InternalTableSessionPool extends SessionPool implements ITableSessionPool { + + InternalTableSessionPool(final InternalTableSessionPoolBuilder builder) { + super(builder); + } + + @Override + public ITableSession getSession() throws IoTDBConnectionException { + return getPooledTableSession(); + } + + @Override + protected Session.Builder createSessionBuilder() { + return new InternalTableSessionBuilder(); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java new file mode 100644 index 0000000000000..85c681badd5c0 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java @@ -0,0 +1,45 @@ +/* + * 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.commons.client.internal.session; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.isession.pool.ITableSessionPool; +import org.apache.iotdb.session.i18n.SessionMessages; +import org.apache.iotdb.session.pool.TableSessionPoolBuilder; + +import java.util.Objects; + +/** Server-side builder for table-model internal session pools. */ +public class InternalTableSessionPoolBuilder extends TableSessionPoolBuilder { + + public InternalTableSessionPoolBuilder() {} + + public InternalTableSessionPoolBuilder(final INodeSupplier nodeSupplier) { + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + } + + @Override + protected ITableSessionPool createSessionPool() { + return new InternalTableSessionPool(this); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java index 24798528ddca2..68e3117f65591 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java @@ -19,7 +19,6 @@ package org.apache.iotdb.commons.client.internal.subscription; -import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.session.i18n.SessionMessages; import org.apache.iotdb.session.subscription.consumer.ISubscriptionTablePullConsumer; @@ -44,16 +43,6 @@ public InternalSubscriptionTablePullConsumerBuilder(final INodeSupplier nodeSupp SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); } - public InternalSubscriptionTablePullConsumerBuilder endPoint(final TEndPoint endPoint) { - final TEndPoint endpoint = - Objects.requireNonNull( - endPoint, SessionMessages.EXCEPTION_END_POINT_MUST_NOT_BE_NULL_908A4B27); - this.host = endpoint.getIp(); - this.port = endpoint.getPort(); - this.nodeUrls = null; - return this; - } - @Override protected ISubscriptionTablePullConsumer createConsumer() { return new InternalSubscriptionTablePullConsumer(this, nodeSupplier); diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java index 9f5e6f20aff48..0a3f477284b4f 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java @@ -19,7 +19,6 @@ package org.apache.iotdb.commons.client.internal.subscription; -import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.session.i18n.SessionMessages; import org.apache.iotdb.session.subscription.ISubscriptionTableSession; @@ -44,16 +43,6 @@ public InternalSubscriptionTableSessionBuilder(final INodeSupplier nodeSupplier) SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); } - public InternalSubscriptionTableSessionBuilder endPoint(final TEndPoint endPoint) { - final TEndPoint endpoint = - Objects.requireNonNull( - endPoint, SessionMessages.EXCEPTION_END_POINT_MUST_NOT_BE_NULL_908A4B27); - this.host = endpoint.getIp(); - this.rpcPort = endpoint.getPort(); - this.nodeUrls = null; - return this; - } - public InternalSubscriptionTableSessionBuilder nodeUrls(final List nodeUrls) { this.nodeUrls = nodeUrls; return this; From 568497b06b7a0562ed50a7bdbdf8220ff245c68f Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 15:39:41 +0800 Subject: [PATCH 21/32] Restore table session pool inheritance for internal pools --- .../session/InternalTableSessionPool.java | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java index 38c6309ed167b..7536644242eb9 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java @@ -19,25 +19,25 @@ package org.apache.iotdb.commons.client.internal.session; -import org.apache.iotdb.isession.ITableSession; -import org.apache.iotdb.isession.pool.ITableSessionPool; -import org.apache.iotdb.rpc.IoTDBConnectionException; import org.apache.iotdb.session.Session; import org.apache.iotdb.session.pool.SessionPool; +import org.apache.iotdb.session.pool.TableSessionPool; -final class InternalTableSessionPool extends SessionPool implements ITableSessionPool { +final class InternalTableSessionPool extends TableSessionPool { InternalTableSessionPool(final InternalTableSessionPoolBuilder builder) { - super(builder); + super(new InternalTableSessionPoolDelegate(builder)); } - @Override - public ITableSession getSession() throws IoTDBConnectionException { - return getPooledTableSession(); - } + private static final class InternalTableSessionPoolDelegate extends SessionPool { + + private InternalTableSessionPoolDelegate(final InternalTableSessionPoolBuilder builder) { + super(builder); + } - @Override - protected Session.Builder createSessionBuilder() { - return new InternalTableSessionBuilder(); + @Override + protected Session.Builder createSessionBuilder() { + return new InternalTableSessionBuilder(); + } } } From 1dd4f91b00bb85532990ba533fd68cf2adbf1029 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 15:41:26 +0800 Subject: [PATCH 22/32] Use tree session builder for internal table pool sessions --- .../client/internal/session/InternalTableSessionPool.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java index 7536644242eb9..8fa2a214155b9 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java @@ -37,7 +37,7 @@ private InternalTableSessionPoolDelegate(final InternalTableSessionPoolBuilder b @Override protected Session.Builder createSessionBuilder() { - return new InternalTableSessionBuilder(); + return new InternalSessionBuilder(); } } } From 858fd6e62d40c1819c8c8048f1565e38c9e1e834 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 16:20:26 +0800 Subject: [PATCH 23/32] Use internal RPC endpoints for redirects --- .../db/protocol/session/IClientSession.java | 22 +++++ .../session/InternalClientSession.java | 2 + .../db/protocol/session/SessionManager.java | 73 ++++++++------- .../thrift/impl/ClientRPCServiceImpl.java | 9 +- .../impl/DataNodeInternalRPCServiceImpl.java | 3 +- .../queryengine/common/DataNodeEndPoints.java | 11 +++ .../queryengine/common/MPPQueryContext.java | 8 +- .../db/queryengine/plan/Coordinator.java | 14 +-- .../db/queryengine/plan/analyze/Analysis.java | 22 ++++- .../plan/analyze/AnalyzeVisitor.java | 90 ++++++++++--------- .../db/queryengine/plan/analyze/Analyzer.java | 5 +- .../queryengine/plan/analyze/IAnalysis.java | 5 +- .../plan/analyze/load/LoadTsFileAnalyzer.java | 24 ++--- .../plan/execution/QueryExecution.java | 4 +- .../plan/node/write/InsertRowNode.java | 5 +- .../plan/node/write/InsertRowsNode.java | 8 +- .../node/write/InsertRowsOfOneDeviceNode.java | 5 +- .../plan/node/write/InsertTabletNode.java | 8 +- .../node/write/RelationalInsertRowsNode.java | 8 +- .../write/RelationalInsertTabletNode.java | 19 ++-- .../plan/relational/analyzer/Analysis.java | 26 +++++- .../plan/relational/analyzer/Analyzer.java | 3 +- .../metrics/IoTDBInternalLocalReporter.java | 3 +- .../planner/TableModelPlannerTest.java | 13 ++- .../queryengine/common/SessionInfo.java | 23 +++-- 25 files changed, 265 insertions(+), 148 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java index aaf73db5c4dd6..3b8fde376ef42 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java @@ -53,6 +53,12 @@ public abstract class IClientSession { private boolean login = false; + // Session opened by the database itself for internal operations. + private boolean databaseInternalSession = false; + + // Session uses the internal RPC endpoint family for redirects. + private boolean internalRpcSession = false; + private long logInTime; private SqlDialect sqlDialect = SqlDialect.TREE; @@ -122,6 +128,22 @@ public void setLogin(boolean login) { this.login = login; } + public boolean isDatabaseInternalSession() { + return databaseInternalSession; + } + + public void setDatabaseInternalSession(final boolean databaseInternalSession) { + this.databaseInternalSession = databaseInternalSession; + } + + public boolean isInternalRpcSession() { + return internalRpcSession; + } + + public void setInternalRpcSession(final boolean internalRpcSession) { + this.internalRpcSession = internalRpcSession; + } + public void setLogInTime(long logInTime) { this.logInTime = logInTime; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java index 8bbb248e9897f..6dfc7bd5f9f0e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java @@ -39,6 +39,8 @@ public class InternalClientSession extends IClientSession { public InternalClientSession(String clientID) { this.clientID = clientID; + setDatabaseInternalSession(true); + setInternalRpcSession(true); } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java index 4495c0d296153..181c666deb6eb 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java @@ -265,6 +265,7 @@ public BasicOpenSessionResp loginInternal( return openSessionResp; } + session.setInternalRpcSession(true); session.setSqlDialect(sqlDialect); supplySession(session, userId, username, ZoneId.of(zoneId), clientVersion); final String logInMessage = DataNodeMiscMessages.EXCEPTION_LOGIN_SUCCESSFULLY_1C2EF670; @@ -584,55 +585,65 @@ public static SessionManager getInstance() { public SessionInfo getSessionInfo(IClientSession session) { return new SessionInfo( - session.getId(), - new UserEntity(session.getUserId(), session.getUsername(), session.getClientAddress()), - session.getZoneId(), - session.getClientVersion(), - session.getDatabaseName(), - session.getSqlDialect(), - session.isEnableAutoCreateSchema()); + session.getId(), + new UserEntity(session.getUserId(), session.getUsername(), session.getClientAddress()), + session.getZoneId(), + session.getClientVersion(), + session.getDatabaseName(), + session.getSqlDialect(), + session.isEnableAutoCreateSchema()) + .setDatabaseInternalSession(session.isDatabaseInternalSession()) + .setInternalRpcSession(session.isInternalRpcSession()); } // Sometimes we need to switch from table model to tree model, // e.g., when loading a tree model TsFile under table model dialect. public SessionInfo copySessionInfoForTreeModel(final SessionInfo sessionInfo) { return new SessionInfo( - sessionInfo.getSessionId(), - sessionInfo.getUserEntity(), - ZoneId.systemDefault(), - sessionInfo.getVersion(), - sessionInfo.getDatabaseName().orElse(null), - SqlDialect.TREE); + sessionInfo.getSessionId(), + sessionInfo.getUserEntity(), + ZoneId.systemDefault(), + sessionInfo.getVersion(), + sessionInfo.getDatabaseName().orElse(null), + SqlDialect.TREE) + .setDatabaseInternalSession(sessionInfo.isDatabaseInternalSession()) + .setInternalRpcSession(sessionInfo.isInternalRpcSession()); } public SessionInfo getSessionInfoOfTreeModel(IClientSession session) { return new SessionInfo( - session.getId(), - new UserEntity(session.getUserId(), session.getUsername(), session.getClientAddress()), - ZoneId.systemDefault(), - session.getClientVersion(), - session.getDatabaseName(), - SqlDialect.TREE); + session.getId(), + new UserEntity(session.getUserId(), session.getUsername(), session.getClientAddress()), + ZoneId.systemDefault(), + session.getClientVersion(), + session.getDatabaseName(), + SqlDialect.TREE) + .setDatabaseInternalSession(session.isDatabaseInternalSession()) + .setInternalRpcSession(session.isInternalRpcSession()); } public SessionInfo getSessionInfoOfTableModel(IClientSession session) { return new SessionInfo( - session.getId(), - new UserEntity(session.getUserId(), session.getUsername(), session.getClientAddress()), - ZoneId.systemDefault(), - session.getClientVersion(), - session.getDatabaseName(), - SqlDialect.TABLE); + session.getId(), + new UserEntity(session.getUserId(), session.getUsername(), session.getClientAddress()), + ZoneId.systemDefault(), + session.getClientVersion(), + session.getDatabaseName(), + SqlDialect.TABLE) + .setDatabaseInternalSession(session.isDatabaseInternalSession()) + .setInternalRpcSession(session.isInternalRpcSession()); } public SessionInfo getSessionInfoOfPipeReceiver(IClientSession session, String databaseName) { return new SessionInfo( - session.getId(), - new UserEntity(session.getUserId(), session.getUsername(), session.getClientAddress()), - ZoneId.systemDefault(), - session.getClientVersion(), - databaseName, - SqlDialect.TABLE); + session.getId(), + new UserEntity(session.getUserId(), session.getUsername(), session.getClientAddress()), + ZoneId.systemDefault(), + session.getClientVersion(), + databaseName, + SqlDialect.TABLE) + .setDatabaseInternalSession(session.isDatabaseInternalSession()) + .setInternalRpcSession(session.isInternalRpcSession()); } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java index db832dcbcebe4..bbaba9f7d256a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java @@ -67,6 +67,7 @@ import org.apache.iotdb.db.protocol.session.IClientSession; import org.apache.iotdb.db.protocol.session.SessionManager; import org.apache.iotdb.db.protocol.thrift.OperationType; +import org.apache.iotdb.db.queryengine.common.DataNodeEndPoints; import org.apache.iotdb.db.queryengine.common.header.DatasetHeader; import org.apache.iotdb.db.queryengine.common.header.DatasetHeaderFactory; import org.apache.iotdb.db.queryengine.execution.aggregation.AccumulatorFactory; @@ -1449,11 +1450,9 @@ public TSExecuteStatementResp executeFastLastDataQueryForOneDeviceV2( resp = createResponse(queryExecution.getDatasetHeader(), queryId); TSStatus tsstatus = new TSStatus(TSStatusCode.REDIRECTION_RECOMMEND.getStatusCode()); tsstatus.setRedirectNode( - regionReplicaSets - .get(regionReplicaSets.size() - 1) - .dataNodeLocations - .get(0) - .clientRpcEndPoint); + DataNodeEndPoints.getRedirectEndPoint( + regionReplicaSets.get(regionReplicaSets.size() - 1).dataNodeLocations.get(0), + clientSession.isInternalRpcSession())); resp.setStatus(tsstatus); finished = SELECT_RESULT.apply(resp, queryExecution, req.fetchSize); resp.setMoreData(!finished); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java index 4b1e05f4e857d..e8d7c4af9755d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java @@ -3573,7 +3573,8 @@ public TSStatus insertRecord(TSInsertRecordReq req) throws TException { AuthorityChecker.SUPER_USER, IoTDBDescriptor.getInstance().getConfig().getInternalAddress()), ZoneId.systemDefault()) - .setInternal(true); + .setDatabaseInternalSession(true) + .setInternalRpcSession(true); long queryId = SESSION_MANAGER.requestQueryId(); ExecutionResult result = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/DataNodeEndPoints.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/DataNodeEndPoints.java index fcc877f3d280c..7a132888611a9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/DataNodeEndPoints.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/DataNodeEndPoints.java @@ -64,6 +64,17 @@ public static final TDataNodeLocation getLocalDataNodeLocation() { LOCAL_SCHEMA_REGION_CONSENSUS_CLIENT_ENDPOINT); } + public static TEndPoint getRedirectEndPoint( + final TDataNodeLocation dataNodeLocation, final boolean internalRpcSession) { + return internalRpcSession + ? dataNodeLocation.getInternalEndPoint() + : dataNodeLocation.getClientRpcEndPoint(); + } + + public static TEndPoint getLocalRedirectEndPoint(final boolean internalRpcSession) { + return internalRpcSession ? LOCAL_HOST_INTERNAL_ENDPOINT : LOCAL_CLIENT_RPC_ENDPOINT; + } + public static boolean isSameNode(TEndPoint endPoint) { return endPoint.equals(LOCAL_HOST_DATA_BLOCK_ENDPOINT); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java index d609c5c8e028e..46d35b6f7190f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java @@ -1087,8 +1087,12 @@ public String getUsername() { } @Override - public boolean isInternal() { - return session != null && session.isInternal(); + public boolean isDatabaseInternalSession() { + return session != null && session.isDatabaseInternalSession(); + } + + public boolean isInternalRpcSession() { + return session != null && session.isInternalRpcSession(); } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java index eb3d0e6869871..d7fed651103ac 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java @@ -439,12 +439,14 @@ private IQueryExecution createQueryExecutionForTreeModel( final SessionInfo session = queryContext.getSession(); queryContext.setSession( new SessionInfo( - session.getSessionId(), - session.getUserEntity(), - session.getZoneId(), - session.getVersion(), - session.getDatabaseName().orElse(null), - SqlDialect.TABLE)); + session.getSessionId(), + session.getUserEntity(), + session.getZoneId(), + session.getVersion(), + session.getDatabaseName().orElse(null), + SqlDialect.TABLE) + .setDatabaseInternalSession(session.isDatabaseInternalSession()) + .setInternalRpcSession(session.isInternalRpcSession())); return createQueryExecutionForTableModel( statement, new SqlParser(), diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analysis.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analysis.java index f1d3fdd887327..a594224d88a05 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analysis.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analysis.java @@ -32,6 +32,7 @@ import org.apache.iotdb.commons.queryengine.common.NodeRef; import org.apache.iotdb.commons.schema.template.Template; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.common.DataNodeEndPoints; import org.apache.iotdb.db.queryengine.common.DeviceContext; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; import org.apache.iotdb.db.queryengine.common.TimeseriesContext; @@ -100,6 +101,8 @@ public class Analysis implements IAnalysis { private List redirectNodeList; + private final boolean internalRpcSession; + // map from output column name (for every node) to its datatype private final Map, TSDataType> expressionTypes = new LinkedHashMap<>(); @@ -349,6 +352,11 @@ public void setDeviceToTimeseriesSchemas( private final Set> fromWhereFilterNodes = new HashSet<>(); public Analysis() { + this(false); + } + + public Analysis(final boolean internalRpcSession) { + this.internalRpcSession = internalRpcSession; this.finishQueryAfterAnalyze = false; } @@ -430,16 +438,22 @@ public List getRedirectNodeList() { } @Override - public void setRedirectNodeList(List redirectNodeList) { - this.redirectNodeList = redirectNodeList; + public void setRedirectNodeListFromDataNodeLocations( + final List dataNodeLocations) { + redirectNodeList = new ArrayList<>(dataNodeLocations.size()); + for (final TDataNodeLocation dataNodeLocation : dataNodeLocations) { + redirectNodeList.add( + DataNodeEndPoints.getRedirectEndPoint(dataNodeLocation, internalRpcSession)); + } } @Override - public void addEndPointToRedirectNodeList(TEndPoint endPoint) { + public void addDataNodeLocationToRedirectNodeList(final TDataNodeLocation dataNodeLocation) { if (redirectNodeList == null) { redirectNodeList = new ArrayList<>(); } - redirectNodeList.add(endPoint); + redirectNodeList.add( + DataNodeEndPoints.getRedirectEndPoint(dataNodeLocation, internalRpcSession)); } public Expression getGlobalTimePredicate() { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java index 7aca24384e239..9f1b5e8c19541 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java @@ -235,15 +235,25 @@ public class AnalyzeVisitor extends StatementVisitor private final ISchemaFetcher schemaFetcher; private final IModelFetcher modelFetcher; private final TreeAnalysisMutationJournal mutationJournal; + private final boolean internalRpcSession; public AnalyzeVisitor( IPartitionFetcher partitionFetcher, ISchemaFetcher schemaFetcher, TreeAnalysisMutationJournal mutationJournal) { + this(partitionFetcher, schemaFetcher, mutationJournal, false); + } + + public AnalyzeVisitor( + IPartitionFetcher partitionFetcher, + ISchemaFetcher schemaFetcher, + TreeAnalysisMutationJournal mutationJournal, + final boolean internalRpcSession) { this.partitionFetcher = partitionFetcher; this.schemaFetcher = schemaFetcher; this.modelFetcher = ModelFetcher.getInstance(); this.mutationJournal = mutationJournal; + this.internalRpcSession = internalRpcSession; } @Override @@ -282,7 +292,7 @@ public Analysis visitExplainAnalyze( @Override public Analysis visitQuery(QueryStatement queryStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setLastLevelUseWildcard(queryStatement.isLastLevelUseWildcard()); try { @@ -2651,7 +2661,7 @@ public Analysis visitCreateTimeseries( } } - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(createTimeSeriesStatement); checkIsTemplateCompatible( @@ -2749,7 +2759,7 @@ public Analysis visitCreateAlignedTimeseries( .MEASUREMENT_UNDER_AN_ALIGNED_DEVICE_IS_NOT_ALLOWED_TO_HAVE_THE_SAME_MEASUREMENT_NAME); } - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(createAlignedTimeSeriesStatement); checkIsTemplateCompatible( @@ -2777,7 +2787,7 @@ public Analysis visitInternalCreateTimeseries( MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(internalCreateTimeSeriesStatement); checkIsTemplateCompatible( internalCreateTimeSeriesStatement.getDevicePath(), @@ -2805,7 +2815,7 @@ public Analysis visitInternalCreateMultiTimeSeries( final MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - final Analysis analysis = new Analysis(); + final Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(internalCreateMultiTimeSeriesStatement); final PathPatternTree pathPatternTree = new PathPatternTree(); @@ -2829,7 +2839,7 @@ public Analysis visitCreateMultiTimeSeries( final CreateMultiTimeSeriesStatement createMultiTimeSeriesStatement, final MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - final Analysis analysis = new Analysis(); + final Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(createMultiTimeSeriesStatement); analyzeSchemaProps(createMultiTimeSeriesStatement.getPropsList()); @@ -2857,7 +2867,7 @@ public Analysis visitCreateMultiTimeSeries( public Analysis visitAlterTimeSeries( AlterTimeSeriesStatement alterTimeSeriesStatement, MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(alterTimeSeriesStatement); Pair templateInfo = @@ -2884,7 +2894,7 @@ public Analysis visitAlterTimeSeries( @Override public Analysis visitInsertTablet( InsertTabletStatement insertTabletStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); AnalyzeUtils.analyzeInsert( context, insertTabletStatement, @@ -2900,7 +2910,7 @@ public Analysis visitInsertTablet( public Analysis visitInsertRow(InsertRowStatement insertRowStatement, MPPQueryContext context) { context.setQueryType(QueryType.WRITE); insertRowStatement.semanticCheck(); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); validateSchema( analysis, insertRowStatement, @@ -2940,7 +2950,7 @@ public Analysis visitInsertRows( InsertRowsStatement insertRowsStatement, MPPQueryContext context) { context.setQueryType(QueryType.WRITE); insertRowsStatement.semanticCheck(); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); validateSchema( analysis, insertRowsStatement, @@ -2965,7 +2975,7 @@ public Analysis visitInsertMultiTablets( InsertMultiTabletsStatement insertMultiTabletsStatement, MPPQueryContext context) { context.setQueryType(QueryType.WRITE); insertMultiTabletsStatement.semanticCheck(); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); validateSchema( analysis, insertMultiTabletsStatement, @@ -2991,7 +3001,7 @@ public Analysis visitInsertRowsOfOneDevice( InsertRowsOfOneDeviceStatement insertRowsOfOneDeviceStatement, MPPQueryContext context) { context.setQueryType(QueryType.WRITE); insertRowsOfOneDeviceStatement.semanticCheck(); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); validateSchema( analysis, insertRowsOfOneDeviceStatement, @@ -3050,7 +3060,7 @@ public Analysis visitLoadFile(LoadTsFileStatement loadTsFileStatement, MPPQueryC try (final LoadTsFileAnalyzer loadTsFileAnalyzer = new LoadTsFileAnalyzer( loadTsFileStatement, loadTsFileStatement.isGeneratedByPipe(), context)) { - return (Analysis) loadTsFileAnalyzer.analyzeFileByFile(new Analysis()); + return (Analysis) loadTsFileAnalyzer.analyzeFileByFile(new Analysis(internalRpcSession)); } catch (final Exception e) { final String exceptionMessage = String.format( @@ -3059,7 +3069,7 @@ public Analysis visitLoadFile(LoadTsFileStatement loadTsFileStatement, MPPQueryC loadTsFileStatement, e.getMessage() == null ? e.getClass().getName() : e.getMessage()); logger.warn(exceptionMessage, e); - final Analysis analysis = new Analysis(); + final Analysis analysis = new Analysis(internalRpcSession); analysis.setFinishQueryAfterAnalyze(true); analysis.setFailStatus(RpcUtils.getStatus(TSStatusCode.LOAD_FILE_ERROR, exceptionMessage)); return analysis; @@ -3281,7 +3291,7 @@ private void addPhysicalTimeseriesForActiveCount( @Override public Analysis visitShowTimeSeries( ShowTimeSeriesStatement showTimeSeriesStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showTimeSeriesStatement); if (showTimeSeriesStatement.getPathPattern().getNodeLength() <= 1) { @@ -3350,7 +3360,7 @@ public Analysis visitShowTimeSeries( @Override public Analysis visitShowTTL(ShowTTLStatement showTTLStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showTTLStatement); analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowTTLHeader()); return analysis; @@ -3422,7 +3432,7 @@ private void analyzeDeviceRegionScan( @Override public Analysis visitShowDevices( ShowDevicesStatement showDevicesStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showDevicesStatement); if (showDevicesStatement.hasTimeCondition()) { @@ -3451,7 +3461,7 @@ public Analysis visitShowDevices( @Override public Analysis visitShowCluster( ShowClusterStatement showClusterStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showClusterStatement); if (showClusterStatement.isDetails()) { analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowClusterDetailsHeader()); @@ -3464,7 +3474,7 @@ public Analysis visitShowCluster( @Override public Analysis visitCountDatabase( CountDatabaseStatement countDatabaseStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(countDatabaseStatement); analysis.setRespDatasetHeader(DatasetHeaderFactory.getCountStorageGroupHeader()); return analysis; @@ -3473,7 +3483,7 @@ public Analysis visitCountDatabase( @Override public Analysis visitSeriesSchemaFetch( SeriesSchemaFetchStatement seriesSchemaFetchStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(seriesSchemaFetchStatement); SchemaPartition schemaPartition = @@ -3493,7 +3503,7 @@ public Analysis visitSeriesSchemaFetch( @Override public Analysis visitDeviceSchemaFetch( DeviceSchemaFetchStatement deviceSchemaFetchStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(deviceSchemaFetchStatement); PathPatternTree patternTree = new PathPatternTree(); @@ -3516,7 +3526,7 @@ public Analysis visitDeviceSchemaFetch( @Override public Analysis visitCountDevices( CountDevicesStatement countDevicesStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(countDevicesStatement); if (countDevicesStatement.hasTimeCondition()) { @@ -3543,7 +3553,7 @@ public Analysis visitCountDevices( @Override public Analysis visitCountTimeSeries( CountTimeSeriesStatement countTimeSeriesStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(countTimeSeriesStatement); PathPatternTree patternTree = new PathPatternTree(); @@ -3584,7 +3594,7 @@ public Analysis visitCountTimeSeries( @Override public Analysis visitCountLevelTimeSeries( CountLevelTimeSeriesStatement countLevelTimeSeriesStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(countLevelTimeSeriesStatement); PathPatternTree patternTree = new PathPatternTree(); @@ -3603,7 +3613,7 @@ public Analysis visitCountLevelTimeSeries( @Override public Analysis visitCountNodes(CountNodesStatement countStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(countStatement); PathPatternTree patternTree = new PathPatternTree(); @@ -3651,7 +3661,7 @@ public Analysis visitShowChildNodes( @Override public Analysis visitShowVersion( ShowVersionStatement showVersionStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showVersionStatement); analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowVersionHeader()); analysis.setFinishQueryAfterAnalyze(true); @@ -3660,7 +3670,7 @@ public Analysis visitShowVersion( private Analysis visitSchemaNodeManagementPartition( ShowStatement statement, PartialPath path, PathPatternTree scope, DatasetHeader header) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(statement); PathPatternTree patternTree = new PathPatternTree(); @@ -3686,7 +3696,7 @@ private Analysis visitSchemaNodeManagementPartition( public Analysis visitDeleteData( DeleteDataStatement deleteDataStatement, MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(deleteDataStatement); PathPatternTree patternTree = new PathPatternTree(); @@ -3782,7 +3792,7 @@ public Analysis visitCreateSchemaTemplate( DataNodeQueryMessages .MEASUREMENT_UNDER_TEMPLATE_IS_NOT_ALLOWED_TO_HAVE_THE_SAME_MEASUREMENT_NAME); } - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(createTemplateStatement); return analysis; } @@ -3791,7 +3801,7 @@ public Analysis visitCreateSchemaTemplate( public Analysis visitSetSchemaTemplate( SetSchemaTemplateStatement setSchemaTemplateStatement, MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(setSchemaTemplateStatement); return analysis; } @@ -3800,7 +3810,7 @@ public Analysis visitSetSchemaTemplate( public Analysis visitActivateTemplate( ActivateTemplateStatement activateTemplateStatement, MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(activateTemplateStatement); PartialPath activatePath = activateTemplateStatement.getPath(); @@ -3831,7 +3841,7 @@ public Analysis visitActivateTemplate( public Analysis visitBatchActivateTemplate( BatchActivateTemplateStatement batchActivateTemplateStatement, MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(batchActivateTemplateStatement); Map> deviceTemplateSetInfoMap = @@ -3868,7 +3878,7 @@ public Analysis visitInternalBatchActivateTemplate( InternalBatchActivateTemplateStatement internalBatchActivateTemplateStatement, MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(internalBatchActivateTemplateStatement); PathPatternTree patternTree = new PathPatternTree(); @@ -3889,7 +3899,7 @@ public Analysis visitInternalBatchActivateTemplate( @Override public Analysis visitShowPathsUsingTemplate( ShowPathsUsingTemplateStatement showPathsUsingTemplateStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showPathsUsingTemplateStatement); analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowPathsUsingTemplateHeader()); @@ -3936,7 +3946,7 @@ public Analysis visitShowPathsUsingTemplate( @Override public Analysis visitShowQueries( ShowQueriesStatement showQueriesStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showQueriesStatement); analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowQueriesHeader()); analysis.setVirtualSource(true); @@ -3973,7 +3983,7 @@ public Analysis visitShowQueries( @Override public Analysis visitShowReceivers( ShowReceiversStatement showReceiversStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showReceiversStatement); analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowReceiversHeader()); analysis.setVirtualSource(true); @@ -4001,7 +4011,7 @@ public Analysis visitShowReceivers( @Override public Analysis visitShowDiskUsage( ShowDiskUsageStatement showDiskUsageStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showDiskUsageStatement); analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowDiskUsageHeader()); analysis.setVirtualSource(true); @@ -4078,7 +4088,7 @@ private boolean analyzeWhereForShowStatements( @Override public Analysis visitCreateLogicalView( CreateLogicalViewStatement createLogicalViewStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); context.setQueryType(QueryType.OTHER); analysis.setRealStatement(createLogicalViewStatement); @@ -4357,7 +4367,7 @@ private void checkTargetPathsInCreateLogicalView( public Analysis visitShowLogicalView( final ShowLogicalViewStatement showLogicalViewStatement, final MPPQueryContext context) { context.setQueryType(QueryType.READ); - final Analysis analysis = new Analysis(); + final Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showLogicalViewStatement); final PathPatternTree patternTree = new PathPatternTree(); @@ -4374,7 +4384,7 @@ public Analysis visitShowLogicalView( @Override public Analysis visitShowCurrentTimestamp( ShowCurrentTimestampStatement showCurrentTimestampStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showCurrentTimestampStatement); analysis.setFinishQueryAfterAnalyze(true); analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowCurrentTimestampHeader()); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analyzer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analyzer.java index 4ef11925f26ce..e8d124024626b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analyzer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analyzer.java @@ -54,7 +54,9 @@ public Analyzer( public Analysis analyze(Statement statement) { long startTime = System.nanoTime(); mutationJournal.prepareForAnalyze(); - AnalyzeVisitor visitor = new AnalyzeVisitor(partitionFetcher, schemaFetcher, mutationJournal); + AnalyzeVisitor visitor = + new AnalyzeVisitor( + partitionFetcher, schemaFetcher, mutationJournal, context.isInternalRpcSession()); Analysis analysis = null; context.setReserveMemoryForSchemaTreeFunc( mem -> { @@ -78,7 +80,6 @@ public Analysis analyze(Statement statement) { // for test compatibility analysis.setDatabaseName(context.getDatabaseName().orElse(null)); } - if (statement.isQuery()) { long analyzeCost = System.nanoTime() diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IAnalysis.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IAnalysis.java index 4f48961691d74..1e62896ceeb1b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IAnalysis.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IAnalysis.java @@ -19,6 +19,7 @@ package org.apache.iotdb.db.queryengine.plan.analyze; +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.partition.DataPartition; @@ -68,9 +69,9 @@ default void setRealStatement(Statement realStatement) {} List getRedirectNodeList(); - void setRedirectNodeList(List redirectNodeList); + void setRedirectNodeListFromDataNodeLocations(List dataNodeLocations); - void addEndPointToRedirectNodeList(TEndPoint endPoint); + void addDataNodeLocationToRedirectNodeList(TDataNodeLocation dataNodeLocation); TimePredicate getConvertedTimePredicate(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/load/LoadTsFileAnalyzer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/load/LoadTsFileAnalyzer.java index 802e4a1697e34..d06aff0dbd461 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/load/LoadTsFileAnalyzer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/load/LoadTsFileAnalyzer.java @@ -475,21 +475,25 @@ && handleSingleMiniFile(i)) { if (isTableModelFile) { final SessionInfo newSessionInfo = new SessionInfo( - sessionInfo.getSessionId(), - sessionInfo.getUserEntity(), - sessionInfo.getZoneId(), - sessionInfo.getDatabaseName().orElse(null), - SqlDialect.TABLE); + sessionInfo.getSessionId(), + sessionInfo.getUserEntity(), + sessionInfo.getZoneId(), + sessionInfo.getDatabaseName().orElse(null), + SqlDialect.TABLE) + .setDatabaseInternalSession(sessionInfo.isDatabaseInternalSession()) + .setInternalRpcSession(sessionInfo.isInternalRpcSession()); context.setSession(newSessionInfo); doAnalyzeSingleTableFile(tsFile, reader, timeseriesMetadataIterator, tableSchemaMap); } else { final SessionInfo newSessionInfo = new SessionInfo( - sessionInfo.getSessionId(), - sessionInfo.getUserEntity(), - sessionInfo.getZoneId(), - sessionInfo.getDatabaseName().orElse(null), - SqlDialect.TREE); + sessionInfo.getSessionId(), + sessionInfo.getUserEntity(), + sessionInfo.getZoneId(), + sessionInfo.getDatabaseName().orElse(null), + SqlDialect.TREE) + .setDatabaseInternalSession(sessionInfo.isDatabaseInternalSession()) + .setInternalRpcSession(sessionInfo.isInternalRpcSession()); context.setSession(newSessionInfo); doAnalyzeSingleTreeFile(tsFile, reader, timeseriesMetadataIterator); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java index 2324bea787789..d77d816960edf 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java @@ -75,6 +75,7 @@ import static com.google.common.base.Throwables.throwIfUnchecked; import static org.apache.iotdb.calc.metric.QueryExecutionMetricSet.WAIT_FOR_RESULT; import static org.apache.iotdb.commons.utils.ErrorHandlingCommonUtils.getRootCause; +import static org.apache.iotdb.db.queryengine.common.DataNodeEndPoints.getLocalRedirectEndPoint; import static org.apache.iotdb.db.queryengine.common.DataNodeEndPoints.isSameNode; import static org.apache.iotdb.rpc.TSStatusCode.DATE_OUT_OF_RANGE; @@ -749,7 +750,8 @@ private ExecutionResult getExecutionResult(QueryState state) { // info to client if (!CONFIG.isEnable13DataInsertAdapt() || IoTDBConstant.ClientVersion.V_1_0.equals(context.getSession().getVersion())) { - planner.setRedirectInfo(analysis, CONFIG.getAddressAndPort(), tsstatus); + planner.setRedirectInfo( + analysis, getLocalRedirectEndPoint(context.isInternalRpcSession()), tsstatus); } return new ExecutionResult( context.getQueryId(), tsstatus, stateMachine.isFailureCausedByInterruption(tsstatus)); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java index ecf76892e4cd9..896bb1505bd6a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java @@ -135,9 +135,8 @@ public List splitByPartition(IAnalysis analysis) { .getDataRegionReplicaSetForWriting( getDeviceID(), timePartitionSlot, analysis.getDatabaseName()); // collect redirectInfo - analysis.setRedirectNodeList( - Collections.singletonList( - dataRegionReplicaSet.getDataNodeLocations().get(0).getClientRpcEndPoint())); + analysis.setRedirectNodeListFromDataNodeLocations( + Collections.singletonList(dataRegionReplicaSet.getDataNodeLocations().get(0))); return Collections.singletonList(this); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsNode.java index c2ba7c04791ec..e58834119df3b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsNode.java @@ -19,7 +19,7 @@ package org.apache.iotdb.db.queryengine.plan.planner.plan.node.write; -import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.consensus.index.ProgressIndex; @@ -304,7 +304,7 @@ public void markAsGeneratedByRemoteConsensusLeader() { @Override public List splitByPartition(IAnalysis analysis) { Map splitMap = new HashMap<>(); - List redirectInfo = new ArrayList<>(); + List redirectInfo = new ArrayList<>(); for (int i = 0; i < insertRowNodeList.size(); i++) { InsertRowNode insertRowNode = insertRowNodeList.get(i); // Data region for insert row node @@ -317,7 +317,7 @@ public List splitByPartition(IAnalysis analysis) { TimePartitionUtils.getTimePartitionSlot(insertRowNode.getTime()), null); // Collect redirectInfo - redirectInfo.add(dataRegionReplicaSet.getDataNodeLocations().get(0).getClientRpcEndPoint()); + redirectInfo.add(dataRegionReplicaSet.getDataNodeLocations().get(0)); InsertRowsNode tmpNode = splitMap.get(dataRegionReplicaSet); if (tmpNode != null) { tmpNode.addOneInsertRowNode(insertRowNode, i); @@ -331,7 +331,7 @@ public List splitByPartition(IAnalysis analysis) { splitMap.put(dataRegionReplicaSet, tmpNode); } } - analysis.setRedirectNodeList(redirectInfo); + analysis.setRedirectNodeListFromDataNodeLocations(redirectInfo); return new ArrayList<>(splitMap.values()); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsOfOneDeviceNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsOfOneDeviceNode.java index cee1e325df228..a0a8a112b297e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsOfOneDeviceNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsOfOneDeviceNode.java @@ -212,9 +212,8 @@ public List splitByPartition(IAnalysis analysis) { tmpIndexList.add(insertRowNodeIndexList.get(i)); if (i == insertRowNodeList.size() - 1) { - analysis.setRedirectNodeList( - Collections.singletonList( - dataRegionReplicaSet.getDataNodeLocations().get(0).getClientRpcEndPoint())); + analysis.setRedirectNodeListFromDataNodeLocations( + Collections.singletonList(dataRegionReplicaSet.getDataNodeLocations().get(0))); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertTabletNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertTabletNode.java index 5c91d0af4eaaf..29e42812c2c36 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertTabletNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertTabletNode.java @@ -294,12 +294,8 @@ protected Map> splitByReplicaSet( deviceID, splitInfo.timePartitionSlots, analysis.getDatabaseName()); splitInfo.replicaSets = replicaSets; // collect redirectInfo - analysis.addEndPointToRedirectNodeList( - replicaSets - .get(replicaSets.size() - 1) - .getDataNodeLocations() - .get(0) - .getClientRpcEndPoint()); + analysis.addDataNodeLocationToRedirectNodeList( + replicaSets.get(replicaSets.size() - 1).getDataNodeLocations().get(0)); for (int i = 0; i < replicaSets.size(); i++) { List subRanges = splitMap.computeIfAbsent(replicaSets.get(i), x -> new ArrayList<>()); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowsNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowsNode.java index 4a98819c6487f..234762d599a59 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowsNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowsNode.java @@ -19,7 +19,7 @@ package org.apache.iotdb.db.queryengine.plan.planner.plan.node.write; -import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.IPlanVisitor; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; @@ -196,7 +196,7 @@ public String getTableName() { public List splitByPartition(IAnalysis analysis) { List writePlanNodeList = new ArrayList<>(); Map splitMap = new HashMap<>(); - List redirectInfo = new ArrayList<>(); + List redirectInfo = new ArrayList<>(); for (int i = 0; i < getInsertRowNodeList().size(); i++) { InsertRowNode insertRowNode = getInsertRowNodeList().get(i); // Data region for insert row node @@ -210,7 +210,7 @@ public List splitByPartition(IAnalysis analysis) { analysis.getDatabaseName()); // Collect redirectInfo - redirectInfo.add(dataRegionReplicaSet.getDataNodeLocations().get(0).getClientRpcEndPoint()); + redirectInfo.add(dataRegionReplicaSet.getDataNodeLocations().get(0)); RelationalInsertRowsNode tmpNode = splitMap.get(dataRegionReplicaSet); if (tmpNode != null) { tmpNode.addOneInsertRowNode(insertRowNode, i); @@ -224,7 +224,7 @@ public List splitByPartition(IAnalysis analysis) { splitMap.put(dataRegionReplicaSet, tmpNode); } } - analysis.setRedirectNodeList(redirectInfo); + analysis.setRedirectNodeListFromDataNodeLocations(redirectInfo); writePlanNodeList.addAll(splitMap.values()); return writePlanNodeList; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertTabletNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertTabletNode.java index 029c55499bf9d..f2c7fef47b8c2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertTabletNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertTabletNode.java @@ -19,7 +19,7 @@ package org.apache.iotdb.db.queryengine.plan.planner.plan.node.write; -import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.path.PartialPath; @@ -266,7 +266,7 @@ protected InsertTabletNode getEmptySplit(int count) { protected Map> splitByReplicaSet( final Map deviceIDSplitInfoMap, final IAnalysis analysis) { final Map> splitMap = new HashMap<>(); - final Map endPointMap = new HashMap<>(); + final Map dataNodeLocationMap = new HashMap<>(); for (final Map.Entry entry : deviceIDSplitInfoMap.entrySet()) { final IDeviceID deviceID = entry.getKey(); @@ -278,13 +278,8 @@ protected Map> splitByReplicaSet( deviceID, splitInfo.timePartitionSlots, analysis.getDatabaseName()); splitInfo.replicaSets = replicaSets; // collect redirectInfo - endPointMap.put( - deviceID, - replicaSets - .get(replicaSets.size() - 1) - .getDataNodeLocations() - .get(0) - .getClientRpcEndPoint()); + dataNodeLocationMap.put( + deviceID, replicaSets.get(replicaSets.size() - 1).getDataNodeLocations().get(0)); for (int i = 0; i < replicaSets.size(); i++) { final List subRanges = splitMap.computeIfAbsent(replicaSets.get(i), x -> new ArrayList<>()); @@ -292,12 +287,12 @@ protected Map> splitByReplicaSet( subRanges.add(splitInfo.ranges.get(2 * i + 1)); } } - final List redirectNodeList = new ArrayList<>(times.length); + final List redirectNodeList = new ArrayList<>(times.length); for (int i = 0; i < times.length; i++) { final IDeviceID deviceId = getDeviceID(i); - redirectNodeList.add(endPointMap.get(deviceId)); + redirectNodeList.add(dataNodeLocationMap.get(deviceId)); } - analysis.setRedirectNodeList(redirectNodeList); + analysis.setRedirectNodeListFromDataNodeLocations(redirectNodeList); return splitMap; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analysis.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analysis.java index 884d5cbd4e74a..af8aa3fd402d7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analysis.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analysis.java @@ -19,6 +19,7 @@ package org.apache.iotdb.db.queryengine.plan.relational.analyzer; +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.partition.DataPartition; @@ -62,6 +63,7 @@ import org.apache.iotdb.commons.queryengine.plan.statement.component.FillPolicy; import org.apache.iotdb.commons.schema.table.InformationSchema; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.common.DataNodeEndPoints; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; import org.apache.iotdb.db.queryengine.common.header.DatasetHeader; import org.apache.iotdb.db.queryengine.plan.analyze.IAnalysis; @@ -118,6 +120,8 @@ public class Analysis implements IAnalysis { private String databaseName; private List redirectNodeList; + private final boolean internalRpcSession; + @Nullable private Statement root; private boolean needSetHighestPriority; @@ -266,6 +270,14 @@ public class Analysis implements IAnalysis { private SqlParser sqlParser; public Analysis(@Nullable Statement root, Map, Expression> parameters) { + this(root, parameters, false); + } + + public Analysis( + @Nullable Statement root, + Map, Expression> parameters, + final boolean internalRpcSession) { + this.internalRpcSession = internalRpcSession; this.root = root; this.parameters = ImmutableMap.copyOf( @@ -1050,16 +1062,22 @@ public List getRedirectNodeList() { } @Override - public void setRedirectNodeList(final List redirectNodeList) { - this.redirectNodeList = redirectNodeList; + public void setRedirectNodeListFromDataNodeLocations( + final List dataNodeLocations) { + redirectNodeList = new ArrayList<>(dataNodeLocations.size()); + for (final TDataNodeLocation dataNodeLocation : dataNodeLocations) { + redirectNodeList.add( + DataNodeEndPoints.getRedirectEndPoint(dataNodeLocation, internalRpcSession)); + } } @Override - public void addEndPointToRedirectNodeList(final TEndPoint endPoint) { + public void addDataNodeLocationToRedirectNodeList(final TDataNodeLocation dataNodeLocation) { if (redirectNodeList == null) { redirectNodeList = new ArrayList<>(); } - redirectNodeList.add(endPoint); + redirectNodeList.add( + DataNodeEndPoints.getRedirectEndPoint(dataNodeLocation, internalRpcSession)); } public void setTableFunctionAnalysis( diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analyzer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analyzer.java index 8b46055d02a8f..ac8dd2ca84fb0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analyzer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analyzer.java @@ -86,7 +86,8 @@ public Analysis analyze(Statement statement) { parameterLookup, warningCollector); - Analysis analysis = new Analysis(rewrittenStatement, parameterLookup); + Analysis analysis = + new Analysis(rewrittenStatement, parameterLookup, context.isInternalRpcSession()); analysis.setSqlParser(statementAnalyzerFactory.getSqlParser()); // Register CTE passed by parent query. context diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/metrics/IoTDBInternalLocalReporter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/metrics/IoTDBInternalLocalReporter.java index 5749d99b66ff7..c4bf9ee2929c6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/metrics/IoTDBInternalLocalReporter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/metrics/IoTDBInternalLocalReporter.java @@ -103,7 +103,8 @@ public IoTDBInternalLocalReporter() { AuthorityChecker.SUPER_USER, IoTDBDescriptor.getInstance().getConfig().getInternalAddress()), ZoneId.systemDefault()) - .setInternal(true), + .setDatabaseInternalSession(true) + .setInternalRpcSession(true), ClusterPartitionFetcher.getInstance(), ClusterSchemaFetcher.getInstance()); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlannerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlannerTest.java index 99c0ed023576e..9bff113a90e00 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlannerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlannerTest.java @@ -19,6 +19,7 @@ package org.apache.iotdb.db.queryengine.plan.relational.planner; +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.db.queryengine.plan.relational.analyzer.Analysis; @@ -44,7 +45,17 @@ public void testSetRedirectInfoForInsertRows() { new Analysis(new InsertRows(insertRowsStatement, null), Collections.emptyMap()); final TEndPoint localEndPoint = new TEndPoint("127.0.0.1", 6667); final TEndPoint remoteEndPoint = new TEndPoint("127.0.0.2", 6667); - analysis.setRedirectNodeList(Arrays.asList(localEndPoint, remoteEndPoint)); + analysis.setRedirectNodeListFromDataNodeLocations( + Arrays.asList( + new TDataNodeLocation( + 0, localEndPoint, localEndPoint, localEndPoint, localEndPoint, localEndPoint), + new TDataNodeLocation( + 1, + remoteEndPoint, + remoteEndPoint, + remoteEndPoint, + remoteEndPoint, + remoteEndPoint))); final TSStatus status = RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS); createPlanner().setRedirectInfo(analysis, localEndPoint, status); diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/common/SessionInfo.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/common/SessionInfo.java index d6a0f4120abb7..1dbc6012c5df9 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/common/SessionInfo.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/common/SessionInfo.java @@ -52,7 +52,11 @@ public class SessionInfo { // Whether the session is opened by the database itself. Like enableAutoCreateSchema, this is // only consulted during analysis and therefore does not need serialization. - private transient boolean internal; + private transient boolean databaseInternalSession; + + // Whether the session uses the internal RPC endpoint family. Used to select internal redirect + // endpoints without changing the meaning of internal database sessions. + private transient boolean internalRpcSession; public SessionInfo(long sessionId, UserEntity userEntity, ZoneId zoneId) { this.sessionId = sessionId; @@ -169,12 +173,21 @@ public boolean isEnableAutoCreateSchema() { return enableAutoCreateSchema; } - public boolean isInternal() { - return internal; + public boolean isDatabaseInternalSession() { + return databaseInternalSession; + } + + public SessionInfo setDatabaseInternalSession(final boolean databaseInternalSession) { + this.databaseInternalSession = databaseInternalSession; + return this; + } + + public boolean isInternalRpcSession() { + return internalRpcSession; } - public SessionInfo setInternal(final boolean internal) { - this.internal = internal; + public SessionInfo setInternalRpcSession(final boolean internalRpcSession) { + this.internalRpcSession = internalRpcSession; return this; } From e22fd5b316da782a6886f00981b75143521f1361 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 16:32:29 +0800 Subject: [PATCH 24/32] Refine internal session auto-fetch requirements --- .../apache/iotdb/session/i18n/SessionMessages.java | 3 +++ .../apache/iotdb/session/i18n/SessionMessages.java | 3 +++ .../iotdb/db/queryengine/plan/analyze/Analysis.java | 5 ++--- .../db/queryengine/plan/analyze/IAnalysis.java | 4 ++-- .../plan/planner/plan/node/write/InsertRowNode.java | 2 +- .../planner/plan/node/write/InsertRowsNode.java | 2 +- .../plan/node/write/InsertRowsOfOneDeviceNode.java | 2 +- .../planner/plan/node/write/InsertTabletNode.java | 2 +- .../plan/node/write/RelationalInsertRowsNode.java | 2 +- .../plan/node/write/RelationalInsertTabletNode.java | 2 +- .../plan/relational/analyzer/Analysis.java | 5 ++--- .../relational/planner/TableModelPlannerTest.java | 2 +- .../client/internal/session/InternalSession.java | 13 ++++++++++--- .../internal/session/InternalSessionBuilder.java | 10 +++++++++- .../internal/session/InternalSessionPool.java | 2 +- .../session/InternalSessionPoolBuilder.java | 4 ++++ .../session/InternalTableSessionBuilder.java | 10 +++++++++- .../internal/session/InternalTableSessionPool.java | 2 +- .../session/InternalTableSessionPoolBuilder.java | 4 ++++ 19 files changed, 57 insertions(+), 22 deletions(-) diff --git a/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java b/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java index 341aae9e3de62..df11f6bc33dc9 100644 --- a/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java +++ b/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java @@ -27,6 +27,9 @@ public final class SessionMessages { public static final String EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392 = "Internal node supplier must not be null"; + public static final String EXCEPTION_INTERNAL_AUTO_FETCH_REQUIRES_NODE_SUPPLIER_9CB6BA02 = + "An internal node supplier is required when auto-fetch is enabled"; + // Session public static final String NODE_URLS_EMPTY = "nodeUrls shouldn't be empty."; public static final String REDIRECT_TWICE_LOG = "redirect twice"; diff --git a/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java b/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java index ed04bf7126f02..c2f954b5cc07e 100644 --- a/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java +++ b/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java @@ -27,6 +27,9 @@ public final class SessionMessages { public static final String EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392 = "内部节点 supplier 不能为空"; + public static final String EXCEPTION_INTERNAL_AUTO_FETCH_REQUIRES_NODE_SUPPLIER_9CB6BA02 = + "启用自动发现时，内部节点 supplier 不能为空"; + // Session public static final String NODE_URLS_EMPTY = "nodeUrls 不能为空。"; public static final String REDIRECT_TWICE_LOG = "重定向了两次"; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analysis.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analysis.java index a594224d88a05..ac9ec086375b8 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analysis.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analysis.java @@ -438,8 +438,7 @@ public List getRedirectNodeList() { } @Override - public void setRedirectNodeListFromDataNodeLocations( - final List dataNodeLocations) { + public void setRedirectNodeList(final List dataNodeLocations) { redirectNodeList = new ArrayList<>(dataNodeLocations.size()); for (final TDataNodeLocation dataNodeLocation : dataNodeLocations) { redirectNodeList.add( @@ -448,7 +447,7 @@ public void setRedirectNodeListFromDataNodeLocations( } @Override - public void addDataNodeLocationToRedirectNodeList(final TDataNodeLocation dataNodeLocation) { + public void addEndPointToRedirectNodeList(final TDataNodeLocation dataNodeLocation) { if (redirectNodeList == null) { redirectNodeList = new ArrayList<>(); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IAnalysis.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IAnalysis.java index 1e62896ceeb1b..3aee6015c63ab 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IAnalysis.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IAnalysis.java @@ -69,9 +69,9 @@ default void setRealStatement(Statement realStatement) {} List getRedirectNodeList(); - void setRedirectNodeListFromDataNodeLocations(List dataNodeLocations); + void setRedirectNodeList(List dataNodeLocations); - void addDataNodeLocationToRedirectNodeList(TDataNodeLocation dataNodeLocation); + void addEndPointToRedirectNodeList(TDataNodeLocation dataNodeLocation); TimePredicate getConvertedTimePredicate(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java index 896bb1505bd6a..1230400e31045 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java @@ -135,7 +135,7 @@ public List splitByPartition(IAnalysis analysis) { .getDataRegionReplicaSetForWriting( getDeviceID(), timePartitionSlot, analysis.getDatabaseName()); // collect redirectInfo - analysis.setRedirectNodeListFromDataNodeLocations( + analysis.setRedirectNodeList( Collections.singletonList(dataRegionReplicaSet.getDataNodeLocations().get(0))); return Collections.singletonList(this); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsNode.java index e58834119df3b..44481da77c30f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsNode.java @@ -331,7 +331,7 @@ public List splitByPartition(IAnalysis analysis) { splitMap.put(dataRegionReplicaSet, tmpNode); } } - analysis.setRedirectNodeListFromDataNodeLocations(redirectInfo); + analysis.setRedirectNodeList(redirectInfo); return new ArrayList<>(splitMap.values()); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsOfOneDeviceNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsOfOneDeviceNode.java index a0a8a112b297e..ce78aea859bd6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsOfOneDeviceNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsOfOneDeviceNode.java @@ -212,7 +212,7 @@ public List splitByPartition(IAnalysis analysis) { tmpIndexList.add(insertRowNodeIndexList.get(i)); if (i == insertRowNodeList.size() - 1) { - analysis.setRedirectNodeListFromDataNodeLocations( + analysis.setRedirectNodeList( Collections.singletonList(dataRegionReplicaSet.getDataNodeLocations().get(0))); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertTabletNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertTabletNode.java index 29e42812c2c36..bd197ce6168da 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertTabletNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertTabletNode.java @@ -294,7 +294,7 @@ protected Map> splitByReplicaSet( deviceID, splitInfo.timePartitionSlots, analysis.getDatabaseName()); splitInfo.replicaSets = replicaSets; // collect redirectInfo - analysis.addDataNodeLocationToRedirectNodeList( + analysis.addEndPointToRedirectNodeList( replicaSets.get(replicaSets.size() - 1).getDataNodeLocations().get(0)); for (int i = 0; i < replicaSets.size(); i++) { List subRanges = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowsNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowsNode.java index 234762d599a59..5fff5ad440919 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowsNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowsNode.java @@ -224,7 +224,7 @@ public List splitByPartition(IAnalysis analysis) { splitMap.put(dataRegionReplicaSet, tmpNode); } } - analysis.setRedirectNodeListFromDataNodeLocations(redirectInfo); + analysis.setRedirectNodeList(redirectInfo); writePlanNodeList.addAll(splitMap.values()); return writePlanNodeList; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertTabletNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertTabletNode.java index f2c7fef47b8c2..50924da4447b9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertTabletNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertTabletNode.java @@ -292,7 +292,7 @@ protected Map> splitByReplicaSet( final IDeviceID deviceId = getDeviceID(i); redirectNodeList.add(dataNodeLocationMap.get(deviceId)); } - analysis.setRedirectNodeListFromDataNodeLocations(redirectNodeList); + analysis.setRedirectNodeList(redirectNodeList); return splitMap; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analysis.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analysis.java index af8aa3fd402d7..375c8d59be283 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analysis.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analysis.java @@ -1062,8 +1062,7 @@ public List getRedirectNodeList() { } @Override - public void setRedirectNodeListFromDataNodeLocations( - final List dataNodeLocations) { + public void setRedirectNodeList(final List dataNodeLocations) { redirectNodeList = new ArrayList<>(dataNodeLocations.size()); for (final TDataNodeLocation dataNodeLocation : dataNodeLocations) { redirectNodeList.add( @@ -1072,7 +1071,7 @@ public void setRedirectNodeListFromDataNodeLocations( } @Override - public void addDataNodeLocationToRedirectNodeList(final TDataNodeLocation dataNodeLocation) { + public void addEndPointToRedirectNodeList(final TDataNodeLocation dataNodeLocation) { if (redirectNodeList == null) { redirectNodeList = new ArrayList<>(); } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlannerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlannerTest.java index 9bff113a90e00..67f5ff1083d63 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlannerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlannerTest.java @@ -45,7 +45,7 @@ public void testSetRedirectInfoForInsertRows() { new Analysis(new InsertRows(insertRowsStatement, null), Collections.emptyMap()); final TEndPoint localEndPoint = new TEndPoint("127.0.0.1", 6667); final TEndPoint remoteEndPoint = new TEndPoint("127.0.0.2", 6667); - analysis.setRedirectNodeListFromDataNodeLocations( + analysis.setRedirectNodeList( Arrays.asList( new TDataNodeLocation( 0, localEndPoint, localEndPoint, localEndPoint, localEndPoint, localEndPoint), diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java index f08455bc00df2..84d14e20fd378 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java @@ -25,12 +25,19 @@ import org.apache.iotdb.session.Session; /** Base session for trusted server-to-server client RPC connections. */ -public class InternalSession extends Session { +final class InternalSession extends Session { private final INodeSupplier nodeSupplier; - protected InternalSession( - final AbstractSessionBuilder builder, final INodeSupplier nodeSupplier) { + InternalSession(final InternalSessionBuilder builder) { + this(builder, builder.getNodeSupplier()); + } + + InternalSession(final InternalTableSessionBuilder builder) { + this(builder, builder.getNodeSupplier()); + } + + private InternalSession(final AbstractSessionBuilder builder, final INodeSupplier nodeSupplier) { super(builder); this.nodeSupplier = nodeSupplier; } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java index 0c81de3594172..5b3d36e7f3a02 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java @@ -41,8 +41,16 @@ public InternalSessionBuilder(final INodeSupplier nodeSupplier) { SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); } + INodeSupplier getNodeSupplier() { + return nodeSupplier; + } + @Override protected Session createSession() { - return new InternalSession(this, nodeSupplier); + if (enableAutoFetch && nodeSupplier == null) { + throw new IllegalArgumentException( + SessionMessages.EXCEPTION_INTERNAL_AUTO_FETCH_REQUIRES_NODE_SUPPLIER_9CB6BA02); + } + return new InternalSession(this); } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java index ea4a27a9ab6f6..1e65d06143293 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java @@ -30,6 +30,6 @@ public final class InternalSessionPool extends SessionPool { @Override protected Session.Builder createSessionBuilder() { - return new InternalSessionBuilder(); + return new InternalSessionBuilder().enableAutoFetch(false); } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java index 6f032cdaad1da..8ba30a18dab62 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java @@ -39,6 +39,10 @@ public InternalSessionPoolBuilder(final INodeSupplier nodeSupplier) { @Override public InternalSessionPool build() { + if (enableAutoFetch && nodeSupplier == null) { + throw new IllegalArgumentException( + SessionMessages.EXCEPTION_INTERNAL_AUTO_FETCH_REQUIRES_NODE_SUPPLIER_9CB6BA02); + } return new InternalSessionPool(this); } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java index 161f16eba8e60..132f944212c40 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java @@ -42,8 +42,16 @@ public InternalTableSessionBuilder(final INodeSupplier nodeSupplier) { SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); } + INodeSupplier getNodeSupplier() { + return nodeSupplier; + } + @Override protected Session createSession() { - return new InternalSession(this, nodeSupplier); + if (enableAutoFetch && nodeSupplier == null) { + throw new IllegalArgumentException( + SessionMessages.EXCEPTION_INTERNAL_AUTO_FETCH_REQUIRES_NODE_SUPPLIER_9CB6BA02); + } + return new InternalSession(this); } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java index 8fa2a214155b9..8091b8a7a7517 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java @@ -37,7 +37,7 @@ private InternalTableSessionPoolDelegate(final InternalTableSessionPoolBuilder b @Override protected Session.Builder createSessionBuilder() { - return new InternalSessionBuilder(); + return new InternalSessionBuilder().enableAutoFetch(false); } } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java index 85c681badd5c0..6c1d9b04253e3 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java @@ -40,6 +40,10 @@ public InternalTableSessionPoolBuilder(final INodeSupplier nodeSupplier) { @Override protected ITableSessionPool createSessionPool() { + if (enableAutoFetch && nodeSupplier == null) { + throw new IllegalArgumentException( + SessionMessages.EXCEPTION_INTERNAL_AUTO_FETCH_REQUIRES_NODE_SUPPLIER_9CB6BA02); + } return new InternalTableSessionPool(this); } } From 1ec246e34da2c0a777172bafe89b8d92a3fc3268 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 17:09:00 +0800 Subject: [PATCH 25/32] Refine internal session lifecycle and authorization --- .../it/utils/InternalSessionITUtils.java | 158 ++++++++++++++++ .../it/session/InternalTableSessionIT.java | 174 ++++++++++++++++++ .../InternalTableSessionRedirectIT.java | 138 ++++++++++++++ .../iotdb/session/it/InternalSessionIT.java | 88 ++++++++- .../session/it/InternalSessionRedirectIT.java | 147 +++++++++++++++ .../apache/iotdb/isession/INodeSupplier.java | 8 + .../iotdb/session/DummyNodesSupplier.java | 5 + .../org/apache/iotdb/session/Session.java | 8 +- .../iotdb/session/pool/SessionPool.java | 4 +- .../iotdb/db/auth/AuthorityChecker.java | 2 +- .../security/TreeAccessCheckContext.java | 10 +- .../security/TreeAccessCheckVisitor.java | 6 +- .../apache/iotdb/db/auth/TreeAccessTest.java | 2 +- .../iotdb/commons/audit/IAuditEntity.java | 4 +- .../session/InternalSessionBuilder.java | 8 +- .../session/InternalSessionPoolBuilder.java | 8 +- .../session/InternalTableSessionBuilder.java | 8 +- .../InternalTableSessionPoolBuilder.java | 8 +- .../InternalSubscriptionSessionWrapper.java | 20 ++ ...lSubscriptionTablePullConsumerBuilder.java | 7 +- ...ternalSubscriptionTableSessionBuilder.java | 8 +- .../session/InternalBuildersTest.java | 47 +++++ .../session/InternalSessionProtocolTest.java | 29 ++- 23 files changed, 859 insertions(+), 38 deletions(-) create mode 100644 integration-test/src/test/java/org/apache/iotdb/it/utils/InternalSessionITUtils.java create mode 100644 integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionIT.java create mode 100644 integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionRedirectIT.java create mode 100644 integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionRedirectIT.java create mode 100644 iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalBuildersTest.java diff --git a/integration-test/src/test/java/org/apache/iotdb/it/utils/InternalSessionITUtils.java b/integration-test/src/test/java/org/apache/iotdb/it/utils/InternalSessionITUtils.java new file mode 100644 index 0000000000000..74001c0afdbad --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/it/utils/InternalSessionITUtils.java @@ -0,0 +1,158 @@ +/* + * 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.it.utils; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.schema.column.ColumnHeaderConstant; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.isession.ITableSession; +import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper; +import org.apache.iotdb.rpc.UrlUtils; +import org.apache.iotdb.session.DummyNodesSupplier; +import org.apache.iotdb.session.Session; + +import org.apache.tsfile.read.common.Field; +import org.apache.tsfile.read.common.RowRecord; +import org.awaitility.Awaitility; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.Assert.assertFalse; + +public final class InternalSessionITUtils { + + private static final String DATA_REGION = "DataRegion"; + private static final String LEADER = "Leader"; + + private InternalSessionITUtils() {} + + public static TEndPoint getInternalEndPoint(final DataNodeWrapper dataNode) { + return new TEndPoint(dataNode.getInternalAddress(), dataNode.getInternalPort()); + } + + public static String getEndPointUrl(final TEndPoint endPoint) { + return UrlUtils.convertTEndPointIpv4AndIpv6Url(endPoint); + } + + public static List getInternalEndPoints() { + final List endPoints = new ArrayList<>(); + for (final DataNodeWrapper dataNode : EnvFactory.getEnv().getDataNodeWrapperList()) { + endPoints.add(getInternalEndPoint(dataNode)); + } + return endPoints; + } + + public static INodeSupplier createInternalNodeSupplier() { + return new DummyNodesSupplier(getInternalEndPoints()); + } + + public static TEndPoint findNonLeaderInternalEndPoint( + final Set leaderInternalEndPoints) { + for (final TEndPoint endPoint : getInternalEndPoints()) { + if (!leaderInternalEndPoints.contains(endPoint)) { + return endPoint; + } + } + throw new IllegalStateException("No non-leader DataNode is available"); + } + + public static Set awaitDataRegionLeaderInternalEndPoints( + final Session session, final String database) { + final AtomicReference> leaderEndPoints = + new AtomicReference<>(Collections.emptySet()); + Awaitility.await() + .atMost(30, TimeUnit.SECONDS) + .untilAsserted( + () -> { + final Set currentLeaderEndPoints = + queryDataRegionLeaderInternalEndPoints(session, database); + assertFalse(currentLeaderEndPoints.isEmpty()); + leaderEndPoints.set(currentLeaderEndPoints); + }); + return leaderEndPoints.get(); + } + + public static Session getUnderlyingSession(final ITableSession tableSession) + throws ReflectiveOperationException { + Class currentClass = tableSession.getClass(); + while (currentClass != null) { + for (final java.lang.reflect.Field field : currentClass.getDeclaredFields()) { + if (Session.class.equals(field.getType())) { + field.setAccessible(true); + return (Session) field.get(tableSession); + } + } + currentClass = currentClass.getSuperclass(); + } + throw new NoSuchFieldException("Session field not found in " + tableSession.getClass()); + } + + @SuppressWarnings("unchecked") + public static T getSessionField(final Session session, final String fieldName) + throws ReflectiveOperationException { + final java.lang.reflect.Field field = Session.class.getDeclaredField(fieldName); + field.setAccessible(true); + return (T) field.get(session); + } + + private static Set queryDataRegionLeaderInternalEndPoints( + final Session session, final String database) throws Exception { + final Set leaderEndPoints = new HashSet<>(); + try (SessionDataSet dataSet = session.executeQueryStatement("SHOW REGIONS")) { + final List columnNames = dataSet.getColumnNames(); + final int typeIndex = columnNames.indexOf(ColumnHeaderConstant.TYPE); + final int databaseIndex = columnNames.indexOf(ColumnHeaderConstant.DATABASE); + final int rpcAddressIndex = columnNames.indexOf(ColumnHeaderConstant.RPC_ADDRESS); + final int rpcPortIndex = columnNames.indexOf(ColumnHeaderConstant.RPC_PORT); + final int roleIndex = columnNames.indexOf(ColumnHeaderConstant.ROLE); + + while (dataSet.hasNext()) { + final RowRecord rowRecord = dataSet.next(); + if (!DATA_REGION.equals(getString(rowRecord, typeIndex)) + || !database.equals(getString(rowRecord, databaseIndex)) + || !LEADER.equalsIgnoreCase(getString(rowRecord, roleIndex))) { + continue; + } + final String rpcAddress = getString(rowRecord, rpcAddressIndex); + final int rpcPort = rowRecord.getFields().get(rpcPortIndex).getIntV(); + for (final DataNodeWrapper dataNode : EnvFactory.getEnv().getDataNodeWrapperList()) { + if (dataNode.getIp().equals(rpcAddress) && dataNode.getPort() == rpcPort) { + leaderEndPoints.add(getInternalEndPoint(dataNode)); + break; + } + } + } + } + return leaderEndPoints; + } + + private static String getString(final RowRecord rowRecord, final int index) { + final Field field = rowRecord.getFields().get(index); + return field == null ? null : field.getStringValue(); + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionIT.java new file mode 100644 index 0000000000000..76624becc4c23 --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionIT.java @@ -0,0 +1,174 @@ +/* + * 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.relational.it.session; + +import org.apache.iotdb.commons.client.internal.session.InternalTableSessionBuilder; +import org.apache.iotdb.commons.client.internal.session.InternalTableSessionPoolBuilder; +import org.apache.iotdb.isession.ITableSession; +import org.apache.iotdb.isession.SessionConfig; +import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.isession.pool.ITableSessionPool; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.it.utils.InternalSessionITUtils; +import org.apache.iotdb.itbase.category.TableClusterIT; +import org.apache.iotdb.itbase.category.TableLocalStandaloneIT; + +import org.apache.tsfile.enums.ColumnCategory; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.IMeasurementSchema; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +@RunWith(IoTDBTestRunner.class) +@Category({TableLocalStandaloneIT.class, TableClusterIT.class}) +public class InternalTableSessionIT { + + private static final String DATABASE = "internal_table_session_it"; + private static final String TABLE = "t"; + + @Before + public void setUp() throws Exception { + EnvFactory.getEnv().initClusterEnvironment(); + } + + @After + public void tearDown() throws Exception { + EnvFactory.getEnv().cleanClusterEnvironment(); + } + + @Test + public void testInternalTableSessionUsesInternalPort() throws Exception { + final DataNodeWrapper dataNode = EnvFactory.getEnv().getDataNodeWrapper(0); + try (ITableSession session = + new InternalTableSessionBuilder(InternalSessionITUtils.createInternalNodeSupplier()) + .nodeUrls( + Collections.singletonList( + InternalSessionITUtils.getEndPointUrl( + InternalSessionITUtils.getInternalEndPoint(dataNode)))) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .build()) { + session.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + try { + createTableAndInsert(session); + assertTableValue(session); + } finally { + session.executeNonQueryStatement("DROP DATABASE " + DATABASE); + } + } + } + + @Test + public void testInternalTableSessionPoolUsesInternalPort() throws Exception { + final DataNodeWrapper dataNode = EnvFactory.getEnv().getDataNodeWrapper(0); + try (ITableSessionPool sessionPool = + new InternalTableSessionPoolBuilder(InternalSessionITUtils.createInternalNodeSupplier()) + .nodeUrls( + Collections.singletonList( + InternalSessionITUtils.getEndPointUrl( + InternalSessionITUtils.getInternalEndPoint(dataNode)))) + .user(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .maxSize(1) + .build()) { + try (ITableSession session = sessionPool.getSession()) { + session.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + try { + createTableAndInsert(session); + assertTableValue(session); + } finally { + session.executeNonQueryStatement("DROP DATABASE " + DATABASE); + } + } + } + } + + @Test + public void testInternalTableSessionWithoutAutoFetchUsesConfiguredInternalEndpoint() + throws Exception { + final DataNodeWrapper dataNode = EnvFactory.getEnv().getDataNodeWrapper(0); + try (ITableSession session = + new InternalTableSessionBuilder() + .enableAutoFetch(false) + .nodeUrls( + Collections.singletonList( + InternalSessionITUtils.getEndPointUrl( + InternalSessionITUtils.getInternalEndPoint(dataNode)))) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .build()) { + session.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + try { + session.executeNonQueryStatement( + "CREATE TABLE " + DATABASE + "." + TABLE + " (tag1 STRING TAG, value INT64 FIELD)"); + } finally { + session.executeNonQueryStatement("DROP DATABASE " + DATABASE); + } + } + } + + private static void createTableAndInsert(final ITableSession session) throws Exception { + session.executeNonQueryStatement("USE " + DATABASE); + session.executeNonQueryStatement( + "CREATE TABLE " + TABLE + " (tag1 STRING TAG, value INT64 FIELD)"); + session.insert(createTablet()); + session.insert(createTablet()); + } + + private static void assertTableValue(final ITableSession session) throws Exception { + try (SessionDataSet dataSet = + session.executeQueryStatement("SELECT value FROM " + DATABASE + "." + TABLE)) { + assertTrue(dataSet.hasNext()); + assertEquals(42L, dataSet.next().getFields().get(0).getLongV()); + } + } + + private static Tablet createTablet() { + final List schemas = + Arrays.asList( + new MeasurementSchema("tag1", TSDataType.STRING), + new MeasurementSchema("value", TSDataType.INT64)); + final Tablet tablet = + new Tablet( + TABLE, + IMeasurementSchema.getMeasurementNameList(schemas), + IMeasurementSchema.getDataTypeList(schemas), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD), + 1); + tablet.addTimestamp(0, 1L); + tablet.addValue("tag1", 0, "d1"); + tablet.addValue("value", 0, 42L); + return tablet; + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionRedirectIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionRedirectIT.java new file mode 100644 index 0000000000000..f5bf9f64c6ac4 --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionRedirectIT.java @@ -0,0 +1,138 @@ +/* + * 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.relational.it.session; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.client.internal.session.InternalTableSessionBuilder; +import org.apache.iotdb.consensus.ConsensusFactory; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.isession.ITableSession; +import org.apache.iotdb.isession.SessionConfig; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.it.utils.InternalSessionITUtils; +import org.apache.iotdb.itbase.category.TableClusterIT; +import org.apache.iotdb.session.Session; + +import org.apache.tsfile.enums.ColumnCategory; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.IMeasurementSchema; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +@RunWith(IoTDBTestRunner.class) +@Category(TableClusterIT.class) +public class InternalTableSessionRedirectIT { + + private static final String DATABASE = "internal_table_session_redirect_it"; + private static final String TABLE = "t"; + + @Before + public void setUp() throws Exception { + EnvFactory.getEnv() + .getConfig() + .getCommonConfig() + .setSchemaReplicationFactor(1) + .setDataReplicationFactor(1) + .setDataRegionGroupExtensionPolicy("CUSTOM") + .setDefaultDataRegionGroupNumPerDatabase(1) + .setSchemaRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS) + .setDataRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS); + EnvFactory.getEnv().initClusterEnvironment(1, 3); + } + + @After + public void tearDown() throws Exception { + EnvFactory.getEnv().cleanClusterEnvironment(); + } + + @Test + public void testTabletInsertRedirectsToInternalLeaderEndpoint() throws Exception { + final INodeSupplier nodeSupplier = InternalSessionITUtils.createInternalNodeSupplier(); + final TEndPoint initialEndPoint = InternalSessionITUtils.getInternalEndPoints().get(0); + final Set leaderEndPoints; + + try (ITableSession setupSession = + new InternalTableSessionBuilder(nodeSupplier) + .nodeUrls( + Collections.singletonList(InternalSessionITUtils.getEndPointUrl(initialEndPoint))) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .build()) { + setupSession.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + setupSession.executeNonQueryStatement( + "CREATE TABLE " + DATABASE + "." + TABLE + " (tag1 STRING TAG, value INT64 FIELD)"); + leaderEndPoints = + InternalSessionITUtils.awaitDataRegionLeaderInternalEndPoints( + InternalSessionITUtils.getUnderlyingSession(setupSession), DATABASE); + } + final TEndPoint nonLeaderEndPoint = + InternalSessionITUtils.findNonLeaderInternalEndPoint(leaderEndPoints); + + try (ITableSession session = + new InternalTableSessionBuilder(nodeSupplier) + .nodeUrls( + Collections.singletonList(InternalSessionITUtils.getEndPointUrl(nonLeaderEndPoint))) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .database(DATABASE) + .build()) { + session.insert(createTablet()); + + final Session underlyingSession = InternalSessionITUtils.getUnderlyingSession(session); + final Map tableModelDeviceIdToEndpoint = + InternalSessionITUtils.getSessionField(underlyingSession, "tableModelDeviceIdToEndpoint"); + assertEquals(1, tableModelDeviceIdToEndpoint.size()); + assertTrue(leaderEndPoints.contains(tableModelDeviceIdToEndpoint.values().iterator().next())); + } + } + + private static Tablet createTablet() { + final List schemas = + Arrays.asList( + new MeasurementSchema("tag1", TSDataType.STRING), + new MeasurementSchema("value", TSDataType.INT64)); + final Tablet tablet = + new Tablet( + TABLE, + IMeasurementSchema.getMeasurementNameList(schemas), + IMeasurementSchema.getDataTypeList(schemas), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD), + 1); + tablet.addTimestamp(0, 1L); + tablet.addValue("tag1", 0, "d1"); + tablet.addValue("value", 0, 42L); + return tablet; + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionIT.java b/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionIT.java index 5f8c9bfdcef4a..aee7733b16c7c 100644 --- a/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionIT.java @@ -19,17 +19,20 @@ package org.apache.iotdb.session.it; -import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.commons.client.internal.session.InternalSessionBuilder; +import org.apache.iotdb.commons.client.internal.session.InternalSessionPoolBuilder; import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.isession.pool.SessionDataSetWrapper; import org.apache.iotdb.it.env.EnvFactory; import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper; import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.it.utils.InternalSessionITUtils; import org.apache.iotdb.itbase.category.ClusterIT; import org.apache.iotdb.itbase.category.LocalStandaloneIT; -import org.apache.iotdb.session.DummyNodesSupplier; import org.apache.iotdb.session.Session; +import org.apache.iotdb.session.pool.SessionPool; import org.apache.tsfile.enums.TSDataType; import org.junit.After; @@ -62,16 +65,14 @@ public void tearDown() throws Exception { @Test public void testInternalSessionUsesInternalPort() throws Exception { final DataNodeWrapper dataNode = EnvFactory.getEnv().getDataNodeWrapper(0); - final INodeSupplier nodeSupplier = - new DummyNodesSupplier( - Collections.singletonList( - new TEndPoint(dataNode.getInternalAddress(), dataNode.getInternalPort()))); + final INodeSupplier nodeSupplier = InternalSessionITUtils.createInternalNodeSupplier(); try (Session session = new InternalSessionBuilder(nodeSupplier) .host(dataNode.getInternalAddress()) .port(dataNode.getInternalPort()) - .username("root") + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) .build()) { session.open(); try { @@ -84,6 +85,12 @@ public void testInternalSessionUsesInternalPort() throws Exception { Collections.singletonList("s1"), Collections.singletonList(TSDataType.INT64), Collections.singletonList(1L)); + session.insertRecord( + DATABASE + ".d1", + 1L, + Collections.singletonList("s1"), + Collections.singletonList(TSDataType.INT64), + Collections.singletonList(1L)); try (SessionDataSet dataSet = session.executeQueryStatement("SELECT s1 FROM " + DATABASE + ".d1")) { @@ -95,4 +102,71 @@ public void testInternalSessionUsesInternalPort() throws Exception { } } } + + @Test + public void testInternalSessionPoolUsesInternalPort() throws Exception { + final DataNodeWrapper dataNode = EnvFactory.getEnv().getDataNodeWrapper(0); + final INodeSupplier nodeSupplier = InternalSessionITUtils.createInternalNodeSupplier(); + + final SessionPool sessionPool = + new InternalSessionPoolBuilder(nodeSupplier) + .host(dataNode.getInternalAddress()) + .port(dataNode.getInternalPort()) + .user(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .maxSize(1) + .build(); + try { + sessionPool.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + try { + sessionPool.executeNonQueryStatement( + "CREATE TIMESERIES " + DATABASE + ".d1.s1 WITH DATATYPE=INT64, ENCODING=RLE"); + sessionPool.insertRecord( + DATABASE + ".d1", + 1L, + Collections.singletonList("s1"), + Collections.singletonList(TSDataType.INT64), + Collections.singletonList(1L)); + sessionPool.insertRecord( + DATABASE + ".d1", + 1L, + Collections.singletonList("s1"), + Collections.singletonList(TSDataType.INT64), + Collections.singletonList(1L)); + + try (SessionDataSetWrapper dataSet = + sessionPool.executeQueryStatement("SELECT s1 FROM " + DATABASE + ".d1")) { + assertTrue(dataSet.hasNext()); + assertEquals(1L, dataSet.next().getFields().get(0).getLongV()); + } + } finally { + sessionPool.executeNonQueryStatement("DELETE DATABASE " + DATABASE); + } + } finally { + sessionPool.close(); + } + } + + @Test + public void testInternalSessionWithoutAutoFetchUsesConfiguredInternalEndpoint() throws Exception { + final DataNodeWrapper dataNode = EnvFactory.getEnv().getDataNodeWrapper(0); + + try (Session session = + new InternalSessionBuilder() + .enableAutoFetch(false) + .host(dataNode.getInternalAddress()) + .port(dataNode.getInternalPort()) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .build()) { + session.open(); + session.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + try { + session.executeNonQueryStatement( + "CREATE TIMESERIES " + DATABASE + ".d1.s1 WITH DATATYPE=INT64, ENCODING=RLE"); + } finally { + session.executeNonQueryStatement("DELETE DATABASE " + DATABASE); + } + } + } } diff --git a/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionRedirectIT.java b/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionRedirectIT.java new file mode 100644 index 0000000000000..2c786619bc499 --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionRedirectIT.java @@ -0,0 +1,147 @@ +/* + * 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.it; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.client.internal.session.InternalSessionBuilder; +import org.apache.iotdb.consensus.ConsensusFactory; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.isession.SessionConfig; +import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.it.utils.InternalSessionITUtils; +import org.apache.iotdb.itbase.category.ClusterIT; +import org.apache.iotdb.session.Session; + +import org.apache.tsfile.enums.TSDataType; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.util.Collections; +import java.util.Map; +import java.util.Set; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +@RunWith(IoTDBTestRunner.class) +@Category(ClusterIT.class) +public class InternalSessionRedirectIT { + + private static final String DATABASE = "root.internal_session_redirect_it"; + private static final String DEVICE = DATABASE + ".d1"; + private static final String TIMESERIES = DEVICE + ".s1"; + + @Before + public void setUp() throws Exception { + EnvFactory.getEnv() + .getConfig() + .getCommonConfig() + .setSchemaReplicationFactor(1) + .setDataReplicationFactor(1) + .setDataRegionGroupExtensionPolicy("CUSTOM") + .setDefaultDataRegionGroupNumPerDatabase(1) + .setSchemaRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS) + .setDataRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS); + EnvFactory.getEnv().initClusterEnvironment(1, 3); + } + + @After + public void tearDown() throws Exception { + EnvFactory.getEnv().cleanClusterEnvironment(); + } + + @Test + public void testLastDataQueryRedirectsToInternalLeaderEndpoint() throws Exception { + final INodeSupplier nodeSupplier = InternalSessionITUtils.createInternalNodeSupplier(); + final TEndPoint initialEndPoint = InternalSessionITUtils.getInternalEndPoints().get(0); + + try (Session setupSession = + new InternalSessionBuilder(nodeSupplier) + .host(initialEndPoint.getIp()) + .port(initialEndPoint.getPort()) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .build()) { + setupSession.open(); + setupSession.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + setupSession.executeNonQueryStatement( + "CREATE TIMESERIES " + TIMESERIES + " WITH DATATYPE=INT64, ENCODING=RLE"); + setupSession.insertRecord( + DEVICE, + 1L, + Collections.singletonList("s1"), + Collections.singletonList(TSDataType.INT64), + Collections.singletonList(42L)); + // The first write only caches the leader endpoint; retry once to persist the setup value. + setupSession.insertRecord( + DEVICE, + 1L, + Collections.singletonList("s1"), + Collections.singletonList(TSDataType.INT64), + Collections.singletonList(42L)); + } + + final Set leaderEndPoints; + try (Session discoverySession = + new InternalSessionBuilder(nodeSupplier) + .host(initialEndPoint.getIp()) + .port(initialEndPoint.getPort()) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .build()) { + discoverySession.open(); + leaderEndPoints = + InternalSessionITUtils.awaitDataRegionLeaderInternalEndPoints(discoverySession, DATABASE); + } + final TEndPoint nonLeaderEndPoint = + InternalSessionITUtils.findNonLeaderInternalEndPoint(leaderEndPoints); + + try (Session session = + new InternalSessionBuilder(nodeSupplier) + .host(nonLeaderEndPoint.getIp()) + .port(nonLeaderEndPoint.getPort()) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .build()) { + session.open(); + session.insertRecord( + DEVICE, + 2L, + Collections.singletonList("s1"), + Collections.singletonList(TSDataType.INT64), + Collections.singletonList(43L)); + try (SessionDataSet dataSet = + session.executeLastDataQuery(Collections.singletonList(TIMESERIES))) { + assertTrue(dataSet.hasNext()); + dataSet.next(); + } + + final Map deviceIdToEndpoint = + InternalSessionITUtils.getSessionField(session, "deviceIdToEndpoint"); + assertEquals(1, deviceIdToEndpoint.size()); + assertTrue(leaderEndPoints.contains(deviceIdToEndpoint.get(DEVICE))); + } + } +} 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 index 9c8d32f946120..553de87cc5ecd 100644 --- 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 @@ -41,4 +41,12 @@ default List getSeedNodes() { default void requestRefresh() { // Suppliers without automatic discovery have nothing to refresh. } + + /** + * Whether this supplier's lifecycle is managed externally and must not be closed by IoTDB + * clients. + */ + default boolean isExternalLifeCycleNodeSupplier() { + return false; + } } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/DummyNodesSupplier.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/DummyNodesSupplier.java index bb3186d42007b..3f59f26a6ed32 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/DummyNodesSupplier.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/DummyNodesSupplier.java @@ -41,6 +41,11 @@ public void close() { // do nothing } + @Override + public boolean isExternalLifeCycleNodeSupplier() { + return true; + } + @Override public Optional getQueryEndPoint() { if (availableNodes == null || availableNodes.isEmpty()) { 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 b478e3311e24a..1d6bca912a735 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 @@ -545,7 +545,9 @@ public synchronized void open(boolean enableRPCCompaction, int connectionTimeout this.executorService = null; } if (this.availableNodes != null) { - this.availableNodes.close(); + if (!this.availableNodes.isExternalLifeCycleNodeSupplier()) { + this.availableNodes.close(); + } this.availableNodes = null; } @@ -698,7 +700,9 @@ public synchronized void close() throws IoTDBConnectionException { if (this.executorService != null) { this.executorService.shutdown(); this.executorService = null; - this.availableNodes.close(); + if (!this.availableNodes.isExternalLifeCycleNodeSupplier()) { + this.availableNodes.close(); + } this.availableNodes = null; } isClosed = true; 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 35dd03886d932..eeacd841114c4 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 @@ -853,7 +853,9 @@ public synchronized void close() { this.executorService = null; } if (availableNodes != null) { - this.availableNodes.close(); + if (!this.availableNodes.isExternalLifeCycleNodeSupplier()) { + this.availableNodes.close(); + } this.availableNodes = null; } LOGGER.info(SessionMessages.CLOSING_SESSION_POOL); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/auth/AuthorityChecker.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/auth/AuthorityChecker.java index c91d554125c41..b9579dc241fb3 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/auth/AuthorityChecker.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/auth/AuthorityChecker.java @@ -229,7 +229,7 @@ public static TSStatus checkAuthority(Statement statement, IAuditEntity auditEnt auditEntity.getUserId(), auditEntity.getUsername(), auditEntity.getCliHostname()) - .setInternal(auditEntity.isInternal()) + .setDatabaseInternalSession(auditEntity.isDatabaseInternalSession()) .setSqlString(auditEntity.getSqlString())); } return status; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckContext.java index ac09431fc2caf..346b6d6028aae 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckContext.java @@ -32,7 +32,7 @@ public class TreeAccessCheckContext implements IAuditEntity { private final long userId; private final String username; private final String cliHostname; - private boolean internal; + private boolean databaseInternalSession; public TreeAccessCheckContext(long userId, String username, String cliHostname) { this.userId = userId; @@ -56,12 +56,12 @@ public String getCliHostname() { } @Override - public boolean isInternal() { - return internal; + public boolean isDatabaseInternalSession() { + return databaseInternalSession; } - public TreeAccessCheckContext setInternal(final boolean internal) { - this.internal = internal; + public TreeAccessCheckContext setDatabaseInternalSession(final boolean databaseInternalSession) { + this.databaseInternalSession = databaseInternalSession; return this; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java index aa1cab64b0f05..06ea582a6afec 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java @@ -1136,7 +1136,7 @@ && isAuditTreeDatabase(statement.getDatabasePath())) { .setMessage(getReservedDatabaseNameErrorMsg(TREE_MODEL_AUDIT_DATABASE)); } // The system database is created by the database itself, not through SQL. - if (!context.isInternal() + if (!context.isDatabaseInternalSession() && InternalDatabaseUtils.isSystemDatabasePath(statement.getDatabasePath())) { context .setDatabase(statement.getDatabasePath().getFullPath()) @@ -1209,7 +1209,7 @@ public TSStatus visitDeleteDatabase( context .setAuditLogOperation(AuditLogOperation.DDL) .setDatabase(statement.getPrefixPath().toString()); - if (!context.isInternal() + if (!context.isDatabaseInternalSession() && SchemaConstant.SYSTEM_DATABASE.equals(String.join(".", statement.getPrefixPath()))) { AUDIT_LOGGER.recordObjectAuthenticationAuditLog( context.setPrivilegeType(PrivilegeType.MANAGE_DATABASE).setResult(false), @@ -1291,7 +1291,7 @@ private TSStatus checkShowOrCountDatabasePermission( */ public static TSStatus checkSystemDatabaseWrite( final IAuditEntity context, final Collection checkedPaths) { - if (context.isInternal()) { + if (context.isDatabaseInternalSession()) { return null; } for (PartialPath path : checkedPaths) { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/auth/TreeAccessTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/auth/TreeAccessTest.java index 1bd9ffb2f7a71..998a766fd0712 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/auth/TreeAccessTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/auth/TreeAccessTest.java @@ -254,7 +254,7 @@ public void testSystemDatabaseWriteRequiresInternalCaller() throws Exception { new TreeAccessCheckContext(AuthorityChecker.SUPER_USER_ID, AuthorityChecker.SUPER_USER, ""); final TreeAccessCheckContext internalCaller = new TreeAccessCheckContext(AuthorityChecker.SUPER_USER_ID, AuthorityChecker.SUPER_USER, "") - .setInternal(true); + .setDatabaseInternalSession(true); final InsertStatement systemInsert = new InsertStatement(); systemInsert.setDevice(new PartialPath("root.__system.d")); diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/audit/IAuditEntity.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/audit/IAuditEntity.java index ad4c85f18d35d..ffc5effea4d14 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/audit/IAuditEntity.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/audit/IAuditEntity.java @@ -59,8 +59,8 @@ public interface IAuditEntity { IAuditEntity setSqlString(String sqlString); - /** Whether the caller is a dedicated internal writer of the database itself. */ - default boolean isInternal() { + /** Whether the caller is a database-internal session used for dedicated internal writes. */ + default boolean isDatabaseInternalSession() { return false; } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java index 5b3d36e7f3a02..d1346d4c672bd 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java @@ -25,7 +25,13 @@ import java.util.Objects; -/** Server-side builder for tree-model internal sessions. */ +/** + * Server-side builder for tree-model internal sessions. + * + *

Callers sharing an external {@link INodeSupplier} must return {@code true} from {@link + * INodeSupplier#isExternalLifeCycleNodeSupplier()}; otherwise the created session may close it + * during its lifecycle, including when the session is reopened. + */ public class InternalSessionBuilder extends Session.Builder { private final INodeSupplier nodeSupplier; diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java index 8ba30a18dab62..42be53fc49a89 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java @@ -25,7 +25,13 @@ import java.util.Objects; -/** Server-side builder for tree-model internal session pools. */ +/** + * Server-side builder for tree-model internal session pools. + * + *

Callers sharing an external {@link INodeSupplier} must return {@code true} from {@link + * INodeSupplier#isExternalLifeCycleNodeSupplier()}; otherwise the created pool may close it when + * the pool is closed. + */ public class InternalSessionPoolBuilder extends SessionPool.Builder { public InternalSessionPoolBuilder() {} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java index 132f944212c40..044abea3a638b 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java @@ -26,7 +26,13 @@ import java.util.Objects; -/** Server-side builder for table-model internal sessions. */ +/** + * Server-side builder for table-model internal sessions. + * + *

Callers sharing an external {@link INodeSupplier} must return {@code true} from {@link + * INodeSupplier#isExternalLifeCycleNodeSupplier()}; otherwise the created session may close it + * during its lifecycle, including when the session is reopened. + */ public class InternalTableSessionBuilder extends TableSessionBuilder { private final INodeSupplier nodeSupplier; diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java index 6c1d9b04253e3..6c53a222a80b3 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java @@ -26,7 +26,13 @@ import java.util.Objects; -/** Server-side builder for table-model internal session pools. */ +/** + * Server-side builder for table-model internal session pools. + * + *

Callers sharing an external {@link INodeSupplier} must return {@code true} from {@link + * INodeSupplier#isExternalLifeCycleNodeSupplier()}; otherwise the created pool may close it when + * the pool is closed. + */ public class InternalTableSessionPoolBuilder extends TableSessionPoolBuilder { public InternalTableSessionPoolBuilder() {} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java index d88e449630863..d436c99dcf737 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java @@ -20,6 +20,8 @@ package org.apache.iotdb.commons.client.internal.subscription; import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.isession.SessionConfig; +import org.apache.iotdb.rpc.IoTDBConnectionException; import org.apache.iotdb.rpc.RpcServiceConstants; import org.apache.iotdb.session.AbstractSessionBuilder; import org.apache.iotdb.session.subscription.SubscriptionSessionWrapper; @@ -27,11 +29,16 @@ final class InternalSubscriptionSessionWrapper extends SubscriptionSessionWrapper { private final INodeSupplier nodeSupplier; + private final boolean endpointConfigured; InternalSubscriptionSessionWrapper( final AbstractSessionBuilder builder, final INodeSupplier nodeSupplier) { super(builder); this.nodeSupplier = nodeSupplier; + this.endpointConfigured = + builder.nodeUrls != null + || !SessionConfig.DEFAULT_HOST.equals(builder.host) + || builder.rpcPort != SessionConfig.DEFAULT_PORT; } @Override @@ -39,6 +46,19 @@ protected String getRpcServiceName() { return RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE; } + @Override + public void open() throws IoTDBConnectionException { + if (nodeSupplier != null) { + enableAutoFetch = true; + if (!endpointConfigured) { + nodeSupplier + .getQueryEndPoint() + .ifPresent(endpoint -> defaultEndPoint = endpoint.deepCopy()); + } + } + super.open(); + } + @Override protected void initThreadPool() { if (nodeSupplier == null) { diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java index 68e3117f65591..bb32d6224b358 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java @@ -26,7 +26,12 @@ import java.util.Objects; -/** Server-side builder for table subscription pull consumers. */ +/** + * Server-side builder for table subscription pull consumers. + * + *

Callers sharing an external {@link INodeSupplier} must return {@code true} from {@link + * INodeSupplier#isExternalLifeCycleNodeSupplier()}. + */ public class InternalSubscriptionTablePullConsumerBuilder extends SubscriptionTablePullConsumerBuilder { diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java index 0a3f477284b4f..9005773782423 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java @@ -27,7 +27,13 @@ import java.util.List; import java.util.Objects; -/** Server-side builder for table subscription-management sessions. */ +/** + * Server-side builder for table subscription-management sessions. + * + *

Callers sharing an external {@link INodeSupplier} must return {@code true} from {@link + * INodeSupplier#isExternalLifeCycleNodeSupplier()}; otherwise the created session may close it + * during its lifecycle, including when the session is reopened. + */ public class InternalSubscriptionTableSessionBuilder extends SubscriptionTableSessionBuilder { private final INodeSupplier nodeSupplier; diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalBuildersTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalBuildersTest.java new file mode 100644 index 0000000000000..fbe4bb969d8a1 --- /dev/null +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalBuildersTest.java @@ -0,0 +1,47 @@ +/* + * 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.commons.client.internal.session; + +import org.apache.iotdb.session.i18n.SessionMessages; + +import org.junit.Test; +import org.junit.function.ThrowingRunnable; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +public class InternalBuildersTest { + + @Test + public void testInternalBuildersRejectAutoFetchWithoutNodeSupplier() { + assertAutoFetchRequiresNodeSupplier(() -> new InternalSessionBuilder().build()); + assertAutoFetchRequiresNodeSupplier(() -> new InternalTableSessionBuilder().build()); + assertAutoFetchRequiresNodeSupplier(() -> new InternalSessionPoolBuilder().build()); + assertAutoFetchRequiresNodeSupplier(() -> new InternalTableSessionPoolBuilder().build()); + } + + private static void assertAutoFetchRequiresNodeSupplier(final ThrowingRunnable builder) { + final IllegalArgumentException exception = + assertThrows(IllegalArgumentException.class, builder); + assertEquals( + SessionMessages.EXCEPTION_INTERNAL_AUTO_FETCH_REQUIRES_NODE_SUPPLIER_9CB6BA02, + exception.getMessage()); + } +} diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalSessionProtocolTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalSessionProtocolTest.java index 996a47d902d28..104e56a1c4500 100644 --- a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalSessionProtocolTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalSessionProtocolTest.java @@ -21,19 +21,23 @@ import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.client.internal.subscription.InternalSubscriptionTableSessionBuilder; import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.rpc.DeepCopyRpcTransportFactory; import org.apache.iotdb.rpc.RpcServiceConstants; import org.apache.iotdb.rpc.TSStatusCode; 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.iotdb.service.rpc.thrift.TSProtocolVersion; import org.apache.iotdb.session.DummyNodesSupplier; import org.apache.iotdb.session.Session; -import org.apache.iotdb.session.SessionConnection; +import org.apache.iotdb.session.subscription.ISubscriptionTableSession; +import org.apache.thrift.TException; import org.apache.thrift.TMultiplexedProcessor; import org.apache.thrift.protocol.TBinaryProtocol; import org.apache.thrift.server.TThreadPoolServer; @@ -42,7 +46,6 @@ import org.junit.Before; import org.junit.Test; -import java.lang.reflect.Field; import java.net.ServerSocket; import java.util.Collections; import java.util.concurrent.CountDownLatch; @@ -83,18 +86,19 @@ public void testReconnectRetainsMultiplexedClientServiceName() throws Exception .username("root") .build()) { session.open(); - - final SessionConnection connection = - (SessionConnection) getField(session, "defaultSessionConnection"); - assertTrue(connection.reconnect()); + session.executeNonQueryStatement("SET CONFIGURATION 'force_reconnect'='true'"); assertEquals(2, openSessionCalls.get()); } } - private static Object getField(final Object target, final String fieldName) throws Exception { - final Field field = Session.class.getDeclaredField(fieldName); - field.setAccessible(true); - return field.get(target); + @Test + public void testInternalSubscriptionSessionUsesSuppliedEndpoint() throws Exception { + final INodeSupplier nodeSupplier = + new DummyNodesSupplier(Collections.singletonList(server.endpoint)); + try (ISubscriptionTableSession session = + new InternalSubscriptionTableSessionBuilder(nodeSupplier).username("root").build()) { + assertEquals(1, openSessionCalls.get()); + } } private static final class TestServer implements AutoCloseable { @@ -121,6 +125,11 @@ private TestServer(final AtomicInteger openSessionCalls) throws Exception { }); when(handler.closeSession(any(TSCloseSessionReq.class))) .thenReturn(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode())); + when(handler.executeUpdateStatementV2(any(TSExecuteStatementReq.class))) + .thenThrow(new TException("force reconnect")) + .thenReturn( + new TSExecuteStatementResp( + new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()))); final TMultiplexedProcessor processor = new TMultiplexedProcessor(); processor.registerProcessor( From 2c33cab76e8bfcc143524366632a2e37abacc7b8 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 17:17:59 +0800 Subject: [PATCH 26/32] Narrow database internal session propagation --- .../iotdb/db/protocol/session/IClientSession.java | 11 ----------- .../db/protocol/session/InternalClientSession.java | 1 - .../iotdb/db/protocol/session/SessionManager.java | 4 ---- 3 files changed, 16 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java index 3b8fde376ef42..4c10f5b0eef77 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java @@ -53,9 +53,6 @@ public abstract class IClientSession { private boolean login = false; - // Session opened by the database itself for internal operations. - private boolean databaseInternalSession = false; - // Session uses the internal RPC endpoint family for redirects. private boolean internalRpcSession = false; @@ -128,14 +125,6 @@ public void setLogin(boolean login) { this.login = login; } - public boolean isDatabaseInternalSession() { - return databaseInternalSession; - } - - public void setDatabaseInternalSession(final boolean databaseInternalSession) { - this.databaseInternalSession = databaseInternalSession; - } - public boolean isInternalRpcSession() { return internalRpcSession; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java index 6dfc7bd5f9f0e..22f1a3b051bd2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java @@ -39,7 +39,6 @@ public class InternalClientSession extends IClientSession { public InternalClientSession(String clientID) { this.clientID = clientID; - setDatabaseInternalSession(true); setInternalRpcSession(true); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java index 181c666deb6eb..b3c3f143b76bb 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java @@ -592,7 +592,6 @@ public SessionInfo getSessionInfo(IClientSession session) { session.getDatabaseName(), session.getSqlDialect(), session.isEnableAutoCreateSchema()) - .setDatabaseInternalSession(session.isDatabaseInternalSession()) .setInternalRpcSession(session.isInternalRpcSession()); } @@ -618,7 +617,6 @@ public SessionInfo getSessionInfoOfTreeModel(IClientSession session) { session.getClientVersion(), session.getDatabaseName(), SqlDialect.TREE) - .setDatabaseInternalSession(session.isDatabaseInternalSession()) .setInternalRpcSession(session.isInternalRpcSession()); } @@ -630,7 +628,6 @@ public SessionInfo getSessionInfoOfTableModel(IClientSession session) { session.getClientVersion(), session.getDatabaseName(), SqlDialect.TABLE) - .setDatabaseInternalSession(session.isDatabaseInternalSession()) .setInternalRpcSession(session.isInternalRpcSession()); } @@ -642,7 +639,6 @@ public SessionInfo getSessionInfoOfPipeReceiver(IClientSession session, String d session.getClientVersion(), databaseName, SqlDialect.TABLE) - .setDatabaseInternalSession(session.isDatabaseInternalSession()) .setInternalRpcSession(session.isInternalRpcSession()); } From 415d5e68875a1b59151bc6a41b69881b30aff61b Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 17:23:36 +0800 Subject: [PATCH 27/32] Preserve internal session flags during load conversion --- .../LoadTsFileDataTypeConverter.java | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTsFileDataTypeConverter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTsFileDataTypeConverter.java index f78d6c5c53715..a341d6597b4bd 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTsFileDataTypeConverter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTsFileDataTypeConverter.java @@ -265,7 +265,9 @@ private TSStatus executeForTableModel( session.getClientVersion(), databaseName, SqlDialect.TABLE, - allowAutoCreateSchema && session.isEnableAutoCreateSchema()), + allowAutoCreateSchema && session.isEnableAutoCreateSchema()) + .setDatabaseInternalSession(isDatabaseInternalSessionContext()) + .setInternalRpcSession(isInternalRpcSessionContext()), "", LocalExecutionPlanner.getInstance().metadata, IoTDBDescriptor.getInstance().getConfig().getQueryTimeoutThreshold()) @@ -334,7 +336,9 @@ private TSStatus executeForTreeModel( && session.isEnableAutoCreateSchema() && (Objects.isNull(context) || Objects.isNull(context.getSession()) - || context.getSession().isEnableAutoCreateSchema())), + || context.getSession().isEnableAutoCreateSchema())) + .setDatabaseInternalSession(isDatabaseInternalSessionContext()) + .setInternalRpcSession(isInternalRpcSessionContext()), "", ClusterPartitionFetcher.getInstance(), ClusterSchemaFetcher.getInstance(), @@ -355,4 +359,16 @@ public boolean isSuccessful(final TSStatus status) { || status.getCode() == TSStatusCode.REDIRECTION_RECOMMEND.getStatusCode() || status.getCode() == TSStatusCode.LOAD_IDEMPOTENT_CONFLICT_EXCEPTION.getStatusCode()); } + + private boolean isDatabaseInternalSessionContext() { + return Objects.nonNull(context) + && Objects.nonNull(context.getSession()) + && context.getSession().isDatabaseInternalSession(); + } + + private boolean isInternalRpcSessionContext() { + return Objects.nonNull(context) + && Objects.nonNull(context.getSession()) + && context.getSession().isInternalRpcSession(); + } } From d608659e684bfd0deb4e9b531fe00c4588066a0c Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 17:25:33 +0800 Subject: [PATCH 28/32] Remove unused session message constants --- .../en/org/apache/iotdb/session/i18n/SessionMessages.java | 6 ------ .../zh/org/apache/iotdb/session/i18n/SessionMessages.java | 5 ----- 2 files changed, 11 deletions(-) diff --git a/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java b/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java index df11f6bc33dc9..101cb74be6f7d 100644 --- a/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java +++ b/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java @@ -21,9 +21,6 @@ public final class SessionMessages { - public static final String EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081 = - "No DataNode internal endpoint is available"; - public static final String EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392 = "Internal node supplier must not be null"; @@ -312,9 +309,6 @@ private SessionMessages() {} public static final String EXCEPTION_SESSION_CAN_NOT_CONNECT_TO_ARG_226FBA66 = "Session can not connect to {}"; - public static final String EXCEPTION_END_POINT_MUST_NOT_BE_NULL_908A4B27 = - "Endpoint must not be null."; - public static final String EXCEPTION_ALL_VALUES_ARE_NULL_AND_THIS_SUBMISSION_IS_IGNORED_DEVICEID_IS_ARG_TIME_IS_ARG_MEASUREMENTS_IS_ARG_5524221C = "All values are null and this submission is ignored,deviceId is [{}],time is" diff --git a/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java b/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java index c2f954b5cc07e..870ca233acd62 100644 --- a/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java +++ b/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java @@ -21,9 +21,6 @@ public final class SessionMessages { - public static final String EXCEPTION_NO_DATANODE_INTERNAL_ENDPOINT_IS_AVAILABLE_249D1081 = - "没有可用的 DataNode 内部端点"; - public static final String EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392 = "内部节点 supplier 不能为空"; @@ -309,8 +306,6 @@ private SessionMessages() {} public static final String EXCEPTION_SESSION_CAN_NOT_CONNECT_TO_ARG_226FBA66 = "Session 无法连接到 {}"; - public static final String EXCEPTION_END_POINT_MUST_NOT_BE_NULL_908A4B27 = "Endpoint 不能为空。"; - public static final String EXCEPTION_ALL_VALUES_ARE_NULL_AND_THIS_SUBMISSION_IS_IGNORED_DEVICEID_IS_ARG_TIME_IS_ARG_MEASUREMENTS_IS_ARG_5524221C = "所有值均为空，本次提交被忽略，deviceId 为 [{}]，time 为 [{}]，measurements 为 [{}]"; From db18fe0545382bac41ea5348963412c9a3f52c6a Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 17:30:21 +0800 Subject: [PATCH 29/32] spotless --- .../plan/analyze/load/LoadTsFileAnalyzer.java | 28 +++++------ .../LoadTsFileDataTypeConverter.java | 48 +++++++++---------- 2 files changed, 38 insertions(+), 38 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/load/LoadTsFileAnalyzer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/load/LoadTsFileAnalyzer.java index dd5ac8bdcbca3..e5e2e5244c890 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/load/LoadTsFileAnalyzer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/load/LoadTsFileAnalyzer.java @@ -486,13 +486,13 @@ && handleSingleMiniFile(i)) { if (isTableModelFile) { final SessionInfo newSessionInfo = new SessionInfo( - sessionInfo.getSessionId(), - sessionInfo.getUserEntity(), - sessionInfo.getZoneId(), - sessionInfo.getVersion(), - sessionInfo.getDatabaseName().orElse(null), - SqlDialect.TABLE, - sessionInfo.isEnableAutoCreateSchema()) + sessionInfo.getSessionId(), + sessionInfo.getUserEntity(), + sessionInfo.getZoneId(), + sessionInfo.getVersion(), + sessionInfo.getDatabaseName().orElse(null), + SqlDialect.TABLE, + sessionInfo.isEnableAutoCreateSchema()) .setDatabaseInternalSession(sessionInfo.isDatabaseInternalSession()) .setInternalRpcSession(sessionInfo.isInternalRpcSession()); context.setSession(newSessionInfo); @@ -500,13 +500,13 @@ && handleSingleMiniFile(i)) { } else { final SessionInfo newSessionInfo = new SessionInfo( - sessionInfo.getSessionId(), - sessionInfo.getUserEntity(), - sessionInfo.getZoneId(), - sessionInfo.getVersion(), - sessionInfo.getDatabaseName().orElse(null), - SqlDialect.TREE, - sessionInfo.isEnableAutoCreateSchema()) + sessionInfo.getSessionId(), + sessionInfo.getUserEntity(), + sessionInfo.getZoneId(), + sessionInfo.getVersion(), + sessionInfo.getDatabaseName().orElse(null), + SqlDialect.TREE, + sessionInfo.isEnableAutoCreateSchema()) .setDatabaseInternalSession(sessionInfo.isDatabaseInternalSession()) .setInternalRpcSession(sessionInfo.isInternalRpcSession()); context.setSession(newSessionInfo); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTsFileDataTypeConverter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTsFileDataTypeConverter.java index a341d6597b4bd..523661170f885 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTsFileDataTypeConverter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTsFileDataTypeConverter.java @@ -258,16 +258,16 @@ private TSStatus executeForTableModel( session, SESSION_MANAGER.requestQueryId(), new SessionInfo( - session.getId(), - new UserEntity( - session.getUserId(), session.getUsername(), session.getClientAddress()), - ZoneId.systemDefault(), - session.getClientVersion(), - databaseName, - SqlDialect.TABLE, - allowAutoCreateSchema && session.isEnableAutoCreateSchema()) - .setDatabaseInternalSession(isDatabaseInternalSessionContext()) - .setInternalRpcSession(isInternalRpcSessionContext()), + session.getId(), + new UserEntity( + session.getUserId(), session.getUsername(), session.getClientAddress()), + ZoneId.systemDefault(), + session.getClientVersion(), + databaseName, + SqlDialect.TABLE, + allowAutoCreateSchema && session.isEnableAutoCreateSchema()) + .setDatabaseInternalSession(isDatabaseInternalSessionContext()) + .setInternalRpcSession(isInternalRpcSessionContext()), "", LocalExecutionPlanner.getInstance().metadata, IoTDBDescriptor.getInstance().getConfig().getQueryTimeoutThreshold()) @@ -325,20 +325,20 @@ private TSStatus executeForTreeModel( isGeneratedByPipe ? new PipeEnrichedStatement(statement) : statement, SESSION_MANAGER.requestQueryId(), new SessionInfo( - session.getId(), - new UserEntity( - session.getUserId(), session.getUsername(), session.getClientAddress()), - ZoneId.systemDefault(), - session.getClientVersion(), - session.getDatabaseName(), - SqlDialect.TREE, - allowAutoCreateSchema - && session.isEnableAutoCreateSchema() - && (Objects.isNull(context) - || Objects.isNull(context.getSession()) - || context.getSession().isEnableAutoCreateSchema())) - .setDatabaseInternalSession(isDatabaseInternalSessionContext()) - .setInternalRpcSession(isInternalRpcSessionContext()), + session.getId(), + new UserEntity( + session.getUserId(), session.getUsername(), session.getClientAddress()), + ZoneId.systemDefault(), + session.getClientVersion(), + session.getDatabaseName(), + SqlDialect.TREE, + allowAutoCreateSchema + && session.isEnableAutoCreateSchema() + && (Objects.isNull(context) + || Objects.isNull(context.getSession()) + || context.getSession().isEnableAutoCreateSchema())) + .setDatabaseInternalSession(isDatabaseInternalSessionContext()) + .setInternalRpcSession(isInternalRpcSessionContext()), "", ClusterPartitionFetcher.getInstance(), ClusterSchemaFetcher.getInstance(), From 9824dfa319fbaaaeb4bb2798c6fdc31fc05d2f78 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 17:56:25 +0800 Subject: [PATCH 30/32] Declare isession dependency in node commons --- iotdb-core/node-commons/pom.xml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/iotdb-core/node-commons/pom.xml b/iotdb-core/node-commons/pom.xml index 3e8a5b87b2d3c..18623c2ca20c9 100644 --- a/iotdb-core/node-commons/pom.xml +++ b/iotdb-core/node-commons/pom.xml @@ -40,6 +40,11 @@ service-rpc 2.0.11-SNAPSHOT + + org.apache.iotdb + isession + 2.0.11-SNAPSHOT + org.apache.iotdb iotdb-subscription From 1ff0e4f1d0132604764990c017c282525da11e2d Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 18:08:58 +0800 Subject: [PATCH 31/32] Remove internal RPC flag from local client sessions --- .../apache/iotdb/db/protocol/session/InternalClientSession.java | 1 - 1 file changed, 1 deletion(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java index 22f1a3b051bd2..8bbb248e9897f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java @@ -39,7 +39,6 @@ public class InternalClientSession extends IClientSession { public InternalClientSession(String clientID) { this.clientID = clientID; - setInternalRpcSession(true); } @Override From e76cc51c47dd2984bb24b41533f7de394b7c03df Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Fri, 6 Nov 2026 18:58:48 +0800 Subject: [PATCH 32/32] Seed table region before redirect test discovery --- .../it/session/InternalTableSessionRedirectIT.java | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionRedirectIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionRedirectIT.java index f5bf9f64c6ac4..870c95594ff2d 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionRedirectIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionRedirectIT.java @@ -93,6 +93,8 @@ public void testTabletInsertRedirectsToInternalLeaderEndpoint() throws Exception setupSession.executeNonQueryStatement("CREATE DATABASE " + DATABASE); setupSession.executeNonQueryStatement( "CREATE TABLE " + DATABASE + "." + TABLE + " (tag1 STRING TAG, value INT64 FIELD)"); + // DataRegionGroups are allocated on the first data write, not when the database is created. + setupSession.insert(createTablet(1L)); leaderEndPoints = InternalSessionITUtils.awaitDataRegionLeaderInternalEndPoints( InternalSessionITUtils.getUnderlyingSession(setupSession), DATABASE); @@ -108,7 +110,7 @@ public void testTabletInsertRedirectsToInternalLeaderEndpoint() throws Exception .password(SessionConfig.DEFAULT_PASSWORD) .database(DATABASE) .build()) { - session.insert(createTablet()); + session.insert(createTablet(2L)); final Session underlyingSession = InternalSessionITUtils.getUnderlyingSession(session); final Map tableModelDeviceIdToEndpoint = @@ -118,7 +120,7 @@ public void testTabletInsertRedirectsToInternalLeaderEndpoint() throws Exception } } - private static Tablet createTablet() { + private static Tablet createTablet(final long timestamp) { final List schemas = Arrays.asList( new MeasurementSchema("tag1", TSDataType.STRING), @@ -130,7 +132,7 @@ private static Tablet createTablet() { IMeasurementSchema.getDataTypeList(schemas), Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD), 1); - tablet.addTimestamp(0, 1L); + tablet.addTimestamp(0, timestamp); tablet.addValue("tag1", 0, "d1"); tablet.addValue("value", 0, 42L); return tablet;