From 5e9353f6bfea79b0be3916919b03da786e20824e Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Wed, 4 Nov 2026 15:41:24 +0800 Subject: [PATCH 01/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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) {} } }