diff --git a/integration-test/src/test/java/org/apache/iotdb/it/utils/InternalSessionITUtils.java b/integration-test/src/test/java/org/apache/iotdb/it/utils/InternalSessionITUtils.java new file mode 100644 index 0000000000000..74001c0afdbad --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/it/utils/InternalSessionITUtils.java @@ -0,0 +1,158 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.it.utils; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.schema.column.ColumnHeaderConstant; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.isession.ITableSession; +import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper; +import org.apache.iotdb.rpc.UrlUtils; +import org.apache.iotdb.session.DummyNodesSupplier; +import org.apache.iotdb.session.Session; + +import org.apache.tsfile.read.common.Field; +import org.apache.tsfile.read.common.RowRecord; +import org.awaitility.Awaitility; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.Assert.assertFalse; + +public final class InternalSessionITUtils { + + private static final String DATA_REGION = "DataRegion"; + private static final String LEADER = "Leader"; + + private InternalSessionITUtils() {} + + public static TEndPoint getInternalEndPoint(final DataNodeWrapper dataNode) { + return new TEndPoint(dataNode.getInternalAddress(), dataNode.getInternalPort()); + } + + public static String getEndPointUrl(final TEndPoint endPoint) { + return UrlUtils.convertTEndPointIpv4AndIpv6Url(endPoint); + } + + public static List getInternalEndPoints() { + final List endPoints = new ArrayList<>(); + for (final DataNodeWrapper dataNode : EnvFactory.getEnv().getDataNodeWrapperList()) { + endPoints.add(getInternalEndPoint(dataNode)); + } + return endPoints; + } + + public static INodeSupplier createInternalNodeSupplier() { + return new DummyNodesSupplier(getInternalEndPoints()); + } + + public static TEndPoint findNonLeaderInternalEndPoint( + final Set leaderInternalEndPoints) { + for (final TEndPoint endPoint : getInternalEndPoints()) { + if (!leaderInternalEndPoints.contains(endPoint)) { + return endPoint; + } + } + throw new IllegalStateException("No non-leader DataNode is available"); + } + + public static Set awaitDataRegionLeaderInternalEndPoints( + final Session session, final String database) { + final AtomicReference> leaderEndPoints = + new AtomicReference<>(Collections.emptySet()); + Awaitility.await() + .atMost(30, TimeUnit.SECONDS) + .untilAsserted( + () -> { + final Set currentLeaderEndPoints = + queryDataRegionLeaderInternalEndPoints(session, database); + assertFalse(currentLeaderEndPoints.isEmpty()); + leaderEndPoints.set(currentLeaderEndPoints); + }); + return leaderEndPoints.get(); + } + + public static Session getUnderlyingSession(final ITableSession tableSession) + throws ReflectiveOperationException { + Class currentClass = tableSession.getClass(); + while (currentClass != null) { + for (final java.lang.reflect.Field field : currentClass.getDeclaredFields()) { + if (Session.class.equals(field.getType())) { + field.setAccessible(true); + return (Session) field.get(tableSession); + } + } + currentClass = currentClass.getSuperclass(); + } + throw new NoSuchFieldException("Session field not found in " + tableSession.getClass()); + } + + @SuppressWarnings("unchecked") + public static T getSessionField(final Session session, final String fieldName) + throws ReflectiveOperationException { + final java.lang.reflect.Field field = Session.class.getDeclaredField(fieldName); + field.setAccessible(true); + return (T) field.get(session); + } + + private static Set queryDataRegionLeaderInternalEndPoints( + final Session session, final String database) throws Exception { + final Set leaderEndPoints = new HashSet<>(); + try (SessionDataSet dataSet = session.executeQueryStatement("SHOW REGIONS")) { + final List columnNames = dataSet.getColumnNames(); + final int typeIndex = columnNames.indexOf(ColumnHeaderConstant.TYPE); + final int databaseIndex = columnNames.indexOf(ColumnHeaderConstant.DATABASE); + final int rpcAddressIndex = columnNames.indexOf(ColumnHeaderConstant.RPC_ADDRESS); + final int rpcPortIndex = columnNames.indexOf(ColumnHeaderConstant.RPC_PORT); + final int roleIndex = columnNames.indexOf(ColumnHeaderConstant.ROLE); + + while (dataSet.hasNext()) { + final RowRecord rowRecord = dataSet.next(); + if (!DATA_REGION.equals(getString(rowRecord, typeIndex)) + || !database.equals(getString(rowRecord, databaseIndex)) + || !LEADER.equalsIgnoreCase(getString(rowRecord, roleIndex))) { + continue; + } + final String rpcAddress = getString(rowRecord, rpcAddressIndex); + final int rpcPort = rowRecord.getFields().get(rpcPortIndex).getIntV(); + for (final DataNodeWrapper dataNode : EnvFactory.getEnv().getDataNodeWrapperList()) { + if (dataNode.getIp().equals(rpcAddress) && dataNode.getPort() == rpcPort) { + leaderEndPoints.add(getInternalEndPoint(dataNode)); + break; + } + } + } + } + return leaderEndPoints; + } + + private static String getString(final RowRecord rowRecord, final int index) { + final Field field = rowRecord.getFields().get(index); + return field == null ? null : field.getStringValue(); + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionIT.java new file mode 100644 index 0000000000000..76624becc4c23 --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionIT.java @@ -0,0 +1,174 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.relational.it.session; + +import org.apache.iotdb.commons.client.internal.session.InternalTableSessionBuilder; +import org.apache.iotdb.commons.client.internal.session.InternalTableSessionPoolBuilder; +import org.apache.iotdb.isession.ITableSession; +import org.apache.iotdb.isession.SessionConfig; +import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.isession.pool.ITableSessionPool; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.it.utils.InternalSessionITUtils; +import org.apache.iotdb.itbase.category.TableClusterIT; +import org.apache.iotdb.itbase.category.TableLocalStandaloneIT; + +import org.apache.tsfile.enums.ColumnCategory; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.IMeasurementSchema; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +@RunWith(IoTDBTestRunner.class) +@Category({TableLocalStandaloneIT.class, TableClusterIT.class}) +public class InternalTableSessionIT { + + private static final String DATABASE = "internal_table_session_it"; + private static final String TABLE = "t"; + + @Before + public void setUp() throws Exception { + EnvFactory.getEnv().initClusterEnvironment(); + } + + @After + public void tearDown() throws Exception { + EnvFactory.getEnv().cleanClusterEnvironment(); + } + + @Test + public void testInternalTableSessionUsesInternalPort() throws Exception { + final DataNodeWrapper dataNode = EnvFactory.getEnv().getDataNodeWrapper(0); + try (ITableSession session = + new InternalTableSessionBuilder(InternalSessionITUtils.createInternalNodeSupplier()) + .nodeUrls( + Collections.singletonList( + InternalSessionITUtils.getEndPointUrl( + InternalSessionITUtils.getInternalEndPoint(dataNode)))) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .build()) { + session.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + try { + createTableAndInsert(session); + assertTableValue(session); + } finally { + session.executeNonQueryStatement("DROP DATABASE " + DATABASE); + } + } + } + + @Test + public void testInternalTableSessionPoolUsesInternalPort() throws Exception { + final DataNodeWrapper dataNode = EnvFactory.getEnv().getDataNodeWrapper(0); + try (ITableSessionPool sessionPool = + new InternalTableSessionPoolBuilder(InternalSessionITUtils.createInternalNodeSupplier()) + .nodeUrls( + Collections.singletonList( + InternalSessionITUtils.getEndPointUrl( + InternalSessionITUtils.getInternalEndPoint(dataNode)))) + .user(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .maxSize(1) + .build()) { + try (ITableSession session = sessionPool.getSession()) { + session.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + try { + createTableAndInsert(session); + assertTableValue(session); + } finally { + session.executeNonQueryStatement("DROP DATABASE " + DATABASE); + } + } + } + } + + @Test + public void testInternalTableSessionWithoutAutoFetchUsesConfiguredInternalEndpoint() + throws Exception { + final DataNodeWrapper dataNode = EnvFactory.getEnv().getDataNodeWrapper(0); + try (ITableSession session = + new InternalTableSessionBuilder() + .enableAutoFetch(false) + .nodeUrls( + Collections.singletonList( + InternalSessionITUtils.getEndPointUrl( + InternalSessionITUtils.getInternalEndPoint(dataNode)))) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .build()) { + session.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + try { + session.executeNonQueryStatement( + "CREATE TABLE " + DATABASE + "." + TABLE + " (tag1 STRING TAG, value INT64 FIELD)"); + } finally { + session.executeNonQueryStatement("DROP DATABASE " + DATABASE); + } + } + } + + private static void createTableAndInsert(final ITableSession session) throws Exception { + session.executeNonQueryStatement("USE " + DATABASE); + session.executeNonQueryStatement( + "CREATE TABLE " + TABLE + " (tag1 STRING TAG, value INT64 FIELD)"); + session.insert(createTablet()); + session.insert(createTablet()); + } + + private static void assertTableValue(final ITableSession session) throws Exception { + try (SessionDataSet dataSet = + session.executeQueryStatement("SELECT value FROM " + DATABASE + "." + TABLE)) { + assertTrue(dataSet.hasNext()); + assertEquals(42L, dataSet.next().getFields().get(0).getLongV()); + } + } + + private static Tablet createTablet() { + final List schemas = + Arrays.asList( + new MeasurementSchema("tag1", TSDataType.STRING), + new MeasurementSchema("value", TSDataType.INT64)); + final Tablet tablet = + new Tablet( + TABLE, + IMeasurementSchema.getMeasurementNameList(schemas), + IMeasurementSchema.getDataTypeList(schemas), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD), + 1); + tablet.addTimestamp(0, 1L); + tablet.addValue("tag1", 0, "d1"); + tablet.addValue("value", 0, 42L); + return tablet; + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionRedirectIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionRedirectIT.java new file mode 100644 index 0000000000000..870c95594ff2d --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/session/InternalTableSessionRedirectIT.java @@ -0,0 +1,140 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.relational.it.session; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.client.internal.session.InternalTableSessionBuilder; +import org.apache.iotdb.consensus.ConsensusFactory; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.isession.ITableSession; +import org.apache.iotdb.isession.SessionConfig; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.it.utils.InternalSessionITUtils; +import org.apache.iotdb.itbase.category.TableClusterIT; +import org.apache.iotdb.session.Session; + +import org.apache.tsfile.enums.ColumnCategory; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.IMeasurementSchema; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +@RunWith(IoTDBTestRunner.class) +@Category(TableClusterIT.class) +public class InternalTableSessionRedirectIT { + + private static final String DATABASE = "internal_table_session_redirect_it"; + private static final String TABLE = "t"; + + @Before + public void setUp() throws Exception { + EnvFactory.getEnv() + .getConfig() + .getCommonConfig() + .setSchemaReplicationFactor(1) + .setDataReplicationFactor(1) + .setDataRegionGroupExtensionPolicy("CUSTOM") + .setDefaultDataRegionGroupNumPerDatabase(1) + .setSchemaRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS) + .setDataRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS); + EnvFactory.getEnv().initClusterEnvironment(1, 3); + } + + @After + public void tearDown() throws Exception { + EnvFactory.getEnv().cleanClusterEnvironment(); + } + + @Test + public void testTabletInsertRedirectsToInternalLeaderEndpoint() throws Exception { + final INodeSupplier nodeSupplier = InternalSessionITUtils.createInternalNodeSupplier(); + final TEndPoint initialEndPoint = InternalSessionITUtils.getInternalEndPoints().get(0); + final Set leaderEndPoints; + + try (ITableSession setupSession = + new InternalTableSessionBuilder(nodeSupplier) + .nodeUrls( + Collections.singletonList(InternalSessionITUtils.getEndPointUrl(initialEndPoint))) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .build()) { + setupSession.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + setupSession.executeNonQueryStatement( + "CREATE TABLE " + DATABASE + "." + TABLE + " (tag1 STRING TAG, value INT64 FIELD)"); + // DataRegionGroups are allocated on the first data write, not when the database is created. + setupSession.insert(createTablet(1L)); + leaderEndPoints = + InternalSessionITUtils.awaitDataRegionLeaderInternalEndPoints( + InternalSessionITUtils.getUnderlyingSession(setupSession), DATABASE); + } + final TEndPoint nonLeaderEndPoint = + InternalSessionITUtils.findNonLeaderInternalEndPoint(leaderEndPoints); + + try (ITableSession session = + new InternalTableSessionBuilder(nodeSupplier) + .nodeUrls( + Collections.singletonList(InternalSessionITUtils.getEndPointUrl(nonLeaderEndPoint))) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .database(DATABASE) + .build()) { + session.insert(createTablet(2L)); + + final Session underlyingSession = InternalSessionITUtils.getUnderlyingSession(session); + final Map tableModelDeviceIdToEndpoint = + InternalSessionITUtils.getSessionField(underlyingSession, "tableModelDeviceIdToEndpoint"); + assertEquals(1, tableModelDeviceIdToEndpoint.size()); + assertTrue(leaderEndPoints.contains(tableModelDeviceIdToEndpoint.values().iterator().next())); + } + } + + private static Tablet createTablet(final long timestamp) { + final List schemas = + Arrays.asList( + new MeasurementSchema("tag1", TSDataType.STRING), + new MeasurementSchema("value", TSDataType.INT64)); + final Tablet tablet = + new Tablet( + TABLE, + IMeasurementSchema.getMeasurementNameList(schemas), + IMeasurementSchema.getDataTypeList(schemas), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD), + 1); + tablet.addTimestamp(0, timestamp); + tablet.addValue("tag1", 0, "d1"); + tablet.addValue("value", 0, 42L); + return tablet; + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionIT.java b/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionIT.java new file mode 100644 index 0000000000000..aee7733b16c7c --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionIT.java @@ -0,0 +1,172 @@ +/* + * 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.commons.client.internal.session.InternalSessionBuilder; +import org.apache.iotdb.commons.client.internal.session.InternalSessionPoolBuilder; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.isession.SessionConfig; +import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.isession.pool.SessionDataSetWrapper; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.it.utils.InternalSessionITUtils; +import org.apache.iotdb.itbase.category.ClusterIT; +import org.apache.iotdb.itbase.category.LocalStandaloneIT; +import org.apache.iotdb.session.Session; +import org.apache.iotdb.session.pool.SessionPool; + +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 = InternalSessionITUtils.createInternalNodeSupplier(); + + try (Session session = + new InternalSessionBuilder(nodeSupplier) + .host(dataNode.getInternalAddress()) + .port(dataNode.getInternalPort()) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .build()) { + 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)); + 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); + } + } + } + + @Test + public void testInternalSessionPoolUsesInternalPort() throws Exception { + final DataNodeWrapper dataNode = EnvFactory.getEnv().getDataNodeWrapper(0); + final INodeSupplier nodeSupplier = InternalSessionITUtils.createInternalNodeSupplier(); + + final SessionPool sessionPool = + new InternalSessionPoolBuilder(nodeSupplier) + .host(dataNode.getInternalAddress()) + .port(dataNode.getInternalPort()) + .user(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .maxSize(1) + .build(); + try { + sessionPool.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + try { + sessionPool.executeNonQueryStatement( + "CREATE TIMESERIES " + DATABASE + ".d1.s1 WITH DATATYPE=INT64, ENCODING=RLE"); + sessionPool.insertRecord( + DATABASE + ".d1", + 1L, + Collections.singletonList("s1"), + Collections.singletonList(TSDataType.INT64), + Collections.singletonList(1L)); + sessionPool.insertRecord( + DATABASE + ".d1", + 1L, + Collections.singletonList("s1"), + Collections.singletonList(TSDataType.INT64), + Collections.singletonList(1L)); + + try (SessionDataSetWrapper dataSet = + sessionPool.executeQueryStatement("SELECT s1 FROM " + DATABASE + ".d1")) { + assertTrue(dataSet.hasNext()); + assertEquals(1L, dataSet.next().getFields().get(0).getLongV()); + } + } finally { + sessionPool.executeNonQueryStatement("DELETE DATABASE " + DATABASE); + } + } finally { + sessionPool.close(); + } + } + + @Test + public void testInternalSessionWithoutAutoFetchUsesConfiguredInternalEndpoint() throws Exception { + final DataNodeWrapper dataNode = EnvFactory.getEnv().getDataNodeWrapper(0); + + try (Session session = + new InternalSessionBuilder() + .enableAutoFetch(false) + .host(dataNode.getInternalAddress()) + .port(dataNode.getInternalPort()) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .build()) { + session.open(); + session.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + try { + session.executeNonQueryStatement( + "CREATE TIMESERIES " + DATABASE + ".d1.s1 WITH DATATYPE=INT64, ENCODING=RLE"); + } finally { + session.executeNonQueryStatement("DELETE DATABASE " + DATABASE); + } + } + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionRedirectIT.java b/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionRedirectIT.java new file mode 100644 index 0000000000000..2c786619bc499 --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/session/it/InternalSessionRedirectIT.java @@ -0,0 +1,147 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.session.it; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.client.internal.session.InternalSessionBuilder; +import org.apache.iotdb.consensus.ConsensusFactory; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.isession.SessionConfig; +import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.it.utils.InternalSessionITUtils; +import org.apache.iotdb.itbase.category.ClusterIT; +import org.apache.iotdb.session.Session; + +import org.apache.tsfile.enums.TSDataType; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.util.Collections; +import java.util.Map; +import java.util.Set; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +@RunWith(IoTDBTestRunner.class) +@Category(ClusterIT.class) +public class InternalSessionRedirectIT { + + private static final String DATABASE = "root.internal_session_redirect_it"; + private static final String DEVICE = DATABASE + ".d1"; + private static final String TIMESERIES = DEVICE + ".s1"; + + @Before + public void setUp() throws Exception { + EnvFactory.getEnv() + .getConfig() + .getCommonConfig() + .setSchemaReplicationFactor(1) + .setDataReplicationFactor(1) + .setDataRegionGroupExtensionPolicy("CUSTOM") + .setDefaultDataRegionGroupNumPerDatabase(1) + .setSchemaRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS) + .setDataRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS); + EnvFactory.getEnv().initClusterEnvironment(1, 3); + } + + @After + public void tearDown() throws Exception { + EnvFactory.getEnv().cleanClusterEnvironment(); + } + + @Test + public void testLastDataQueryRedirectsToInternalLeaderEndpoint() throws Exception { + final INodeSupplier nodeSupplier = InternalSessionITUtils.createInternalNodeSupplier(); + final TEndPoint initialEndPoint = InternalSessionITUtils.getInternalEndPoints().get(0); + + try (Session setupSession = + new InternalSessionBuilder(nodeSupplier) + .host(initialEndPoint.getIp()) + .port(initialEndPoint.getPort()) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .build()) { + setupSession.open(); + setupSession.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + setupSession.executeNonQueryStatement( + "CREATE TIMESERIES " + TIMESERIES + " WITH DATATYPE=INT64, ENCODING=RLE"); + setupSession.insertRecord( + DEVICE, + 1L, + Collections.singletonList("s1"), + Collections.singletonList(TSDataType.INT64), + Collections.singletonList(42L)); + // The first write only caches the leader endpoint; retry once to persist the setup value. + setupSession.insertRecord( + DEVICE, + 1L, + Collections.singletonList("s1"), + Collections.singletonList(TSDataType.INT64), + Collections.singletonList(42L)); + } + + final Set leaderEndPoints; + try (Session discoverySession = + new InternalSessionBuilder(nodeSupplier) + .host(initialEndPoint.getIp()) + .port(initialEndPoint.getPort()) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .build()) { + discoverySession.open(); + leaderEndPoints = + InternalSessionITUtils.awaitDataRegionLeaderInternalEndPoints(discoverySession, DATABASE); + } + final TEndPoint nonLeaderEndPoint = + InternalSessionITUtils.findNonLeaderInternalEndPoint(leaderEndPoints); + + try (Session session = + new InternalSessionBuilder(nodeSupplier) + .host(nonLeaderEndPoint.getIp()) + .port(nonLeaderEndPoint.getPort()) + .username(SessionConfig.DEFAULT_USER) + .password(SessionConfig.DEFAULT_PASSWORD) + .build()) { + session.open(); + session.insertRecord( + DEVICE, + 2L, + Collections.singletonList("s1"), + Collections.singletonList(TSDataType.INT64), + Collections.singletonList(43L)); + try (SessionDataSet dataSet = + session.executeLastDataQuery(Collections.singletonList(TIMESERIES))) { + assertTrue(dataSet.hasNext()); + dataSet.next(); + } + + final Map deviceIdToEndpoint = + InternalSessionITUtils.getSessionField(session, "deviceIdToEndpoint"); + assertEquals(1, deviceIdToEndpoint.size()); + assertTrue(leaderEndPoints.contains(deviceIdToEndpoint.get(DEVICE))); + } + } +} diff --git a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/INodeSupplier.java b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/INodeSupplier.java index 9c8d32f946120..553de87cc5ecd 100644 --- a/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/INodeSupplier.java +++ b/iotdb-client/isession/src/main/java/org/apache/iotdb/isession/INodeSupplier.java @@ -41,4 +41,12 @@ default List getSeedNodes() { default void requestRefresh() { // Suppliers without automatic discovery have nothing to refresh. } + + /** + * Whether this supplier's lifecycle is managed externally and must not be closed by IoTDB + * clients. + */ + default boolean isExternalLifeCycleNodeSupplier() { + return false; + } } diff --git a/iotdb-client/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..101cb74be6f7d 100644 --- a/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java +++ b/iotdb-client/session/src/main/i18n/en/org/apache/iotdb/session/i18n/SessionMessages.java @@ -21,6 +21,12 @@ public final class SessionMessages { + public static final String EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392 = + "Internal node supplier must not be null"; + + public static final String EXCEPTION_INTERNAL_AUTO_FETCH_REQUIRES_NODE_SUPPLIER_9CB6BA02 = + "An internal node supplier is required when auto-fetch is enabled"; + // Session public static final String NODE_URLS_EMPTY = "nodeUrls shouldn't be empty."; public static final String REDIRECT_TWICE_LOG = "redirect twice"; diff --git a/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java b/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java index 838a7e2148329..870ca233acd62 100644 --- a/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java +++ b/iotdb-client/session/src/main/i18n/zh/org/apache/iotdb/session/i18n/SessionMessages.java @@ -21,6 +21,12 @@ public final class SessionMessages { + public static final String EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392 = + "内部节点 supplier 不能为空"; + + public static final String EXCEPTION_INTERNAL_AUTO_FETCH_REQUIRES_NODE_SUPPLIER_9CB6BA02 = + "启用自动发现时，内部节点 supplier 不能为空"; + // Session public static final String NODE_URLS_EMPTY = "nodeUrls 不能为空。"; public static final String REDIRECT_TWICE_LOG = "重定向了两次"; diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/DummyNodesSupplier.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/DummyNodesSupplier.java index bb3186d42007b..3f59f26a6ed32 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/DummyNodesSupplier.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/DummyNodesSupplier.java @@ -41,6 +41,11 @@ public void close() { // do nothing } + @Override + public boolean isExternalLifeCycleNodeSupplier() { + return true; + } + @Override public Optional getQueryEndPoint() { if (availableNodes == null || availableNodes.isEmpty()) { diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/Session.java index 89b3e348b163a..579564f3f2b4a 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 @@ -548,32 +548,15 @@ public synchronized void open(boolean enableRPCCompaction, int connectionTimeout this.executorService = null; } if (this.availableNodes != null) { - this.availableNodes.close(); + if (!this.availableNodes.isExternalLifeCycleNodeSupplier()) { + this.availableNodes.close(); + } this.availableNodes = null; } 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(), - thriftMaxReusableBufferSize); + initAvailableNodes(enableRPCCompaction, connectionTimeoutInMs); } else { // A static supplier has no separate seed tier. Keep every configured address after reopen. Set staticNodes = new LinkedHashSet<>(getNodeUrls()); @@ -593,7 +576,7 @@ public synchronized void open(boolean enableRPCCompaction, int connectionTimeout } } - private void initThreadPool() { + protected void initThreadPool() { this.executorService = Executors.newSingleThreadScheduledExecutor( r -> { @@ -610,6 +593,30 @@ private void initThreadPool() { }); } + protected void initAvailableNodes( + final boolean enableRPCCompaction, final int connectionTimeoutInMs) { + this.availableNodes = + NodesSupplier.createNodeSupplier( + getNodeUrls(), + SessionUtils.parseSeedNodeUrls(initialNodeUrls), + executorService, + username, + password, + zoneId, + thriftDefaultBufferSize, + thriftMaxFrameSize, + connectionTimeoutInMs, + useSSL, + trustStore, + trustStorePwd, + keyStore, + keyStorePwd, + sslProtocol, + enableRPCCompaction, + version.toString(), + thriftMaxReusableBufferSize); + } + private static List shuffleNodeUrls(List endPoints) { try { Collections.shuffle(endPoints); @@ -726,7 +733,9 @@ public synchronized void close() throws IoTDBConnectionException { if (this.executorService != null) { this.executorService.shutdown(); this.executorService = null; - this.availableNodes.close(); + if (!this.availableNodes.isExternalLifeCycleNodeSupplier()) { + this.availableNodes.close(); + } this.availableNodes = null; } isClosed = true; @@ -4477,6 +4486,10 @@ protected void setDefaultSessionConnection(SessionConnection defaultSessionConne this.defaultSessionConnection = defaultSessionConnection; } + protected String getRpcServiceName() { + return null; + } + public boolean isEnableThriftRpcCompaction() { return enableThriftRpcCompaction; } @@ -4644,6 +4657,10 @@ public Session build() { SessionMessages .EXCEPTION_YOU_SHOULD_SPECIFY_EITHER_NODEURLS_HOST_RPCPORT_BUT_NOT_BOTH_77E7B084); } + return createSession(); + } + + protected Session createSession() { return new Session(this); } } 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 f2158ef0ffacb..ca8862cf1c989 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; @@ -233,11 +235,15 @@ 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); + final String rpcServiceName = session.getRpcServiceName(); + if (rpcServiceName != null) { + protocol = new TMultiplexedProtocol(protocol, 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 b4be67b11b62c..3922316ac9ae3 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 @@ -451,7 +451,7 @@ public ITableSession build() throws IoTDBConnectionException { Collections.singletonList(SessionConfig.DEFAULT_HOST + ":" + SessionConfig.DEFAULT_PORT); } this.sqlDialect = TABLE; - Session newSession = new Session(this); + Session newSession = createSession(); newSession.enableIoTDBRpcCompression = isIoTDBRpcCompressionEnabled; newSession.tabletCompressionMinRowSize = tabletCompressionMinRowSize; @@ -463,4 +463,8 @@ public ITableSession build() throws IoTDBConnectionException { } return new TableSession(newSession); } + + protected Session createSession() { + return new Session(this); + } } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/AbstractSessionPoolBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/AbstractSessionPoolBuilder.java index 91c93d8dd63b8..f8ce0c3d54c5b 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/AbstractSessionPoolBuilder.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/AbstractSessionPoolBuilder.java @@ -19,6 +19,7 @@ package org.apache.iotdb.session.pool; +import org.apache.iotdb.isession.INodeSupplier; import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.session.AbstractSessionBuilder; @@ -26,4 +27,5 @@ public class AbstractSessionPoolBuilder extends AbstractSessionBuilder { int maxSize = SessionConfig.DEFAULT_SESSION_POOL_MAX_SIZE; long waitToGetSessionTimeoutInMs = 60_000; int connectionTimeoutInMs = SessionConfig.DEFAULT_CONNECTION_TIMEOUT_MS; + protected INodeSupplier nodeSupplier; } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/SessionPool.java index a3e5b6d473b7f..eedd53f89bcf7 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 @@ -560,9 +560,11 @@ public SessionPool(AbstractSessionPoolBuilder builder) { this.queryTimeoutInMs = builder.timeOut; this.enableAutoCreateSchema = builder.enableAutoCreateSchema; - if (enableAutoFetch) { + final INodeSupplier suppliedNodes = builder.nodeSupplier; + if (enableAutoFetch && suppliedNodes == null) { initThreadPool(); } + final List endPointList; if (builder.nodeUrls != null) { if (builder.nodeUrls.isEmpty()) { throw new IllegalArgumentException(SessionMessages.NODE_URLS_EMPTY); @@ -571,23 +573,21 @@ public SessionPool(AbstractSessionPoolBuilder builder) { this.host = null; this.port = -1; this.formattedNodeUrls = builder.nodeUrls.toString(); - if (enableAutoFetch) { - initAvailableNodes(SessionUtils.parseSeedNodeUrls(nodeUrls)); - } else { - this.availableNodes = new DummyNodesSupplier(SessionUtils.parseSeedNodeUrls(nodeUrls)); - } - + endPointList = SessionUtils.parseSeedNodeUrls(nodeUrls); } else { this.host = builder.host; this.port = builder.rpcPort; this.nodeUrls = null; this.formattedNodeUrls = UrlUtils.formatTEndPointIpv4AndIpv6Url(host, port); - if (enableAutoFetch) { - initAvailableNodes(Collections.singletonList(new TEndPoint(host, port))); - } else { - this.availableNodes = - new DummyNodesSupplier(Collections.singletonList(new TEndPoint(host, port))); - } + endPointList = Collections.singletonList(new TEndPoint(host, port)); + } + + if (suppliedNodes != null) { + this.availableNodes = suppliedNodes; + } else if (enableAutoFetch) { + initAvailableNodes(endPointList); + } else { + this.availableNodes = new DummyNodesSupplier(endPointList); } } @@ -596,7 +596,7 @@ private Session constructNewSession() { if (nodeUrls == null) { // Construct custom Session session = - new Session.Builder() + createSessionBuilder() .host(host) .port(port) .username(user) @@ -627,7 +627,7 @@ private Session constructNewSession() { } else { // Construct redirect-able Session session = - new Session.Builder() + createSessionBuilder() .nodeUrls(nodeUrls) .username(user) .password(password) @@ -659,6 +659,10 @@ private Session constructNewSession() { return session; } + protected Session.Builder createSessionBuilder() { + return new Session.Builder(); + } + private void initThreadPool() { this.executorService = Executors.newSingleThreadScheduledExecutor( @@ -855,7 +859,9 @@ public synchronized void close() { this.executorService = null; } if (availableNodes != null) { - this.availableNodes.close(); + if (!this.availableNodes.isExternalLifeCycleNodeSupplier()) { + this.availableNodes.close(); + } this.availableNodes = null; } LOGGER.info(SessionMessages.CLOSING_SESSION_POOL); diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPool.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPool.java index 4e08f202a14ad..5e3a53c4462ff 100644 --- a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPool.java +++ b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPool.java @@ -27,7 +27,7 @@ public class TableSessionPool implements ITableSessionPool { private final SessionPool sessionPool; - TableSessionPool(SessionPool sessionPool) { + protected TableSessionPool(SessionPool sessionPool) { this.sessionPool = sessionPool; } diff --git a/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java b/iotdb-client/session/src/main/java/org/apache/iotdb/session/pool/TableSessionPoolBuilder.java index 8104a774306fc..fbd071a092b01 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 @@ -358,7 +358,10 @@ public ITableSessionPool build() { Collections.singletonList(SessionConfig.DEFAULT_HOST + ":" + SessionConfig.DEFAULT_PORT); } this.sqlDialect = "table"; - SessionPool sessionPool = new SessionPool(this); - return new TableSessionPool(sessionPool); + return createSessionPool(); + } + + protected ITableSessionPool createSessionPool() { + return new TableSessionPool(new SessionPool(this)); } } 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/SubscriptionSessionWrapper.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java index d7e715fa62762..a14b8eb5aace1 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionSessionWrapper.java @@ -30,7 +30,9 @@ import java.time.ZoneId; import java.util.Objects; -public final class SubscriptionSessionWrapper extends Session { +public class SubscriptionSessionWrapper extends Session { + + private String rpcServiceName; public SubscriptionSessionWrapper(final AbstractSessionBuilder builder) { super(builder); @@ -70,4 +72,28 @@ public SubscriptionSessionConnection getSessionConnection() throws IoTDBConnecti public int getThriftMaxFrameSize() { return thriftMaxFrameSize; } + + public void setRpcServiceName(final String rpcServiceName) { + this.rpcServiceName = rpcServiceName; + } + + @Override + protected String getRpcServiceName() { + return 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..5a74a564d02a1 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/SubscriptionTableSessionBuilder.java @@ -63,14 +63,48 @@ 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; } public ISubscriptionTableSession build() throws IoTDBConnectionException { - final ISubscriptionTableSession session = new SubscriptionTableSession(this); + final ISubscriptionTableSession session = createSubscriptionSession(); session.open(); return session; } + + protected ISubscriptionTableSession createSubscriptionSession() { + return new SubscriptionTableSession(this); + } } 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 bdbd1a6b9ad0b..677f536553b4c 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, @@ -585,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 = @@ -600,7 +643,7 @@ AbstractSubscriptionProvider constructProviderAndHandshake(final TEndPoint endPo this.thriftMaxFrameSize, this.heartbeatIntervalMs, this.connectionTimeoutInMs); - provider.setConsumerInstanceId(consumerInstanceId); + configureProvider(provider); 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 ae41d537d4484..de8112fc8e65b 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 955eec7050fcb..ad35e59d32bac 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/SubscriptionTablePullConsumerBuilder.java b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java index efc99debb4977..dbb1e020ab431 100644 --- a/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java +++ b/iotdb-client/subscription/src/main/java/org/apache/iotdb/session/subscription/consumer/table/SubscriptionTablePullConsumerBuilder.java @@ -130,6 +130,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); @@ -150,6 +186,10 @@ public SubscriptionTablePullConsumerBuilder autoCommitIntervalMs( } public ISubscriptionTablePullConsumer build() { + return createConsumer(); + } + + protected ISubscriptionTablePullConsumer createConsumer() { return new SubscriptionTablePullConsumer(this); } } 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-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..a29657ff0c344 --- /dev/null +++ b/iotdb-client/subscription/src/test/java/org/apache/iotdb/session/subscription/consumer/base/SubscriptionSslConfigTest.java @@ -0,0 +1,157 @@ +/* + * 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.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.rpc.subscription.config.ConsumerConstant; +import org.apache.iotdb.session.AbstractSessionBuilder; +import org.apache.iotdb.session.subscription.SubscriptionTreeSessionBuilder; +import org.apache.iotdb.session.subscription.consumer.table.SubscriptionTablePullConsumerBuilder; +import org.apache.iotdb.session.subscription.consumer.tree.SubscriptionTreePullConsumer; + +import org.junit.Test; + +import java.util.Properties; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +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 RecordingProvider provider = new RecordingProvider(); + + consumer.configureProvider(provider); + + assertSslConfig(provider); + } + + @Test + public void testConsumerPropertiesPropagateSslConfigToProvider() { + 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); + final RecordingProvider provider = new RecordingProvider(); + + consumer.configureProvider(provider); + + assertSslConfig(provider); + } + + private static void assertSslConfig(final RecordingProvider provider) { + assertTrue(provider.useSSL); + assertEquals("trust-store", provider.trustStore); + assertEquals("trust-password", provider.trustStorePwd); + assertEquals("key-store", provider.keyStore); + assertEquals("key-password", provider.keyStorePwd); + assertEquals("TLSv1.2", provider.sslProtocol); + } + + private static final class RecordingProvider extends AbstractSubscriptionProvider { + + 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 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) {} + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/auth/AuthorityChecker.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/auth/AuthorityChecker.java index c91d554125c41..b9579dc241fb3 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/auth/AuthorityChecker.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/auth/AuthorityChecker.java @@ -229,7 +229,7 @@ public static TSStatus checkAuthority(Statement statement, IAuditEntity auditEnt auditEntity.getUserId(), auditEntity.getUsername(), auditEntity.getCliHostname()) - .setInternal(auditEntity.isInternal()) + .setDatabaseInternalSession(auditEntity.isDatabaseInternalSession()) .setSqlString(auditEntity.getSqlString())); } return status; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java index aaf73db5c4dd6..4c10f5b0eef77 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java @@ -53,6 +53,9 @@ public abstract class IClientSession { private boolean login = false; + // Session uses the internal RPC endpoint family for redirects. + private boolean internalRpcSession = false; + private long logInTime; private SqlDialect sqlDialect = SqlDialect.TREE; @@ -122,6 +125,14 @@ public void setLogin(boolean login) { this.login = login; } + public boolean isInternalRpcSession() { + return internalRpcSession; + } + + public void setInternalRpcSession(final boolean internalRpcSession) { + this.internalRpcSession = internalRpcSession; + } + public void setLogInTime(long logInTime) { this.logInTime = logInTime; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java index 85d9acae276ec..b3c3f143b76bb 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java @@ -226,6 +226,80 @@ 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.setInternalRpcSession(true); + 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); } @@ -511,55 +585,61 @@ public static SessionManager getInstance() { public SessionInfo getSessionInfo(IClientSession session) { return new SessionInfo( - session.getId(), - new UserEntity(session.getUserId(), session.getUsername(), session.getClientAddress()), - session.getZoneId(), - session.getClientVersion(), - session.getDatabaseName(), - session.getSqlDialect(), - session.isEnableAutoCreateSchema()); + session.getId(), + new UserEntity(session.getUserId(), session.getUsername(), session.getClientAddress()), + session.getZoneId(), + session.getClientVersion(), + session.getDatabaseName(), + session.getSqlDialect(), + session.isEnableAutoCreateSchema()) + .setInternalRpcSession(session.isInternalRpcSession()); } // Sometimes we need to switch from table model to tree model, // e.g., when loading a tree model TsFile under table model dialect. public SessionInfo copySessionInfoForTreeModel(final SessionInfo sessionInfo) { return new SessionInfo( - sessionInfo.getSessionId(), - sessionInfo.getUserEntity(), - ZoneId.systemDefault(), - sessionInfo.getVersion(), - sessionInfo.getDatabaseName().orElse(null), - SqlDialect.TREE); + sessionInfo.getSessionId(), + sessionInfo.getUserEntity(), + ZoneId.systemDefault(), + sessionInfo.getVersion(), + sessionInfo.getDatabaseName().orElse(null), + SqlDialect.TREE) + .setDatabaseInternalSession(sessionInfo.isDatabaseInternalSession()) + .setInternalRpcSession(sessionInfo.isInternalRpcSession()); } public SessionInfo getSessionInfoOfTreeModel(IClientSession session) { return new SessionInfo( - session.getId(), - new UserEntity(session.getUserId(), session.getUsername(), session.getClientAddress()), - ZoneId.systemDefault(), - session.getClientVersion(), - session.getDatabaseName(), - SqlDialect.TREE); + session.getId(), + new UserEntity(session.getUserId(), session.getUsername(), session.getClientAddress()), + ZoneId.systemDefault(), + session.getClientVersion(), + session.getDatabaseName(), + SqlDialect.TREE) + .setInternalRpcSession(session.isInternalRpcSession()); } public SessionInfo getSessionInfoOfTableModel(IClientSession session) { return new SessionInfo( - session.getId(), - new UserEntity(session.getUserId(), session.getUsername(), session.getClientAddress()), - ZoneId.systemDefault(), - session.getClientVersion(), - session.getDatabaseName(), - SqlDialect.TABLE); + session.getId(), + new UserEntity(session.getUserId(), session.getUsername(), session.getClientAddress()), + ZoneId.systemDefault(), + session.getClientVersion(), + session.getDatabaseName(), + SqlDialect.TABLE) + .setInternalRpcSession(session.isInternalRpcSession()); } public SessionInfo getSessionInfoOfPipeReceiver(IClientSession session, String databaseName) { return new SessionInfo( - session.getId(), - new UserEntity(session.getUserId(), session.getUsername(), session.getClientAddress()), - ZoneId.systemDefault(), - session.getClientVersion(), - databaseName, - SqlDialect.TABLE); + session.getId(), + new UserEntity(session.getUserId(), session.getUsername(), session.getClientAddress()), + ZoneId.systemDefault(), + session.getClientVersion(), + databaseName, + SqlDialect.TABLE) + .setInternalRpcSession(session.isInternalRpcSession()); } @Override 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 796957d17352e..e65a48205ffec 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 @@ -68,6 +68,7 @@ import org.apache.iotdb.db.protocol.session.IClientSession; import org.apache.iotdb.db.protocol.session.SessionManager; import org.apache.iotdb.db.protocol.thrift.OperationType; +import org.apache.iotdb.db.queryengine.common.DataNodeEndPoints; import org.apache.iotdb.db.queryengine.common.header.DatasetHeader; import org.apache.iotdb.db.queryengine.common.header.DatasetHeaderFactory; import org.apache.iotdb.db.queryengine.execution.aggregation.AccumulatorFactory; @@ -1450,11 +1451,9 @@ public TSExecuteStatementResp executeFastLastDataQueryForOneDeviceV2( resp = createResponse(queryExecution.getDatasetHeader(), queryId); TSStatus tsstatus = new TSStatus(TSStatusCode.REDIRECTION_RECOMMEND.getStatusCode()); tsstatus.setRedirectNode( - regionReplicaSets - .get(regionReplicaSets.size() - 1) - .dataNodeLocations - .get(0) - .clientRpcEndPoint); + DataNodeEndPoints.getRedirectEndPoint( + regionReplicaSets.get(regionReplicaSets.size() - 1).dataNodeLocations.get(0), + clientSession.isInternalRpcSession())); resp.setStatus(tsstatus); finished = SELECT_RESULT.apply(resp, queryExecution, req.fetchSize); resp.setMoreData(!finished); @@ -1627,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()) { @@ -1652,6 +1642,30 @@ 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, + 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/DataNodeInternalRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java index 836f8d47f2983..743404a3978f8 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 @@ -3572,7 +3572,8 @@ public TSStatus insertRecord(TSInsertRecordReq req) throws TException { AuthorityChecker.SUPER_USER, IoTDBDescriptor.getInstance().getConfig().getInternalAddress()), ZoneId.systemDefault()) - .setInternal(true); + .setDatabaseInternalSession(true) + .setInternalRpcSession(true); long queryId = SESSION_MANAGER.requestQueryId(); ExecutionResult result = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/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/queryengine/common/DataNodeEndPoints.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/DataNodeEndPoints.java index fcc877f3d280c..7a132888611a9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/DataNodeEndPoints.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/DataNodeEndPoints.java @@ -64,6 +64,17 @@ public static final TDataNodeLocation getLocalDataNodeLocation() { LOCAL_SCHEMA_REGION_CONSENSUS_CLIENT_ENDPOINT); } + public static TEndPoint getRedirectEndPoint( + final TDataNodeLocation dataNodeLocation, final boolean internalRpcSession) { + return internalRpcSession + ? dataNodeLocation.getInternalEndPoint() + : dataNodeLocation.getClientRpcEndPoint(); + } + + public static TEndPoint getLocalRedirectEndPoint(final boolean internalRpcSession) { + return internalRpcSession ? LOCAL_HOST_INTERNAL_ENDPOINT : LOCAL_CLIENT_RPC_ENDPOINT; + } + public static boolean isSameNode(TEndPoint endPoint) { return endPoint.equals(LOCAL_HOST_DATA_BLOCK_ENDPOINT); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java index d609c5c8e028e..46d35b6f7190f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java @@ -1087,8 +1087,12 @@ public String getUsername() { } @Override - public boolean isInternal() { - return session != null && session.isInternal(); + public boolean isDatabaseInternalSession() { + return session != null && session.isDatabaseInternalSession(); + } + + public boolean isInternalRpcSession() { + return session != null && session.isInternalRpcSession(); } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java index eb3d0e6869871..d7fed651103ac 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java @@ -439,12 +439,14 @@ private IQueryExecution createQueryExecutionForTreeModel( final SessionInfo session = queryContext.getSession(); queryContext.setSession( new SessionInfo( - session.getSessionId(), - session.getUserEntity(), - session.getZoneId(), - session.getVersion(), - session.getDatabaseName().orElse(null), - SqlDialect.TABLE)); + session.getSessionId(), + session.getUserEntity(), + session.getZoneId(), + session.getVersion(), + session.getDatabaseName().orElse(null), + SqlDialect.TABLE) + .setDatabaseInternalSession(session.isDatabaseInternalSession()) + .setInternalRpcSession(session.isInternalRpcSession())); return createQueryExecutionForTableModel( statement, new SqlParser(), diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analysis.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analysis.java index f1d3fdd887327..ac9ec086375b8 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analysis.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analysis.java @@ -32,6 +32,7 @@ import org.apache.iotdb.commons.queryengine.common.NodeRef; import org.apache.iotdb.commons.schema.template.Template; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.common.DataNodeEndPoints; import org.apache.iotdb.db.queryengine.common.DeviceContext; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; import org.apache.iotdb.db.queryengine.common.TimeseriesContext; @@ -100,6 +101,8 @@ public class Analysis implements IAnalysis { private List redirectNodeList; + private final boolean internalRpcSession; + // map from output column name (for every node) to its datatype private final Map, TSDataType> expressionTypes = new LinkedHashMap<>(); @@ -349,6 +352,11 @@ public void setDeviceToTimeseriesSchemas( private final Set> fromWhereFilterNodes = new HashSet<>(); public Analysis() { + this(false); + } + + public Analysis(final boolean internalRpcSession) { + this.internalRpcSession = internalRpcSession; this.finishQueryAfterAnalyze = false; } @@ -430,16 +438,21 @@ public List getRedirectNodeList() { } @Override - public void setRedirectNodeList(List redirectNodeList) { - this.redirectNodeList = redirectNodeList; + public void setRedirectNodeList(final List dataNodeLocations) { + redirectNodeList = new ArrayList<>(dataNodeLocations.size()); + for (final TDataNodeLocation dataNodeLocation : dataNodeLocations) { + redirectNodeList.add( + DataNodeEndPoints.getRedirectEndPoint(dataNodeLocation, internalRpcSession)); + } } @Override - public void addEndPointToRedirectNodeList(TEndPoint endPoint) { + public void addEndPointToRedirectNodeList(final TDataNodeLocation dataNodeLocation) { if (redirectNodeList == null) { redirectNodeList = new ArrayList<>(); } - redirectNodeList.add(endPoint); + redirectNodeList.add( + DataNodeEndPoints.getRedirectEndPoint(dataNodeLocation, internalRpcSession)); } public Expression getGlobalTimePredicate() { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java index 7aca24384e239..9f1b5e8c19541 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/AnalyzeVisitor.java @@ -235,15 +235,25 @@ public class AnalyzeVisitor extends StatementVisitor private final ISchemaFetcher schemaFetcher; private final IModelFetcher modelFetcher; private final TreeAnalysisMutationJournal mutationJournal; + private final boolean internalRpcSession; public AnalyzeVisitor( IPartitionFetcher partitionFetcher, ISchemaFetcher schemaFetcher, TreeAnalysisMutationJournal mutationJournal) { + this(partitionFetcher, schemaFetcher, mutationJournal, false); + } + + public AnalyzeVisitor( + IPartitionFetcher partitionFetcher, + ISchemaFetcher schemaFetcher, + TreeAnalysisMutationJournal mutationJournal, + final boolean internalRpcSession) { this.partitionFetcher = partitionFetcher; this.schemaFetcher = schemaFetcher; this.modelFetcher = ModelFetcher.getInstance(); this.mutationJournal = mutationJournal; + this.internalRpcSession = internalRpcSession; } @Override @@ -282,7 +292,7 @@ public Analysis visitExplainAnalyze( @Override public Analysis visitQuery(QueryStatement queryStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setLastLevelUseWildcard(queryStatement.isLastLevelUseWildcard()); try { @@ -2651,7 +2661,7 @@ public Analysis visitCreateTimeseries( } } - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(createTimeSeriesStatement); checkIsTemplateCompatible( @@ -2749,7 +2759,7 @@ public Analysis visitCreateAlignedTimeseries( .MEASUREMENT_UNDER_AN_ALIGNED_DEVICE_IS_NOT_ALLOWED_TO_HAVE_THE_SAME_MEASUREMENT_NAME); } - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(createAlignedTimeSeriesStatement); checkIsTemplateCompatible( @@ -2777,7 +2787,7 @@ public Analysis visitInternalCreateTimeseries( MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(internalCreateTimeSeriesStatement); checkIsTemplateCompatible( internalCreateTimeSeriesStatement.getDevicePath(), @@ -2805,7 +2815,7 @@ public Analysis visitInternalCreateMultiTimeSeries( final MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - final Analysis analysis = new Analysis(); + final Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(internalCreateMultiTimeSeriesStatement); final PathPatternTree pathPatternTree = new PathPatternTree(); @@ -2829,7 +2839,7 @@ public Analysis visitCreateMultiTimeSeries( final CreateMultiTimeSeriesStatement createMultiTimeSeriesStatement, final MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - final Analysis analysis = new Analysis(); + final Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(createMultiTimeSeriesStatement); analyzeSchemaProps(createMultiTimeSeriesStatement.getPropsList()); @@ -2857,7 +2867,7 @@ public Analysis visitCreateMultiTimeSeries( public Analysis visitAlterTimeSeries( AlterTimeSeriesStatement alterTimeSeriesStatement, MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(alterTimeSeriesStatement); Pair templateInfo = @@ -2884,7 +2894,7 @@ public Analysis visitAlterTimeSeries( @Override public Analysis visitInsertTablet( InsertTabletStatement insertTabletStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); AnalyzeUtils.analyzeInsert( context, insertTabletStatement, @@ -2900,7 +2910,7 @@ public Analysis visitInsertTablet( public Analysis visitInsertRow(InsertRowStatement insertRowStatement, MPPQueryContext context) { context.setQueryType(QueryType.WRITE); insertRowStatement.semanticCheck(); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); validateSchema( analysis, insertRowStatement, @@ -2940,7 +2950,7 @@ public Analysis visitInsertRows( InsertRowsStatement insertRowsStatement, MPPQueryContext context) { context.setQueryType(QueryType.WRITE); insertRowsStatement.semanticCheck(); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); validateSchema( analysis, insertRowsStatement, @@ -2965,7 +2975,7 @@ public Analysis visitInsertMultiTablets( InsertMultiTabletsStatement insertMultiTabletsStatement, MPPQueryContext context) { context.setQueryType(QueryType.WRITE); insertMultiTabletsStatement.semanticCheck(); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); validateSchema( analysis, insertMultiTabletsStatement, @@ -2991,7 +3001,7 @@ public Analysis visitInsertRowsOfOneDevice( InsertRowsOfOneDeviceStatement insertRowsOfOneDeviceStatement, MPPQueryContext context) { context.setQueryType(QueryType.WRITE); insertRowsOfOneDeviceStatement.semanticCheck(); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); validateSchema( analysis, insertRowsOfOneDeviceStatement, @@ -3050,7 +3060,7 @@ public Analysis visitLoadFile(LoadTsFileStatement loadTsFileStatement, MPPQueryC try (final LoadTsFileAnalyzer loadTsFileAnalyzer = new LoadTsFileAnalyzer( loadTsFileStatement, loadTsFileStatement.isGeneratedByPipe(), context)) { - return (Analysis) loadTsFileAnalyzer.analyzeFileByFile(new Analysis()); + return (Analysis) loadTsFileAnalyzer.analyzeFileByFile(new Analysis(internalRpcSession)); } catch (final Exception e) { final String exceptionMessage = String.format( @@ -3059,7 +3069,7 @@ public Analysis visitLoadFile(LoadTsFileStatement loadTsFileStatement, MPPQueryC loadTsFileStatement, e.getMessage() == null ? e.getClass().getName() : e.getMessage()); logger.warn(exceptionMessage, e); - final Analysis analysis = new Analysis(); + final Analysis analysis = new Analysis(internalRpcSession); analysis.setFinishQueryAfterAnalyze(true); analysis.setFailStatus(RpcUtils.getStatus(TSStatusCode.LOAD_FILE_ERROR, exceptionMessage)); return analysis; @@ -3281,7 +3291,7 @@ private void addPhysicalTimeseriesForActiveCount( @Override public Analysis visitShowTimeSeries( ShowTimeSeriesStatement showTimeSeriesStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showTimeSeriesStatement); if (showTimeSeriesStatement.getPathPattern().getNodeLength() <= 1) { @@ -3350,7 +3360,7 @@ public Analysis visitShowTimeSeries( @Override public Analysis visitShowTTL(ShowTTLStatement showTTLStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showTTLStatement); analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowTTLHeader()); return analysis; @@ -3422,7 +3432,7 @@ private void analyzeDeviceRegionScan( @Override public Analysis visitShowDevices( ShowDevicesStatement showDevicesStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showDevicesStatement); if (showDevicesStatement.hasTimeCondition()) { @@ -3451,7 +3461,7 @@ public Analysis visitShowDevices( @Override public Analysis visitShowCluster( ShowClusterStatement showClusterStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showClusterStatement); if (showClusterStatement.isDetails()) { analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowClusterDetailsHeader()); @@ -3464,7 +3474,7 @@ public Analysis visitShowCluster( @Override public Analysis visitCountDatabase( CountDatabaseStatement countDatabaseStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(countDatabaseStatement); analysis.setRespDatasetHeader(DatasetHeaderFactory.getCountStorageGroupHeader()); return analysis; @@ -3473,7 +3483,7 @@ public Analysis visitCountDatabase( @Override public Analysis visitSeriesSchemaFetch( SeriesSchemaFetchStatement seriesSchemaFetchStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(seriesSchemaFetchStatement); SchemaPartition schemaPartition = @@ -3493,7 +3503,7 @@ public Analysis visitSeriesSchemaFetch( @Override public Analysis visitDeviceSchemaFetch( DeviceSchemaFetchStatement deviceSchemaFetchStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(deviceSchemaFetchStatement); PathPatternTree patternTree = new PathPatternTree(); @@ -3516,7 +3526,7 @@ public Analysis visitDeviceSchemaFetch( @Override public Analysis visitCountDevices( CountDevicesStatement countDevicesStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(countDevicesStatement); if (countDevicesStatement.hasTimeCondition()) { @@ -3543,7 +3553,7 @@ public Analysis visitCountDevices( @Override public Analysis visitCountTimeSeries( CountTimeSeriesStatement countTimeSeriesStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(countTimeSeriesStatement); PathPatternTree patternTree = new PathPatternTree(); @@ -3584,7 +3594,7 @@ public Analysis visitCountTimeSeries( @Override public Analysis visitCountLevelTimeSeries( CountLevelTimeSeriesStatement countLevelTimeSeriesStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(countLevelTimeSeriesStatement); PathPatternTree patternTree = new PathPatternTree(); @@ -3603,7 +3613,7 @@ public Analysis visitCountLevelTimeSeries( @Override public Analysis visitCountNodes(CountNodesStatement countStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(countStatement); PathPatternTree patternTree = new PathPatternTree(); @@ -3651,7 +3661,7 @@ public Analysis visitShowChildNodes( @Override public Analysis visitShowVersion( ShowVersionStatement showVersionStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showVersionStatement); analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowVersionHeader()); analysis.setFinishQueryAfterAnalyze(true); @@ -3660,7 +3670,7 @@ public Analysis visitShowVersion( private Analysis visitSchemaNodeManagementPartition( ShowStatement statement, PartialPath path, PathPatternTree scope, DatasetHeader header) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(statement); PathPatternTree patternTree = new PathPatternTree(); @@ -3686,7 +3696,7 @@ private Analysis visitSchemaNodeManagementPartition( public Analysis visitDeleteData( DeleteDataStatement deleteDataStatement, MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(deleteDataStatement); PathPatternTree patternTree = new PathPatternTree(); @@ -3782,7 +3792,7 @@ public Analysis visitCreateSchemaTemplate( DataNodeQueryMessages .MEASUREMENT_UNDER_TEMPLATE_IS_NOT_ALLOWED_TO_HAVE_THE_SAME_MEASUREMENT_NAME); } - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(createTemplateStatement); return analysis; } @@ -3791,7 +3801,7 @@ public Analysis visitCreateSchemaTemplate( public Analysis visitSetSchemaTemplate( SetSchemaTemplateStatement setSchemaTemplateStatement, MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(setSchemaTemplateStatement); return analysis; } @@ -3800,7 +3810,7 @@ public Analysis visitSetSchemaTemplate( public Analysis visitActivateTemplate( ActivateTemplateStatement activateTemplateStatement, MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(activateTemplateStatement); PartialPath activatePath = activateTemplateStatement.getPath(); @@ -3831,7 +3841,7 @@ public Analysis visitActivateTemplate( public Analysis visitBatchActivateTemplate( BatchActivateTemplateStatement batchActivateTemplateStatement, MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(batchActivateTemplateStatement); Map> deviceTemplateSetInfoMap = @@ -3868,7 +3878,7 @@ public Analysis visitInternalBatchActivateTemplate( InternalBatchActivateTemplateStatement internalBatchActivateTemplateStatement, MPPQueryContext context) { context.setQueryType(QueryType.OTHER); - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(internalBatchActivateTemplateStatement); PathPatternTree patternTree = new PathPatternTree(); @@ -3889,7 +3899,7 @@ public Analysis visitInternalBatchActivateTemplate( @Override public Analysis visitShowPathsUsingTemplate( ShowPathsUsingTemplateStatement showPathsUsingTemplateStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showPathsUsingTemplateStatement); analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowPathsUsingTemplateHeader()); @@ -3936,7 +3946,7 @@ public Analysis visitShowPathsUsingTemplate( @Override public Analysis visitShowQueries( ShowQueriesStatement showQueriesStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showQueriesStatement); analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowQueriesHeader()); analysis.setVirtualSource(true); @@ -3973,7 +3983,7 @@ public Analysis visitShowQueries( @Override public Analysis visitShowReceivers( ShowReceiversStatement showReceiversStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showReceiversStatement); analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowReceiversHeader()); analysis.setVirtualSource(true); @@ -4001,7 +4011,7 @@ public Analysis visitShowReceivers( @Override public Analysis visitShowDiskUsage( ShowDiskUsageStatement showDiskUsageStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showDiskUsageStatement); analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowDiskUsageHeader()); analysis.setVirtualSource(true); @@ -4078,7 +4088,7 @@ private boolean analyzeWhereForShowStatements( @Override public Analysis visitCreateLogicalView( CreateLogicalViewStatement createLogicalViewStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); context.setQueryType(QueryType.OTHER); analysis.setRealStatement(createLogicalViewStatement); @@ -4357,7 +4367,7 @@ private void checkTargetPathsInCreateLogicalView( public Analysis visitShowLogicalView( final ShowLogicalViewStatement showLogicalViewStatement, final MPPQueryContext context) { context.setQueryType(QueryType.READ); - final Analysis analysis = new Analysis(); + final Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showLogicalViewStatement); final PathPatternTree patternTree = new PathPatternTree(); @@ -4374,7 +4384,7 @@ public Analysis visitShowLogicalView( @Override public Analysis visitShowCurrentTimestamp( ShowCurrentTimestampStatement showCurrentTimestampStatement, MPPQueryContext context) { - Analysis analysis = new Analysis(); + Analysis analysis = new Analysis(internalRpcSession); analysis.setRealStatement(showCurrentTimestampStatement); analysis.setFinishQueryAfterAnalyze(true); analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowCurrentTimestampHeader()); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analyzer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analyzer.java index 4ef11925f26ce..e8d124024626b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analyzer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/Analyzer.java @@ -54,7 +54,9 @@ public Analyzer( public Analysis analyze(Statement statement) { long startTime = System.nanoTime(); mutationJournal.prepareForAnalyze(); - AnalyzeVisitor visitor = new AnalyzeVisitor(partitionFetcher, schemaFetcher, mutationJournal); + AnalyzeVisitor visitor = + new AnalyzeVisitor( + partitionFetcher, schemaFetcher, mutationJournal, context.isInternalRpcSession()); Analysis analysis = null; context.setReserveMemoryForSchemaTreeFunc( mem -> { @@ -78,7 +80,6 @@ public Analysis analyze(Statement statement) { // for test compatibility analysis.setDatabaseName(context.getDatabaseName().orElse(null)); } - if (statement.isQuery()) { long analyzeCost = System.nanoTime() diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IAnalysis.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IAnalysis.java index 4f48961691d74..3aee6015c63ab 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IAnalysis.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/IAnalysis.java @@ -19,6 +19,7 @@ package org.apache.iotdb.db.queryengine.plan.analyze; +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.partition.DataPartition; @@ -68,9 +69,9 @@ default void setRealStatement(Statement realStatement) {} List getRedirectNodeList(); - void setRedirectNodeList(List redirectNodeList); + void setRedirectNodeList(List dataNodeLocations); - void addEndPointToRedirectNodeList(TEndPoint endPoint); + void addEndPointToRedirectNodeList(TDataNodeLocation dataNodeLocation); TimePredicate getConvertedTimePredicate(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/load/LoadTsFileAnalyzer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/load/LoadTsFileAnalyzer.java index bdabeffe235f1..e5e2e5244c890 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/load/LoadTsFileAnalyzer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/load/LoadTsFileAnalyzer.java @@ -486,25 +486,29 @@ && handleSingleMiniFile(i)) { if (isTableModelFile) { final SessionInfo newSessionInfo = new SessionInfo( - sessionInfo.getSessionId(), - sessionInfo.getUserEntity(), - sessionInfo.getZoneId(), - sessionInfo.getVersion(), - sessionInfo.getDatabaseName().orElse(null), - SqlDialect.TABLE, - sessionInfo.isEnableAutoCreateSchema()); + sessionInfo.getSessionId(), + sessionInfo.getUserEntity(), + sessionInfo.getZoneId(), + sessionInfo.getVersion(), + sessionInfo.getDatabaseName().orElse(null), + SqlDialect.TABLE, + sessionInfo.isEnableAutoCreateSchema()) + .setDatabaseInternalSession(sessionInfo.isDatabaseInternalSession()) + .setInternalRpcSession(sessionInfo.isInternalRpcSession()); context.setSession(newSessionInfo); doAnalyzeSingleTableFile(tsFile, reader, timeseriesMetadataIterator, tableSchemaMap); } else { final SessionInfo newSessionInfo = new SessionInfo( - sessionInfo.getSessionId(), - sessionInfo.getUserEntity(), - sessionInfo.getZoneId(), - sessionInfo.getVersion(), - sessionInfo.getDatabaseName().orElse(null), - SqlDialect.TREE, - sessionInfo.isEnableAutoCreateSchema()); + sessionInfo.getSessionId(), + sessionInfo.getUserEntity(), + sessionInfo.getZoneId(), + sessionInfo.getVersion(), + sessionInfo.getDatabaseName().orElse(null), + SqlDialect.TREE, + sessionInfo.isEnableAutoCreateSchema()) + .setDatabaseInternalSession(sessionInfo.isDatabaseInternalSession()) + .setInternalRpcSession(sessionInfo.isInternalRpcSession()); context.setSession(newSessionInfo); doAnalyzeSingleTreeFile(tsFile, reader, timeseriesMetadataIterator); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java index 2324bea787789..d77d816960edf 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java @@ -75,6 +75,7 @@ import static com.google.common.base.Throwables.throwIfUnchecked; import static org.apache.iotdb.calc.metric.QueryExecutionMetricSet.WAIT_FOR_RESULT; import static org.apache.iotdb.commons.utils.ErrorHandlingCommonUtils.getRootCause; +import static org.apache.iotdb.db.queryengine.common.DataNodeEndPoints.getLocalRedirectEndPoint; import static org.apache.iotdb.db.queryengine.common.DataNodeEndPoints.isSameNode; import static org.apache.iotdb.rpc.TSStatusCode.DATE_OUT_OF_RANGE; @@ -749,7 +750,8 @@ private ExecutionResult getExecutionResult(QueryState state) { // info to client if (!CONFIG.isEnable13DataInsertAdapt() || IoTDBConstant.ClientVersion.V_1_0.equals(context.getSession().getVersion())) { - planner.setRedirectInfo(analysis, CONFIG.getAddressAndPort(), tsstatus); + planner.setRedirectInfo( + analysis, getLocalRedirectEndPoint(context.isInternalRpcSession()), tsstatus); } return new ExecutionResult( context.getQueryId(), tsstatus, stateMachine.isFailureCausedByInterruption(tsstatus)); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java index ecf76892e4cd9..1230400e31045 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowNode.java @@ -136,8 +136,7 @@ public List splitByPartition(IAnalysis analysis) { getDeviceID(), timePartitionSlot, analysis.getDatabaseName()); // collect redirectInfo analysis.setRedirectNodeList( - Collections.singletonList( - dataRegionReplicaSet.getDataNodeLocations().get(0).getClientRpcEndPoint())); + Collections.singletonList(dataRegionReplicaSet.getDataNodeLocations().get(0))); return Collections.singletonList(this); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsNode.java index c2ba7c04791ec..44481da77c30f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsNode.java @@ -19,7 +19,7 @@ package org.apache.iotdb.db.queryengine.plan.planner.plan.node.write; -import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.consensus.index.ProgressIndex; @@ -304,7 +304,7 @@ public void markAsGeneratedByRemoteConsensusLeader() { @Override public List splitByPartition(IAnalysis analysis) { Map splitMap = new HashMap<>(); - List redirectInfo = new ArrayList<>(); + List redirectInfo = new ArrayList<>(); for (int i = 0; i < insertRowNodeList.size(); i++) { InsertRowNode insertRowNode = insertRowNodeList.get(i); // Data region for insert row node @@ -317,7 +317,7 @@ public List splitByPartition(IAnalysis analysis) { TimePartitionUtils.getTimePartitionSlot(insertRowNode.getTime()), null); // Collect redirectInfo - redirectInfo.add(dataRegionReplicaSet.getDataNodeLocations().get(0).getClientRpcEndPoint()); + redirectInfo.add(dataRegionReplicaSet.getDataNodeLocations().get(0)); InsertRowsNode tmpNode = splitMap.get(dataRegionReplicaSet); if (tmpNode != null) { tmpNode.addOneInsertRowNode(insertRowNode, i); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsOfOneDeviceNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsOfOneDeviceNode.java index cee1e325df228..ce78aea859bd6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsOfOneDeviceNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertRowsOfOneDeviceNode.java @@ -213,8 +213,7 @@ public List splitByPartition(IAnalysis analysis) { if (i == insertRowNodeList.size() - 1) { analysis.setRedirectNodeList( - Collections.singletonList( - dataRegionReplicaSet.getDataNodeLocations().get(0).getClientRpcEndPoint())); + Collections.singletonList(dataRegionReplicaSet.getDataNodeLocations().get(0))); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertTabletNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertTabletNode.java index 5c91d0af4eaaf..bd197ce6168da 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertTabletNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/InsertTabletNode.java @@ -295,11 +295,7 @@ protected Map> splitByReplicaSet( splitInfo.replicaSets = replicaSets; // collect redirectInfo analysis.addEndPointToRedirectNodeList( - replicaSets - .get(replicaSets.size() - 1) - .getDataNodeLocations() - .get(0) - .getClientRpcEndPoint()); + replicaSets.get(replicaSets.size() - 1).getDataNodeLocations().get(0)); for (int i = 0; i < replicaSets.size(); i++) { List subRanges = splitMap.computeIfAbsent(replicaSets.get(i), x -> new ArrayList<>()); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowsNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowsNode.java index 4a98819c6487f..5fff5ad440919 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowsNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertRowsNode.java @@ -19,7 +19,7 @@ package org.apache.iotdb.db.queryengine.plan.planner.plan.node.write; -import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.IPlanVisitor; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; @@ -196,7 +196,7 @@ public String getTableName() { public List splitByPartition(IAnalysis analysis) { List writePlanNodeList = new ArrayList<>(); Map splitMap = new HashMap<>(); - List redirectInfo = new ArrayList<>(); + List redirectInfo = new ArrayList<>(); for (int i = 0; i < getInsertRowNodeList().size(); i++) { InsertRowNode insertRowNode = getInsertRowNodeList().get(i); // Data region for insert row node @@ -210,7 +210,7 @@ public List splitByPartition(IAnalysis analysis) { analysis.getDatabaseName()); // Collect redirectInfo - redirectInfo.add(dataRegionReplicaSet.getDataNodeLocations().get(0).getClientRpcEndPoint()); + redirectInfo.add(dataRegionReplicaSet.getDataNodeLocations().get(0)); RelationalInsertRowsNode tmpNode = splitMap.get(dataRegionReplicaSet); if (tmpNode != null) { tmpNode.addOneInsertRowNode(insertRowNode, i); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertTabletNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertTabletNode.java index 029c55499bf9d..50924da4447b9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertTabletNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/write/RelationalInsertTabletNode.java @@ -19,7 +19,7 @@ package org.apache.iotdb.db.queryengine.plan.planner.plan.node.write; -import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.path.PartialPath; @@ -266,7 +266,7 @@ protected InsertTabletNode getEmptySplit(int count) { protected Map> splitByReplicaSet( final Map deviceIDSplitInfoMap, final IAnalysis analysis) { final Map> splitMap = new HashMap<>(); - final Map endPointMap = new HashMap<>(); + final Map dataNodeLocationMap = new HashMap<>(); for (final Map.Entry entry : deviceIDSplitInfoMap.entrySet()) { final IDeviceID deviceID = entry.getKey(); @@ -278,13 +278,8 @@ protected Map> splitByReplicaSet( deviceID, splitInfo.timePartitionSlots, analysis.getDatabaseName()); splitInfo.replicaSets = replicaSets; // collect redirectInfo - endPointMap.put( - deviceID, - replicaSets - .get(replicaSets.size() - 1) - .getDataNodeLocations() - .get(0) - .getClientRpcEndPoint()); + dataNodeLocationMap.put( + deviceID, replicaSets.get(replicaSets.size() - 1).getDataNodeLocations().get(0)); for (int i = 0; i < replicaSets.size(); i++) { final List subRanges = splitMap.computeIfAbsent(replicaSets.get(i), x -> new ArrayList<>()); @@ -292,10 +287,10 @@ protected Map> splitByReplicaSet( subRanges.add(splitInfo.ranges.get(2 * i + 1)); } } - final List redirectNodeList = new ArrayList<>(times.length); + final List redirectNodeList = new ArrayList<>(times.length); for (int i = 0; i < times.length; i++) { final IDeviceID deviceId = getDeviceID(i); - redirectNodeList.add(endPointMap.get(deviceId)); + redirectNodeList.add(dataNodeLocationMap.get(deviceId)); } analysis.setRedirectNodeList(redirectNodeList); return splitMap; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analysis.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analysis.java index 884d5cbd4e74a..375c8d59be283 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analysis.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analysis.java @@ -19,6 +19,7 @@ package org.apache.iotdb.db.queryengine.plan.relational.analyzer; +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.partition.DataPartition; @@ -62,6 +63,7 @@ import org.apache.iotdb.commons.queryengine.plan.statement.component.FillPolicy; import org.apache.iotdb.commons.schema.table.InformationSchema; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.common.DataNodeEndPoints; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; import org.apache.iotdb.db.queryengine.common.header.DatasetHeader; import org.apache.iotdb.db.queryengine.plan.analyze.IAnalysis; @@ -118,6 +120,8 @@ public class Analysis implements IAnalysis { private String databaseName; private List redirectNodeList; + private final boolean internalRpcSession; + @Nullable private Statement root; private boolean needSetHighestPriority; @@ -266,6 +270,14 @@ public class Analysis implements IAnalysis { private SqlParser sqlParser; public Analysis(@Nullable Statement root, Map, Expression> parameters) { + this(root, parameters, false); + } + + public Analysis( + @Nullable Statement root, + Map, Expression> parameters, + final boolean internalRpcSession) { + this.internalRpcSession = internalRpcSession; this.root = root; this.parameters = ImmutableMap.copyOf( @@ -1050,16 +1062,21 @@ public List getRedirectNodeList() { } @Override - public void setRedirectNodeList(final List redirectNodeList) { - this.redirectNodeList = redirectNodeList; + public void setRedirectNodeList(final List dataNodeLocations) { + redirectNodeList = new ArrayList<>(dataNodeLocations.size()); + for (final TDataNodeLocation dataNodeLocation : dataNodeLocations) { + redirectNodeList.add( + DataNodeEndPoints.getRedirectEndPoint(dataNodeLocation, internalRpcSession)); + } } @Override - public void addEndPointToRedirectNodeList(final TEndPoint endPoint) { + public void addEndPointToRedirectNodeList(final TDataNodeLocation dataNodeLocation) { if (redirectNodeList == null) { redirectNodeList = new ArrayList<>(); } - redirectNodeList.add(endPoint); + redirectNodeList.add( + DataNodeEndPoints.getRedirectEndPoint(dataNodeLocation, internalRpcSession)); } public void setTableFunctionAnalysis( diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analyzer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analyzer.java index 8b46055d02a8f..ac8dd2ca84fb0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analyzer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analyzer.java @@ -86,7 +86,8 @@ public Analysis analyze(Statement statement) { parameterLookup, warningCollector); - Analysis analysis = new Analysis(rewrittenStatement, parameterLookup); + Analysis analysis = + new Analysis(rewrittenStatement, parameterLookup, context.isInternalRpcSession()); analysis.setSqlParser(statementAnalyzerFactory.getSqlParser()); // Register CTE passed by parent query. context diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckContext.java index ac09431fc2caf..346b6d6028aae 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckContext.java @@ -32,7 +32,7 @@ public class TreeAccessCheckContext implements IAuditEntity { private final long userId; private final String username; private final String cliHostname; - private boolean internal; + private boolean databaseInternalSession; public TreeAccessCheckContext(long userId, String username, String cliHostname) { this.userId = userId; @@ -56,12 +56,12 @@ public String getCliHostname() { } @Override - public boolean isInternal() { - return internal; + public boolean isDatabaseInternalSession() { + return databaseInternalSession; } - public TreeAccessCheckContext setInternal(final boolean internal) { - this.internal = internal; + public TreeAccessCheckContext setDatabaseInternalSession(final boolean databaseInternalSession) { + this.databaseInternalSession = databaseInternalSession; return this; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java index aa1cab64b0f05..06ea582a6afec 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java @@ -1136,7 +1136,7 @@ && isAuditTreeDatabase(statement.getDatabasePath())) { .setMessage(getReservedDatabaseNameErrorMsg(TREE_MODEL_AUDIT_DATABASE)); } // The system database is created by the database itself, not through SQL. - if (!context.isInternal() + if (!context.isDatabaseInternalSession() && InternalDatabaseUtils.isSystemDatabasePath(statement.getDatabasePath())) { context .setDatabase(statement.getDatabasePath().getFullPath()) @@ -1209,7 +1209,7 @@ public TSStatus visitDeleteDatabase( context .setAuditLogOperation(AuditLogOperation.DDL) .setDatabase(statement.getPrefixPath().toString()); - if (!context.isInternal() + if (!context.isDatabaseInternalSession() && SchemaConstant.SYSTEM_DATABASE.equals(String.join(".", statement.getPrefixPath()))) { AUDIT_LOGGER.recordObjectAuthenticationAuditLog( context.setPrivilegeType(PrivilegeType.MANAGE_DATABASE).setResult(false), @@ -1291,7 +1291,7 @@ private TSStatus checkShowOrCountDatabasePermission( */ public static TSStatus checkSystemDatabaseWrite( final IAuditEntity context, final Collection checkedPaths) { - if (context.isInternal()) { + if (context.isDatabaseInternalSession()) { return null; } for (PartialPath path : checkedPaths) { diff --git a/iotdb-core/datanode/src/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..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 @@ -32,12 +32,19 @@ 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 org.apache.thrift.TProcessor; import java.util.concurrent.atomic.AtomicReference; @@ -48,9 +55,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 +69,11 @@ public ServiceType getID() { @Override public void initTProcessor() { DataNodeInternalRPCServiceImpl service = getImpl(); + IClientRPCServiceWithHandler clientService = getClientImpl(); initSyncedServiceImpl(null); - processor = new Processor<>(service); + processor = + createMultiplexedProcessor( + new Processor<>(service), new InternalClientRPCServiceProcessor(clientService)); } @Override @@ -78,7 +90,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 +107,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 +133,20 @@ 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(); + } + private static class DataNodeInternalRPCServiceHolder { private static final DataNodeInternalRPCService INSTANCE = new DataNodeInternalRPCService(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/metrics/IoTDBInternalLocalReporter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/metrics/IoTDBInternalLocalReporter.java index 5749d99b66ff7..c4bf9ee2929c6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/metrics/IoTDBInternalLocalReporter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/metrics/IoTDBInternalLocalReporter.java @@ -103,7 +103,8 @@ public IoTDBInternalLocalReporter() { AuthorityChecker.SUPER_USER, IoTDBDescriptor.getInstance().getConfig().getInternalAddress()), ZoneId.systemDefault()) - .setInternal(true), + .setDatabaseInternalSession(true) + .setInternalRpcSession(true), ClusterPartitionFetcher.getInstance(), ClusterSchemaFetcher.getInstance()); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTsFileDataTypeConverter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTsFileDataTypeConverter.java index f78d6c5c53715..523661170f885 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTsFileDataTypeConverter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTsFileDataTypeConverter.java @@ -258,14 +258,16 @@ private TSStatus executeForTableModel( session, SESSION_MANAGER.requestQueryId(), new SessionInfo( - session.getId(), - new UserEntity( - session.getUserId(), session.getUsername(), session.getClientAddress()), - ZoneId.systemDefault(), - session.getClientVersion(), - databaseName, - SqlDialect.TABLE, - allowAutoCreateSchema && session.isEnableAutoCreateSchema()), + session.getId(), + new UserEntity( + session.getUserId(), session.getUsername(), session.getClientAddress()), + ZoneId.systemDefault(), + session.getClientVersion(), + databaseName, + SqlDialect.TABLE, + allowAutoCreateSchema && session.isEnableAutoCreateSchema()) + .setDatabaseInternalSession(isDatabaseInternalSessionContext()) + .setInternalRpcSession(isInternalRpcSessionContext()), "", LocalExecutionPlanner.getInstance().metadata, IoTDBDescriptor.getInstance().getConfig().getQueryTimeoutThreshold()) @@ -323,18 +325,20 @@ private TSStatus executeForTreeModel( isGeneratedByPipe ? new PipeEnrichedStatement(statement) : statement, SESSION_MANAGER.requestQueryId(), new SessionInfo( - session.getId(), - new UserEntity( - session.getUserId(), session.getUsername(), session.getClientAddress()), - ZoneId.systemDefault(), - session.getClientVersion(), - session.getDatabaseName(), - SqlDialect.TREE, - allowAutoCreateSchema - && session.isEnableAutoCreateSchema() - && (Objects.isNull(context) - || Objects.isNull(context.getSession()) - || context.getSession().isEnableAutoCreateSchema())), + session.getId(), + new UserEntity( + session.getUserId(), session.getUsername(), session.getClientAddress()), + ZoneId.systemDefault(), + session.getClientVersion(), + session.getDatabaseName(), + SqlDialect.TREE, + allowAutoCreateSchema + && session.isEnableAutoCreateSchema() + && (Objects.isNull(context) + || Objects.isNull(context.getSession()) + || context.getSession().isEnableAutoCreateSchema())) + .setDatabaseInternalSession(isDatabaseInternalSessionContext()) + .setInternalRpcSession(isInternalRpcSessionContext()), "", ClusterPartitionFetcher.getInstance(), ClusterSchemaFetcher.getInstance(), @@ -355,4 +359,16 @@ public boolean isSuccessful(final TSStatus status) { || status.getCode() == TSStatusCode.REDIRECTION_RECOMMEND.getStatusCode() || status.getCode() == TSStatusCode.LOAD_IDEMPOTENT_CONFLICT_EXCEPTION.getStatusCode()); } + + private boolean isDatabaseInternalSessionContext() { + return Objects.nonNull(context) + && Objects.nonNull(context.getSession()) + && context.getSession().isDatabaseInternalSession(); + } + + private boolean isInternalRpcSessionContext() { + return Objects.nonNull(context) + && Objects.nonNull(context.getSession()) + && context.getSession().isInternalRpcSession(); + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/auth/TreeAccessTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/auth/TreeAccessTest.java index 1bd9ffb2f7a71..998a766fd0712 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/auth/TreeAccessTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/auth/TreeAccessTest.java @@ -254,7 +254,7 @@ public void testSystemDatabaseWriteRequiresInternalCaller() throws Exception { new TreeAccessCheckContext(AuthorityChecker.SUPER_USER_ID, AuthorityChecker.SUPER_USER, ""); final TreeAccessCheckContext internalCaller = new TreeAccessCheckContext(AuthorityChecker.SUPER_USER_ID, AuthorityChecker.SUPER_USER, "") - .setInternal(true); + .setDatabaseInternalSession(true); final InsertStatement systemInsert = new InsertStatement(); systemInsert.setDevice(new PartialPath("root.__system.d")); diff --git a/iotdb-core/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(); + } + } + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlannerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlannerTest.java index 99c0ed023576e..67f5ff1083d63 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlannerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TableModelPlannerTest.java @@ -19,6 +19,7 @@ package org.apache.iotdb.db.queryengine.plan.relational.planner; +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.db.queryengine.plan.relational.analyzer.Analysis; @@ -44,7 +45,17 @@ public void testSetRedirectInfoForInsertRows() { new Analysis(new InsertRows(insertRowsStatement, null), Collections.emptyMap()); final TEndPoint localEndPoint = new TEndPoint("127.0.0.1", 6667); final TEndPoint remoteEndPoint = new TEndPoint("127.0.0.2", 6667); - analysis.setRedirectNodeList(Arrays.asList(localEndPoint, remoteEndPoint)); + analysis.setRedirectNodeList( + Arrays.asList( + new TDataNodeLocation( + 0, localEndPoint, localEndPoint, localEndPoint, localEndPoint, localEndPoint), + new TDataNodeLocation( + 1, + remoteEndPoint, + remoteEndPoint, + remoteEndPoint, + remoteEndPoint, + remoteEndPoint))); final TSStatus status = RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS); createPlanner().setRedirectInfo(analysis, localEndPoint, status); diff --git a/iotdb-core/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(); + } +} diff --git a/iotdb-core/node-commons/pom.xml b/iotdb-core/node-commons/pom.xml index 3bb68a29562c3..18623c2ca20c9 100644 --- a/iotdb-core/node-commons/pom.xml +++ b/iotdb-core/node-commons/pom.xml @@ -40,11 +40,21 @@ service-rpc 2.0.11-SNAPSHOT + + org.apache.iotdb + isession + 2.0.11-SNAPSHOT + org.apache.iotdb iotdb-subscription 2.0.11-SNAPSHOT + + org.apache.iotdb + iotdb-session + 2.0.11-SNAPSHOT + org.apache.tsfile common diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/audit/IAuditEntity.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/audit/IAuditEntity.java index ad4c85f18d35d..ffc5effea4d14 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/audit/IAuditEntity.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/audit/IAuditEntity.java @@ -59,8 +59,8 @@ public interface IAuditEntity { IAuditEntity setSqlString(String sqlString); - /** Whether the caller is a dedicated internal writer of the database itself. */ - default boolean isInternal() { + /** Whether the caller is a database-internal session used for dedicated internal writes. */ + default boolean isDatabaseInternalSession() { return false; } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java new file mode 100644 index 0000000000000..84d14e20fd378 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSession.java @@ -0,0 +1,66 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.commons.client.internal.session; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.rpc.RpcServiceConstants; +import org.apache.iotdb.session.AbstractSessionBuilder; +import org.apache.iotdb.session.Session; + +/** Base session for trusted server-to-server client RPC connections. */ +final class InternalSession extends Session { + + private final INodeSupplier nodeSupplier; + + InternalSession(final InternalSessionBuilder builder) { + this(builder, builder.getNodeSupplier()); + } + + InternalSession(final InternalTableSessionBuilder builder) { + this(builder, builder.getNodeSupplier()); + } + + private InternalSession(final AbstractSessionBuilder builder, final INodeSupplier nodeSupplier) { + super(builder); + this.nodeSupplier = nodeSupplier; + } + + @Override + protected String getRpcServiceName() { + return RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE; + } + + @Override + protected void initThreadPool() { + if (nodeSupplier == null) { + super.initThreadPool(); + } + } + + @Override + protected void initAvailableNodes( + final boolean enableRPCCompaction, final int connectionTimeoutInMs) { + if (nodeSupplier == null) { + super.initAvailableNodes(enableRPCCompaction, connectionTimeoutInMs); + } else { + this.availableNodes = nodeSupplier; + } + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java new file mode 100644 index 0000000000000..d1346d4c672bd --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionBuilder.java @@ -0,0 +1,62 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.commons.client.internal.session; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.session.Session; +import org.apache.iotdb.session.i18n.SessionMessages; + +import java.util.Objects; + +/** + * Server-side builder for tree-model internal sessions. + * + *

Callers sharing an external {@link INodeSupplier} must return {@code true} from {@link + * INodeSupplier#isExternalLifeCycleNodeSupplier()}; otherwise the created session may close it + * during its lifecycle, including when the session is reopened. + */ +public class InternalSessionBuilder extends Session.Builder { + + private final INodeSupplier nodeSupplier; + + public InternalSessionBuilder() { + this.nodeSupplier = null; + } + + public InternalSessionBuilder(final INodeSupplier nodeSupplier) { + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + } + + INodeSupplier getNodeSupplier() { + return nodeSupplier; + } + + @Override + protected Session createSession() { + if (enableAutoFetch && nodeSupplier == null) { + throw new IllegalArgumentException( + SessionMessages.EXCEPTION_INTERNAL_AUTO_FETCH_REQUIRES_NODE_SUPPLIER_9CB6BA02); + } + return new InternalSession(this); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java new file mode 100644 index 0000000000000..1e65d06143293 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPool.java @@ -0,0 +1,35 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.commons.client.internal.session; + +import org.apache.iotdb.session.Session; +import org.apache.iotdb.session.pool.SessionPool; + +public final class InternalSessionPool extends SessionPool { + + InternalSessionPool(final InternalSessionPoolBuilder builder) { + super(builder); + } + + @Override + protected Session.Builder createSessionBuilder() { + return new InternalSessionBuilder().enableAutoFetch(false); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java new file mode 100644 index 0000000000000..42be53fc49a89 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalSessionPoolBuilder.java @@ -0,0 +1,54 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.commons.client.internal.session; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.session.i18n.SessionMessages; +import org.apache.iotdb.session.pool.SessionPool; + +import java.util.Objects; + +/** + * Server-side builder for tree-model internal session pools. + * + *

Callers sharing an external {@link INodeSupplier} must return {@code true} from {@link + * INodeSupplier#isExternalLifeCycleNodeSupplier()}; otherwise the created pool may close it when + * the pool is closed. + */ +public class InternalSessionPoolBuilder extends SessionPool.Builder { + + public InternalSessionPoolBuilder() {} + + public InternalSessionPoolBuilder(final INodeSupplier nodeSupplier) { + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + } + + @Override + public InternalSessionPool build() { + if (enableAutoFetch && nodeSupplier == null) { + throw new IllegalArgumentException( + SessionMessages.EXCEPTION_INTERNAL_AUTO_FETCH_REQUIRES_NODE_SUPPLIER_9CB6BA02); + } + return new InternalSessionPool(this); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.java new file mode 100644 index 0000000000000..044abea3a638b --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionBuilder.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.commons.client.internal.session; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.session.Session; +import org.apache.iotdb.session.TableSessionBuilder; +import org.apache.iotdb.session.i18n.SessionMessages; + +import java.util.Objects; + +/** + * Server-side builder for table-model internal sessions. + * + *

Callers sharing an external {@link INodeSupplier} must return {@code true} from {@link + * INodeSupplier#isExternalLifeCycleNodeSupplier()}; otherwise the created session may close it + * during its lifecycle, including when the session is reopened. + */ +public class InternalTableSessionBuilder extends TableSessionBuilder { + + private final INodeSupplier nodeSupplier; + + public InternalTableSessionBuilder() { + this.nodeSupplier = null; + } + + public InternalTableSessionBuilder(final INodeSupplier nodeSupplier) { + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + } + + INodeSupplier getNodeSupplier() { + return nodeSupplier; + } + + @Override + protected Session createSession() { + if (enableAutoFetch && nodeSupplier == null) { + throw new IllegalArgumentException( + SessionMessages.EXCEPTION_INTERNAL_AUTO_FETCH_REQUIRES_NODE_SUPPLIER_9CB6BA02); + } + return new InternalSession(this); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java new file mode 100644 index 0000000000000..8091b8a7a7517 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPool.java @@ -0,0 +1,43 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.commons.client.internal.session; + +import org.apache.iotdb.session.Session; +import org.apache.iotdb.session.pool.SessionPool; +import org.apache.iotdb.session.pool.TableSessionPool; + +final class InternalTableSessionPool extends TableSessionPool { + + InternalTableSessionPool(final InternalTableSessionPoolBuilder builder) { + super(new InternalTableSessionPoolDelegate(builder)); + } + + private static final class InternalTableSessionPoolDelegate extends SessionPool { + + private InternalTableSessionPoolDelegate(final InternalTableSessionPoolBuilder builder) { + super(builder); + } + + @Override + protected Session.Builder createSessionBuilder() { + return new InternalSessionBuilder().enableAutoFetch(false); + } + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java new file mode 100644 index 0000000000000..6c53a222a80b3 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/session/InternalTableSessionPoolBuilder.java @@ -0,0 +1,55 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.commons.client.internal.session; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.isession.pool.ITableSessionPool; +import org.apache.iotdb.session.i18n.SessionMessages; +import org.apache.iotdb.session.pool.TableSessionPoolBuilder; + +import java.util.Objects; + +/** + * Server-side builder for table-model internal session pools. + * + *

Callers sharing an external {@link INodeSupplier} must return {@code true} from {@link + * INodeSupplier#isExternalLifeCycleNodeSupplier()}; otherwise the created pool may close it when + * the pool is closed. + */ +public class InternalTableSessionPoolBuilder extends TableSessionPoolBuilder { + + public InternalTableSessionPoolBuilder() {} + + public InternalTableSessionPoolBuilder(final INodeSupplier nodeSupplier) { + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + } + + @Override + protected ITableSessionPool createSessionPool() { + if (enableAutoFetch && nodeSupplier == null) { + throw new IllegalArgumentException( + SessionMessages.EXCEPTION_INTERNAL_AUTO_FETCH_REQUIRES_NODE_SUPPLIER_9CB6BA02); + } + return new InternalTableSessionPool(this); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java new file mode 100644 index 0000000000000..d436c99dcf737 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionSessionWrapper.java @@ -0,0 +1,78 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.commons.client.internal.subscription; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.isession.SessionConfig; +import org.apache.iotdb.rpc.IoTDBConnectionException; +import org.apache.iotdb.rpc.RpcServiceConstants; +import org.apache.iotdb.session.AbstractSessionBuilder; +import org.apache.iotdb.session.subscription.SubscriptionSessionWrapper; + +final class InternalSubscriptionSessionWrapper extends SubscriptionSessionWrapper { + + private final INodeSupplier nodeSupplier; + private final boolean endpointConfigured; + + InternalSubscriptionSessionWrapper( + final AbstractSessionBuilder builder, final INodeSupplier nodeSupplier) { + super(builder); + this.nodeSupplier = nodeSupplier; + this.endpointConfigured = + builder.nodeUrls != null + || !SessionConfig.DEFAULT_HOST.equals(builder.host) + || builder.rpcPort != SessionConfig.DEFAULT_PORT; + } + + @Override + protected String getRpcServiceName() { + return RpcServiceConstants.INTERNAL_CLIENT_RPC_SERVICE; + } + + @Override + public void open() throws IoTDBConnectionException { + if (nodeSupplier != null) { + enableAutoFetch = true; + if (!endpointConfigured) { + nodeSupplier + .getQueryEndPoint() + .ifPresent(endpoint -> defaultEndPoint = endpoint.deepCopy()); + } + } + super.open(); + } + + @Override + protected void initThreadPool() { + if (nodeSupplier == null) { + super.initThreadPool(); + } + } + + @Override + protected void initAvailableNodes( + final boolean enableRPCCompaction, final int connectionTimeoutInMs) { + if (nodeSupplier == null) { + super.initAvailableNodes(enableRPCCompaction, connectionTimeoutInMs); + } else { + this.availableNodes = nodeSupplier; + } + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumer.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumer.java new file mode 100644 index 0000000000000..9269c9232ece8 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumer.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.commons.client.internal.subscription; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.rpc.RpcServiceConstants; +import org.apache.iotdb.session.subscription.consumer.base.AbstractSubscriptionProvider; +import org.apache.iotdb.session.subscription.consumer.table.SubscriptionTablePullConsumer; +import org.apache.iotdb.session.subscription.consumer.table.SubscriptionTablePullConsumerBuilder; + +final class InternalSubscriptionTablePullConsumer extends SubscriptionTablePullConsumer { + + private final INodeSupplier nodeSupplier; + + InternalSubscriptionTablePullConsumer( + final SubscriptionTablePullConsumerBuilder builder, final INodeSupplier nodeSupplier) { + super(builder); + this.nodeSupplier = nodeSupplier; + } + + @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-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java new file mode 100644 index 0000000000000..bb32d6224b358 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTablePullConsumerBuilder.java @@ -0,0 +1,55 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.commons.client.internal.subscription; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.session.i18n.SessionMessages; +import org.apache.iotdb.session.subscription.consumer.ISubscriptionTablePullConsumer; +import org.apache.iotdb.session.subscription.consumer.table.SubscriptionTablePullConsumerBuilder; + +import java.util.Objects; + +/** + * Server-side builder for table subscription pull consumers. + * + *

Callers sharing an external {@link INodeSupplier} must return {@code true} from {@link + * INodeSupplier#isExternalLifeCycleNodeSupplier()}. + */ +public class InternalSubscriptionTablePullConsumerBuilder + extends SubscriptionTablePullConsumerBuilder { + + private final INodeSupplier nodeSupplier; + + public InternalSubscriptionTablePullConsumerBuilder() { + this.nodeSupplier = null; + } + + public InternalSubscriptionTablePullConsumerBuilder(final INodeSupplier nodeSupplier) { + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + } + + @Override + protected ISubscriptionTablePullConsumer createConsumer() { + return new InternalSubscriptionTablePullConsumer(this, nodeSupplier); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSession.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSession.java new file mode 100644 index 0000000000000..f4aa5bbd30ce6 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSession.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.commons.client.internal.subscription; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.session.AbstractSessionBuilder; +import org.apache.iotdb.session.subscription.SubscriptionTableSession; + +final class InternalSubscriptionTableSession extends SubscriptionTableSession { + + InternalSubscriptionTableSession( + final AbstractSessionBuilder builder, final INodeSupplier nodeSupplier) { + super(new InternalSubscriptionSessionWrapper(builder, nodeSupplier)); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java new file mode 100644 index 0000000000000..9005773782423 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/internal/subscription/InternalSubscriptionTableSessionBuilder.java @@ -0,0 +1,61 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.commons.client.internal.subscription; + +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.session.i18n.SessionMessages; +import org.apache.iotdb.session.subscription.ISubscriptionTableSession; +import org.apache.iotdb.session.subscription.SubscriptionTableSessionBuilder; + +import java.util.List; +import java.util.Objects; + +/** + * Server-side builder for table subscription-management sessions. + * + *

Callers sharing an external {@link INodeSupplier} must return {@code true} from {@link + * INodeSupplier#isExternalLifeCycleNodeSupplier()}; otherwise the created session may close it + * during its lifecycle, including when the session is reopened. + */ +public class InternalSubscriptionTableSessionBuilder extends SubscriptionTableSessionBuilder { + + private final INodeSupplier nodeSupplier; + + public InternalSubscriptionTableSessionBuilder() { + this.nodeSupplier = null; + } + + public InternalSubscriptionTableSessionBuilder(final INodeSupplier nodeSupplier) { + this.nodeSupplier = + Objects.requireNonNull( + nodeSupplier, + SessionMessages.EXCEPTION_INTERNAL_NODE_SUPPLIER_MUST_NOT_BE_NULL_C5CAF392); + } + + public InternalSubscriptionTableSessionBuilder nodeUrls(final List nodeUrls) { + this.nodeUrls = nodeUrls; + return this; + } + + @Override + protected ISubscriptionTableSession createSubscriptionSession() { + return new InternalSubscriptionTableSession(this, nodeSupplier); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/common/SessionInfo.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/common/SessionInfo.java index d6a0f4120abb7..1dbc6012c5df9 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/common/SessionInfo.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/common/SessionInfo.java @@ -52,7 +52,11 @@ public class SessionInfo { // Whether the session is opened by the database itself. Like enableAutoCreateSchema, this is // only consulted during analysis and therefore does not need serialization. - private transient boolean internal; + private transient boolean databaseInternalSession; + + // Whether the session uses the internal RPC endpoint family. Used to select internal redirect + // endpoints without changing the meaning of internal database sessions. + private transient boolean internalRpcSession; public SessionInfo(long sessionId, UserEntity userEntity, ZoneId zoneId) { this.sessionId = sessionId; @@ -169,12 +173,21 @@ public boolean isEnableAutoCreateSchema() { return enableAutoCreateSchema; } - public boolean isInternal() { - return internal; + public boolean isDatabaseInternalSession() { + return databaseInternalSession; + } + + public SessionInfo setDatabaseInternalSession(final boolean databaseInternalSession) { + this.databaseInternalSession = databaseInternalSession; + return this; + } + + public boolean isInternalRpcSession() { + return internalRpcSession; } - public SessionInfo setInternal(final boolean internal) { - this.internal = internal; + public SessionInfo setInternalRpcSession(final boolean internalRpcSession) { + this.internalRpcSession = internalRpcSession; return this; } diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalBuildersTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalBuildersTest.java new file mode 100644 index 0000000000000..fbe4bb969d8a1 --- /dev/null +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalBuildersTest.java @@ -0,0 +1,47 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.commons.client.internal.session; + +import org.apache.iotdb.session.i18n.SessionMessages; + +import org.junit.Test; +import org.junit.function.ThrowingRunnable; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +public class InternalBuildersTest { + + @Test + public void testInternalBuildersRejectAutoFetchWithoutNodeSupplier() { + assertAutoFetchRequiresNodeSupplier(() -> new InternalSessionBuilder().build()); + assertAutoFetchRequiresNodeSupplier(() -> new InternalTableSessionBuilder().build()); + assertAutoFetchRequiresNodeSupplier(() -> new InternalSessionPoolBuilder().build()); + assertAutoFetchRequiresNodeSupplier(() -> new InternalTableSessionPoolBuilder().build()); + } + + private static void assertAutoFetchRequiresNodeSupplier(final ThrowingRunnable builder) { + final IllegalArgumentException exception = + assertThrows(IllegalArgumentException.class, builder); + assertEquals( + SessionMessages.EXCEPTION_INTERNAL_AUTO_FETCH_REQUIRES_NODE_SUPPLIER_9CB6BA02, + exception.getMessage()); + } +} diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalSessionProtocolTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalSessionProtocolTest.java new file mode 100644 index 0000000000000..104e56a1c4500 --- /dev/null +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/client/internal/session/InternalSessionProtocolTest.java @@ -0,0 +1,167 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.commons.client.internal.session; + +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.client.internal.subscription.InternalSubscriptionTableSessionBuilder; +import org.apache.iotdb.isession.INodeSupplier; +import org.apache.iotdb.rpc.DeepCopyRpcTransportFactory; +import org.apache.iotdb.rpc.RpcServiceConstants; +import org.apache.iotdb.rpc.TSStatusCode; +import org.apache.iotdb.service.rpc.thrift.IClientRPCService; +import org.apache.iotdb.service.rpc.thrift.TSCloseSessionReq; +import org.apache.iotdb.service.rpc.thrift.TSExecuteStatementReq; +import org.apache.iotdb.service.rpc.thrift.TSExecuteStatementResp; +import org.apache.iotdb.service.rpc.thrift.TSOpenSessionReq; +import org.apache.iotdb.service.rpc.thrift.TSOpenSessionResp; +import org.apache.iotdb.service.rpc.thrift.TSProtocolVersion; +import org.apache.iotdb.session.DummyNodesSupplier; +import org.apache.iotdb.session.Session; +import org.apache.iotdb.session.subscription.ISubscriptionTableSession; + +import org.apache.thrift.TException; +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 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 (Session session = + new InternalSessionBuilder(nodeSupplier) + .host(server.endpoint.getIp()) + .port(server.endpoint.getPort()) + .username("root") + .build()) { + session.open(); + session.executeNonQueryStatement("SET CONFIGURATION 'force_reconnect'='true'"); + assertEquals(2, openSessionCalls.get()); + } + } + + @Test + public void testInternalSubscriptionSessionUsesSuppliedEndpoint() throws Exception { + final INodeSupplier nodeSupplier = + new DummyNodesSupplier(Collections.singletonList(server.endpoint)); + try (ISubscriptionTableSession session = + new InternalSubscriptionTableSessionBuilder(nodeSupplier).username("root").build()) { + assertEquals(1, openSessionCalls.get()); + } + } + + private static final class TestServer implements AutoCloseable { + 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())); + when(handler.executeUpdateStatementV2(any(TSExecuteStatementReq.class))) + .thenThrow(new TException("force reconnect")) + .thenReturn( + new TSExecuteStatementResp( + new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()))); + + final TMultiplexedProcessor processor = new TMultiplexedProcessor(); + processor.registerProcessor( + 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); + } + } +}