diff --git a/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/env/AbstractEnv.java b/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/env/AbstractEnv.java index 342b4aba26d8c..135011406da9c 100644 --- a/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/env/AbstractEnv.java +++ b/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/env/AbstractEnv.java @@ -24,7 +24,6 @@ import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.client.ClientPoolFactory; import org.apache.iotdb.commons.client.IClientManager; -import org.apache.iotdb.commons.client.exception.ClientManagerException; import org.apache.iotdb.commons.client.sync.SyncConfigNodeIServiceClient; import org.apache.iotdb.commons.cluster.NodeStatus; import org.apache.iotdb.commons.exception.PortOccupiedException; @@ -40,7 +39,6 @@ import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.isession.pool.ISessionPool; import org.apache.iotdb.isession.pool.ITableSessionPool; -import org.apache.iotdb.it.env.EnvFactory; import org.apache.iotdb.it.env.cluster.EnvUtils; import org.apache.iotdb.it.env.cluster.config.MppClusterConfig; import org.apache.iotdb.it.env.cluster.config.MppCommonConfig; @@ -1580,7 +1578,7 @@ public void ensureNodeStatus( Throwable lastException = null; for (int i = 0; i < retryCount; i++) { try (final SyncConfigNodeIServiceClient client = - (SyncConfigNodeIServiceClient) EnvFactory.getEnv().getLeaderConfigNodeConnection()) { + (SyncConfigNodeIServiceClient) getLeaderConfigNodeConnection()) { final List errorMessages = new ArrayList<>(nodes.size()); final Map nodeIds = new HashMap<>(nodes.size()); final TShowClusterResp showClusterResp = client.showCluster(); @@ -1641,7 +1639,7 @@ public void ensureNodeStatus( } else { lastException = new IllegalStateException(String.join(". ", errorMessages)); } - } catch (final TException | ClientManagerException | IOException | InterruptedException e) { + } catch (final TException | IOException | InterruptedException e) { lastException = e; } try { diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/auto/basic/IoTDBShowReceiversLifecycleIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/auto/basic/IoTDBShowReceiversLifecycleIT.java new file mode 100644 index 0000000000000..e31b6dc9b9448 --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/auto/basic/IoTDBShowReceiversLifecycleIT.java @@ -0,0 +1,379 @@ +/* + * 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.pipe.it.dual.treemodel.auto.basic; + +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.client.sync.SyncConfigNodeIServiceClient; +import org.apache.iotdb.confignode.rpc.thrift.TCreatePipeReq; +import org.apache.iotdb.db.it.utils.TestUtils; +import org.apache.iotdb.isession.SessionConfig; +import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.itbase.category.MultiClusterIT2DualTreeAutoBasic; +import org.apache.iotdb.itbase.env.BaseEnv; +import org.apache.iotdb.pipe.it.dual.treemodel.auto.AbstractPipeDualTreeModelAutoIT; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.awaitility.Awaitility; +import org.junit.Assert; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.TimeUnit; + +@RunWith(IoTDBTestRunner.class) +@Category({MultiClusterIT2DualTreeAutoBasic.class}) +public class IoTDBShowReceiversLifecycleIT extends AbstractPipeDualTreeModelAutoIT { + + @Override + protected void setupConfig() { + super.setupConfig(); + senderEnv.getConfig().getCommonConfig().setPipeAirGapReceiverEnabled(true); + receiverEnv.getConfig().getCommonConfig().setPipeAirGapReceiverEnabled(true); + } + + @Test + public void testShowReceiversPipeIdsDisappearAfterDropPipe() throws Exception { + final String database = "root.show_receivers_lifecycle"; + final String pipeName = "show_receivers_lifecycle_pipe"; + + createThriftPipe(database, pipeName); + + assertShowReceivers("show receivers", BaseEnv.TREE_SQL_DIALECT, pipeName); + assertShowReceivers( + "select * from information_schema.receivers", BaseEnv.TABLE_SQL_DIALECT, pipeName); + + TestUtils.executeNonQueries( + senderEnv, Collections.singletonList("drop pipe " + pipeName), null); + + assertShowReceiversDoesNotContainPipe("show receivers", BaseEnv.TREE_SQL_DIALECT, pipeName); + assertShowReceiversDoesNotContainPipe( + "select * from information_schema.receivers", BaseEnv.TABLE_SQL_DIALECT, pipeName); + } + + @Test + public void testShowReceiversPipeIdsDisappearAfterStopPipe() throws Exception { + final String database = "root.show_receivers_lifecycle_stop"; + final String pipeName = "show_receivers_lifecycle_stop_pipe"; + + createThriftPipe(database, pipeName); + + assertShowReceivers("show receivers", BaseEnv.TREE_SQL_DIALECT, pipeName); + assertShowReceivers( + "select * from information_schema.receivers", BaseEnv.TABLE_SQL_DIALECT, pipeName); + + TestUtils.executeNonQueries( + senderEnv, Collections.singletonList("stop pipe " + pipeName), null); + + assertShowReceiversDoesNotContainPipe("show receivers", BaseEnv.TREE_SQL_DIALECT, pipeName); + assertShowReceiversDoesNotContainPipe( + "select * from information_schema.receivers", BaseEnv.TABLE_SQL_DIALECT, pipeName); + } + + @Test + public void testShowReceiversIncludesDataNodeAndConfigNodeReceivers() throws Exception { + final String database = "root.show_receivers_config_node"; + final String pipeName = "show_receivers_config_node_pipe"; + + createThriftPipe(database, pipeName, "all"); + + TestUtils.executeNonQueries( + senderEnv, + Arrays.asList( + "create timeseries " + database + ".d1.s2 with datatype=BOOLEAN, encoding=PLAIN", + "insert into " + database + ".d1(time, s1, s2) values (2, 2, true)", + "flush"), + null); + + TestUtils.assertDataEventuallyOnEnv( + receiverEnv, + "count timeseries " + database + ".**", + "count(timeseries),", + Collections.singleton("2,")); + TestUtils.assertDataEventuallyOnEnv( + receiverEnv, + "select count(s1) from " + database + ".d1", + "count(" + database + ".d1.s1),", + Collections.singleton("2,")); + + assertShowReceiversContainDataNodeAndConfigNode( + "show receivers", BaseEnv.TREE_SQL_DIALECT, pipeName); + assertShowReceiversContainDataNodeAndConfigNode( + "select * from information_schema.receivers", BaseEnv.TABLE_SQL_DIALECT, pipeName); + } + + @Test + public void testShowReceiversShowsAirGapProtocol() throws Exception { + final String database = "root.show_receivers_air_gap"; + final String pipeName = "show_receivers_air_gap_pipe"; + + createAirGapPipe(database, pipeName); + + TestUtils.assertDataEventuallyOnEnv( + receiverEnv, + "select count(s1) from " + database + ".d1", + "count(" + database + ".d1.s1),", + Collections.singleton("1,")); + + assertShowReceiversContainProtocol( + "show receivers", BaseEnv.TREE_SQL_DIALECT, pipeName, "air_gap"); + assertShowReceiversContainProtocol( + "select * from information_schema.receivers", + BaseEnv.TABLE_SQL_DIALECT, + pipeName, + "air_gap"); + } + + private void createThriftPipe(final String database, final String pipeName) throws Exception { + createThriftPipe(database, pipeName, "data.insert"); + } + + private void createThriftPipe( + final String database, final String pipeName, final String sourceInclusion) throws Exception { + createPipe(database, pipeName, sourceInclusion, "iotdb-thrift-sink", false); + } + + private void createAirGapPipe(final String database, final String pipeName) throws Exception { + createPipe(database, pipeName, "data.insert", "iotdb-air-gap-sink", true); + } + + private void createPipe( + final String database, + final String pipeName, + final String sourceInclusion, + final String sinkName, + final boolean useAirGapPort) + throws Exception { + TestUtils.executeNonQueries( + senderEnv, + Arrays.asList( + "create database " + database, + "create timeseries " + database + ".d1.s1 with datatype=INT32, encoding=PLAIN", + "insert into " + database + ".d1(time, s1) values (1, 1)"), + null); + awaitUntilFlush(senderEnv); + + final DataNodeWrapper receiverDataNode = receiverEnv.getDataNodeWrapper(0); + + final Map sourceAttributes = new HashMap<>(); + final Map processorAttributes = new HashMap<>(); + final Map sinkAttributes = new HashMap<>(); + + sourceAttributes.put("source.path", database + ".**"); + sourceAttributes.put("source.inclusion", sourceInclusion); + sourceAttributes.put("user", SessionConfig.DEFAULT_USER); + + sinkAttributes.put("sink", sinkName); + sinkAttributes.put("sink.batch.enable", "false"); + sinkAttributes.put("sink.ip", receiverDataNode.getIp()); + sinkAttributes.put( + "sink.port", + Integer.toString( + useAirGapPort + ? receiverDataNode.getPipeAirGapReceiverPort() + : receiverDataNode.getPort())); + + try (final SyncConfigNodeIServiceClient client = + (SyncConfigNodeIServiceClient) senderEnv.getLeaderConfigNodeConnection()) { + TSStatus status = + client.createPipe( + new TCreatePipeReq(pipeName, sinkAttributes) + .setExtractorAttributes(sourceAttributes) + .setProcessorAttributes(processorAttributes)); + Assert.assertEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), status.getCode()); + + status = client.startPipe(pipeName); + Assert.assertEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), status.getCode()); + } + } + + private void assertShowReceivers( + final String sql, final String sqlDialect, final String pipeName) { + Awaitility.await() + .pollInSameThread() + .pollDelay(1L, TimeUnit.SECONDS) + .pollInterval(1L, TimeUnit.SECONDS) + .atMost(60L, TimeUnit.SECONDS) + .untilAsserted(() -> Assert.assertTrue(hasExpectedReceiver(sql, sqlDialect, pipeName))); + } + + private boolean hasExpectedReceiver( + final String sql, final String sqlDialect, final String pipeName) throws SQLException { + try (final Connection connection = + receiverEnv.getConnection( + SessionConfig.DEFAULT_USER, SessionConfig.DEFAULT_PASSWORD, sqlDialect); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery(sql)) { + while (resultSet.next()) { + final String senderClusterId = getString(resultSet, "SenderClusterId", "sender_cluster_id"); + final String senderAddress = getString(resultSet, "SenderAddress", "sender_address"); + final String userName = getString(resultSet, "UserName", "user_name"); + final String senderPorts = getString(resultSet, "SenderPorts", "sender_ports"); + final String pipeIds = getString(resultSet, "PipeIDs", "pipe_ids"); + if ("DataNode".equals(getString(resultSet, "ReceiverNodeType", "receiver_node_type")) + && "thrift".equals(getString(resultSet, "Protocol", "protocol")) + && senderClusterId != null + && !senderClusterId.isEmpty() + && senderAddress != null + && !senderAddress.isEmpty() + && SessionConfig.DEFAULT_USER.equals(userName) + && senderPorts != null + && !senderPorts.isEmpty() + && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 + && getInt(resultSet, "PipeCount", "pipe_count") >= 1 + && pipeIds != null + && pipeIds.contains(pipeName + "@") + && getString(resultSet, "LastHandshakeTime", "last_handshake_time") != null + && getString(resultSet, "LastTransferTime", "last_transfer_time") != null) { + return true; + } + } + return false; + } + } + + private void assertShowReceiversDoesNotContainPipe( + final String sql, final String sqlDialect, final String pipeName) { + Awaitility.await() + .pollInSameThread() + .pollDelay(1L, TimeUnit.SECONDS) + .pollInterval(1L, TimeUnit.SECONDS) + .atMost(60L, TimeUnit.SECONDS) + .untilAsserted(() -> Assert.assertFalse(containsPipe(sql, sqlDialect, pipeName))); + } + + private boolean containsPipe(final String sql, final String sqlDialect, final String pipeName) + throws SQLException { + try (final Connection connection = + receiverEnv.getConnection( + SessionConfig.DEFAULT_USER, SessionConfig.DEFAULT_PASSWORD, sqlDialect); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery(sql)) { + while (resultSet.next()) { + final String pipeIds = getString(resultSet, "PipeIDs", "pipe_ids"); + if (pipeIds != null && pipeIds.contains(pipeName + "@")) { + return true; + } + } + return false; + } + } + + private void assertShowReceiversContainDataNodeAndConfigNode( + final String sql, final String sqlDialect, final String pipeName) { + Awaitility.await() + .pollInSameThread() + .pollDelay(1L, TimeUnit.SECONDS) + .pollInterval(1L, TimeUnit.SECONDS) + .atMost(60L, TimeUnit.SECONDS) + .untilAsserted( + () -> Assert.assertTrue(hasDataNodeAndConfigNodeReceivers(sql, sqlDialect, pipeName))); + } + + private boolean hasDataNodeAndConfigNodeReceivers( + final String sql, final String sqlDialect, final String pipeName) throws SQLException { + boolean hasDataNode = false; + boolean hasConfigNode = false; + try (final Connection connection = + receiverEnv.getConnection( + SessionConfig.DEFAULT_USER, SessionConfig.DEFAULT_PASSWORD, sqlDialect); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery(sql)) { + while (resultSet.next()) { + final String receiverNodeType = + getString(resultSet, "ReceiverNodeType", "receiver_node_type"); + final String protocol = getString(resultSet, "Protocol", "protocol"); + final String pipeIds = getString(resultSet, "PipeIDs", "pipe_ids"); + if (!"thrift".equals(protocol) || pipeIds == null || !pipeIds.contains(pipeName + "@")) { + continue; + } + if ("DataNode".equals(receiverNodeType)) { + hasDataNode = true; + } else if ("ConfigNode".equals(receiverNodeType)) { + hasConfigNode = true; + } + } + } + return hasDataNode && hasConfigNode; + } + + private void assertShowReceiversContainProtocol( + final String sql, final String sqlDialect, final String pipeName, final String protocol) { + Awaitility.await() + .pollInSameThread() + .pollDelay(1L, TimeUnit.SECONDS) + .pollInterval(1L, TimeUnit.SECONDS) + .atMost(60L, TimeUnit.SECONDS) + .untilAsserted( + () -> + Assert.assertTrue( + hasDataNodeReceiverWithProtocol(sql, sqlDialect, pipeName, protocol))); + } + + private boolean hasDataNodeReceiverWithProtocol( + final String sql, final String sqlDialect, final String pipeName, final String protocol) + throws SQLException { + try (final Connection connection = + receiverEnv.getConnection( + SessionConfig.DEFAULT_USER, SessionConfig.DEFAULT_PASSWORD, sqlDialect); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery(sql)) { + while (resultSet.next()) { + final String pipeIds = getString(resultSet, "PipeIDs", "pipe_ids"); + if ("DataNode".equals(getString(resultSet, "ReceiverNodeType", "receiver_node_type")) + && protocol.equals(getString(resultSet, "Protocol", "protocol")) + && pipeIds != null + && pipeIds.contains(pipeName + "@")) { + return true; + } + } + } + return false; + } + + private static String getString( + final ResultSet resultSet, final String treeColumnName, final String tableColumnName) + throws SQLException { + try { + return resultSet.getString(treeColumnName); + } catch (final SQLException ignored) { + return resultSet.getString(tableColumnName); + } + } + + private static int getInt( + final ResultSet resultSet, final String treeColumnName, final String tableColumnName) + throws SQLException { + try { + return resultSet.getInt(treeColumnName); + } catch (final SQLException ignored) { + return resultSet.getInt(tableColumnName); + } + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/single/IoTDBShowReceiversIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/single/IoTDBShowReceiversIT.java new file mode 100644 index 0000000000000..5e765d5677d54 --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/single/IoTDBShowReceiversIT.java @@ -0,0 +1,711 @@ +/* + * 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.pipe.it.single; + +import org.apache.iotdb.commons.client.property.ThriftClientProperty; +import org.apache.iotdb.commons.client.sync.SyncConfigNodeIServiceClient; +import org.apache.iotdb.commons.cluster.NodeStatus; +import org.apache.iotdb.commons.conf.CommonDescriptor; +import org.apache.iotdb.commons.pipe.sink.client.IoTDBSyncClient; +import org.apache.iotdb.confignode.rpc.thrift.TDataNodeInfo; +import org.apache.iotdb.confignode.rpc.thrift.TShowDataNodesResp; +import org.apache.iotdb.db.it.utils.TestUtils; +import org.apache.iotdb.db.pipe.sink.payload.evolvable.request.PipeTransferDataNodeHandshakeV1Req; +import org.apache.iotdb.isession.SessionConfig; +import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.itbase.category.MultiClusterIT1; +import org.apache.iotdb.itbase.env.BaseEnv; +import org.apache.iotdb.rpc.TSStatusCode; +import org.apache.iotdb.service.rpc.thrift.TPipeTransferResp; + +import org.awaitility.Awaitility; +import org.junit.Assert; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.Arrays; +import java.util.Collections; +import java.util.concurrent.TimeUnit; + +@RunWith(IoTDBTestRunner.class) +@Category({MultiClusterIT1.class}) +public class IoTDBShowReceiversIT extends AbstractPipeSingleIT { + + @Test + public void testShowReceiversInTreeAndTableModel() { + createWriteBackPipe("root.show_receivers", "show_receivers_pipe"); + + assertShowReceivers("show receivers", BaseEnv.TREE_SQL_DIALECT); + assertShowReceivers("select * from information_schema.receivers", BaseEnv.TABLE_SQL_DIALECT); + } + + @Test + public void testInformationSchemaReceiversWithProtocolFilter() { + createWriteBackPipe("root.show_receivers_filter", "show_receivers_filter_pipe"); + + assertShowReceivers( + "select receiver_node_type, receiver_node_id, protocol, sender_cluster_id, " + + "sender_address, user_name, sender_ports, connection_count, pipe_count, pipe_ids, " + + "last_handshake_time, last_transfer_time " + + "from information_schema.receivers where protocol = 'writeback'", + BaseEnv.TABLE_SQL_DIALECT, + "show_receivers_filter_pipe"); + } + + @Test + public void testInformationSchemaReceiversProjectedColumns() { + final String pipeName = "show_receivers_project_pipe"; + createWriteBackPipe("root.show_receivers_project", pipeName); + + assertProjectedShowReceivers(pipeName); + } + + @Test + public void testShowReceiversAggregatesMultiplePipesFromSameSender() { + final String pipeName1 = "show_receivers_aggregate_pipe_1"; + final String pipeName2 = "show_receivers_aggregate_pipe_2"; + createTwoWriteBackPipes("root.show_receivers_aggregate", pipeName1, pipeName2); + + assertAggregatedShowReceivers("show receivers", BaseEnv.TREE_SQL_DIALECT, pipeName1, pipeName2); + assertAggregatedShowReceivers( + "select * from information_schema.receivers", + BaseEnv.TABLE_SQL_DIALECT, + pipeName1, + pipeName2); + } + + @Test + public void testShowReceiversWithStoppedDataNode() throws Exception { + Assert.assertTrue(env.getDataNodeWrapperList().size() >= 3); + createWriteBackPipe("root.show_receivers_ha", "show_receivers_ha_pipe"); + + assertShowReceivers("show receivers", BaseEnv.TREE_SQL_DIALECT, "show_receivers_ha_pipe"); + assertShowReceivers( + "select * from information_schema.receivers", + BaseEnv.TABLE_SQL_DIALECT, + "show_receivers_ha_pipe"); + + final int stoppedDataNodeIndex = 0; + final DataNodeWrapper stoppedDataNode = env.getDataNodeWrapper(stoppedDataNodeIndex); + final int stoppedDataNodeId = getDataNodeId(stoppedDataNode); + final DataNodeWrapper queryDataNode = env.getDataNodeWrapper(1); + + env.shutdownDataNode(stoppedDataNodeIndex); + env.ensureNodeStatus( + Collections.singletonList(stoppedDataNode), Collections.singletonList(NodeStatus.Unknown)); + + assertShowReceiversWithoutDataNode( + "show receivers", BaseEnv.TREE_SQL_DIALECT, queryDataNode, stoppedDataNodeId); + assertShowReceiversWithoutDataNode( + "select * from information_schema.receivers", + BaseEnv.TABLE_SQL_DIALECT, + queryDataNode, + stoppedDataNodeId); + } + + @Test + public void testShowReceiversPermissionInTreeAndTableModel() { + final String password = "Passwd123456@"; + final String normalUser = "show_receiver_normal"; + final String systemUser = "show_receiver_system"; + final String pipeName = "show_receivers_auth_pipe"; + + TestUtils.executeNonQueries( + env, + Arrays.asList( + "create user " + normalUser + " '" + password + "'", + "create user " + systemUser + " '" + password + "'", + "grant system on root.** to user " + systemUser), + null); + createWriteBackPipe("root.show_receivers_auth", pipeName); + + assertShowReceivers("show receivers", BaseEnv.TREE_SQL_DIALECT, pipeName); + assertShowReceivers( + "select * from information_schema.receivers", BaseEnv.TABLE_SQL_DIALECT, pipeName); + + assertUserCannotSeeReceivers("show receivers", BaseEnv.TREE_SQL_DIALECT, normalUser, password); + assertUserCannotSeeReceivers( + "select * from information_schema.receivers", + BaseEnv.TABLE_SQL_DIALECT, + normalUser, + password); + + assertShowReceiversAsUser( + "show receivers", BaseEnv.TREE_SQL_DIALECT, pipeName, systemUser, password); + assertShowReceiversAsUser( + "select * from information_schema.receivers", + BaseEnv.TABLE_SQL_DIALECT, + pipeName, + systemUser, + password); + } + + @Test + public void testNormalUserCanSeeOwnReceiverInTreeAndTableModel() { + final String password = "Passwd123456@"; + final String sinkUser = "show_receiver_sink_user"; + final String otherUser = "show_receiver_other_user"; + final String pipeName = "show_receivers_own_user_pipe"; + + TestUtils.executeNonQueries( + env, + Arrays.asList( + "create user " + sinkUser + " '" + password + "'", + "create user " + otherUser + " '" + password + "'", + "grant write on root.** to user " + sinkUser), + null); + createWriteBackPipeWithSinkUser("root.show_receivers_own_user", pipeName, sinkUser, password); + + assertShowReceiversAsUser( + "show receivers", BaseEnv.TREE_SQL_DIALECT, pipeName, sinkUser, password, sinkUser); + assertShowReceiversAsUser( + "select * from information_schema.receivers", + BaseEnv.TABLE_SQL_DIALECT, + pipeName, + sinkUser, + password, + sinkUser); + + assertUserCannotSeeReceivers("show receivers", BaseEnv.TREE_SQL_DIALECT, otherUser, password); + assertUserCannotSeeReceivers( + "select * from information_schema.receivers", + BaseEnv.TABLE_SQL_DIALECT, + otherUser, + password); + } + + @Test + public void testOldSenderRuntimeShowsUnknownPipeIdsAndDisappearsAfterClientClose() + throws Exception { + final DataNodeWrapper receiverDataNode = env.getDataNodeWrapper(0); + final int receiverDataNodeId = getDataNodeId(receiverDataNode); + + final IoTDBSyncClient client = createPipeTransferClient(receiverDataNode); + try { + sendOldDataNodeHandshake(client); + + assertOldSenderReceiverRuntime( + "show receivers", BaseEnv.TREE_SQL_DIALECT, receiverDataNodeId); + assertOldSenderReceiverRuntime( + "select * from information_schema.receivers", + BaseEnv.TABLE_SQL_DIALECT, + receiverDataNodeId); + } finally { + client.close(); + } + + assertNoThriftReceiverRuntimeOnDataNode( + "show receivers", BaseEnv.TREE_SQL_DIALECT, receiverDataNodeId); + assertNoThriftReceiverRuntimeOnDataNode( + "select * from information_schema.receivers", + BaseEnv.TABLE_SQL_DIALECT, + receiverDataNodeId); + } + + @Test + public void testReceiverRuntimeClearedAfterDataNodeRestartAndCanReconnect() throws Exception { + Assert.assertTrue(env.getDataNodeWrapperList().size() >= 2); + final int restartedDataNodeIndex = 0; + final DataNodeWrapper restartedDataNode = env.getDataNodeWrapper(restartedDataNodeIndex); + final int restartedDataNodeId = getDataNodeId(restartedDataNode); + + IoTDBSyncClient client = createPipeTransferClient(restartedDataNode); + try { + sendOldDataNodeHandshake(client); + assertOldSenderReceiverRuntime( + "show receivers", BaseEnv.TREE_SQL_DIALECT, restartedDataNodeId); + + env.shutdownDataNode(restartedDataNodeIndex); + } finally { + client.close(); + } + + env.ensureNodeStatus( + Collections.singletonList(restartedDataNode), + Collections.singletonList(NodeStatus.Unknown)); + assertNoThriftReceiverRuntimeOnDataNode( + "show receivers", BaseEnv.TREE_SQL_DIALECT, restartedDataNodeId); + + env.startDataNode(restartedDataNodeIndex); + env.ensureNodeStatus( + Collections.singletonList(restartedDataNode), + Collections.singletonList(NodeStatus.Running)); + + client = createPipeTransferClient(restartedDataNode); + try { + sendOldDataNodeHandshake(client); + assertOldSenderReceiverRuntime( + "select * from information_schema.receivers", + BaseEnv.TABLE_SQL_DIALECT, + restartedDataNodeId); + } finally { + client.close(); + } + } + + private void createWriteBackPipe(final String database, final String pipeName) { + TestUtils.executeNonQueries( + env, + Arrays.asList( + "create database " + database, + "create timeseries " + database + ".d1.s1 with datatype=INT32, encoding=PLAIN", + "create pipe " + + pipeName + + " with source ('pattern'='" + + database + + "') with sink ('sink'='write-back-sink')", + "insert into " + database + ".d1(time, s1) values (1, 1)", + "flush"), + null); + } + + private void createWriteBackPipeWithSinkUser( + final String database, + final String pipeName, + final String sinkUser, + final String sinkPassword) { + TestUtils.executeNonQueries( + env, + Arrays.asList( + "create database " + database, + "create timeseries " + database + ".d1.s1 with datatype=INT32, encoding=PLAIN", + "create pipe " + + pipeName + + " with source ('pattern'='" + + database + + "') with sink ('sink'='write-back-sink', 'user'='" + + sinkUser + + "', 'password'='" + + sinkPassword + + "')", + "insert into " + database + ".d1(time, s1) values (1, 1)", + "flush"), + null); + } + + private void createTwoWriteBackPipes( + final String database, final String firstPipeName, final String secondPipeName) { + TestUtils.executeNonQueries( + env, + Arrays.asList( + "create database " + database, + "create timeseries " + database + ".d1.s1 with datatype=INT32, encoding=PLAIN", + "create pipe " + + firstPipeName + + " with source ('pattern'='" + + database + + "') with sink ('sink'='write-back-sink')", + "create pipe " + + secondPipeName + + " with source ('pattern'='" + + database + + "') with sink ('sink'='write-back-sink')", + "insert into " + database + ".d1(time, s1) values (1, 1)", + "flush"), + null); + } + + private void assertShowReceivers(final String sql, final String sqlDialect) { + assertShowReceivers(sql, sqlDialect, "show_receivers_pipe"); + } + + private void assertShowReceivers( + final String sql, final String sqlDialect, final String pipeName) { + Awaitility.await() + .pollInSameThread() + .pollDelay(1L, TimeUnit.SECONDS) + .pollInterval(1L, TimeUnit.SECONDS) + .atMost(60L, TimeUnit.SECONDS) + .untilAsserted(() -> Assert.assertTrue(hasExpectedReceiver(sql, sqlDialect, pipeName))); + } + + private void assertShowReceiversAsUser( + final String sql, + final String sqlDialect, + final String pipeName, + final String userName, + final String password) { + assertShowReceiversAsUser(sql, sqlDialect, pipeName, userName, password, "root"); + } + + private void assertShowReceiversAsUser( + final String sql, + final String sqlDialect, + final String pipeName, + final String userName, + final String password, + final String expectedReceiverUserName) { + Awaitility.await() + .pollInSameThread() + .pollDelay(1L, TimeUnit.SECONDS) + .pollInterval(1L, TimeUnit.SECONDS) + .atMost(60L, TimeUnit.SECONDS) + .untilAsserted( + () -> + Assert.assertTrue( + hasExpectedReceiver( + sql, sqlDialect, pipeName, userName, password, expectedReceiverUserName))); + } + + private void assertAggregatedShowReceivers( + final String sql, + final String sqlDialect, + final String firstPipeName, + final String secondPipeName) { + Awaitility.await() + .pollInSameThread() + .pollDelay(1L, TimeUnit.SECONDS) + .pollInterval(1L, TimeUnit.SECONDS) + .atMost(60L, TimeUnit.SECONDS) + .untilAsserted( + () -> + Assert.assertTrue( + hasAggregatedReceiver(sql, sqlDialect, firstPipeName, secondPipeName))); + } + + private void assertProjectedShowReceivers(final String pipeName) { + Awaitility.await() + .pollInSameThread() + .pollDelay(1L, TimeUnit.SECONDS) + .pollInterval(1L, TimeUnit.SECONDS) + .atMost(60L, TimeUnit.SECONDS) + .untilAsserted(() -> Assert.assertTrue(hasProjectedReceiver(pipeName))); + } + + private void assertUserCannotSeeReceivers( + final String sql, final String sqlDialect, final String userName, final String password) { + Awaitility.await() + .pollInSameThread() + .pollDelay(1L, TimeUnit.SECONDS) + .pollInterval(1L, TimeUnit.SECONDS) + .atMost(60L, TimeUnit.SECONDS) + .untilAsserted( + () -> Assert.assertFalse(hasAnyReceiver(sql, sqlDialect, userName, password))); + } + + private IoTDBSyncClient createPipeTransferClient(final DataNodeWrapper dataNodeWrapper) + throws Exception { + return new IoTDBSyncClient( + new ThriftClientProperty.Builder().build(), + dataNodeWrapper.getIp(), + dataNodeWrapper.getPort(), + false, + null, + null); + } + + private void sendOldDataNodeHandshake(final IoTDBSyncClient client) throws Exception { + final TPipeTransferResp resp = + client.pipeTransfer( + PipeTransferDataNodeHandshakeV1Req.toTPipeTransferReq( + CommonDescriptor.getInstance().getConfig().getTimestampPrecision())); + Assert.assertEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), resp.getStatus().getCode()); + } + + private void assertOldSenderReceiverRuntime( + final String sql, final String sqlDialect, final int receiverDataNodeId) { + Awaitility.await() + .pollInSameThread() + .pollDelay(1L, TimeUnit.SECONDS) + .pollInterval(1L, TimeUnit.SECONDS) + .atMost(60L, TimeUnit.SECONDS) + .untilAsserted( + () -> + Assert.assertTrue( + hasOldSenderReceiverRuntime(sql, sqlDialect, receiverDataNodeId))); + } + + private boolean hasOldSenderReceiverRuntime( + final String sql, final String sqlDialect, final int receiverDataNodeId) throws SQLException { + try (final Connection connection = + getAvailableReceiverQueryConnection( + SessionConfig.DEFAULT_USER, SessionConfig.DEFAULT_PASSWORD, sqlDialect); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery(sql)) { + while (resultSet.next()) { + if ("DataNode".equals(getString(resultSet, "ReceiverNodeType", "receiver_node_type")) + && receiverDataNodeId == getInt(resultSet, "ReceiverNodeId", "receiver_node_id") + && "thrift".equals(getString(resultSet, "Protocol", "protocol")) + && "Unknown".equals(getString(resultSet, "SenderClusterId", "sender_cluster_id")) + && SessionConfig.DEFAULT_USER.equals(getString(resultSet, "UserName", "user_name")) + && getString(resultSet, "SenderAddress", "sender_address") != null + && getString(resultSet, "SenderPorts", "sender_ports") != null + && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 + && getInt(resultSet, "PipeCount", "pipe_count") == 0 + && "Unknown".equals(getString(resultSet, "PipeIDs", "pipe_ids")) + && getString(resultSet, "LastHandshakeTime", "last_handshake_time") != null + && getString(resultSet, "LastTransferTime", "last_transfer_time") != null) { + return true; + } + } + return false; + } + } + + private void assertNoThriftReceiverRuntimeOnDataNode( + final String sql, final String sqlDialect, final int receiverDataNodeId) { + Awaitility.await() + .pollInSameThread() + .pollDelay(1L, TimeUnit.SECONDS) + .pollInterval(1L, TimeUnit.SECONDS) + .atMost(60L, TimeUnit.SECONDS) + .untilAsserted( + () -> + Assert.assertFalse( + hasThriftReceiverRuntimeOnDataNode(sql, sqlDialect, receiverDataNodeId))); + } + + private boolean hasThriftReceiverRuntimeOnDataNode( + final String sql, final String sqlDialect, final int receiverDataNodeId) throws SQLException { + try (final Connection connection = + getAvailableReceiverQueryConnection( + SessionConfig.DEFAULT_USER, SessionConfig.DEFAULT_PASSWORD, sqlDialect); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery(sql)) { + while (resultSet.next()) { + if ("DataNode".equals(getString(resultSet, "ReceiverNodeType", "receiver_node_type")) + && receiverDataNodeId == getInt(resultSet, "ReceiverNodeId", "receiver_node_id") + && "thrift".equals(getString(resultSet, "Protocol", "protocol"))) { + return true; + } + } + return false; + } + } + + private boolean hasExpectedReceiver( + final String sql, final String sqlDialect, final String pipeName) throws SQLException { + return hasExpectedReceiver( + sql, sqlDialect, pipeName, SessionConfig.DEFAULT_USER, SessionConfig.DEFAULT_PASSWORD); + } + + private boolean hasExpectedReceiver( + final String sql, + final String sqlDialect, + final String pipeName, + final String userName, + final String password) + throws SQLException { + return hasExpectedReceiver(sql, sqlDialect, pipeName, userName, password, "root"); + } + + private boolean hasExpectedReceiver( + final String sql, + final String sqlDialect, + final String pipeName, + final String userName, + final String password, + final String expectedReceiverUserName) + throws SQLException { + try (final Connection connection = getReceiverQueryConnection(userName, password, sqlDialect); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery(sql)) { + while (resultSet.next()) { + final String senderClusterId = getString(resultSet, "SenderClusterId", "sender_cluster_id"); + final String senderAddress = getString(resultSet, "SenderAddress", "sender_address"); + final String userNameInRow = getString(resultSet, "UserName", "user_name"); + final String senderPorts = getString(resultSet, "SenderPorts", "sender_ports"); + final String pipeIds = getString(resultSet, "PipeIDs", "pipe_ids"); + if ("DataNode".equals(getString(resultSet, "ReceiverNodeType", "receiver_node_type")) + && "writeback".equals(getString(resultSet, "Protocol", "protocol")) + && senderClusterId != null + && !senderClusterId.isEmpty() + && senderAddress != null + && !senderAddress.isEmpty() + && expectedReceiverUserName.equals(userNameInRow) + && senderPorts != null + && !senderPorts.isEmpty() + && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 + && getInt(resultSet, "PipeCount", "pipe_count") >= 1 + && pipeIds != null + && pipeIds.contains(pipeName + "@") + && getString(resultSet, "LastHandshakeTime", "last_handshake_time") != null + && getString(resultSet, "LastTransferTime", "last_transfer_time") != null) { + return true; + } + } + return false; + } + } + + private boolean hasAggregatedReceiver( + final String sql, + final String sqlDialect, + final String firstPipeName, + final String secondPipeName) + throws SQLException { + try (final Connection connection = + getReceiverQueryConnection( + SessionConfig.DEFAULT_USER, SessionConfig.DEFAULT_PASSWORD, sqlDialect); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery(sql)) { + while (resultSet.next()) { + final String pipeIds = getString(resultSet, "PipeIDs", "pipe_ids"); + if ("DataNode".equals(getString(resultSet, "ReceiverNodeType", "receiver_node_type")) + && "writeback".equals(getString(resultSet, "Protocol", "protocol")) + && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 + && getInt(resultSet, "PipeCount", "pipe_count") >= 2 + && pipeIds != null + && pipeIds.contains(firstPipeName + "@") + && pipeIds.contains(secondPipeName + "@")) { + return true; + } + } + return false; + } + } + + private boolean hasProjectedReceiver(final String pipeName) throws SQLException { + try (final Connection connection = + getReceiverQueryConnection( + SessionConfig.DEFAULT_USER, + SessionConfig.DEFAULT_PASSWORD, + BaseEnv.TABLE_SQL_DIALECT); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = + statement.executeQuery( + "select receiver_node_type, sender_cluster_id, sender_address, " + + "connection_count, pipe_ids " + + "from information_schema.receivers where protocol = 'writeback'")) { + while (resultSet.next()) { + if ("DataNode".equals(resultSet.getString(1)) + && resultSet.getString(2) != null + && !resultSet.getString(2).isEmpty() + && resultSet.getString(3) != null + && !resultSet.getString(3).isEmpty() + && resultSet.getInt(4) >= 1 + && resultSet.getString(5) != null + && resultSet.getString(5).contains(pipeName + "@")) { + return true; + } + } + return false; + } + } + + private boolean hasAnyReceiver( + final String sql, final String sqlDialect, final String userName, final String password) + throws SQLException { + try (final Connection connection = getReceiverQueryConnection(userName, password, sqlDialect); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery(sql)) { + return resultSet.next(); + } + } + + private Connection getReceiverQueryConnection( + final String userName, final String password, final String sqlDialect) throws SQLException { + // Receiver runtime rows contain volatile timestamps. A single coordinator still executes the + // product cluster-level plan, but avoids the test wrapper comparing independently collected + // snapshots from multiple client entry points. + return env.getConnection(env.getDataNodeWrapper(0), userName, password, sqlDialect); + } + + private Connection getAvailableReceiverQueryConnection( + final String userName, final String password, final String sqlDialect) throws SQLException { + for (final DataNodeWrapper dataNodeWrapper : env.getDataNodeWrapperList()) { + if (dataNodeWrapper.isAlive()) { + return env.getConnection(dataNodeWrapper, userName, password, sqlDialect); + } + } + return getReceiverQueryConnection(userName, password, sqlDialect); + } + + private int getDataNodeId(final DataNodeWrapper targetDataNode) throws Exception { + try (final SyncConfigNodeIServiceClient client = + (SyncConfigNodeIServiceClient) env.getLeaderConfigNodeConnection()) { + final TShowDataNodesResp response = client.showDataNodes(); + Assert.assertEquals( + TSStatusCode.SUCCESS_STATUS.getStatusCode(), response.getStatus().getCode()); + for (final TDataNodeInfo dataNodeInfo : response.getDataNodesInfoList()) { + if (targetDataNode.getIp().equals(dataNodeInfo.getRpcAddresss()) + && targetDataNode.getPort() == dataNodeInfo.getRpcPort()) { + return dataNodeInfo.getDataNodeId(); + } + } + } + throw new AssertionError("Cannot find DataNodeId for " + targetDataNode.getIpAndPortString()); + } + + private static String getString( + final ResultSet resultSet, final String treeColumnName, final String tableColumnName) + throws SQLException { + try { + return resultSet.getString(treeColumnName); + } catch (final SQLException ignored) { + return resultSet.getString(tableColumnName); + } + } + + private static int getInt( + final ResultSet resultSet, final String treeColumnName, final String tableColumnName) + throws SQLException { + try { + return resultSet.getInt(treeColumnName); + } catch (final SQLException ignored) { + return resultSet.getInt(tableColumnName); + } + } + + private void assertShowReceiversWithoutDataNode( + final String sql, + final String sqlDialect, + final DataNodeWrapper queryDataNode, + final int excludedDataNodeId) { + Awaitility.await() + .pollInSameThread() + .pollDelay(1L, TimeUnit.SECONDS) + .pollInterval(1L, TimeUnit.SECONDS) + .atMost(60L, TimeUnit.SECONDS) + .untilAsserted( + () -> + assertQueryResultDoesNotContainDataNode( + sql, sqlDialect, queryDataNode, excludedDataNodeId)); + } + + private void assertQueryResultDoesNotContainDataNode( + final String sql, + final String sqlDialect, + final DataNodeWrapper queryDataNode, + final int excludedDataNodeId) + throws SQLException { + try (final Connection connection = + env.getConnection( + queryDataNode, + SessionConfig.DEFAULT_USER, + SessionConfig.DEFAULT_PASSWORD, + sqlDialect); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery(sql)) { + while (resultSet.next()) { + final int receiverDataNodeId = resultSet.getInt(2); + if (!resultSet.wasNull()) { + Assert.assertNotEquals(excludedDataNodeId, receiverDataNodeId); + } + } + } + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java new file mode 100644 index 0000000000000..0712387775b9d --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/IoTDBShowReceiversTripleIT.java @@ -0,0 +1,287 @@ +/* + * 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.pipe.it.triple; + +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.client.sync.SyncConfigNodeIServiceClient; +import org.apache.iotdb.confignode.rpc.thrift.TCreatePipeReq; +import org.apache.iotdb.confignode.rpc.thrift.TDataNodeInfo; +import org.apache.iotdb.confignode.rpc.thrift.TShowDataNodesResp; +import org.apache.iotdb.consensus.ConsensusFactory; +import org.apache.iotdb.db.it.utils.TestUtils; +import org.apache.iotdb.isession.SessionConfig; +import org.apache.iotdb.it.env.MultiEnvFactory; +import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.itbase.category.MultiClusterIT3; +import org.apache.iotdb.itbase.env.BaseEnv; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.awaitility.Awaitility; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.TimeUnit; + +@RunWith(IoTDBTestRunner.class) +@Category({MultiClusterIT3.class}) +public class IoTDBShowReceiversTripleIT { + + private BaseEnv senderEnv1; + private BaseEnv senderEnv2; + private BaseEnv receiverEnv; + + @Before + public void setUp() { + MultiEnvFactory.createEnv(3); + senderEnv1 = MultiEnvFactory.getEnv(0); + senderEnv2 = MultiEnvFactory.getEnv(1); + receiverEnv = MultiEnvFactory.getEnv(2); + + setupConfig(senderEnv1); + setupConfig(senderEnv2); + setupConfig(receiverEnv); + + senderEnv1.initClusterEnvironment(1, 1); + senderEnv2.initClusterEnvironment(1, 1); + receiverEnv.initClusterEnvironment(1, 3); + } + + private static void setupConfig(final BaseEnv env) { + env.getConfig() + .getCommonConfig() + .setAutoCreateSchemaEnabled(true) + .setConfigNodeConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS) + .setSchemaRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS) + .setEnforceStrongPassword(false) + .setPipeMemoryManagementEnabled(false) + .setIsPipeEnableMemoryCheck(false) + .setPipeAutoSplitFullEnabled(false); + env.getConfig().getCommonConfig().setDnConnectionTimeoutMs(600000); + } + + @After + public void tearDown() { + senderEnv1.cleanClusterEnvironment(); + senderEnv2.cleanClusterEnvironment(); + receiverEnv.cleanClusterEnvironment(); + } + + @Test + public void testShowReceiversAggregatesMultipleSendersOnMultipleDataNodes() throws Exception { + Assert.assertTrue(receiverEnv.getDataNodeWrapperList().size() >= 2); + final DataNodeWrapper receiverDataNode1 = receiverEnv.getDataNodeWrapper(0); + final DataNodeWrapper receiverDataNode2 = receiverEnv.getDataNodeWrapper(1); + final int receiverDataNodeId1 = getDataNodeId(receiverDataNode1); + final int receiverDataNodeId2 = getDataNodeId(receiverDataNode2); + final String pipeName1 = "show_receivers_sender_a_pipe"; + final String pipeName2 = "show_receivers_sender_b_pipe"; + final String database1 = "root.show_receivers_sender_a"; + final String database2 = "root.show_receivers_sender_b"; + + createThriftPipe(senderEnv1, receiverDataNode1, database1, pipeName1); + createThriftPipe(senderEnv2, receiverDataNode2, database2, pipeName2); + + TestUtils.assertDataEventuallyOnEnv( + receiverEnv, + "select count(s1) from " + database1 + ".d1", + "count(" + database1 + ".d1.s1),", + Collections.singleton("1,")); + TestUtils.assertDataEventuallyOnEnv( + receiverEnv, + "select count(s1) from " + database2 + ".d1", + "count(" + database2 + ".d1.s1),", + Collections.singleton("1,")); + + assertShowReceiversContainPipesOnDifferentDataNodes( + "show receivers", + BaseEnv.TREE_SQL_DIALECT, + pipeName1, + receiverDataNodeId1, + pipeName2, + receiverDataNodeId2); + assertShowReceiversContainPipesOnDifferentDataNodes( + "select * from information_schema.receivers", + BaseEnv.TABLE_SQL_DIALECT, + pipeName1, + receiverDataNodeId1, + pipeName2, + receiverDataNodeId2); + } + + private void createThriftPipe( + final BaseEnv senderEnv, + final DataNodeWrapper receiverDataNode, + final String database, + final String pipeName) + throws Exception { + TestUtils.executeNonQueries( + senderEnv, + Arrays.asList( + "create database " + database, + "create timeseries " + database + ".d1.s1 with datatype=INT32, encoding=PLAIN", + "insert into " + database + ".d1(time, s1) values (1, 1)", + "flush"), + null); + + final Map sourceAttributes = new HashMap<>(); + final Map processorAttributes = new HashMap<>(); + final Map sinkAttributes = new HashMap<>(); + + sourceAttributes.put("source.path", database + ".**"); + sourceAttributes.put("source.inclusion", "data.insert"); + sourceAttributes.put("user", SessionConfig.DEFAULT_USER); + + sinkAttributes.put("sink", "iotdb-thrift-sink"); + sinkAttributes.put("sink.batch.enable", "false"); + sinkAttributes.put("sink.ip", receiverDataNode.getIp()); + sinkAttributes.put("sink.port", Integer.toString(receiverDataNode.getPort())); + + try (final SyncConfigNodeIServiceClient client = + (SyncConfigNodeIServiceClient) senderEnv.getLeaderConfigNodeConnection()) { + TSStatus status = + client.createPipe( + new TCreatePipeReq(pipeName, sinkAttributes) + .setExtractorAttributes(sourceAttributes) + .setProcessorAttributes(processorAttributes)); + Assert.assertEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), status.getCode()); + + status = client.startPipe(pipeName); + Assert.assertEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), status.getCode()); + } + } + + private void assertShowReceiversContainPipesOnDifferentDataNodes( + final String sql, + final String sqlDialect, + final String firstPipeName, + final int firstReceiverDataNodeId, + final String secondPipeName, + final int secondReceiverDataNodeId) { + Awaitility.await() + .pollInSameThread() + .pollDelay(1L, TimeUnit.SECONDS) + .pollInterval(1L, TimeUnit.SECONDS) + .atMost(60L, TimeUnit.SECONDS) + .untilAsserted( + () -> + Assert.assertTrue( + hasReceiversOnDifferentDataNodes( + sql, + sqlDialect, + firstPipeName, + firstReceiverDataNodeId, + secondPipeName, + secondReceiverDataNodeId))); + } + + private boolean hasReceiversOnDifferentDataNodes( + final String sql, + final String sqlDialect, + final String firstPipeName, + final int firstReceiverDataNodeId, + final String secondPipeName, + final int secondReceiverDataNodeId) + throws SQLException { + boolean hasFirstPipe = false; + boolean hasSecondPipe = false; + final Set receiverDataNodeIds = new HashSet<>(); + try (final Connection connection = + receiverEnv.getConnection( + SessionConfig.DEFAULT_USER, SessionConfig.DEFAULT_PASSWORD, sqlDialect); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery(sql)) { + while (resultSet.next()) { + final String pipeIds = getString(resultSet, "PipeIDs", "pipe_ids"); + if (!"DataNode".equals(getString(resultSet, "ReceiverNodeType", "receiver_node_type")) + || !"thrift".equals(getString(resultSet, "Protocol", "protocol")) + || pipeIds == null) { + continue; + } + final int receiverDataNodeId = getInt(resultSet, "ReceiverNodeId", "receiver_node_id"); + if (receiverDataNodeId == firstReceiverDataNodeId + && pipeIds.contains(firstPipeName + "@") + && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 + && getInt(resultSet, "PipeCount", "pipe_count") >= 1) { + hasFirstPipe = true; + receiverDataNodeIds.add(receiverDataNodeId); + } + if (receiverDataNodeId == secondReceiverDataNodeId + && pipeIds.contains(secondPipeName + "@") + && getInt(resultSet, "ConnectionCount", "connection_count") >= 1 + && getInt(resultSet, "PipeCount", "pipe_count") >= 1) { + hasSecondPipe = true; + receiverDataNodeIds.add(receiverDataNodeId); + } + } + } + return hasFirstPipe && hasSecondPipe && receiverDataNodeIds.size() >= 2; + } + + private int getDataNodeId(final DataNodeWrapper targetDataNode) throws Exception { + try (final SyncConfigNodeIServiceClient client = + (SyncConfigNodeIServiceClient) receiverEnv.getLeaderConfigNodeConnection()) { + final TShowDataNodesResp response = client.showDataNodes(); + Assert.assertEquals( + TSStatusCode.SUCCESS_STATUS.getStatusCode(), response.getStatus().getCode()); + for (final TDataNodeInfo dataNodeInfo : response.getDataNodesInfoList()) { + if (targetDataNode.getIp().equals(dataNodeInfo.getRpcAddresss()) + && targetDataNode.getPort() == dataNodeInfo.getRpcPort()) { + return dataNodeInfo.getDataNodeId(); + } + } + } + throw new AssertionError("Cannot find DataNodeId for " + targetDataNode.getIpAndPortString()); + } + + private static String getString( + final ResultSet resultSet, final String treeColumnName, final String tableColumnName) + throws SQLException { + try { + return resultSet.getString(treeColumnName); + } catch (final SQLException ignored) { + return resultSet.getString(tableColumnName); + } + } + + private static int getInt( + final ResultSet resultSet, final String treeColumnName, final String tableColumnName) + throws SQLException { + try { + return resultSet.getInt(treeColumnName); + } catch (final SQLException ignored) { + return resultSet.getInt(tableColumnName); + } + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/schema/IoTDBDatabaseIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/schema/IoTDBDatabaseIT.java index dac15d8b66bf4..01f2d47e8559e 100644 --- a/integration-test/src/test/java/org/apache/iotdb/relational/it/schema/IoTDBDatabaseIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/schema/IoTDBDatabaseIT.java @@ -419,6 +419,7 @@ public void testInformationSchema() throws SQLException { "pipes,INF,", "queries,INF,", "queries_costs_histogram,INF,", + "receivers,INF,", "regions,INF,", "services,INF,", "subscriptions,INF,", @@ -492,6 +493,23 @@ public void testInformationSchema() throws SQLException { "exception_message,STRING,ATTRIBUTE,", "remaining_event_count,INT64,ATTRIBUTE,", "estimated_remaining_seconds,DOUBLE,ATTRIBUTE,"))); + TestUtils.assertResultSetEqual( + statement.executeQuery("desc receivers"), + "ColumnName,DataType,Category,", + new HashSet<>( + Arrays.asList( + "receiver_node_type,STRING,TAG,", + "receiver_node_id,INT32,TAG,", + "protocol,STRING,TAG,", + "sender_address,STRING,TAG,", + "sender_ports,STRING,ATTRIBUTE,", + "connection_count,INT32,ATTRIBUTE,", + "pipe_count,INT32,ATTRIBUTE,", + "pipe_ids,STRING,ATTRIBUTE,", + "user_name,STRING,TAG,", + "sender_cluster_id,STRING,TAG,", + "last_handshake_time,TIMESTAMP,ATTRIBUTE,", + "last_transfer_time,TIMESTAMP,ATTRIBUTE,"))); TestUtils.assertResultSetEqual( statement.executeQuery("desc pipe_plugins"), "ColumnName,DataType,Category,", @@ -690,6 +708,7 @@ public void testInformationSchema() throws SQLException { "information_schema,topics,INF,USING,null,SYSTEM VIEW,", "information_schema,pipe_plugins,INF,USING,null,SYSTEM VIEW,", "information_schema,pipes,INF,USING,null,SYSTEM VIEW,", + "information_schema,receivers,INF,USING,null,SYSTEM VIEW,", "information_schema,services,INF,USING,null,SYSTEM VIEW,", "information_schema,subscriptions,INF,USING,null,SYSTEM VIEW,", "information_schema,views,INF,USING,null,SYSTEM VIEW,", @@ -708,7 +727,7 @@ public void testInformationSchema() throws SQLException { TestUtils.assertResultSetEqual( statement.executeQuery("count devices from tables where status = 'USING'"), "count(devices),", - Collections.singleton("23,")); + Collections.singleton("24,")); TestUtils.assertResultSetEqual( statement.executeQuery( "select * from columns where table_name = 'queries' or database = 'test'"), diff --git a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IdentifierParser.g4 b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IdentifierParser.g4 index e96d84e892c8c..1477a67166483 100644 --- a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IdentifierParser.g4 +++ b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IdentifierParser.g4 @@ -196,6 +196,7 @@ keyWords | QUERY | QUERYID | QUOTA + | RECEIVERS | RANGE | READONLY | READ @@ -298,4 +299,4 @@ keyWords | OPTION | INF | CURRENT_TIMESTAMP - ; \ No newline at end of file + ; diff --git a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4 b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4 index 2f91b573aef0e..1a9c7278cca30 100644 --- a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4 +++ b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4 @@ -55,7 +55,7 @@ ddlStatement // ExternalService | createService | startService | stopService | dropService | showService // Pipe Task - | createPipe | alterPipe | dropPipe | startPipe | stopPipe | showPipes + | createPipe | alterPipe | dropPipe | startPipe | stopPipe | showPipes | showReceivers // Pipe Plugin | createPipePlugin | dropPipePlugin | showPipePlugins // Subscription @@ -701,6 +701,10 @@ showPipes : SHOW ((PIPE pipeName=identifier) | PIPES (WHERE (CONNECTOR | SINK) USED BY pipeName=identifier)?) ; +showReceivers + : SHOW RECEIVERS + ; + // Pipe Plugin ========================================================================================= createPipePlugin : CREATE PIPEPLUGIN (IF NOT EXISTS)? pluginName=identifier AS className=STRING_LITERAL uriClause diff --git a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/SqlLexer.g4 b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/SqlLexer.g4 index 3950de50ea477..1a440e45e0dc4 100644 --- a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/SqlLexer.g4 +++ b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/SqlLexer.g4 @@ -723,6 +723,10 @@ QUOTA : Q U O T A ; +RECEIVERS + : R E C E I V E R S + ; + RANGE : R A N G E ; @@ -1411,4 +1415,4 @@ fragment V: [vV]; fragment W: [wW]; fragment X: [xX]; fragment Y: [yY]; -fragment Z: [zZ]; \ No newline at end of file +fragment Z: [zZ]; diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/receiver/protocol/IoTDBConfigNodeReceiver.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/receiver/protocol/IoTDBConfigNodeReceiver.java index 5f414f2cd69fd..a6e1da1954188 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/receiver/protocol/IoTDBConfigNodeReceiver.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/receiver/protocol/IoTDBConfigNodeReceiver.java @@ -35,11 +35,13 @@ import org.apache.iotdb.commons.pipe.datastructure.pattern.TreePattern; import org.apache.iotdb.commons.pipe.receiver.IoTDBFileReceiver; import org.apache.iotdb.commons.pipe.receiver.PipeReceiverStatusHandler; +import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeRegistry; import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapPseudoTPipeTransferRequest; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeRequestType; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferCompressedReq; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferFileSealReqV1; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferFileSealReqV2; +import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferPipeReceiverRuntimeInfoCleanupReq; import org.apache.iotdb.commons.schema.column.ColumnHeaderConstant; import org.apache.iotdb.commons.schema.table.Audit; import org.apache.iotdb.commons.schema.table.TreeViewSchema; @@ -194,20 +196,22 @@ public TPipeTransferResp receive(final TPipeTransferReq req) { PipeTransferConfigNodeHandshakeV1Req.fromTPipeTransferReq(req)); PipeConfigNodeReceiverMetrics.getInstance() .recordHandshakeConfigNodeV1Timer(System.nanoTime() - startTime); - return resp; + return recordConfigNodeHandshakeIfSuccess(resp, req); case HANDSHAKE_CONFIGNODE_V2: resp = handleTransferHandshakeV2( PipeTransferConfigNodeHandshakeV2Req.fromTPipeTransferReq(req)); - userEntity.setAuditLogOperation(AuditLogOperation.DDL); + if (Objects.nonNull(userEntity)) { + userEntity.setAuditLogOperation(AuditLogOperation.DDL); + } PipeConfigNodeReceiverMetrics.getInstance() .recordHandshakeConfigNodeV2Timer(System.nanoTime() - startTime); - return resp; + return recordConfigNodeHandshakeIfSuccess(resp, req); case TRANSFER_CONFIG_PLAN: resp = handleTransferConfigPlan(PipeTransferConfigPlanReq.fromTPipeTransferReq(req)); PipeConfigNodeReceiverMetrics.getInstance() .recordTransferConfigPlanTimer(System.nanoTime() - startTime); - return resp; + return recordConfigNodeTransferIfSuccess(resp); case TRANSFER_CONFIG_SNAPSHOT_PIECE: resp = handleTransferFilePiece( @@ -216,16 +220,23 @@ public TPipeTransferResp receive(final TPipeTransferReq req) { false); PipeConfigNodeReceiverMetrics.getInstance() .recordTransferConfigSnapshotPieceTimer(System.nanoTime() - startTime); - return resp; + return recordConfigNodeTransferIfSuccess(resp); case TRANSFER_CONFIG_SNAPSHOT_SEAL: resp = handleTransferFileSealV2( PipeTransferConfigSnapshotSealReq.fromTPipeTransferReq(req)); PipeConfigNodeReceiverMetrics.getInstance() .recordTransferConfigSnapshotSealTimer(System.nanoTime() - startTime); - return resp; + return recordConfigNodeTransferIfSuccess(resp); case TRANSFER_COMPRESSED: return receive(PipeTransferCompressedReq.fromTPipeTransferReq(req)); + case TRANSFER_PIPE_RECEIVER_RUNTIME_INFO_CLEANUP: + final PipeTransferPipeReceiverRuntimeInfoCleanupReq cleanupReq = + PipeTransferPipeReceiverRuntimeInfoCleanupReq.fromTPipeTransferReq(req); + PipeReceiverRuntimeRegistry.getInstance() + .removePipeFromAllSessions( + cleanupReq.getPipeName(), cleanupReq.getPipeCreationTime()); + return new TPipeTransferResp(RpcUtils.SUCCESS_STATUS); default: break; } @@ -259,7 +270,32 @@ public TPipeTransferResp receive(final TPipeTransferReq req) { private boolean needHandshake(final PipeRequestType type) { return Objects.isNull(receiverFileDirWithIdSuffix.get()) && type != PipeRequestType.HANDSHAKE_CONFIGNODE_V1 - && type != PipeRequestType.HANDSHAKE_CONFIGNODE_V2; + && type != PipeRequestType.HANDSHAKE_CONFIGNODE_V2 + && type != PipeRequestType.TRANSFER_PIPE_RECEIVER_RUNTIME_INFO_CLEANUP; + } + + private TPipeTransferResp recordConfigNodeHandshakeIfSuccess( + final TPipeTransferResp resp, final TPipeTransferReq req) { + if (isSuccess(resp)) { + recordPipeReceiverHandshake( + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + ConfigNodeDescriptor.getInstance().getConf().getConfigNodeId(), + getProtocol(req)); + } + return resp; + } + + private TPipeTransferResp recordConfigNodeTransferIfSuccess(final TPipeTransferResp resp) { + if (isSuccess(resp)) { + recordPipeReceiverTransfer(); + } + return resp; + } + + private static String getProtocol(final TPipeTransferReq req) { + return req instanceof AirGapPseudoTPipeTransferRequest + ? PipeReceiverRuntimeRegistry.PROTOCOL_AIR_GAP + : PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT; } private TPipeTransferResp handleTransferConfigPlan(final PipeTransferConfigPlanReq req) diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/sink/protocol/IoTDBConfigRegionAirGapSink.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/sink/protocol/IoTDBConfigRegionAirGapSink.java index c8c2890f5be3e..f8e91993eeb90 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/sink/protocol/IoTDBConfigRegionAirGapSink.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/sink/protocol/IoTDBConfigRegionAirGapSink.java @@ -89,6 +89,7 @@ protected byte[] generateHandShakeV2Payload() throws IOException { Boolean.toString(shouldMarkAsPipeRequest)); params.put( PipeTransferHandshakeConstant.HANDSHAKE_KEY_SKIP_IF, Boolean.toString(skipIfNoPrivileges)); + appendPipeInfoToHandshakeParams(params); return PipeTransferConfigNodeHandshakeV2Req.toTPipeTransferBytes(params); } diff --git a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/pipe/sink/PipeConfigNodeThriftRequestTest.java b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/pipe/sink/PipeConfigNodeThriftRequestTest.java index dd948a451cf90..bd77d72d1bef1 100644 --- a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/pipe/sink/PipeConfigNodeThriftRequestTest.java +++ b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/pipe/sink/PipeConfigNodeThriftRequestTest.java @@ -55,11 +55,13 @@ public class PipeConfigNodeThriftRequestTest { public void testPipeTransferConfigHandshakeReq() throws IOException { PipeTransferConfigNodeHandshakeV1Req req = PipeTransferConfigNodeHandshakeV1Req.toTPipeTransferReq(TIME_PRECISION); + final int originalBodyPosition = req.body.position(); PipeTransferConfigNodeHandshakeV1Req deserializeReq = PipeTransferConfigNodeHandshakeV1Req.fromTPipeTransferReq(req); Assert.assertEquals(req.getVersion(), deserializeReq.getVersion()); Assert.assertEquals(req.getType(), deserializeReq.getType()); + Assert.assertEquals(originalBodyPosition, req.body.position()); Assert.assertEquals(req.getTimestampPrecision(), deserializeReq.getTimestampPrecision()); } @@ -83,17 +85,21 @@ public void testPipeTransferConfigHandshakeReqFromLegacyV13Body() throws IOExcep public void testPipeTransferConfigHandshakeV2Req() throws IOException { final Map params = new HashMap<>(); params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_TIME_PRECISION, TIME_PRECISION); - params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_CLUSTER_ID, "cluster"); + params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_CLUSTER_ID, "cluster-a"); params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_USERNAME, "root"); params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_PASSWORD, "root"); + params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_PIPE_NAME, "pipe-a"); + params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_PIPE_CREATION_TIME, "1"); final PipeTransferConfigNodeHandshakeV2Req req = PipeTransferConfigNodeHandshakeV2Req.toTPipeTransferReq(params); + final int originalBodyPosition = req.body.position(); final PipeTransferConfigNodeHandshakeV2Req deserializeReq = PipeTransferConfigNodeHandshakeV2Req.fromTPipeTransferReq(req); Assert.assertEquals(req.getVersion(), deserializeReq.getVersion()); Assert.assertEquals(req.getType(), deserializeReq.getType()); + Assert.assertEquals(originalBodyPosition, req.body.position()); Assert.assertEquals(params, deserializeReq.getParams()); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtask.java index 10b746778e35e..6ae72fc7724bc 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtask.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtask.java @@ -126,6 +126,9 @@ protected boolean executeOnce() { try { if (Objects.isNull(event)) { + if (shouldStopSubmittingSelf.get()) { + return false; + } transferHeartbeatEvent(CRON_HEARTBEAT_EVENT); return false; } @@ -272,6 +275,12 @@ && isEventFromPipe((EnrichedEvent) lastExceptionEvent, committerKey)) { } } + public void discardReceiverRuntimeSessions() { + if (outputPipeSink instanceof IoTDBSink) { + ((IoTDBSink) outputPipeSink).discardReceiverRuntimeSessions(); + } + } + private static boolean isEventFromPipe( final EnrichedEvent event, final CommitterKey committerKey) { return committerKey.getPipeName().equals(event.getPipeName()) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtaskLifeCycle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtaskLifeCycle.java index b49332346591d..d93021f68714f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtaskLifeCycle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/agent/task/subtask/sink/PipeSinkSubtaskLifeCycle.java @@ -147,6 +147,7 @@ public synchronized void stop() { if (runningTaskCount == 1) { executor.stop(subtask.getTaskID()); + subtask.discardReceiverRuntimeSessions(); } runningTaskCount--; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java index 978d924c05651..f1d3a90e11ef8 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java @@ -37,14 +37,17 @@ import org.apache.iotdb.commons.pipe.datastructure.pattern.TreePattern; import org.apache.iotdb.commons.pipe.receiver.IoTDBFileReceiver; import org.apache.iotdb.commons.pipe.receiver.PipeReceiverStatusHandler; +import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeRegistry; import org.apache.iotdb.commons.pipe.resource.log.PipeLogger; import org.apache.iotdb.commons.pipe.resource.log.PipePeriodicalLogReducer; import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapPseudoTPipeTransferRequest; +import org.apache.iotdb.commons.pipe.sink.payload.thrift.common.PipeTransferHandshakeConstant; import org.apache.iotdb.commons.pipe.sink.payload.thrift.common.PipeTransferSliceReqHandler; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeRequestType; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferCompressedReq; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferFileSealReqV1; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferFileSealReqV2; +import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferPipeReceiverRuntimeInfoCleanupReq; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferSliceReq; import org.apache.iotdb.commons.schema.column.ColumnHeaderConstant; import org.apache.iotdb.commons.utils.PathUtils; @@ -121,16 +124,19 @@ import com.google.common.util.concurrent.ListenableFuture; import org.apache.tsfile.utils.Pair; +import org.apache.tsfile.utils.ReadWriteIOUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.FileNotFoundException; import java.io.IOException; +import java.nio.ByteBuffer; import java.nio.file.Paths; import java.time.ZoneId; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.Locale; import java.util.Map; @@ -182,6 +188,8 @@ public class IoTDBDataNodeReceiver extends IoTDBFileReceiver { // datanode (cluster B). private static final AtomicLong CONFIG_RECEIVER_ID_GENERATOR = new AtomicLong(0); protected final AtomicReference configReceiverId = new AtomicReference<>(); + private final AtomicReference configPipeReceiverRuntimeSessionKey = + new AtomicReference<>(); private final PipeTransferSliceReqHandler sliceReqHandler = new PipeTransferSliceReqHandler(); @@ -233,8 +241,10 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { TSStatusCode.PIPE_HANDSHAKE_ERROR.getStatusCode(), "The receiver memory is not enough to handle the handshake request from datanode.")); } - return handleTransferHandshakeV1( - PipeTransferDataNodeHandshakeV1Req.fromTPipeTransferReq(req)); + return recordDataNodeHandshakeIfSuccess( + handleTransferHandshakeV1( + PipeTransferDataNodeHandshakeV1Req.fromTPipeTransferReq(req)), + req); } finally { PipeDataNodeReceiverMetrics.getInstance() .recordHandshakeDatanodeV1Timer(System.nanoTime() - startTime); @@ -251,8 +261,10 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { TSStatusCode.PIPE_HANDSHAKE_ERROR.getStatusCode(), "The receiver memory is not enough to handle the handshake request from datanode.")); } - return handleTransferHandshakeV2( - PipeTransferDataNodeHandshakeV2Req.fromTPipeTransferReq(req)); + return recordDataNodeHandshakeIfSuccess( + handleTransferHandshakeV2( + PipeTransferDataNodeHandshakeV2Req.fromTPipeTransferReq(req)), + req); } finally { PipeDataNodeReceiverMetrics.getInstance() .recordHandshakeDatanodeV2Timer(System.nanoTime() - startTime); @@ -261,8 +273,9 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { case TRANSFER_TABLET_INSERT_NODE: { try { - return handleTransferTabletInsertNode( - PipeTransferTabletInsertNodeReq.fromTPipeTransferReq(req)); + return recordDataNodeTransferIfSuccess( + handleTransferTabletInsertNode( + PipeTransferTabletInsertNodeReq.fromTPipeTransferReq(req))); } finally { PipeDataNodeReceiverMetrics.getInstance() @@ -272,8 +285,9 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { case TRANSFER_TABLET_INSERT_NODE_V2: { try { - return handleTransferTabletInsertNode( - PipeTransferTabletInsertNodeReqV2.fromTPipeTransferReq(req)); + return recordDataNodeTransferIfSuccess( + handleTransferTabletInsertNode( + PipeTransferTabletInsertNodeReqV2.fromTPipeTransferReq(req))); } finally { PipeDataNodeReceiverMetrics.getInstance() .recordTransferTabletInsertNodeV2Timer(System.nanoTime() - startTime); @@ -282,7 +296,8 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { case TRANSFER_TABLET_RAW: { try { - return handleTransferTabletRaw(PipeTransferTabletRawReq.fromTPipeTransferReq(req)); + return recordDataNodeTransferIfSuccess( + handleTransferTabletRaw(PipeTransferTabletRawReq.fromTPipeTransferReq(req))); } finally { PipeDataNodeReceiverMetrics.getInstance() .recordTransferTabletRawTimer(System.nanoTime() - startTime); @@ -291,8 +306,8 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { case TRANSFER_TABLET_RAW_V2: { try { - return handleTransferTabletRaw( - PipeTransferTabletRawReqV2.fromTPipeTransferReq(req)); + return recordDataNodeTransferIfSuccess( + handleTransferTabletRaw(PipeTransferTabletRawReqV2.fromTPipeTransferReq(req))); } finally { PipeDataNodeReceiverMetrics.getInstance() .recordTransferTabletRawV2Timer(System.nanoTime() - startTime); @@ -301,8 +316,9 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { case TRANSFER_TABLET_BINARY: { try { - return handleTransferTabletBinary( - PipeTransferTabletBinaryReq.fromTPipeTransferReq(req)); + return recordDataNodeTransferIfSuccess( + handleTransferTabletBinary( + PipeTransferTabletBinaryReq.fromTPipeTransferReq(req))); } finally { PipeDataNodeReceiverMetrics.getInstance() .recordTransferTabletBinaryTimer(System.nanoTime() - startTime); @@ -311,8 +327,9 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { case TRANSFER_TABLET_BINARY_V2: { try { - return handleTransferTabletBinary( - PipeTransferTabletBinaryReqV2.fromTPipeTransferReq(req)); + return recordDataNodeTransferIfSuccess( + handleTransferTabletBinary( + PipeTransferTabletBinaryReqV2.fromTPipeTransferReq(req))); } finally { PipeDataNodeReceiverMetrics.getInstance() .recordTransferTabletBinaryV2Timer(System.nanoTime() - startTime); @@ -321,8 +338,9 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { case TRANSFER_TABLET_BATCH: { try { - return handleTransferTabletBatch( - PipeTransferTabletBatchReq.fromTPipeTransferReq(req)); + return recordDataNodeTransferIfSuccess( + handleTransferTabletBatch( + PipeTransferTabletBatchReq.fromTPipeTransferReq(req))); } finally { PipeDataNodeReceiverMetrics.getInstance() .recordTransferTabletBatchTimer(System.nanoTime() - startTime); @@ -331,8 +349,9 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { case TRANSFER_TABLET_BATCH_V2: { try { - return handleTransferTabletBatchV2( - PipeTransferTabletBatchReqV2.fromTPipeTransferReq(req)); + return recordDataNodeTransferIfSuccess( + handleTransferTabletBatchV2( + PipeTransferTabletBatchReqV2.fromTPipeTransferReq(req))); } finally { PipeDataNodeReceiverMetrics.getInstance() .recordTransferTabletBatchV2Timer(System.nanoTime() - startTime); @@ -341,10 +360,11 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { case TRANSFER_TS_FILE_PIECE: { try { - return handleTransferFilePiece( - PipeTransferTsFilePieceReq.fromTPipeTransferReq(req), - req instanceof AirGapPseudoTPipeTransferRequest, - true); + return recordDataNodeTransferIfSuccess( + handleTransferFilePiece( + PipeTransferTsFilePieceReq.fromTPipeTransferReq(req), + req instanceof AirGapPseudoTPipeTransferRequest, + true)); } finally { PipeDataNodeReceiverMetrics.getInstance() .recordTransferTsFilePieceTimer(System.nanoTime() - startTime); @@ -353,8 +373,8 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { case TRANSFER_TS_FILE_SEAL: { try { - return handleTransferFileSealV1( - PipeTransferTsFileSealReq.fromTPipeTransferReq(req)); + return recordDataNodeTransferIfSuccess( + handleTransferFileSealV1(PipeTransferTsFileSealReq.fromTPipeTransferReq(req))); } finally { PipeDataNodeReceiverMetrics.getInstance() .recordTransferTsFileSealTimer(System.nanoTime() - startTime); @@ -363,10 +383,11 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { case TRANSFER_TS_FILE_PIECE_WITH_MOD: { try { - return handleTransferFilePiece( - PipeTransferTsFilePieceWithModReq.fromTPipeTransferReq(req), - req instanceof AirGapPseudoTPipeTransferRequest, - false); + return recordDataNodeTransferIfSuccess( + handleTransferFilePiece( + PipeTransferTsFilePieceWithModReq.fromTPipeTransferReq(req), + req instanceof AirGapPseudoTPipeTransferRequest, + false)); } finally { PipeDataNodeReceiverMetrics.getInstance() @@ -376,8 +397,9 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { case TRANSFER_TS_FILE_SEAL_WITH_MOD: { try { - return handleTransferFileSealV2( - PipeTransferTsFileSealWithModReq.fromTPipeTransferReq(req)); + return recordDataNodeTransferIfSuccess( + handleTransferFileSealV2( + PipeTransferTsFileSealWithModReq.fromTPipeTransferReq(req))); } finally { PipeDataNodeReceiverMetrics.getInstance() .recordTransferTsFileSealWithModTimer(System.nanoTime() - startTime); @@ -386,7 +408,8 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { case TRANSFER_PLAN_NODE: { try { - return handleTransferSchemaPlan(PipeTransferPlanNodeReq.fromTPipeTransferReq(req)); + return recordDataNodeTransferIfSuccess( + handleTransferSchemaPlan(PipeTransferPlanNodeReq.fromTPipeTransferReq(req))); } finally { PipeDataNodeReceiverMetrics.getInstance() .recordTransferSchemaPlanTimer(System.nanoTime() - startTime); @@ -395,10 +418,11 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { case TRANSFER_SCHEMA_SNAPSHOT_PIECE: { try { - return handleTransferFilePiece( - PipeTransferSchemaSnapshotPieceReq.fromTPipeTransferReq(req), - req instanceof AirGapPseudoTPipeTransferRequest, - false); + return recordDataNodeTransferIfSuccess( + handleTransferFilePiece( + PipeTransferSchemaSnapshotPieceReq.fromTPipeTransferReq(req), + req instanceof AirGapPseudoTPipeTransferRequest, + false)); } finally { PipeDataNodeReceiverMetrics.getInstance() @@ -408,8 +432,9 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { case TRANSFER_SCHEMA_SNAPSHOT_SEAL: { try { - return handleTransferFileSealV2( - PipeTransferSchemaSnapshotSealReq.fromTPipeTransferReq(req)); + return recordDataNodeTransferIfSuccess( + handleTransferFileSealV2( + PipeTransferSchemaSnapshotSealReq.fromTPipeTransferReq(req))); } finally { PipeDataNodeReceiverMetrics.getInstance() @@ -425,7 +450,7 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { try { // Config requests will first be received by the DataNode receiver, // then transferred to ConfigNode receiver to execute. - return handleTransferConfigPlan(req); + return recordConfigNodeReceiverRuntimeIfSuccess(handleTransferConfigPlan(req), req); } finally { PipeDataNodeReceiverMetrics.getInstance() .recordTransferConfigPlanTimer(System.nanoTime() - startTime); @@ -449,6 +474,15 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { .recordTransferCompressedTimer(System.nanoTime() - startTime); } } + case TRANSFER_PIPE_RECEIVER_RUNTIME_INFO_CLEANUP: + { + final PipeTransferPipeReceiverRuntimeInfoCleanupReq cleanupReq = + PipeTransferPipeReceiverRuntimeInfoCleanupReq.fromTPipeTransferReq(req); + PipeReceiverRuntimeRegistry.getInstance() + .removePipeFromAllSessions( + cleanupReq.getPipeName(), cleanupReq.getPipeCreationTime()); + return new TPipeTransferResp(RpcUtils.SUCCESS_STATUS); + } default: break; } @@ -804,9 +838,9 @@ private TPipeTransferResp handleTransferSchemaPlan(final PipeTransferPlanNodeReq null)); } - private TPipeTransferResp handleTransferConfigPlan(final TPipeTransferReq req) { + private Pair handleTransferConfigPlan(final TPipeTransferReq req) { return ClusterConfigTaskExecutor.getInstance() - .handleTransferConfigPlan(getConfigReceiverId(), req); + .handleTransferConfigPlanAndGetReceiverNodeId(getConfigReceiverId(), req); } /** Used to identify the sender client */ @@ -841,6 +875,122 @@ private TPipeTransferResp handleTransferSlice(final PipeTransferSliceReq pipeTra return receive(req.get()); } + private TPipeTransferResp recordDataNodeHandshakeIfSuccess( + final TPipeTransferResp resp, final TPipeTransferReq req) { + if (isSuccess(resp)) { + recordPipeReceiverHandshake( + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + IOTDB_CONFIG.getDataNodeId(), + getProtocol(req)); + } + return resp; + } + + private TPipeTransferResp recordDataNodeTransferIfSuccess(final TPipeTransferResp resp) { + if (isSuccess(resp)) { + recordPipeReceiverTransfer(); + } + return resp; + } + + private TPipeTransferResp recordConfigNodeReceiverRuntimeIfSuccess( + final Pair respWithReceiverNodeId, final TPipeTransferReq req) { + final TPipeTransferResp resp = respWithReceiverNodeId.left; + if (!PipeRequestType.isValidatedRequestType(req.getType())) { + return resp; + } + + final PipeRequestType requestType = PipeRequestType.valueOf(req.getType()); + if (requestType == PipeRequestType.HANDSHAKE_CONFIGNODE_V1 + || requestType == PipeRequestType.HANDSHAKE_CONFIGNODE_V2) { + if (isSuccess(resp)) { + recordConfigNodeHandshake(req, requestType, respWithReceiverNodeId.right); + } + } else { + if (isSuccess(resp)) { + PipeReceiverRuntimeRegistry.getInstance() + .markTransfer(configPipeReceiverRuntimeSessionKey.get(), System.currentTimeMillis()); + } + } + return resp; + } + + private void recordConfigNodeHandshake( + final TPipeTransferReq req, final PipeRequestType requestType, final int receiverNodeId) { + final String protocol = getProtocol(req); + final String sessionKey = + String.format( + "%s-%s-%s-%s", + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + receiverNodeId, + protocol, + getConfigReceiverId()); + final String oldSessionKey = configPipeReceiverRuntimeSessionKey.getAndSet(sessionKey); + if (!Objects.equals(oldSessionKey, sessionKey)) { + PipeReceiverRuntimeRegistry.getInstance().deregister(oldSessionKey); + } + + final Map params = + requestType == PipeRequestType.HANDSHAKE_CONFIGNODE_V2 + ? parseHandshakeV2Params(req) + : Collections.emptyMap(); + PipeReceiverRuntimeRegistry.getInstance() + .registerOrUpdateSession( + sessionKey, + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + receiverNodeId, + protocol, + getSenderHost(), + parseSenderPort(getSenderPort()), + params.getOrDefault(PipeTransferHandshakeConstant.HANDSHAKE_KEY_USERNAME, username), + params.getOrDefault( + PipeTransferHandshakeConstant.HANDSHAKE_KEY_CLUSTER_ID, + PipeReceiverRuntimeRegistry.UNKNOWN), + params.get(PipeTransferHandshakeConstant.HANDSHAKE_KEY_PIPE_NAME), + parsePipeCreationTime( + params.get(PipeTransferHandshakeConstant.HANDSHAKE_KEY_PIPE_CREATION_TIME)), + System.currentTimeMillis()); + } + + private static Map parseHandshakeV2Params(final TPipeTransferReq req) { + final Map params = new HashMap<>(); + if (req.getBody() == null) { + return params; + } + final ByteBuffer body = req.body.duplicate(); + body.rewind(); + final int size = ReadWriteIOUtils.readInt(body); + for (int i = 0; i < size; ++i) { + params.put(ReadWriteIOUtils.readString(body), ReadWriteIOUtils.readString(body)); + } + return params; + } + + private static String getProtocol(final TPipeTransferReq req) { + return req instanceof AirGapPseudoTPipeTransferRequest + ? PipeReceiverRuntimeRegistry.PROTOCOL_AIR_GAP + : PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT; + } + + private static int parseSenderPort(final String senderPort) { + try { + return Integer.parseInt(senderPort); + } catch (final Exception e) { + return -1; + } + } + + private static long parsePipeCreationTime(final String pipeCreationTime) { + if (pipeCreationTime == null) { + return Long.MIN_VALUE; + } + try { + return Long.parseLong(pipeCreationTime); + } catch (final NumberFormatException e) { + return Long.MIN_VALUE; + } + } + /** * For {@link InsertRowsStatement} and {@link InsertMultiTabletsStatement}, the returned {@link * TSStatus} will use sub-status to record the endpoint for redirection. Each sub-status records @@ -1388,6 +1538,9 @@ private TSStatus executeStatementForTableModelWithPermissionCheck( @Override public synchronized void handleExit() { + PipeReceiverRuntimeRegistry.getInstance() + .deregister(configPipeReceiverRuntimeSessionKey.getAndSet(null)); + if (Objects.nonNull(configReceiverId.get())) { try { ClusterConfigTaskExecutor.getInstance().handlePipeConfigClientExit(configReceiverId.get()); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/client/IoTDBDataNodeAsyncClientManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/client/IoTDBDataNodeAsyncClientManager.java index 5c38c3a85402f..bb4404709e01e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/client/IoTDBDataNodeAsyncClientManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/client/IoTDBDataNodeAsyncClientManager.java @@ -319,6 +319,7 @@ public void onError(final Exception e) { params.put( PipeTransferHandshakeConstant.HANDSHAKE_KEY_SKIP_IF, Boolean.toString(skipIfNoPrivileges)); + appendPipeInfoToHandshakeParams(params); client.setTimeoutDynamically(PipeConfig.getInstance().getPipeSinkHandshakeTimeoutMs()); client.pipeTransfer(PipeTransferDataNodeHandshakeV2Req.toTPipeTransferReq(params), callback); @@ -418,6 +419,12 @@ public ExecutorService getExecutor() { return executor; } + public void discardReceiverRuntimeSessions() { + for (final TEndPoint endPoint : new HashSet<>(endPointSet)) { + endPoint2Client.clear(endPoint); + } + } + public void close() { isClosed = true; synchronized (IoTDBDataNodeAsyncClientManager.class) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataNodeAirGapSink.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataNodeAirGapSink.java index b593df5661206..9e05c240a155e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataNodeAirGapSink.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataNodeAirGapSink.java @@ -72,6 +72,7 @@ protected byte[] generateHandShakeV2Payload() throws IOException { Boolean.toString(shouldMarkAsPipeRequest)); params.put( PipeTransferHandshakeConstant.HANDSHAKE_KEY_SKIP_IF, Boolean.toString(skipIfNoPrivileges)); + appendPipeInfoToHandshakeParams(params); return PipeTransferDataNodeHandshakeV2Req.toTPipeTransferBytes(params); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java index 7c85268e6d3a9..b5af3199e8ea0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/IoTDBDataRegionAsyncSink.java @@ -177,6 +177,7 @@ public void customize( shouldMarkAsPipeRequest, false, skipIfNoPrivileges); + clientManager.setPipeInfo(pipeName, creationTime); transferTsFileClientManager = new IoTDBDataNodeAsyncClientManager( @@ -193,6 +194,7 @@ public void customize( shouldMarkAsPipeRequest, isSplitTSFileBatchModeEnabled, skipIfNoPrivileges); + transferTsFileClientManager.setPipeInfo(pipeName, creationTime); if (isTabletBatchModeEnabled) { tabletBatchBuilder = new PipeTransferBatchReqBuilder(parameters); @@ -917,6 +919,19 @@ public synchronized void close() { super.close(); } + @Override + public synchronized void discardReceiverRuntimeSessions() { + syncSink.discardReceiverRuntimeSessions(); + + if (clientManager != null) { + clientManager.discardReceiverRuntimeSessions(); + } + + if (transferTsFileClientManager != null) { + transferTsFileClientManager.discardReceiverRuntimeSessions(); + } + } + public synchronized void clearRetryEventsReferenceCount() { while (!retryEventQueue.isEmpty() || !retryTsFileQueue.isEmpty()) { final Event event = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/writeback/WriteBackSink.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/writeback/WriteBackSink.java index 9bec114ede1c0..67fab340dfc22 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/writeback/WriteBackSink.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/writeback/WriteBackSink.java @@ -25,11 +25,14 @@ import org.apache.iotdb.commons.exception.IoTDBRuntimeException; import org.apache.iotdb.commons.exception.auth.AccessDeniedException; import org.apache.iotdb.commons.exception.pipe.PipeRuntimeSinkNonReportTimeConfigurableException; +import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeRegistry; import org.apache.iotdb.commons.pipe.resource.log.PipeLogger; +import org.apache.iotdb.commons.pipe.sink.protocol.PipeConnectorWithEventDiscard; import org.apache.iotdb.commons.queryengine.common.SqlDialect; import org.apache.iotdb.commons.utils.StatusUtils; import org.apache.iotdb.confignode.rpc.thrift.TDatabaseSchema; import org.apache.iotdb.db.auth.AuthorityChecker; +import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.i18n.DataNodePipeMessages; import org.apache.iotdb.db.pipe.event.common.statement.PipeStatementInsertionEvent; @@ -82,6 +85,7 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_IOTDB_CLI_HOSTNAME; @@ -106,15 +110,17 @@ @TreeModel @TableModel -public class WriteBackSink implements PipeConnector { +public class WriteBackSink implements PipeConnector, PipeConnectorWithEventDiscard { private static final Logger LOGGER = LoggerFactory.getLogger(WriteBackSink.class); + private static final IoTDBConfig IOTDB_CONFIG = IoTDBDescriptor.getInstance().getConfig(); // Simulate the behavior of the client-to-server communication // for correctly handling data insertion in IoTDBReceiverAgent#receive method private static final Coordinator COORDINATOR = Coordinator.getInstance(); private static final SessionManager SESSION_MANAGER = SessionManager.getInstance(); public static final AtomicLong id = new AtomicLong(); + private final AtomicReference receiverRuntimeSessionKey = new AtomicReference<>(); private InternalClientSession session; private boolean skipIfNoPrivileges; @@ -139,15 +145,16 @@ public void customize( final PipeParameters parameters, final PipeConnectorRuntimeConfiguration configuration) throws Exception { final PipeRuntimeEnvironment environment = configuration.getRuntimeEnvironment(); - session = - new InternalClientSession( - String.format( - "%s_%s_%s_%s_%s", - WriteBackSink.class.getSimpleName(), - environment.getPipeName(), - environment.getCreationTime(), - environment.getRegionId(), - id.getAndIncrement())); + final long writeBackSinkId = id.getAndIncrement(); + final String sessionId = + String.format( + "%s_%s_%s_%s_%s", + WriteBackSink.class.getSimpleName(), + environment.getPipeName(), + environment.getCreationTime(), + environment.getRegionId(), + writeBackSinkId); + session = new InternalClientSession(sessionId); String userIdString = parameters.getStringOrDefault( @@ -213,6 +220,9 @@ public void customize( throw new PipePasswordCheckException( String.format("Failed to check password for pipe %s.", environment.getPipeName())); } + + recordReceiverRuntimeHandshake( + sessionId, usernameString, environment.getPipeName(), environment.getCreationTime()); } @Override @@ -291,6 +301,10 @@ private void doTransfer( "Write back PipeInsertNodeTabletInsertionEvent %s error, result status %s", pipeInsertNodeTabletInsertionEvent, status)); } + recordReceiverRuntimeTransferIfSuccess( + status, + pipeInsertNodeTabletInsertionEvent.getPipeName(), + pipeInsertNodeTabletInsertionEvent.getCreationTime()); } private void doTransferWrapper(final PipeRawTabletInsertionEvent pipeRawTabletInsertionEvent) @@ -336,6 +350,10 @@ private void doTransfer(final PipeRawTabletInsertionEvent pipeRawTabletInsertion "Write back PipeRawTabletInsertionEvent %s error, result status %s", pipeRawTabletInsertionEvent, status)); } + recordReceiverRuntimeTransferIfSuccess( + status, + pipeRawTabletInsertionEvent.getPipeName(), + pipeRawTabletInsertionEvent.getCreationTime()); } @Override @@ -382,6 +400,10 @@ private void doTransfer(final PipeStatementInsertionEvent pipeStatementInsertion "Write back PipeStatementInsertionEvent %s error, result status %s", pipeStatementInsertionEvent, status)); } + recordReceiverRuntimeTransferIfSuccess( + status, + pipeStatementInsertionEvent.getPipeName(), + pipeStatementInsertionEvent.getCreationTime()); } private static void throwWriteBackExceptionIfNecessary( @@ -395,9 +417,59 @@ private static void throwWriteBackExceptionIfNecessary( @Override public void close() throws Exception { - if (session != null) { - SESSION_MANAGER.closeSession(session, COORDINATOR::cleanupQueryExecution, false); + try { + if (session != null) { + SESSION_MANAGER.closeSession(session, COORDINATOR::cleanupQueryExecution, false); + } + } finally { + PipeReceiverRuntimeRegistry.getInstance() + .deregister(receiverRuntimeSessionKey.getAndSet(null)); + } + } + + private void recordReceiverRuntimeHandshake( + final String sessionId, + final String userName, + final String pipeName, + final long creationTime) { + final String oldSessionKey = receiverRuntimeSessionKey.getAndSet(sessionId); + if (!Objects.equals(oldSessionKey, sessionId)) { + PipeReceiverRuntimeRegistry.getInstance().deregister(oldSessionKey); } + + PipeReceiverRuntimeRegistry.getInstance() + .registerOrUpdateSession( + sessionId, + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + IOTDB_CONFIG.getDataNodeId(), + PipeReceiverRuntimeRegistry.PROTOCOL_WRITEBACK, + IOTDB_CONFIG.getRpcAddress(), + IOTDB_CONFIG.getRpcPort(), + userName, + IOTDB_CONFIG.getClusterId(), + pipeName, + creationTime, + System.currentTimeMillis()); + } + + private void recordReceiverRuntimeTransferIfSuccess( + final TSStatus status, final String pipeName, final long pipeCreationTime) { + if (status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode() + || status.getCode() == TSStatusCode.REDIRECTION_RECOMMEND.getStatusCode()) { + PipeReceiverRuntimeRegistry.getInstance() + .markTransfer( + receiverRuntimeSessionKey.get(), + pipeName, + pipeCreationTime, + System.currentTimeMillis()); + } + } + + @Override + public void discardEventsOfPipe( + final String pipeName, final long creationTime, final int regionId) { + PipeReceiverRuntimeRegistry.getInstance() + .removePipe(receiverRuntimeSessionKey.get(), pipeName, creationTime); } private TSStatus executeStatementForTableModel( diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java index ec55111a200b4..e696f31155e21 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java @@ -363,6 +363,10 @@ public TTransport getTransport() { return transport; } + public TEndPoint getConfigNode() { + return configNode; + } + public void syncLatestConfigNodeList() { configNodes = ConfigNodeInfo.getInstance().getLatestConfigNodes(); cursor = 0; @@ -529,6 +533,7 @@ public TDataNodeRegisterResp registerDataNode(TDataNodeRegisterReq req) throws T for (TConfigNodeLocation configNodeLocation : resp.getConfigNodeList()) { newConfigNodes.add(configNodeLocation.getInternalEndPoint()); } + ConfigNodeInfo.getInstance().updateConfigNodeLocations(resp.getConfigNodeList()); configNodes = newConfigNodes; } } catch (TException e) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeInfo.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeInfo.java index 235a1c9d86e64..a0f9c007434f2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeInfo.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeInfo.java @@ -19,6 +19,7 @@ package org.apache.iotdb.db.protocol.client; +import org.apache.iotdb.common.rpc.thrift.TConfigNodeLocation; import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.commons.consensus.ConfigRegionId; import org.apache.iotdb.commons.exception.BadNodeUrlException; @@ -33,8 +34,10 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.HashMap; import java.util.HashSet; import java.util.List; +import java.util.Map; import java.util.Properties; import java.util.Set; import java.util.concurrent.locks.ReentrantReadWriteLock; @@ -49,6 +52,8 @@ public class ConfigNodeInfo { /** latest config nodes. */ private final Set onlineConfigNodes; + private final Map configNodeIdMap; + public static final ConfigRegionId CONFIG_REGION_ID = new ConfigRegionId(0); SystemPropertiesHandler systemPropertiesHandler = DataNodeSystemPropertiesHandler.getInstance(); @@ -56,6 +61,7 @@ public class ConfigNodeInfo { private ConfigNodeInfo() { this.configNodeInfoReadWriteLock = new ReentrantReadWriteLock(); this.onlineConfigNodes = new HashSet<>(); + this.configNodeIdMap = new HashMap<>(); } public static void reinitializeStatics() { @@ -81,6 +87,7 @@ public boolean updateConfigNodeList(List latestConfigNodes) { try { onlineConfigNodes.clear(); onlineConfigNodes.addAll(latestConfigNodes); + configNodeIdMap.keySet().retainAll(onlineConfigNodes); storeConfigNodeList(); long endTime = System.currentTimeMillis(); logger.info( @@ -96,6 +103,35 @@ public boolean updateConfigNodeList(List latestConfigNodes) { return true; } + public boolean updateConfigNodeLocations(List latestConfigNodeLocations) { + if (latestConfigNodeLocations == null) { + return false; + } + final List latestConfigNodes = new ArrayList<>(); + final Map latestConfigNodeIdMap = new HashMap<>(); + for (TConfigNodeLocation configNodeLocation : latestConfigNodeLocations) { + if (configNodeLocation == null || configNodeLocation.getInternalEndPoint() == null) { + continue; + } + final TEndPoint internalEndPoint = configNodeLocation.getInternalEndPoint(); + latestConfigNodes.add(internalEndPoint); + latestConfigNodeIdMap.put(internalEndPoint, configNodeLocation.getConfigNodeId()); + } + + if (!updateConfigNodeList(latestConfigNodes)) { + return false; + } + + configNodeInfoReadWriteLock.writeLock().lock(); + try { + configNodeIdMap.putAll(latestConfigNodeIdMap); + configNodeIdMap.keySet().retainAll(onlineConfigNodes); + } finally { + configNodeInfoReadWriteLock.writeLock().unlock(); + } + return true; + } + /** * Call this method to store config node list. * @@ -152,6 +188,18 @@ public List getLatestConfigNodes() { return result; } + public int getConfigNodeId(TEndPoint internalEndPoint) { + if (internalEndPoint == null) { + return -1; + } + configNodeInfoReadWriteLock.readLock().lock(); + try { + return configNodeIdMap.getOrDefault(internalEndPoint, -1); + } finally { + configNodeInfoReadWriteLock.readLock().unlock(); + } + } + private static class ConfigNodeInfoHolder { private static ConfigNodeInfo INSTANCE = new ConfigNodeInfo(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/header/DatasetHeaderFactory.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/header/DatasetHeaderFactory.java index 18f15eea8f397..0aebf48d95802 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/header/DatasetHeaderFactory.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/header/DatasetHeaderFactory.java @@ -207,6 +207,10 @@ public static DatasetHeader getShowQueriesHeader() { return new DatasetHeader(ColumnHeaderConstant.showQueriesColumnHeaders, false); } + public static DatasetHeader getShowReceiversHeader() { + return new DatasetHeader(ColumnHeaderConstant.showReceiversColumnHeaders, true); + } + public static DatasetHeader getShowDiskUsageHeader() { return new DatasetHeader(ColumnHeaderConstant.showDiskUsageColumnHeaders, true); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/PipeReceiverRuntimeSnapshotFilter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/PipeReceiverRuntimeSnapshotFilter.java new file mode 100644 index 0000000000000..3300c29ddf674 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/PipeReceiverRuntimeSnapshotFilter.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.db.queryengine.execution.operator.source; + +import org.apache.iotdb.commons.audit.UserEntity; +import org.apache.iotdb.commons.auth.entity.PrivilegeType; +import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeRegistry; +import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeSnapshot; +import org.apache.iotdb.db.auth.AuthorityChecker; + +import java.util.List; +import java.util.stream.Collectors; + +public final class PipeReceiverRuntimeSnapshotFilter { + + private PipeReceiverRuntimeSnapshotFilter() { + // utility class + } + + public static List visibleSnapshots(final UserEntity userEntity) { + final List snapshots = + PipeReceiverRuntimeRegistry.getInstance().snapshot(); + if (canSeeAllReceivers(userEntity)) { + return snapshots; + } + final String userName = userEntity.getUsername(); + return snapshots.stream() + .filter(snapshot -> userName.equals(snapshot.getUserName())) + .collect(Collectors.toList()); + } + + private static boolean canSeeAllReceivers(final UserEntity userEntity) { + if (userEntity == null || userEntity.getUsername() == null) { + return true; + } + if (AuthorityChecker.SUPER_USER_ID == userEntity.getUserId() + || AuthorityChecker.SUPER_USER.equals(userEntity.getUsername())) { + return true; + } + try { + return AuthorityChecker.checkSystemPermission(userEntity.getUsername(), PrivilegeType.SYSTEM); + } catch (final Exception e) { + return false; + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/ShowReceiversOperator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/ShowReceiversOperator.java new file mode 100644 index 0000000000000..6a64a828fbe38 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/ShowReceiversOperator.java @@ -0,0 +1,177 @@ +/* + * 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.queryengine.execution.operator.source; + +import org.apache.iotdb.commons.audit.UserEntity; +import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeRegistry; +import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeSnapshot; +import org.apache.iotdb.commons.queryengine.execution.MemoryEstimationHelper; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.queryengine.utils.DateTimeUtils; +import org.apache.iotdb.commons.queryengine.utils.TimestampPrecisionUtils; +import org.apache.iotdb.db.queryengine.common.header.DatasetHeaderFactory; +import org.apache.iotdb.db.queryengine.execution.operator.OperatorContext; + +import org.apache.tsfile.block.column.ColumnBuilder; +import org.apache.tsfile.common.conf.TSFileDescriptor; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.read.common.block.TsBlock; +import org.apache.tsfile.read.common.block.TsBlockBuilder; +import org.apache.tsfile.read.common.block.column.TimeColumnBuilder; +import org.apache.tsfile.utils.BytesUtils; +import org.apache.tsfile.utils.RamUsageEstimator; + +import java.util.List; +import java.util.concurrent.TimeUnit; + +public class ShowReceiversOperator implements SourceOperator { + + private final OperatorContext operatorContext; + private final PlanNodeId sourceId; + private final UserEntity userEntity; + + private TsBlock tsBlock; + private boolean hasConsumed; + + private static final int DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES = + TSFileDescriptor.getInstance().getConfig().getMaxTsBlockSizeInBytes(); + + private static final long INSTANCE_SIZE = + RamUsageEstimator.shallowSizeOfInstance(ShowReceiversOperator.class); + + public ShowReceiversOperator(OperatorContext operatorContext, PlanNodeId sourceId) { + this(operatorContext, sourceId, null); + } + + public ShowReceiversOperator( + OperatorContext operatorContext, PlanNodeId sourceId, UserEntity userEntity) { + this.operatorContext = operatorContext; + this.sourceId = sourceId; + this.userEntity = userEntity; + } + + @Override + public OperatorContext getOperatorContext() { + return operatorContext; + } + + @Override + public TsBlock next() { + TsBlock res = tsBlock; + hasConsumed = true; + tsBlock = null; + return res; + } + + @Override + public boolean hasNext() { + if (hasConsumed) { + return false; + } + if (tsBlock == null) { + tsBlock = buildTsBlock(); + } + return true; + } + + @Override + public boolean isFinished() { + return hasConsumed; + } + + @Override + public void close() { + // do nothing + } + + @Override + public long calculateMaxPeekMemory() { + return DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES; + } + + @Override + public long calculateMaxReturnSize() { + return DEFAULT_MAX_TSBLOCK_SIZE_IN_BYTES; + } + + @Override + public long calculateRetainedSizeAfterCallingNext() { + return 0; + } + + @Override + public PlanNodeId getSourceId() { + return sourceId; + } + + private TsBlock buildTsBlock() { + final List outputDataTypes = + DatasetHeaderFactory.getShowReceiversHeader().getRespDataTypes(); + final TsBlockBuilder builder = new TsBlockBuilder(outputDataTypes); + final TimeColumnBuilder timeColumnBuilder = builder.getTimeColumnBuilder(); + final ColumnBuilder[] columnBuilders = builder.getValueColumnBuilders(); + final long currentTime = + TimestampPrecisionUtils.convertToCurrPrecision( + System.currentTimeMillis(), TimeUnit.MILLISECONDS); + + for (PipeReceiverRuntimeSnapshot snapshot : + PipeReceiverRuntimeSnapshotFilter.visibleSnapshots(userEntity)) { + timeColumnBuilder.writeLong(currentTime); + columnBuilders[0].writeBinary(BytesUtils.valueOf(snapshot.getReceiverNodeType())); + writeReceiverNodeId(columnBuilders[1], snapshot); + columnBuilders[2].writeBinary(BytesUtils.valueOf(snapshot.getProtocol())); + columnBuilders[3].writeBinary(BytesUtils.valueOf(snapshot.getSenderClusterId())); + columnBuilders[4].writeBinary(BytesUtils.valueOf(snapshot.getSenderAddress())); + columnBuilders[5].writeBinary(BytesUtils.valueOf(snapshot.getUserName())); + columnBuilders[6].writeBinary(BytesUtils.valueOf(snapshot.getSenderPorts())); + columnBuilders[7].writeInt(snapshot.getConnectionCount()); + columnBuilders[8].writeInt(snapshot.getPipeCount()); + columnBuilders[9].writeBinary(BytesUtils.valueOf(snapshot.getPipeIds())); + columnBuilders[10].writeBinary( + BytesUtils.valueOf(formatTime(snapshot.getLastHandshakeTime()))); + columnBuilders[11].writeBinary( + BytesUtils.valueOf(formatTime(snapshot.getLastTransferTime()))); + builder.declarePosition(); + } + return builder.build(); + } + + private static void writeReceiverNodeId( + ColumnBuilder columnBuilder, PipeReceiverRuntimeSnapshot snapshot) { + if (snapshot.isReceiverNodeIdKnown()) { + columnBuilder.writeInt(snapshot.getReceiverNodeId()); + } else { + columnBuilder.appendNull(); + } + } + + private static String formatTime(long timestampInMillis) { + return timestampInMillis <= 0 + ? PipeReceiverRuntimeRegistry.UNKNOWN + : DateTimeUtils.convertLongToDate(timestampInMillis, "ms"); + } + + @Override + public long ramBytesUsed() { + return INSTANCE_SIZE + + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(operatorContext) + + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(sourceId); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/InformationSchemaContentSupplierFactory.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/InformationSchemaContentSupplierFactory.java index a501824850d13..2b23a35b55f5f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/InformationSchemaContentSupplierFactory.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/InformationSchemaContentSupplierFactory.java @@ -34,6 +34,7 @@ import org.apache.iotdb.commons.exception.auth.AccessDeniedException; import org.apache.iotdb.commons.pipe.agent.plugin.builtin.BuiltinPipePlugin; import org.apache.iotdb.commons.pipe.agent.plugin.meta.PipePluginMeta; +import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeSnapshot; import org.apache.iotdb.commons.queryengine.common.ConnectionInfo; import org.apache.iotdb.commons.queryengine.common.SqlDialect; import org.apache.iotdb.commons.queryengine.plan.relational.function.TableBuiltinTableFunction; @@ -82,6 +83,7 @@ import org.apache.iotdb.db.queryengine.common.QueryId; import org.apache.iotdb.db.queryengine.execution.QueryState; import org.apache.iotdb.db.queryengine.execution.operator.OperatorContext; +import org.apache.iotdb.db.queryengine.execution.operator.source.PipeReceiverRuntimeSnapshotFilter; import org.apache.iotdb.db.queryengine.plan.Coordinator; import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution; import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.relational.ShowCreateViewTask; @@ -186,6 +188,8 @@ public static IInformationSchemaContentSupplier getSupplier( return new RegionSupplier(dataTypes, userEntity); case InformationSchema.PIPES: return new PipeSupplier(dataTypes, userEntity.getUsername()); + case InformationSchema.RECEIVERS: + return new ReceiversSupplier(dataTypes, userEntity); case InformationSchema.PIPE_PLUGINS: return new PipePluginSupplier(dataTypes, userEntity); case InformationSchema.TOPICS: @@ -714,6 +718,56 @@ public boolean hasNext() { } } + private static class ReceiversSupplier extends TsBlockSupplier { + private final Iterator iterator; + + private ReceiversSupplier(final List dataTypes, final UserEntity userEntity) { + super(dataTypes); + iterator = PipeReceiverRuntimeSnapshotFilter.visibleSnapshots(userEntity).iterator(); + } + + @Override + protected void constructLine() { + final PipeReceiverRuntimeSnapshot snapshot = iterator.next(); + columnBuilders[0].writeBinary(BytesUtils.valueOf(snapshot.getReceiverNodeType())); + writeReceiverNodeId(columnBuilders[1], snapshot); + columnBuilders[2].writeBinary(BytesUtils.valueOf(snapshot.getProtocol())); + columnBuilders[3].writeBinary(BytesUtils.valueOf(snapshot.getSenderClusterId())); + columnBuilders[4].writeBinary(BytesUtils.valueOf(snapshot.getSenderAddress())); + columnBuilders[5].writeBinary(BytesUtils.valueOf(snapshot.getUserName())); + columnBuilders[6].writeBinary(BytesUtils.valueOf(snapshot.getSenderPorts())); + columnBuilders[7].writeInt(snapshot.getConnectionCount()); + columnBuilders[8].writeInt(snapshot.getPipeCount()); + columnBuilders[9].writeBinary(BytesUtils.valueOf(snapshot.getPipeIds())); + writeTimestamp(columnBuilders[10], snapshot.getLastHandshakeTime()); + writeTimestamp(columnBuilders[11], snapshot.getLastTransferTime()); + resultBuilder.declarePosition(); + } + + private static void writeReceiverNodeId( + ColumnBuilder columnBuilder, PipeReceiverRuntimeSnapshot snapshot) { + if (snapshot.isReceiverNodeIdKnown()) { + columnBuilder.writeInt(snapshot.getReceiverNodeId()); + } else { + columnBuilder.appendNull(); + } + } + + private static void writeTimestamp(ColumnBuilder columnBuilder, long timestampInMillis) { + if (timestampInMillis <= 0) { + columnBuilder.appendNull(); + return; + } + columnBuilder.writeLong( + TimestampPrecisionUtils.convertToCurrPrecision(timestampInMillis, TimeUnit.MILLISECONDS)); + } + + @Override + public boolean hasNext() { + return iterator.hasNext(); + } + } + private static class PipePluginSupplier extends TsBlockSupplier { private final Iterator iterator; 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 be20d218e6a70..d5539bfe93d4c 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 @@ -60,6 +60,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.sys.ExplainAnalyzeStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowDiskUsageStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowQueriesStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowReceiversStatement; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.metadata.IDeviceID; @@ -488,6 +489,7 @@ private boolean hasDataSource() { return (dataPartition != null && !dataPartition.isEmpty()) || (schemaPartition != null && !schemaPartition.isEmpty()) || statement instanceof ShowQueriesStatement + || statement instanceof ShowReceiversStatement || statement instanceof ShowDiskUsageStatement || (statement instanceof QueryStatement && ((QueryStatement) statement).isAggregationQuery()); 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 22d5b0518ac55..27a9778802a8d 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 @@ -145,6 +145,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.sys.ExplainStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowDiskUsageStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowQueriesStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowReceiversStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowVersionStatement; import org.apache.iotdb.db.schemaengine.schemaregion.view.visitor.TransformToExpressionVisitor; import org.apache.iotdb.rpc.RpcUtils; @@ -3822,6 +3823,34 @@ public Analysis visitShowQueries( return analysis; } + @Override + public Analysis visitShowReceivers( + ShowReceiversStatement showReceiversStatement, MPPQueryContext context) { + Analysis analysis = new Analysis(); + analysis.setRealStatement(showReceiversStatement); + analysis.setRespDatasetHeader(DatasetHeaderFactory.getShowReceiversHeader()); + analysis.setVirtualSource(true); + + List allReadableDataNodeLocations = getReadableDataNodeLocations(); + if (allReadableDataNodeLocations.isEmpty()) { + throw new StatementAnalyzeException(DataNodeQueryMessages.NO_RUNNING_DATANODES); + } + analysis.setReadableDataNodeLocations(allReadableDataNodeLocations); + + Set sourceExpressions = new HashSet<>(); + for (ColumnHeader columnHeader : analysis.getRespDatasetHeader().getColumnHeaders()) { + sourceExpressions.add( + TimeSeriesOperand.constructColumnHeaderExpression( + columnHeader.getColumnName(), columnHeader.getColumnType())); + } + analysis.setSourceExpressions(sourceExpressions); + sourceExpressions.forEach(expression -> analyzeExpressionType(analysis, expression)); + + analysis.setMergeOrderParameter(new OrderByParameter(showReceiversStatement.getSortItemList())); + + return analysis; + } + @Override public Analysis visitShowDiskUsage( ShowDiskUsageStatement showDiskUsageStatement, MPPQueryContext context) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java index 6eda402c2eb2c..949c4805213f1 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java @@ -359,6 +359,7 @@ import org.apache.tsfile.external.commons.codec.digest.DigestUtils; import org.apache.tsfile.read.common.block.TsBlockBuilder; import org.apache.tsfile.read.common.block.column.TimeColumnBuilder; +import org.apache.tsfile.utils.Pair; import org.apache.tsfile.utils.ReadWriteIOUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -4127,6 +4128,11 @@ public TSpaceQuotaResp getSpaceQuota() { @Override public TPipeTransferResp handleTransferConfigPlan( final String clientId, final TPipeTransferReq req) { + return handleTransferConfigPlanAndGetReceiverNodeId(clientId, req).left; + } + + public Pair handleTransferConfigPlanAndGetReceiverNodeId( + final String clientId, final TPipeTransferReq req) { final TPipeConfigTransferReq configTransferReq = new TPipeConfigTransferReq( req.version, @@ -4139,18 +4145,23 @@ public TPipeTransferResp handleTransferConfigPlan( CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) { final TPipeConfigTransferResp pipeConfigTransferResp = configNodeClient.handleTransferConfigPlan(configTransferReq); + final int receiverNodeId = + ConfigNodeInfo.getInstance().getConfigNodeId(configNodeClient.getConfigNode()); if (TSStatusCode.SUCCESS_STATUS.getStatusCode() != pipeConfigTransferResp.getStatus().getCode()) { LOGGER.warn( DataNodeQueryMessages.FAILED_TO_HANDLETRANSFERCONFIGPLAN_STATUS_IS, pipeConfigTransferResp); } - return new TPipeTransferResp(pipeConfigTransferResp.status) - .setBody(pipeConfigTransferResp.body); + return new Pair<>( + new TPipeTransferResp(pipeConfigTransferResp.status).setBody(pipeConfigTransferResp.body), + receiverNodeId); } catch (Exception e) { - return new TPipeTransferResp( - new TSStatus(TSStatusCode.INTERNAL_SERVER_ERROR.getStatusCode()) - .setMessage(e.toString())); + return new Pair<>( + new TPipeTransferResp( + new TSStatus(TSStatusCode.INTERNAL_SERVER_ERROR.getStatusCode()) + .setMessage(e.toString())), + -1); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java index b9e2fe21c7a54..5eb9797ea28b5 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java @@ -251,6 +251,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowCurrentUserStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowDiskUsageStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowQueriesStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowReceiversStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowVersionStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.StartRepairDataStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.StopRepairDataStatement; @@ -4413,6 +4414,11 @@ public Statement visitShowPipes(IoTDBSqlParser.ShowPipesContext ctx) { return showPipesStatement; } + @Override + public Statement visitShowReceivers(IoTDBSqlParser.ShowReceiversContext ctx) { + return new ShowReceiversStatement(); + } + @Override public Statement visitCreateTopic(IoTDBSqlParser.CreateTopicContext ctx) { final CreateTopicStatement createTopicStatement = new CreateTopicStatement(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LogicalPlanBuilder.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LogicalPlanBuilder.java index 26d361c1f210f..36e522e5d1368 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LogicalPlanBuilder.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LogicalPlanBuilder.java @@ -90,6 +90,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.SeriesSourceNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.ShowDiskUsageNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.ShowQueriesNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.ShowReceiversNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.TimeseriesRegionScanNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.AggregationDescriptor; import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.AggregationStep; @@ -1328,6 +1329,44 @@ private LogicalPlanBuilder planSingleShowQueries( return this; } + public LogicalPlanBuilder planShowReceivers(Analysis analysis) { + List dataNodeLocations = analysis.getReadableDataNodeLocations(); + if (dataNodeLocations.size() == 1) { + this.root = + planSingleShowReceivers(dataNodeLocations.get(0)) + .planSort(analysis.getMergeOrderParameter()) + .getRoot(); + } else { + List outputColumns = new ArrayList<>(); + MergeSortNode mergeSortNode = + new MergeSortNode( + context.getQueryId().genPlanNodeId(), + analysis.getMergeOrderParameter(), + outputColumns); + + dataNodeLocations.forEach( + dataNodeLocation -> + mergeSortNode.addChild( + this.planSingleShowReceivers(dataNodeLocation) + .planSort(analysis.getMergeOrderParameter()) + .getRoot())); + outputColumns.addAll(mergeSortNode.getChildren().get(0).getOutputColumnNames()); + this.root = mergeSortNode; + } + + ColumnHeaderConstant.showReceiversColumnHeaders.forEach( + columnHeader -> + context + .getTypeProvider() + .setTreeModelType(columnHeader.getColumnName(), columnHeader.getColumnType())); + return this; + } + + private LogicalPlanBuilder planSingleShowReceivers(TDataNodeLocation dataNodeLocation) { + this.root = new ShowReceiversNode(context.getQueryId().genPlanNodeId(), dataNodeLocation); + return this; + } + public LogicalPlanBuilder planShowDiskUsage(Analysis analysis, PartialPath pathPattern) { List dataNodeLocations = analysis.getReadableDataNodeLocations(); if (dataNodeLocations.size() == 1) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LogicalPlanVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LogicalPlanVisitor.java index 2c9304ce0b98c..f04dc1c9263ce 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LogicalPlanVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LogicalPlanVisitor.java @@ -92,6 +92,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.sys.ExplainAnalyzeStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowDiskUsageStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowQueriesStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowReceiversStatement; import org.apache.iotdb.db.schemaengine.SchemaEngineMode; import org.apache.tsfile.enums.TSDataType; @@ -1017,6 +1018,12 @@ public PlanNode visitShowQueries( return planBuilder.getRoot(); } + @Override + public PlanNode visitShowReceivers( + ShowReceiversStatement showReceiversStatement, MPPQueryContext context) { + return new LogicalPlanBuilder(analysis, context).planShowReceivers(analysis).getRoot(); + } + @Override public PlanNode visitShowDiskUsage( ShowDiskUsageStatement showDiskUsageStatement, MPPQueryContext context) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java index 327dbecdda26d..bf5683c7f0530 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/OperatorTreeGenerator.java @@ -40,6 +40,7 @@ import org.apache.iotdb.calc.transformation.dag.column.leaf.LeafColumnTransformer; import org.apache.iotdb.common.rpc.thrift.TAggregationType; import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.audit.UserEntity; import org.apache.iotdb.commons.model.ModelInformation; import org.apache.iotdb.commons.path.AlignedFullPath; import org.apache.iotdb.commons.path.AlignedPath; @@ -140,6 +141,7 @@ import org.apache.iotdb.db.queryengine.execution.operator.source.SeriesScanOperator; import org.apache.iotdb.db.queryengine.execution.operator.source.ShowDiskUsageOperator; import org.apache.iotdb.db.queryengine.execution.operator.source.ShowQueriesOperator; +import org.apache.iotdb.db.queryengine.execution.operator.source.ShowReceiversOperator; import org.apache.iotdb.db.queryengine.execution.operator.window.ConditionWindowParameter; import org.apache.iotdb.db.queryengine.execution.operator.window.CountWindowParameter; import org.apache.iotdb.db.queryengine.execution.operator.window.SessionWindowParameter; @@ -224,6 +226,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.SeriesScanNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.ShowDiskUsageNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.ShowQueriesNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.ShowReceiversNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.TimeseriesRegionScanNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.AggregationDescriptor; import org.apache.iotdb.db.queryengine.plan.planner.plan.parameter.CrossSeriesAggregationDescriptor; @@ -325,6 +328,15 @@ public class OperatorTreeGenerator implements PlanVisitor DESC_BINARY_COMPARATOR = Comparator.reverseOrder(); + private static UserEntity getCurrentUserEntity(final LocalExecutionPlanContext context) { + if (context.getDriverContext() == null + || context.getDriverContext().getFragmentInstanceContext() == null + || context.getDriverContext().getFragmentInstanceContext().getSessionInfo() == null) { + return null; + } + return context.getDriverContext().getFragmentInstanceContext().getSessionInfo().getUserEntity(); + } + @Override public Operator visitPlan(PlanNode node, LocalExecutionPlanContext context) { throw new UnsupportedOperationException( @@ -2478,6 +2490,20 @@ public Operator visitShowQueries(ShowQueriesNode node, LocalExecutionPlanContext node.getAllowedUsername()); } + @Override + public Operator visitShowReceivers(ShowReceiversNode node, LocalExecutionPlanContext context) { + OperatorContext operatorContext = + context + .getDriverContext() + .addOperatorContext( + context.getNextOperatorId(), + node.getPlanNodeId(), + ShowReceiversOperator.class.getSimpleName()); + + return new ShowReceiversOperator( + operatorContext, node.getPlanNodeId(), getCurrentUserEntity(context)); + } + @Override public Operator visitShowDiskUsage(ShowDiskUsageNode node, LocalExecutionPlanContext context) { OperatorContext operatorContext = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/SimpleFragmentParallelPlanner.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/SimpleFragmentParallelPlanner.java index 4e34d36147168..cd9a9fe44c81d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/SimpleFragmentParallelPlanner.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/SimpleFragmentParallelPlanner.java @@ -39,6 +39,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.sys.ExplainAnalyzeStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowDiskUsageStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowQueriesStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowReceiversStatement; import org.apache.tsfile.read.common.Path; import org.apache.tsfile.utils.Pair; @@ -159,6 +160,7 @@ private void produceFragmentInstance(PlanFragment fragment) { if (analysis.getTreeStatement() instanceof QueryStatement || analysis.getTreeStatement() instanceof ExplainAnalyzeStatement || analysis.getTreeStatement() instanceof ShowQueriesStatement + || analysis.getTreeStatement() instanceof ShowReceiversStatement || analysis.getTreeStatement() instanceof ShowDiskUsageStatement || (analysis.getTreeStatement() instanceof ShowTimeSeriesStatement && (((ShowTimeSeriesStatement) analysis.getTreeStatement()).isOrderByHeat() diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/DataNodePlanNodeDeserializer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/DataNodePlanNodeDeserializer.java index 2a9b9b6905c77..59de11d98d5d4 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/DataNodePlanNodeDeserializer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/DataNodePlanNodeDeserializer.java @@ -109,6 +109,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.SeriesScanNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.ShowDiskUsageNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.ShowQueriesNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.ShowReceiversNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.TimeseriesRegionScanNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.ContinuousSameSearchIndexSeparatorNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; @@ -415,6 +416,8 @@ public PlanNode deserialize(ByteBuffer buffer, short nodeType) { return ShowDiskUsageNode.deserialize(buffer); case 108: return CollectNode.deserialize(buffer); + case 110: + return ShowReceiversNode.deserialize(buffer); case 902: return CreateOrUpdateTableDeviceNode.deserialize(buffer); case 903: diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanVisitor.java index 81596f2d38c10..bb33dfac50764 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanVisitor.java @@ -108,6 +108,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.SeriesScanSourceNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.ShowDiskUsageNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.ShowQueriesNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.ShowReceiversNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.TimeseriesRegionScanNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertMultiTabletsNode; @@ -350,6 +351,10 @@ default R visitShowQueries(ShowQueriesNode node, C context) { return visitPlan(node, context); } + default R visitShowReceivers(ShowReceiversNode node, C context) { + return visitPlan(node, context); + } + default R visitShowDiskUsage(ShowDiskUsageNode node, C context) { return visitPlan(node, context); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/source/ShowReceiversNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/source/ShowReceiversNode.java new file mode 100644 index 0000000000000..1cf76707c560b --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/source/ShowReceiversNode.java @@ -0,0 +1,115 @@ +/* + * 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.queryengine.plan.planner.plan.node.source; + +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.IPlanVisitor; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; +import org.apache.iotdb.commons.schema.column.ColumnHeader; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor; + +import com.google.common.collect.ImmutableList; + +import java.io.DataOutputStream; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Objects; + +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.showReceiversColumnHeaders; + +public class ShowReceiversNode extends VirtualSourceNode { + + public static final List SHOW_RECEIVERS_HEADER_COLUMNS = + showReceiversColumnHeaders.stream() + .map(ColumnHeader::getColumnName) + .collect(ImmutableList.toImmutableList()); + + public ShowReceiversNode(PlanNodeId id, TDataNodeLocation dataNodeLocation) { + super(id, dataNodeLocation); + } + + @Override + public List getChildren() { + return ImmutableList.of(); + } + + @Override + public void addChild(PlanNode child) { + throw new UnsupportedOperationException("no child is allowed for ShowReceiversNode"); + } + + @Override + public PlanNodeType getType() { + return PlanNodeType.SHOW_RECEIVERS; + } + + @Override + public PlanNode clone() { + return new ShowReceiversNode(getPlanNodeId(), getDataNodeLocation()); + } + + @Override + public int allowedChildCount() { + return NO_CHILD_ALLOWED; + } + + @Override + public List getOutputColumnNames() { + return SHOW_RECEIVERS_HEADER_COLUMNS; + } + + @Override + public R accept(IPlanVisitor visitor, C context) { + return ((PlanVisitor) visitor).visitShowReceivers(this, context); + } + + @Override + protected void serializeAttributes(ByteBuffer byteBuffer) { + PlanNodeType.SHOW_RECEIVERS.serialize(byteBuffer); + } + + @Override + protected void serializeAttributes(DataOutputStream stream) throws IOException { + PlanNodeType.SHOW_RECEIVERS.serialize(stream); + } + + public static ShowReceiversNode deserialize(ByteBuffer byteBuffer) { + PlanNodeId planNodeId = PlanNodeId.deserialize(byteBuffer); + return new ShowReceiversNode(planNodeId, null); + } + + @Override + public boolean equals(Object o) { + return super.equals(o); + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode()); + } + + @Override + public String toString() { + return "ShowReceiversNode-" + this.getPlanNodeId(); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/DataNodeLocationSupplierFactory.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/DataNodeLocationSupplierFactory.java index d0b3cb330e10c..2d3f6cea91042 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/DataNodeLocationSupplierFactory.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/DataNodeLocationSupplierFactory.java @@ -123,6 +123,7 @@ public List getDataNodeLocations(final String tableName) { case InformationSchema.CONNECTIONS: case InformationSchema.CURRENT_QUERIES: case InformationSchema.QUERIES_COSTS_HISTOGRAM: + case InformationSchema.RECEIVERS: return getReadableDataNodeLocations(); case InformationSchema.DATABASES: case InformationSchema.TABLES: 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 41f91bb461809..89bef654902ec 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 @@ -160,6 +160,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowCurrentUserStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowDiskUsageStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowQueriesStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowReceiversStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowVersionStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.StartRepairDataStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.StopRepairDataStatement; @@ -849,6 +850,16 @@ public TSStatus visitShowPipes(ShowPipesStatement statement, TreeAccessCheckCont return StatusUtils.OK; } + @Override + public TSStatus visitShowReceivers( + ShowReceiversStatement statement, TreeAccessCheckContext context) { + // This query follows SHOW PIPES: it cannot be rejected here. + AUDIT_LOGGER.recordObjectAuthenticationAuditLog( + context.setAuditLogOperation(AuditLogOperation.QUERY).setResult(true), + () -> "SHOW RECEIVERS"); + return StatusUtils.OK; + } + @Override public TSStatus visitDropPipe(DropPipeStatement statement, TreeAccessCheckContext context) { return checkPipeManagement( diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/parser/AstBuilder.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/parser/AstBuilder.java index b3c375129df91..b769a63c12cf1 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/parser/AstBuilder.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/parser/AstBuilder.java @@ -1733,6 +1733,17 @@ public Node visitShowQueriesStatement(RelationalSqlParser.ShowQueriesStatementCo limit); } + @Override + public Node visitShowReceiversStatement(RelationalSqlParser.ShowReceiversStatementContext ctx) { + return new ShowStatement( + getLocation(ctx), + InformationSchema.RECEIVERS, + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty()); + } + @Override public Node visitKillQueryStatement(RelationalSqlParser.KillQueryStatementContext ctx) { if (ctx.queryId == null) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementType.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementType.java index 88a2959f1475b..187a59cd6ffeb 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementType.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementType.java @@ -202,4 +202,5 @@ public enum StatementType { SHOW_EXTERNAL_SERVICE, SHOW_DISK_USAGE, + SHOW_RECEIVERS, } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementVisitor.java index 39af259928b6b..3cd1ad95875c1 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementVisitor.java @@ -148,6 +148,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowCurrentUserStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowDiskUsageStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowQueriesStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowReceiversStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowVersionStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.StartRepairDataStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.StopRepairDataStatement; @@ -544,6 +545,10 @@ public R visitShowQueries(ShowQueriesStatement showQueriesStatement, C context) return visitStatement(showQueriesStatement, context); } + public R visitShowReceivers(ShowReceiversStatement showReceiversStatement, C context) { + return visitStatement(showReceiversStatement, context); + } + public R visitShowDiskUsage(ShowDiskUsageStatement showDiskUsageStatement, C context) { return visitStatement(showDiskUsageStatement, context); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/ShowReceiversStatement.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/ShowReceiversStatement.java new file mode 100644 index 0000000000000..6f396c6948935 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/ShowReceiversStatement.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.db.queryengine.plan.statement.sys; + +import org.apache.iotdb.commons.schema.column.ColumnHeaderConstant; +import org.apache.iotdb.db.queryengine.plan.statement.StatementType; +import org.apache.iotdb.db.queryengine.plan.statement.StatementVisitor; +import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering; +import org.apache.iotdb.db.queryengine.plan.statement.component.SortItem; +import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowStatement; + +import com.google.common.collect.ImmutableList; + +import java.util.List; + +public class ShowReceiversStatement extends ShowStatement { + + private static final List DEFAULT_SORT_ITEMS = + ImmutableList.of( + new SortItem(ColumnHeaderConstant.RECEIVER_NODE_TYPE, Ordering.ASC), + new SortItem(ColumnHeaderConstant.RECEIVER_NODE_ID, Ordering.ASC), + new SortItem(ColumnHeaderConstant.PROTOCOL, Ordering.ASC), + new SortItem(ColumnHeaderConstant.SENDER_CLUSTER_ID, Ordering.ASC), + new SortItem(ColumnHeaderConstant.SENDER_ADDRESS, Ordering.ASC), + new SortItem(ColumnHeaderConstant.RECEIVER_USER_NAME, Ordering.ASC)); + + public ShowReceiversStatement() { + this.statementType = StatementType.SHOW_RECEIVERS; + } + + @Override + public boolean isQuery() { + return true; + } + + @Override + public R accept(StatementVisitor visitor, C context) { + return visitor.visitShowReceivers(this, context); + } + + public List getSortItemList() { + return DEFAULT_SORT_ITEMS; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNode.java index 09a71ad0b8986..3ce09f4047935 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNode.java @@ -487,11 +487,7 @@ protected void storeRuntimeConfigurations( List configNodeLocations, TRuntimeConfiguration runtimeConfiguration) throws StartupException { /* Store ConfigNodeList */ - List configNodeList = new ArrayList<>(); - for (TConfigNodeLocation configNodeLocation : configNodeLocations) { - configNodeList.add(configNodeLocation.getInternalEndPoint()); - } - ConfigNodeInfo.getInstance().updateConfigNodeList(configNodeList); + ConfigNodeInfo.getInstance().updateConfigNodeLocations(configNodeLocations); /* Store templateSetInfo */ ClusterTemplateManager.getInstance() diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java index e23db1f1ca8cd..eb331e50f5913 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java @@ -23,17 +23,22 @@ import org.apache.iotdb.commons.conf.CommonConfig; import org.apache.iotdb.commons.conf.CommonDescriptor; import org.apache.iotdb.commons.pipe.receiver.IoTDBReceiver; +import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeRegistry; import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapELanguageConstant; import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapOneByteResponse; import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapPseudoTPipeTransferRequest; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.IoTDBSinkRequestVersion; +import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; import org.apache.iotdb.db.pipe.receiver.protocol.thrift.IoTDBDataNodeReceiverAgent; +import org.apache.iotdb.db.protocol.session.SessionManager; import org.apache.iotdb.rpc.TSStatusCode; import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq; import org.apache.iotdb.service.rpc.thrift.TPipeTransferResp; import org.apache.tsfile.utils.BytesUtils; +import org.junit.After; import org.junit.Assert; +import org.junit.Before; import org.junit.Test; import java.io.ByteArrayInputStream; @@ -43,11 +48,30 @@ import java.io.OutputStream; import java.lang.reflect.Field; import java.lang.reflect.Method; +import java.net.InetAddress; import java.net.Socket; +import java.net.SocketException; import java.nio.ByteBuffer; +import java.util.concurrent.atomic.AtomicBoolean; public class IoTDBAirGapReceiverTest { + private final PipeReceiverRuntimeRegistry registry = PipeReceiverRuntimeRegistry.getInstance(); + + @Before + public void setUp() { + registry.clear(); + PipeDataNodeAgent.receiver().thrift().handleClientExit(); + SessionManager.getInstance().removeCurrSession(); + } + + @After + public void tearDown() { + registry.clear(); + PipeDataNodeAgent.receiver().thrift().handleClientExit(); + SessionManager.getInstance().removeCurrSession(); + } + @Test public void testRejectOversizedAirGapPayload() throws Exception { final CommonConfig commonConfig = CommonDescriptor.getInstance().getConfig(); @@ -121,6 +145,33 @@ public void testTemporaryUnavailableRetryTimeoutReturnsFail() throws Exception { } } + @Test + public void testAirGapReceiverExitCleansThriftReceiverRuntime() throws Throwable { + final String sessionKey = "DataNode-1-air_gap-4"; + registry.registerOrUpdateSession( + sessionKey, + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_AIR_GAP, + "10.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-air-gap", + 1, + 100); + final RuntimeCleanupReceiver cleanupReceiver = new RuntimeCleanupReceiver(sessionKey); + setThriftReceiverThreadLocal(cleanupReceiver); + + Assert.assertEquals(1, registry.snapshot().size()); + + new IoTDBAirGapReceiver(new EndOfStreamSocket(), 4L).runMayThrow(); + + Assert.assertTrue(cleanupReceiver.isExited()); + Assert.assertTrue(registry.snapshot().isEmpty()); + Assert.assertNull(getThriftReceiverThreadLocalReceiver()); + } + private static void setField(final Object target, final String fieldName, final Object value) throws Exception { final Field field = IoTDBAirGapReceiver.class.getDeclaredField(fieldName); @@ -128,6 +179,23 @@ private static void setField(final Object target, final String fieldName, final field.set(target, value); } + private static void setThriftReceiverThreadLocal(final IoTDBReceiver receiver) throws Exception { + getThriftReceiverThreadLocal().set(receiver); + } + + private static IoTDBReceiver getThriftReceiverThreadLocalReceiver() throws Exception { + return getThriftReceiverThreadLocal().get(); + } + + private static ThreadLocal getThriftReceiverThreadLocal() throws Exception { + final Field field = IoTDBDataNodeReceiverAgent.class.getDeclaredField("receiverThreadLocal"); + field.setAccessible(true); + @SuppressWarnings("unchecked") + final ThreadLocal receiverThreadLocal = + (ThreadLocal) field.get(PipeDataNodeAgent.receiver().thrift()); + return receiverThreadLocal; + } + private static class RecordingSocket extends Socket { private final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); @@ -142,6 +210,46 @@ byte[] getWrittenBytes() { } } + private static class EndOfStreamSocket extends RecordingSocket { + + private boolean closed; + + @Override + public void setSoTimeout(final int timeout) throws SocketException { + // noop for unit test + } + + @Override + public void setKeepAlive(final boolean on) throws SocketException { + // noop for unit test + } + + @Override + public InputStream getInputStream() { + return new ByteArrayInputStream(new byte[0]); + } + + @Override + public InetAddress getInetAddress() { + return InetAddress.getLoopbackAddress(); + } + + @Override + public int getPort() { + return 9001; + } + + @Override + public boolean isClosed() { + return closed; + } + + @Override + public synchronized void close() { + closed = true; + } + } + private static class StubIoTDBDataNodeReceiverAgent extends IoTDBDataNodeReceiverAgent { void setStubReceiver(final IoTDBReceiver receiver) { @@ -172,4 +280,34 @@ public IoTDBSinkRequestVersion getVersion() { return IoTDBSinkRequestVersion.VERSION_1; } } + + private static class RuntimeCleanupReceiver implements IoTDBReceiver { + + private final String sessionKey; + private final AtomicBoolean exited = new AtomicBoolean(false); + + private RuntimeCleanupReceiver(final String sessionKey) { + this.sessionKey = sessionKey; + } + + @Override + public TPipeTransferResp receive(final TPipeTransferReq req) { + return new TPipeTransferResp(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode())); + } + + @Override + public void handleExit() { + PipeReceiverRuntimeRegistry.getInstance().deregister(sessionKey); + exited.set(true); + } + + @Override + public IoTDBSinkRequestVersion getVersion() { + return IoTDBSinkRequestVersion.VERSION_1; + } + + private boolean isExited() { + return exited.get(); + } + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverTest.java index 8f2e86c62d057..156b7d4feeb16 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverTest.java @@ -19,6 +19,8 @@ package org.apache.iotdb.db.pipe.receiver.protocol.thrift; +import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeRegistry; +import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeSnapshot; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertMultiTabletsStatement; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowStatement; @@ -28,17 +30,34 @@ import org.apache.iotdb.db.storageengine.load.active.ActiveLoadPathHelper; import org.apache.iotdb.db.storageengine.load.config.LoadTsFileConfigurator; +import org.junit.After; import org.junit.Assert; +import org.junit.Before; import org.junit.Test; +import java.lang.reflect.Field; import java.nio.file.Files; import java.nio.file.Path; import java.util.Arrays; import java.util.Collections; +import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; public class IoTDBDataNodeReceiverTest { + private final PipeReceiverRuntimeRegistry registry = PipeReceiverRuntimeRegistry.getInstance(); + + @Before + public void setUp() { + registry.clear(); + } + + @After + public void tearDown() { + registry.clear(); + } + @Test public void testLoadTsFileSyncStatementUsesTreeDatabaseLevelFromDatabaseName() throws Exception { final Path tsFile = Files.createTempFile("pipe-load-tree-database-level", ".tsfile"); @@ -157,6 +176,53 @@ public void testLoadTsFileSyncStatementCanSkipVerifySchemaWhenNotConvertingType( } } + @Test + public void testDataNodeReceiverRuntimeIsClearedOnHandleExitAndCanReconnect() { + final TestingDataNodeReceiver receiver = new TestingDataNodeReceiver(); + + receiver.recordDataNodeReceiverRuntime(3, "10.0.0.1", "9001", "root", "cluster-a", "pipe-a", 1); + + List snapshots = registry.snapshot(); + Assert.assertEquals(1, snapshots.size()); + Assert.assertTrue(snapshots.get(0).getPipeIds().contains("pipe-a@")); + + receiver.handleExit(); + Assert.assertTrue(registry.snapshot().isEmpty()); + + receiver.recordDataNodeReceiverRuntime(3, "10.0.0.1", "9002", "root", "cluster-a", "pipe-b", 2); + + snapshots = registry.snapshot(); + Assert.assertEquals(1, snapshots.size()); + Assert.assertEquals("9002", snapshots.get(0).getSenderPorts()); + Assert.assertTrue(snapshots.get(0).getPipeIds().contains("pipe-b@")); + } + + @Test + public void testConfigNodeReceiverRuntimeIsClearedOnHandleExit() throws Exception { + final TestingDataNodeReceiver receiver = new TestingDataNodeReceiver(); + final String configNodeSessionKey = "ConfigNode-7-thrift-test-config-receiver"; + + registry.registerOrUpdateSession( + configNodeSessionKey, + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + 7, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.2", + 9003, + "root", + "cluster-a", + "pipe-config", + 11, + 100); + setConfigPipeReceiverRuntimeSessionKey(receiver, configNodeSessionKey); + + Assert.assertEquals(1, registry.snapshot().size()); + + receiver.handleExit(); + + Assert.assertTrue(registry.snapshot().isEmpty()); + } + @Test public void testClearTreeDatabaseNameForLoadTsFileStatement() throws Exception { final Path tsFile = Files.createTempFile("pipe-load-clear-tree-database", ".tsfile"); @@ -205,4 +271,54 @@ public void testClearTreeDatabaseNameForBatchInsertStatements() { Assert.assertFalse(insertMultiTabletsStatement.getDatabaseName().isPresent()); Assert.assertFalse(tabletStatement.getDatabaseName().isPresent()); } + + @SuppressWarnings("unchecked") + private static void setConfigPipeReceiverRuntimeSessionKey( + final IoTDBDataNodeReceiver receiver, final String sessionKey) throws Exception { + final Field field = + IoTDBDataNodeReceiver.class.getDeclaredField("configPipeReceiverRuntimeSessionKey"); + field.setAccessible(true); + ((AtomicReference) field.get(receiver)).set(sessionKey); + } + + private static class TestingDataNodeReceiver extends IoTDBDataNodeReceiver { + + private String senderHost; + private String senderPort; + + private void recordDataNodeReceiverRuntime( + final int receiverNodeId, + final String senderHost, + final String senderPort, + final String userName, + final String senderClusterId, + final String pipeName, + final long pipeCreationTime) { + this.senderHost = senderHost; + this.senderPort = senderPort; + this.username = userName; + this.senderClusterId = senderClusterId; + this.receiverPipeName = pipeName; + this.receiverPipeCreationTime = pipeCreationTime; + recordPipeReceiverHandshake( + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + receiverNodeId, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT); + } + + @Override + protected String getSenderHost() { + return senderHost; + } + + @Override + protected String getSenderPort() { + return senderPort; + } + + @Override + protected void closeSession() { + // Avoid touching SessionManager in this unit test. + } + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/PipeDataNodeThriftRequestTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/PipeDataNodeThriftRequestTest.java index 7641070800f57..eed6ac39472f2 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/PipeDataNodeThriftRequestTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/PipeDataNodeThriftRequestTest.java @@ -147,11 +147,13 @@ public void testCombineRequestWithUnexpectedStateClassName() throws Exception { public void testPipeTransferDataNodeHandshakeReq() throws IOException { final PipeTransferDataNodeHandshakeV1Req req = PipeTransferDataNodeHandshakeV1Req.toTPipeTransferReq(TIME_PRECISION); + final int originalBodyPosition = req.body.position(); final PipeTransferDataNodeHandshakeV1Req deserializeReq = PipeTransferDataNodeHandshakeV1Req.fromTPipeTransferReq(req); Assert.assertEquals(req.getVersion(), deserializeReq.getVersion()); Assert.assertEquals(req.getType(), deserializeReq.getType()); + Assert.assertEquals(originalBodyPosition, req.body.position()); Assert.assertEquals(req.getTimestampPrecision(), deserializeReq.getTimestampPrecision()); } @@ -174,20 +176,24 @@ public void testPipeTransferDataNodeHandshakeReqFromLegacyV13Body() throws IOExc public void testPipeTransferDataNodeHandshakeV2Req() throws IOException { final Map params = new HashMap<>(); params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_TIME_PRECISION, TIME_PRECISION); - params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_CLUSTER_ID, "cluster"); + params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_CLUSTER_ID, "cluster-a"); params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_USERNAME, "root"); params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_PASSWORD, "root"); + params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_PIPE_NAME, "pipe-a"); + params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_PIPE_CREATION_TIME, "1"); params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_LOAD_TSFILE_STRATEGY, "async"); params.put( PipeTransferHandshakeConstant.HANDSHAKE_KEY_VALIDATE_TSFILE, Boolean.TRUE.toString()); final PipeTransferDataNodeHandshakeV2Req req = PipeTransferDataNodeHandshakeV2Req.toTPipeTransferReq(params); + final int originalBodyPosition = req.body.position(); final PipeTransferDataNodeHandshakeV2Req deserializeReq = PipeTransferDataNodeHandshakeV2Req.fromTPipeTransferReq(req); Assert.assertEquals(req.getVersion(), deserializeReq.getVersion()); Assert.assertEquals(req.getType(), deserializeReq.getType()); + Assert.assertEquals(originalBodyPosition, req.body.position()); Assert.assertEquals(params, deserializeReq.getParams()); } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/PipeSinkTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/PipeSinkTest.java index cf311639ee9a2..8ca88303c16da 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/PipeSinkTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/PipeSinkTest.java @@ -24,6 +24,9 @@ import org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant; import org.apache.iotdb.commons.pipe.config.plugin.configuraion.PipeTaskRuntimeConfiguration; import org.apache.iotdb.commons.pipe.config.plugin.env.PipeTaskSinkRuntimeEnvironment; +import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeRegistry; +import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeSnapshot; +import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent; import org.apache.iotdb.db.pipe.sink.protocol.legacy.IoTDBLegacyPipeSink; import org.apache.iotdb.db.pipe.sink.protocol.opcua.OpcUaSink; @@ -31,6 +34,7 @@ import org.apache.iotdb.db.pipe.sink.protocol.thrift.sync.IoTDBDataRegionSyncSink; import org.apache.iotdb.db.pipe.sink.protocol.websocket.WebSocketConnectorServer; import org.apache.iotdb.db.pipe.sink.protocol.websocket.WebSocketSink; +import org.apache.iotdb.db.pipe.sink.protocol.writeback.WriteBackSink; import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameterValidator; import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters; import org.apache.iotdb.pipe.api.exception.PipeException; @@ -109,6 +113,60 @@ public void testIoTDBThriftAsyncConnectorToOthers() { } } + @Test + public void testWriteBackSinkRegistersReceiverRuntime() throws Exception { + final PipeReceiverRuntimeRegistry registry = PipeReceiverRuntimeRegistry.getInstance(); + final int previousDataNodeId = IoTDBDescriptor.getInstance().getConfig().getDataNodeId(); + registry.clear(); + + try { + IoTDBDescriptor.getInstance().getConfig().setDataNodeId(1); + final String pipeName = "write_back_receiver_runtime_pipe"; + final long creationTime = 1L; + final PipeParameters parameters = + new PipeParameters( + new HashMap() { + { + put( + PipeSinkConstant.SINK_KEY, + BuiltinPipePlugin.WRITE_BACK_SINK.getPipePluginName()); + put(PipeSinkConstant.SINK_IOTDB_USER_ID, "0"); + put(PipeSinkConstant.SINK_IOTDB_USERNAME_KEY, "root"); + } + }); + + try (final WriteBackSink sink = new WriteBackSink()) { + sink.validate(new PipeParameterValidator(parameters)); + sink.customize( + parameters, + new PipeTaskRuntimeConfiguration( + new PipeTaskSinkRuntimeEnvironment(pipeName, creationTime, 1))); + + final List snapshots = registry.snapshot(); + Assert.assertEquals(1, snapshots.size()); + + final PipeReceiverRuntimeSnapshot snapshot = snapshots.get(0); + Assert.assertEquals( + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, snapshot.getReceiverNodeType()); + Assert.assertEquals(PipeReceiverRuntimeRegistry.PROTOCOL_WRITEBACK, snapshot.getProtocol()); + Assert.assertFalse(snapshot.getSenderAddress().isEmpty()); + Assert.assertFalse(snapshot.getSenderPorts().isEmpty()); + Assert.assertEquals(1, snapshot.getConnectionCount()); + Assert.assertEquals(1, snapshot.getPipeCount()); + Assert.assertTrue(snapshot.getPipeIds().contains(pipeName + "@")); + Assert.assertEquals("root", snapshot.getUserName()); + Assert.assertFalse(snapshot.getSenderClusterId().isEmpty()); + Assert.assertTrue(snapshot.getLastHandshakeTime() > 0); + Assert.assertTrue(snapshot.getLastTransferTime() >= snapshot.getLastHandshakeTime()); + } + + Assert.assertTrue(registry.snapshot().isEmpty()); + } finally { + IoTDBDescriptor.getInstance().getConfig().setDataNodeId(previousDataNodeId); + registry.clear(); + } + } + @Test public void testAsyncSinkDropDoesNotRequeueDroppedPipeEvents() throws Exception { try (final IoTDBDataRegionAsyncSink connector = new IoTDBDataRegionAsyncSink()) { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/client/ConfigNodeInfoTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/client/ConfigNodeInfoTest.java new file mode 100644 index 0000000000000..3807856640c49 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/client/ConfigNodeInfoTest.java @@ -0,0 +1,76 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.protocol.client; + +import org.apache.iotdb.common.rpc.thrift.TConfigNodeLocation; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.commons.file.SystemPropertiesHandler; + +import org.junit.Before; +import org.junit.Test; + +import java.util.Arrays; +import java.util.Collections; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +public class ConfigNodeInfoTest { + + @Before + public void setUp() { + ConfigNodeInfo.reinitializeStatics(); + ConfigNodeInfo.getInstance().systemPropertiesHandler = new NoopSystemPropertiesHandler(); + } + + @Test + public void testUpdateConfigNodeLocationsCachesNodeIds() { + final ConfigNodeInfo configNodeInfo = ConfigNodeInfo.getInstance(); + final TEndPoint firstInternalEndPoint = new TEndPoint("127.0.0.1", 10710); + final TEndPoint secondInternalEndPoint = new TEndPoint("127.0.0.2", 10710); + + assertTrue( + configNodeInfo.updateConfigNodeLocations( + Arrays.asList( + new TConfigNodeLocation( + 1, firstInternalEndPoint, new TEndPoint("127.0.0.1", 10720)), + new TConfigNodeLocation( + 2, secondInternalEndPoint, new TEndPoint("127.0.0.2", 10720))))); + assertEquals(1, configNodeInfo.getConfigNodeId(new TEndPoint("127.0.0.1", 10710))); + assertEquals(2, configNodeInfo.getConfigNodeId(new TEndPoint("127.0.0.2", 10710))); + + assertTrue( + configNodeInfo.updateConfigNodeList(Collections.singletonList(firstInternalEndPoint))); + assertEquals(1, configNodeInfo.getConfigNodeId(new TEndPoint("127.0.0.1", 10710))); + assertEquals(-1, configNodeInfo.getConfigNodeId(new TEndPoint("127.0.0.2", 10710))); + } + + private static class NoopSystemPropertiesHandler extends SystemPropertiesHandler { + + private NoopSystemPropertiesHandler() { + super("target/noop-system.properties"); + } + + @Override + public boolean fileExist() { + return false; + } + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/source/ShowReceiversOperatorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/source/ShowReceiversOperatorTest.java new file mode 100644 index 0000000000000..81de47e62eaf4 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/source/ShowReceiversOperatorTest.java @@ -0,0 +1,571 @@ +/* + * 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.queryengine.execution.operator.source; + +import org.apache.iotdb.commons.audit.UserEntity; +import org.apache.iotdb.commons.auth.entity.PrivilegeType; +import org.apache.iotdb.commons.auth.entity.User; +import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeRegistry; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.queryengine.utils.DateTimeUtils; +import org.apache.iotdb.db.auth.AuthorityChecker; +import org.apache.iotdb.db.queryengine.common.header.DatasetHeaderFactory; + +import com.google.common.collect.ImmutableList; +import org.apache.tsfile.read.common.block.TsBlock; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.CONNECTION_COUNT; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.LAST_HANDSHAKE_TIME; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.LAST_TRANSFER_TIME; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.PIPE_COUNT; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.PIPE_IDS; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.PROTOCOL; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.RECEIVER_NODE_ID; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.RECEIVER_NODE_TYPE; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.RECEIVER_USER_NAME; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.SENDER_ADDRESS; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.SENDER_CLUSTER_ID; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.SENDER_PORTS; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +public class ShowReceiversOperatorTest { + + private final PipeReceiverRuntimeRegistry registry = PipeReceiverRuntimeRegistry.getInstance(); + + @Before + public void setUp() { + registry.clear(); + AuthorityChecker.getAuthorityFetcher().getAuthorCache().invalidAllCache(); + } + + @After + public void tearDown() { + registry.clear(); + AuthorityChecker.getAuthorityFetcher().getAuthorCache().invalidAllCache(); + } + + @Test + public void testShowReceiversHeaderColumns() { + assertEquals( + ImmutableList.of( + RECEIVER_NODE_TYPE, + RECEIVER_NODE_ID, + PROTOCOL, + SENDER_CLUSTER_ID, + SENDER_ADDRESS, + RECEIVER_USER_NAME, + SENDER_PORTS, + CONNECTION_COUNT, + PIPE_COUNT, + PIPE_IDS, + LAST_HANDSHAKE_TIME, + LAST_TRANSFER_TIME), + DatasetHeaderFactory.getShowReceiversHeader().getRespColumns()); + assertTrue(DatasetHeaderFactory.getShowReceiversHeader().isIgnoreTimestamp()); + } + + @Test + public void testSingleSenderSingleDataNodeSingleConnectionSinglePipeOutput() { + registry.registerOrUpdateSession( + "data-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + registry.markTransfer("data-1", 200); + + final ShowReceiversOperator operator = + new ShowReceiversOperator(null, new PlanNodeId("show-receivers")); + + assertTrue(operator.hasNext()); + final TsBlock tsBlock = operator.next(); + + assertEquals(1, tsBlock.getPositionCount()); + assertEquals(12, tsBlock.getValueColumnCount()); + assertEquals(PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, getText(tsBlock, 0)); + assertEquals(1, tsBlock.getColumn(1).getInt(0)); + assertEquals(PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, getText(tsBlock, 2)); + assertEquals("cluster-a", getText(tsBlock, 3)); + assertEquals("10.0.0.1", getText(tsBlock, 4)); + assertEquals("root", getText(tsBlock, 5)); + assertEquals("9001", getText(tsBlock, 6)); + assertEquals(1, tsBlock.getColumn(7).getInt(0)); + assertEquals(1, tsBlock.getColumn(8).getInt(0)); + assertTrue(getText(tsBlock, 9).contains("pipe-a@")); + assertEquals(DateTimeUtils.convertLongToDate(100, "ms"), getText(tsBlock, 10)); + assertEquals(DateTimeUtils.convertLongToDate(200, "ms"), getText(tsBlock, 11)); + } + + @Test + public void testMultiplePipesFromSameSenderAreAggregatedInOutput() { + registry.registerOrUpdateSession( + "data-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + registry.registerOrUpdateSession( + "data-2", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9002, + "root", + "cluster-a", + "pipe-b", + 2, + 200); + + final ShowReceiversOperator operator = + new ShowReceiversOperator(null, new PlanNodeId("show-receivers")); + + assertTrue(operator.hasNext()); + final TsBlock tsBlock = operator.next(); + + assertEquals(1, tsBlock.getPositionCount()); + assertEquals("cluster-a", getText(tsBlock, 3)); + assertEquals("10.0.0.1", getText(tsBlock, 4)); + assertEquals("root", getText(tsBlock, 5)); + assertEquals("9001,9002", getText(tsBlock, 6)); + assertEquals(2, tsBlock.getColumn(7).getInt(0)); + assertEquals(2, tsBlock.getColumn(8).getInt(0)); + assertTrue(getText(tsBlock, 9).contains("pipe-a@")); + assertTrue(getText(tsBlock, 9).contains("pipe-b@")); + assertFalse(operator.hasNext()); + } + + @Test + public void testLastTransferTimeUsesHandshakeTimeBeforeTransfer() { + registry.registerOrUpdateSession( + "data-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + + final ShowReceiversOperator operator = + new ShowReceiversOperator(null, new PlanNodeId("show-receivers")); + + assertTrue(operator.hasNext()); + final TsBlock tsBlock = operator.next(); + + assertEquals(1, tsBlock.getPositionCount()); + assertEquals(DateTimeUtils.convertLongToDate(100, "ms"), getText(tsBlock, 10)); + assertEquals(DateTimeUtils.convertLongToDate(100, "ms"), getText(tsBlock, 11)); + } + + @Test + public void testUnknownReceiverNodeIdIsNull() { + registry.registerOrUpdateSession( + "config-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + -1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "127.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + + final ShowReceiversOperator operator = + new ShowReceiversOperator(null, new PlanNodeId("show-receivers")); + + assertTrue(operator.hasNext()); + final TsBlock tsBlock = operator.next(); + + assertEquals(1, tsBlock.getPositionCount()); + assertEquals(12, tsBlock.getValueColumnCount()); + assertTrue(tsBlock.getColumn(1).isNull(0)); + } + + @Test + public void testShowReceiversWritesUnknownPipeIdsAndTimesForUnboundSession() { + registry.registerOrUpdateSession( + "data-legacy", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9001, + "root", + PipeReceiverRuntimeRegistry.UNKNOWN, + null, + Long.MIN_VALUE, + 0); + + final ShowReceiversOperator operator = + new ShowReceiversOperator(null, new PlanNodeId("show-receivers")); + + assertTrue(operator.hasNext()); + final TsBlock tsBlock = operator.next(); + + assertEquals(1, tsBlock.getPositionCount()); + assertEquals(PipeReceiverRuntimeRegistry.UNKNOWN, getText(tsBlock, 3)); + assertEquals("10.0.0.1", getText(tsBlock, 4)); + assertEquals("root", getText(tsBlock, 5)); + assertEquals("9001", getText(tsBlock, 6)); + assertEquals(1, tsBlock.getColumn(7).getInt(0)); + assertEquals(0, tsBlock.getColumn(8).getInt(0)); + assertEquals(PipeReceiverRuntimeRegistry.UNKNOWN, getText(tsBlock, 9)); + assertEquals(PipeReceiverRuntimeRegistry.UNKNOWN, getText(tsBlock, 10)); + assertEquals(PipeReceiverRuntimeRegistry.UNKNOWN, getText(tsBlock, 11)); + assertFalse(operator.hasNext()); + } + + @Test + public void testNormalUserOnlySeesOwnReceiverSnapshots() { + registerUserSession("data-user1", "10.0.0.1", 9001, "user1", "cluster-a", "pipe-a", 1, 100); + registerUserSession("data-user2", "10.0.0.2", 9002, "user2", "cluster-b", "pipe-b", 2, 200); + + final ShowReceiversOperator operator = + new ShowReceiversOperator( + null, new PlanNodeId("show-receivers"), new UserEntity(1L, "user1", "127.0.0.1")); + + assertTrue(operator.hasNext()); + final TsBlock tsBlock = operator.next(); + + assertEquals(1, tsBlock.getPositionCount()); + assertEquals("cluster-a", getText(tsBlock, 3)); + assertEquals("10.0.0.1", getText(tsBlock, 4)); + assertEquals("user1", getText(tsBlock, 5)); + assertFalse(operator.hasNext()); + } + + @Test + public void testRootUserSeesAllReceiverSnapshots() { + registerUserSession("data-user1", "10.0.0.1", 9001, "user1", "cluster-a", "pipe-a", 1, 100); + registerUserSession("data-user2", "10.0.0.2", 9002, "user2", "cluster-b", "pipe-b", 2, 200); + + final ShowReceiversOperator operator = + new ShowReceiversOperator( + null, + new PlanNodeId("show-receivers"), + new UserEntity( + AuthorityChecker.SUPER_USER_ID, AuthorityChecker.SUPER_USER, "127.0.0.1")); + + assertTrue(operator.hasNext()); + final TsBlock tsBlock = operator.next(); + + assertEquals(2, tsBlock.getPositionCount()); + assertEquals("user1", getText(tsBlock, 5, 0)); + assertEquals("user2", getText(tsBlock, 5, 1)); + assertFalse(operator.hasNext()); + } + + @Test + public void testSystemUserSeesAllReceiverSnapshots() { + registerUserSession("data-user1", "10.0.0.1", 9001, "user1", "cluster-a", "pipe-a", 1, 100); + registerUserSession("data-user2", "10.0.0.2", 9002, "user2", "cluster-b", "pipe-b", 2, 200); + + final User systemUser = new User("system_user", "password"); + systemUser.grantSysPrivilege(PrivilegeType.SYSTEM, false); + AuthorityChecker.getAuthorityFetcher() + .getAuthorCache() + .putUserCache(systemUser.getName(), systemUser); + + final ShowReceiversOperator operator = + new ShowReceiversOperator( + null, new PlanNodeId("show-receivers"), new UserEntity(2L, "system_user", "127.0.0.1")); + + assertTrue(operator.hasNext()); + final TsBlock tsBlock = operator.next(); + + assertEquals(2, tsBlock.getPositionCount()); + assertEquals("user1", getText(tsBlock, 5, 0)); + assertEquals("user2", getText(tsBlock, 5, 1)); + assertFalse(operator.hasNext()); + } + + @Test + public void testNullUserEntitySeesAllReceiverSnapshots() { + registerUserSession("data-user1", "10.0.0.1", 9001, "user1", "cluster-a", "pipe-a", 1, 100); + registerUserSession("data-user2", "10.0.0.2", 9002, "user2", "cluster-b", "pipe-b", 2, 200); + + final ShowReceiversOperator operator = + new ShowReceiversOperator(null, new PlanNodeId("show-receivers"), null); + + assertTrue(operator.hasNext()); + final TsBlock tsBlock = operator.next(); + + assertEquals(2, tsBlock.getPositionCount()); + assertEquals("user1", getText(tsBlock, 5, 0)); + assertEquals("user2", getText(tsBlock, 5, 1)); + assertFalse(operator.hasNext()); + } + + @Test + public void testShowReceiversOutputKeepsDefaultSnapshotOrdering() { + registry.registerOrUpdateSession( + "data-cluster-b", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9007, + "root", + "cluster-b", + "pipe-g", + 7, + 700); + registry.registerOrUpdateSession( + "data-user-b", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.2", + 9006, + "bob", + "cluster-a", + "pipe-f", + 6, + 600); + registry.registerOrUpdateSession( + "data-address-a", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9004, + "root", + "cluster-a", + "pipe-d", + 4, + 400); + registry.registerOrUpdateSession( + "config-node-2", + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + 2, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9002, + "root", + "cluster-a", + "pipe-b", + 2, + 200); + registry.registerOrUpdateSession( + "data-user-a", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.2", + 9005, + "alice", + "cluster-a", + "pipe-e", + 5, + 500); + registry.registerOrUpdateSession( + "config-node-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + registry.registerOrUpdateSession( + "data-air-gap", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_AIR_GAP, + "10.0.0.2", + 9003, + "root", + "cluster-a", + "pipe-c", + 3, + 300); + + final ShowReceiversOperator operator = + new ShowReceiversOperator(null, new PlanNodeId("show-receivers")); + + assertTrue(operator.hasNext()); + final TsBlock tsBlock = operator.next(); + + assertEquals(7, tsBlock.getPositionCount()); + assertRowKey( + tsBlock, + 0, + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "cluster-a", + "10.0.0.1", + "root"); + assertRowKey( + tsBlock, + 1, + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + 2, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "cluster-a", + "10.0.0.1", + "root"); + assertRowKey( + tsBlock, + 2, + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_AIR_GAP, + "cluster-a", + "10.0.0.2", + "root"); + assertRowKey( + tsBlock, + 3, + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "cluster-a", + "10.0.0.1", + "root"); + assertRowKey( + tsBlock, + 4, + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "cluster-a", + "10.0.0.2", + "alice"); + assertRowKey( + tsBlock, + 5, + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "cluster-a", + "10.0.0.2", + "bob"); + assertRowKey( + tsBlock, + 6, + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "cluster-b", + "10.0.0.1", + "root"); + assertFalse(operator.hasNext()); + } + + @Test + public void testSameSenderAddressAndUserWithDifferentClustersAreNotAggregatedInOutput() { + registerUserSession("data-cluster-a", "10.0.0.1", 9001, "root", "cluster-a", "pipe-a", 1, 100); + registerUserSession("data-cluster-b", "10.0.0.1", 9002, "root", "cluster-b", "pipe-b", 2, 200); + + final ShowReceiversOperator operator = + new ShowReceiversOperator(null, new PlanNodeId("show-receivers")); + + assertTrue(operator.hasNext()); + final TsBlock tsBlock = operator.next(); + + assertEquals(2, tsBlock.getPositionCount()); + assertEquals("cluster-a", getText(tsBlock, 3, 0)); + assertEquals("9001", getText(tsBlock, 6, 0)); + assertEquals(1, tsBlock.getColumn(7).getInt(0)); + assertEquals(1, tsBlock.getColumn(8).getInt(0)); + assertTrue(getText(tsBlock, 9, 0).contains("pipe-a@")); + assertEquals("cluster-b", getText(tsBlock, 3, 1)); + assertEquals("9002", getText(tsBlock, 6, 1)); + assertEquals(1, tsBlock.getColumn(7).getInt(1)); + assertEquals(1, tsBlock.getColumn(8).getInt(1)); + assertTrue(getText(tsBlock, 9, 1).contains("pipe-b@")); + assertFalse(operator.hasNext()); + } + + private void registerUserSession( + String connectionKey, + String senderAddress, + int senderPort, + String userName, + String senderClusterId, + String pipeName, + long pipeCreationTime, + long handshakeTime) { + registry.registerOrUpdateSession( + connectionKey, + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + senderAddress, + senderPort, + userName, + senderClusterId, + pipeName, + pipeCreationTime, + handshakeTime); + } + + private static String getText(final TsBlock tsBlock, final int columnIndex) { + return getText(tsBlock, columnIndex, 0); + } + + private static String getText(final TsBlock tsBlock, final int columnIndex, final int position) { + return tsBlock.getColumn(columnIndex).getBinary(position).toString(); + } + + private static void assertRowKey( + final TsBlock tsBlock, + final int position, + final String receiverNodeType, + final int receiverNodeId, + final String protocol, + final String senderClusterId, + final String senderAddress, + final String userName) { + assertEquals(receiverNodeType, getText(tsBlock, 0, position)); + assertEquals(receiverNodeId, tsBlock.getColumn(1).getInt(position)); + assertEquals(protocol, getText(tsBlock, 2, position)); + assertEquals(senderClusterId, getText(tsBlock, 3, position)); + assertEquals(senderAddress, getText(tsBlock, 4, position)); + assertEquals(userName, getText(tsBlock, 5, position)); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/InformationSchemaReceiversSupplierTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/InformationSchemaReceiversSupplierTest.java new file mode 100644 index 0000000000000..a206623a51d37 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/InformationSchemaReceiversSupplierTest.java @@ -0,0 +1,406 @@ +/* + * 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.queryengine.execution.operator.source.relational; + +import org.apache.iotdb.commons.audit.UserEntity; +import org.apache.iotdb.commons.auth.entity.PrivilegeType; +import org.apache.iotdb.commons.auth.entity.User; +import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeRegistry; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.queryengine.plan.relational.metadata.ColumnSchema; +import org.apache.iotdb.commons.queryengine.plan.relational.metadata.QualifiedObjectName; +import org.apache.iotdb.commons.queryengine.plan.relational.planner.Symbol; +import org.apache.iotdb.commons.schema.table.InformationSchema; +import org.apache.iotdb.commons.schema.table.TsTable; +import org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory; +import org.apache.iotdb.commons.schema.table.column.TsTableColumnSchema; +import org.apache.iotdb.db.auth.AuthorityChecker; +import org.apache.iotdb.db.queryengine.execution.operator.source.ShowReceiversOperator; +import org.apache.iotdb.db.queryengine.plan.relational.planner.node.InformationSchemaTableScanNode; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.read.common.block.TsBlock; +import org.apache.tsfile.read.common.type.TypeFactory; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +public class InformationSchemaReceiversSupplierTest { + + private final PipeReceiverRuntimeRegistry registry = PipeReceiverRuntimeRegistry.getInstance(); + + @Before + public void setUp() { + registry.clear(); + AuthorityChecker.getAuthorityFetcher().getAuthorCache().invalidAllCache(); + } + + @After + public void tearDown() { + registry.clear(); + AuthorityChecker.getAuthorityFetcher().getAuthorCache().invalidAllCache(); + } + + @Test + public void testReceiversInformationSchemaTableDefinition() { + final List columns = receiversTable().getColumnList(); + + assertEquals(12, columns.size()); + assertColumn( + columns.get(0), "receiver_node_type", TSDataType.STRING, TsTableColumnCategory.TAG); + assertColumn(columns.get(1), "receiver_node_id", TSDataType.INT32, TsTableColumnCategory.TAG); + assertColumn(columns.get(2), "protocol", TSDataType.STRING, TsTableColumnCategory.TAG); + assertColumn(columns.get(3), "sender_cluster_id", TSDataType.STRING, TsTableColumnCategory.TAG); + assertColumn(columns.get(4), "sender_address", TSDataType.STRING, TsTableColumnCategory.TAG); + assertColumn(columns.get(5), "user_name", TSDataType.STRING, TsTableColumnCategory.TAG); + assertColumn( + columns.get(6), "sender_ports", TSDataType.STRING, TsTableColumnCategory.ATTRIBUTE); + assertColumn( + columns.get(7), "connection_count", TSDataType.INT32, TsTableColumnCategory.ATTRIBUTE); + assertColumn(columns.get(8), "pipe_count", TSDataType.INT32, TsTableColumnCategory.ATTRIBUTE); + assertColumn(columns.get(9), "pipe_ids", TSDataType.STRING, TsTableColumnCategory.ATTRIBUTE); + assertColumn( + columns.get(10), + "last_handshake_time", + TSDataType.TIMESTAMP, + TsTableColumnCategory.ATTRIBUTE); + assertColumn( + columns.get(11), + "last_transfer_time", + TSDataType.TIMESTAMP, + TsTableColumnCategory.ATTRIBUTE); + } + + @Test + public void testReceiversSupplierBuildsTableModelRows() { + registry.registerOrUpdateSession( + "data-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + registry.markTransfer("data-1", 200); + + final InformationSchemaContentSupplierFactory.IInformationSchemaContentSupplier supplier = + InformationSchemaContentSupplierFactory.getSupplier( + null, dataTypes(), rootUser(), receiversScanNode()); + + assertTrue(supplier.hasNext()); + final TsBlock tsBlock = supplier.next(); + + assertEquals(1, tsBlock.getPositionCount()); + assertEquals(PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, getText(tsBlock, 0)); + assertEquals(1, tsBlock.getColumn(1).getInt(0)); + assertEquals(PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, getText(tsBlock, 2)); + assertEquals("cluster-a", getText(tsBlock, 3)); + assertEquals("10.0.0.1", getText(tsBlock, 4)); + assertEquals("root", getText(tsBlock, 5)); + assertEquals("9001", getText(tsBlock, 6)); + assertEquals(1, tsBlock.getColumn(7).getInt(0)); + assertEquals(1, tsBlock.getColumn(8).getInt(0)); + assertTrue(getText(tsBlock, 9).contains("pipe-a@")); + assertEquals(100L, tsBlock.getColumn(10).getLong(0)); + assertEquals(200L, tsBlock.getColumn(11).getLong(0)); + assertFalse(supplier.hasNext()); + } + + @Test + public void testReceiversSupplierUsesHandshakeTimeAsInitialTransferTime() { + registry.registerOrUpdateSession( + "data-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + + final InformationSchemaContentSupplierFactory.IInformationSchemaContentSupplier supplier = + InformationSchemaContentSupplierFactory.getSupplier( + null, dataTypes(), rootUser(), receiversScanNode()); + + assertTrue(supplier.hasNext()); + final TsBlock tsBlock = supplier.next(); + + assertEquals(1, tsBlock.getPositionCount()); + assertEquals(100L, tsBlock.getColumn(10).getLong(0)); + assertEquals(100L, tsBlock.getColumn(11).getLong(0)); + assertFalse(supplier.hasNext()); + } + + @Test + public void testReceiversSupplierWritesUnknownFieldsAsNulls() { + registry.registerOrUpdateSession( + "config-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + -1, + PipeReceiverRuntimeRegistry.PROTOCOL_AIR_GAP, + "10.0.0.2", + 9002, + "root", + "cluster-b", + null, + Long.MIN_VALUE, + 0); + + final InformationSchemaContentSupplierFactory.IInformationSchemaContentSupplier supplier = + InformationSchemaContentSupplierFactory.getSupplier( + null, dataTypes(), rootUser(), receiversScanNode()); + + assertTrue(supplier.hasNext()); + final TsBlock tsBlock = supplier.next(); + + assertEquals(1, tsBlock.getPositionCount()); + assertEquals(PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, getText(tsBlock, 0)); + assertTrue(tsBlock.getColumn(1).isNull(0)); + assertEquals(PipeReceiverRuntimeRegistry.PROTOCOL_AIR_GAP, getText(tsBlock, 2)); + assertEquals("cluster-b", getText(tsBlock, 3)); + assertEquals("10.0.0.2", getText(tsBlock, 4)); + assertEquals("root", getText(tsBlock, 5)); + assertEquals(0, tsBlock.getColumn(8).getInt(0)); + assertEquals(PipeReceiverRuntimeRegistry.UNKNOWN, getText(tsBlock, 9)); + assertTrue(tsBlock.getColumn(10).isNull(0)); + assertTrue(tsBlock.getColumn(11).isNull(0)); + assertFalse(supplier.hasNext()); + } + + @Test + public void testReceiversSupplierFiltersByNormalUser() { + registerUserSession("data-user1", "10.0.0.1", 9001, "user1", "cluster-a", "pipe-a", 1, 100); + registerUserSession("data-user2", "10.0.0.2", 9002, "user2", "cluster-b", "pipe-b", 2, 200); + + final InformationSchemaContentSupplierFactory.IInformationSchemaContentSupplier supplier = + InformationSchemaContentSupplierFactory.getSupplier( + null, dataTypes(), new UserEntity(1L, "user1", "127.0.0.1"), receiversScanNode()); + + assertTrue(supplier.hasNext()); + final TsBlock tsBlock = supplier.next(); + + assertEquals(1, tsBlock.getPositionCount()); + assertEquals("cluster-a", getText(tsBlock, 3)); + assertEquals("10.0.0.1", getText(tsBlock, 4)); + assertEquals("user1", getText(tsBlock, 5)); + assertFalse(supplier.hasNext()); + } + + @Test + public void testReceiversSupplierRootSeesAllReceiverSnapshots() { + registerUserSession("data-user1", "10.0.0.1", 9001, "user1", "cluster-a", "pipe-a", 1, 100); + registerUserSession("data-user2", "10.0.0.2", 9002, "user2", "cluster-b", "pipe-b", 2, 200); + + final InformationSchemaContentSupplierFactory.IInformationSchemaContentSupplier supplier = + InformationSchemaContentSupplierFactory.getSupplier( + null, dataTypes(), rootUser(), receiversScanNode()); + + assertTrue(supplier.hasNext()); + final TsBlock tsBlock = supplier.next(); + + assertEquals(2, tsBlock.getPositionCount()); + assertEquals("user1", getText(tsBlock, 5, 0)); + assertEquals("user2", getText(tsBlock, 5, 1)); + assertFalse(supplier.hasNext()); + } + + @Test + public void testReceiversSupplierSystemUserSeesAllReceiverSnapshots() { + registerUserSession("data-user1", "10.0.0.1", 9001, "user1", "cluster-a", "pipe-a", 1, 100); + registerUserSession("data-user2", "10.0.0.2", 9002, "user2", "cluster-b", "pipe-b", 2, 200); + + final User systemUser = new User("system_user", "password"); + systemUser.grantSysPrivilege(PrivilegeType.SYSTEM, false); + AuthorityChecker.getAuthorityFetcher() + .getAuthorCache() + .putUserCache(systemUser.getName(), systemUser); + + final InformationSchemaContentSupplierFactory.IInformationSchemaContentSupplier supplier = + InformationSchemaContentSupplierFactory.getSupplier( + null, dataTypes(), new UserEntity(2L, "system_user", "127.0.0.1"), receiversScanNode()); + + assertTrue(supplier.hasNext()); + final TsBlock tsBlock = supplier.next(); + + assertEquals(2, tsBlock.getPositionCount()); + assertEquals("user1", getText(tsBlock, 5, 0)); + assertEquals("user2", getText(tsBlock, 5, 1)); + assertFalse(supplier.hasNext()); + } + + @Test + public void testReceiversSupplierNullUserEntitySeesAllReceiverSnapshots() { + registerUserSession("data-user1", "10.0.0.1", 9001, "user1", "cluster-a", "pipe-a", 1, 100); + registerUserSession("data-user2", "10.0.0.2", 9002, "user2", "cluster-b", "pipe-b", 2, 200); + + final InformationSchemaContentSupplierFactory.IInformationSchemaContentSupplier supplier = + InformationSchemaContentSupplierFactory.getSupplier( + null, dataTypes(), null, receiversScanNode()); + + assertTrue(supplier.hasNext()); + final TsBlock tsBlock = supplier.next(); + + assertEquals(2, tsBlock.getPositionCount()); + assertEquals("user1", getText(tsBlock, 5, 0)); + assertEquals("user2", getText(tsBlock, 5, 1)); + assertFalse(supplier.hasNext()); + } + + @Test + public void testReceiversSupplierMatchesShowReceiversOutput() { + registry.registerOrUpdateSession( + "data-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + registry.markTransfer("data-1", 200); + + final ShowReceiversOperator showReceiversOperator = + new ShowReceiversOperator(null, new PlanNodeId("show-receivers"), rootUser()); + assertTrue(showReceiversOperator.hasNext()); + final TsBlock showReceiversTsBlock = showReceiversOperator.next(); + final InformationSchemaContentSupplierFactory.IInformationSchemaContentSupplier supplier = + InformationSchemaContentSupplierFactory.getSupplier( + null, dataTypes(), rootUser(), receiversScanNode()); + assertTrue(supplier.hasNext()); + final TsBlock informationSchemaTsBlock = supplier.next(); + + assertEquals( + showReceiversTsBlock.getPositionCount(), informationSchemaTsBlock.getPositionCount()); + assertEquals(1, informationSchemaTsBlock.getPositionCount()); + for (int columnIndex : new int[] {0, 2, 3, 4, 5, 6, 9}) { + assertEquals( + getText(showReceiversTsBlock, columnIndex), + getText(informationSchemaTsBlock, columnIndex)); + } + assertEquals( + showReceiversTsBlock.getColumn(1).getInt(0), + informationSchemaTsBlock.getColumn(1).getInt(0)); + assertEquals( + showReceiversTsBlock.getColumn(7).getInt(0), + informationSchemaTsBlock.getColumn(7).getInt(0)); + assertEquals( + showReceiversTsBlock.getColumn(8).getInt(0), + informationSchemaTsBlock.getColumn(8).getInt(0)); + assertFalse(supplier.hasNext()); + } + + private static InformationSchemaTableScanNode receiversScanNode() { + final List outputSymbols = new ArrayList<>(); + final Map assignments = new LinkedHashMap<>(); + for (TsTableColumnSchema columnSchema : receiversTable().getColumnList()) { + final Symbol symbol = new Symbol(columnSchema.getColumnName()); + outputSymbols.add(symbol); + assignments.put( + symbol, + new ColumnSchema( + columnSchema.getColumnName(), + TypeFactory.getType(columnSchema.getDataType()), + false, + columnSchema.getColumnCategory())); + } + return new InformationSchemaTableScanNode( + new PlanNodeId("information-schema-receivers"), + new QualifiedObjectName( + InformationSchema.INFORMATION_DATABASE, InformationSchema.RECEIVERS), + outputSymbols, + assignments); + } + + private static List dataTypes() { + final List dataTypes = new ArrayList<>(); + for (TsTableColumnSchema columnSchema : receiversTable().getColumnList()) { + dataTypes.add(columnSchema.getDataType()); + } + return dataTypes; + } + + private static TsTable receiversTable() { + return InformationSchema.getSchemaTables().get(InformationSchema.RECEIVERS); + } + + private static void assertColumn( + final TsTableColumnSchema column, + final String columnName, + final TSDataType dataType, + final TsTableColumnCategory columnCategory) { + assertEquals(columnName, column.getColumnName()); + assertEquals(dataType, column.getDataType()); + assertEquals(columnCategory, column.getColumnCategory()); + } + + private void registerUserSession( + String connectionKey, + String senderAddress, + int senderPort, + String userName, + String senderClusterId, + String pipeName, + long pipeCreationTime, + long handshakeTime) { + registry.registerOrUpdateSession( + connectionKey, + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + senderAddress, + senderPort, + userName, + senderClusterId, + pipeName, + pipeCreationTime, + handshakeTime); + } + + private static UserEntity rootUser() { + return new UserEntity(AuthorityChecker.SUPER_USER_ID, AuthorityChecker.SUPER_USER, "127.0.0.1"); + } + + private static String getText(final TsBlock tsBlock, final int columnIndex) { + return getText(tsBlock, columnIndex, 0); + } + + private static String getText(final TsBlock tsBlock, final int columnIndex, final int position) { + return tsBlock.getColumn(columnIndex).getBinary(position).toString(); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGeneratorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGeneratorTest.java index febf29bea1553..cf0ce73308f62 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGeneratorTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGeneratorTest.java @@ -66,6 +66,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.sys.AuthorStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowDiskUsageStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowQueriesStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowReceiversStatement; import org.apache.iotdb.isession.template.TemplateNode; import org.apache.iotdb.rpc.StatementExecutionException; import org.apache.iotdb.service.rpc.thrift.TSAggregationQueryReq; @@ -88,6 +89,7 @@ import org.apache.iotdb.service.rpc.thrift.TSUnsetSchemaTemplateReq; import org.apache.iotdb.session.template.MeasurementNode; +import org.antlr.v4.runtime.misc.ParseCancellationException; import org.apache.tsfile.enums.ColumnCategory; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.metadata.enums.CompressionType; @@ -196,6 +198,32 @@ public void testShowQueries() { "show queries order by a", ZonedDateTime.now().getOffset())); } + @Test + public void testShowReceivers() { + final Statement showReceivers = + StatementGenerator.createStatement("show receivers", ZonedDateTime.now().getOffset()); + Assert.assertTrue(showReceivers instanceof ShowReceiversStatement); + Assert.assertEquals( + Arrays.asList( + new SortItem("ReceiverNodeType", Ordering.ASC), + new SortItem("ReceiverNodeId", Ordering.ASC), + new SortItem("Protocol", Ordering.ASC), + new SortItem("SenderClusterId", Ordering.ASC), + new SortItem("SenderAddress", Ordering.ASC), + new SortItem("UserName", Ordering.ASC)), + ((ShowReceiversStatement) showReceivers).getSortItemList()); + Assert.assertThrows( + ParseCancellationException.class, + () -> + StatementGenerator.createStatement( + "show receivers where protocol = 'thrift'", ZonedDateTime.now().getOffset())); + Assert.assertThrows( + ParseCancellationException.class, + () -> + StatementGenerator.createStatement( + "show receivers order by ReceiverNodeId", ZonedDateTime.now().getOffset())); + } + @Test public void testRawDataQuery() throws IllegalPathException { TSRawDataQueryReq req = diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/source/SourceNodeSerdeTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/source/SourceNodeSerdeTest.java index 5597a6ca8e600..9df412f2b64f9 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/source/SourceNodeSerdeTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/source/SourceNodeSerdeTest.java @@ -31,6 +31,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.LastQueryScanNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.ShowDiskUsageNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.ShowQueriesNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.ShowReceiversNode; import org.apache.iotdb.db.queryengine.plan.relational.planner.node.TableDiskUsageInformationSchemaTableScanNode; import org.apache.tsfile.enums.TSDataType; @@ -117,6 +118,16 @@ public void testShowQueriesNode() throws IllegalPathException { assertEquals(PlanNodeDeserializeHelper.deserialize(byteBuffer), node); } + @Test + public void testShowReceiversNode() throws IllegalPathException { + ShowReceiversNode node = new ShowReceiversNode(new PlanNodeId("test"), null); + + ByteBuffer byteBuffer = ByteBuffer.allocate(2048); + node.serialize(byteBuffer); + byteBuffer.flip(); + assertEquals(PlanNodeDeserializeHelper.deserialize(byteBuffer), node); + } + @Test public void testTableDiskUsageInformationTableScanNode() throws IllegalPathException { List symbols = Arrays.asList(new Symbol("database"), new Symbol("size_in_bytes")); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/AuthTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/AuthTest.java index e758e49962827..a33179a872d00 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/AuthTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/AuthTest.java @@ -19,21 +19,29 @@ package org.apache.iotdb.db.queryengine.plan.relational.analyzer; +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.audit.AuditEventType; +import org.apache.iotdb.commons.audit.AuditLogOperation; import org.apache.iotdb.commons.exception.auth.AccessDeniedException; import org.apache.iotdb.commons.queryengine.common.SessionInfo; import org.apache.iotdb.commons.queryengine.common.SqlDialect; import org.apache.iotdb.commons.queryengine.plan.relational.metadata.QualifiedObjectName; import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Statement; import org.apache.iotdb.commons.queryengine.plan.relational.type.InternalTypeManager; +import org.apache.iotdb.commons.schema.table.InformationSchema; +import org.apache.iotdb.commons.utils.StatusUtils; import org.apache.iotdb.db.protocol.session.IClientSession; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; import org.apache.iotdb.db.queryengine.plan.execution.config.TableConfigTaskVisitor; import org.apache.iotdb.db.queryengine.plan.relational.security.AccessControlImpl; import org.apache.iotdb.db.queryengine.plan.relational.security.ITableAuthChecker; import org.apache.iotdb.db.queryengine.plan.relational.security.TableModelPrivilege; +import org.apache.iotdb.db.queryengine.plan.relational.security.TreeAccessCheckContext; import org.apache.iotdb.db.queryengine.plan.relational.security.TreeAccessCheckVisitor; import org.apache.iotdb.db.queryengine.plan.relational.sql.parser.SqlParser; import org.apache.iotdb.db.queryengine.plan.relational.sql.rewrite.StatementRewrite; +import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.ShowPipesStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowReceiversStatement; import org.junit.Test; import org.mockito.Mockito; @@ -47,6 +55,7 @@ import static org.apache.iotdb.db.queryengine.plan.relational.analyzer.TestUtils.QUERY_ID; import static org.apache.iotdb.db.queryengine.plan.relational.analyzer.TestUtils.TEST_MATADATA; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; @@ -109,6 +118,50 @@ public void testQueryRelatedAuth() { } } + @Test + public void testInformationSchemaReceiversUseSameLocalAuthAsPipes() { + final ITableAuthChecker authChecker = Mockito.mock(ITableAuthChecker.class); + + try { + analyzeSQL( + String.format( + "SELECT * FROM %s.%s", + InformationSchema.INFORMATION_DATABASE, InformationSchema.PIPES), + user1, + authChecker); + analyzeSQL( + String.format( + "SELECT * FROM %s.%s", + InformationSchema.INFORMATION_DATABASE, InformationSchema.RECEIVERS), + user1, + authChecker); + } catch (Exception e) { + fail("Unexpected exception : " + e.getMessage()); + } + + Mockito.verify(authChecker, Mockito.never()).checkGlobalPrivileges(any(), any(), any()); + } + + @Test + public void testTreeShowReceiversUsesSameLocalAuthAsShowPipes() { + final TreeAccessCheckVisitor visitor = new TreeAccessCheckVisitor(); + final TreeAccessCheckContext showPipesContext = + new TreeAccessCheckContext(1, user1, "127.0.0.1"); + final TreeAccessCheckContext showReceiversContext = + new TreeAccessCheckContext(1, user1, "127.0.0.1"); + + final TSStatus showPipesStatus = + visitor.visitShowPipes(new ShowPipesStatement(), showPipesContext); + final TSStatus showReceiversStatus = + visitor.visitShowReceivers(new ShowReceiversStatement(), showReceiversContext); + + assertEquals(StatusUtils.OK.getCode(), showPipesStatus.getCode()); + assertEquals(showPipesStatus.getCode(), showReceiversStatus.getCode()); + assertEquals(AuditLogOperation.QUERY, showReceiversContext.getAuditLogOperation()); + assertEquals(AuditEventType.OBJECT_AUTHENTICATION, showReceiversContext.getAuditEventType()); + assertTrue(showReceiversContext.getResult()); + } + @Test public void testDatabaseManagementRelatedAuth() { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/informationschema/ShowReceiversTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/informationschema/ShowReceiversTest.java new file mode 100644 index 0000000000000..56ffcc06bf6e2 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/informationschema/ShowReceiversTest.java @@ -0,0 +1,128 @@ +/* + * 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.queryengine.plan.relational.planner.informationschema; + +import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.ComparisonExpression; +import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.StringLiteral; +import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.SymbolReference; +import org.apache.iotdb.db.queryengine.plan.planner.plan.LogicalQueryPlan; +import org.apache.iotdb.db.queryengine.plan.relational.planner.PlanTester; + +import com.google.common.collect.ImmutableList; +import org.junit.Test; + +import java.util.Optional; + +import static org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.ComparisonExpression.Operator.EQUAL; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.CONNECTION_COUNT_TABLE_MODEL; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.LAST_HANDSHAKE_TIME_TABLE_MODEL; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.LAST_TRANSFER_TIME_TABLE_MODEL; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.PIPE_COUNT_TABLE_MODEL; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.PIPE_IDS_TABLE_MODEL; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.PROTOCOL_TABLE_MODEL; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.RECEIVER_NODE_ID_TABLE_MODEL; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.RECEIVER_NODE_TYPE_TABLE_MODEL; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.RECEIVER_USER_NAME_TABLE_MODEL; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.SENDER_ADDRESS_TABLE_MODEL; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.SENDER_CLUSTER_ID_TABLE_MODEL; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.SENDER_PORTS_TABLE_MODEL; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanAssert.assertPlan; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.filter; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.infoSchemaTableScan; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.output; +import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.project; + +public class ShowReceiversTest { + + private static final ImmutableList RECEIVERS_COLUMNS = + ImmutableList.of( + RECEIVER_NODE_TYPE_TABLE_MODEL, + RECEIVER_NODE_ID_TABLE_MODEL, + PROTOCOL_TABLE_MODEL, + SENDER_CLUSTER_ID_TABLE_MODEL, + SENDER_ADDRESS_TABLE_MODEL, + RECEIVER_USER_NAME_TABLE_MODEL, + SENDER_PORTS_TABLE_MODEL, + CONNECTION_COUNT_TABLE_MODEL, + PIPE_COUNT_TABLE_MODEL, + PIPE_IDS_TABLE_MODEL, + LAST_HANDSHAKE_TIME_TABLE_MODEL, + LAST_TRANSFER_TIME_TABLE_MODEL); + + private final PlanTester planTester = new PlanTester(); + + @Test + public void testShowReceiversRewrite() { + final LogicalQueryPlan logicalQueryPlan = planTester.createPlan("show receivers"); + assertPlan( + logicalQueryPlan, + output( + infoSchemaTableScan( + "information_schema.receivers", Optional.empty(), RECEIVERS_COLUMNS))); + } + + @Test + public void testSelectReceivers() { + final LogicalQueryPlan logicalQueryPlan = + planTester.createPlan("select * from information_schema.receivers"); + assertPlan( + logicalQueryPlan, + output( + infoSchemaTableScan( + "information_schema.receivers", Optional.empty(), RECEIVERS_COLUMNS))); + } + + @Test + public void testSelectReceiversWithProtocolFilter() { + final LogicalQueryPlan logicalQueryPlan = + planTester.createPlan("select * from information_schema.receivers where protocol='thrift'"); + assertPlan( + logicalQueryPlan, + output( + filter( + new ComparisonExpression( + EQUAL, new SymbolReference(PROTOCOL_TABLE_MODEL), new StringLiteral("thrift")), + infoSchemaTableScan( + "information_schema.receivers", Optional.empty(), RECEIVERS_COLUMNS)))); + } + + @Test + public void testSelectReceiverColumns() { + // Optimizer column-prune for InformationSchemaTableScanNode is not supported now. + final LogicalQueryPlan logicalQueryPlan = + planTester.createPlan( + "select receiver_node_type, sender_address, connection_count, pipe_ids " + + "from information_schema.receivers where protocol='thrift'"); + assertPlan( + logicalQueryPlan, + output( + project( + filter( + new ComparisonExpression( + EQUAL, + new SymbolReference(PROTOCOL_TABLE_MODEL), + new StringLiteral("thrift")), + project( + infoSchemaTableScan( + "information_schema.receivers", + Optional.empty(), + RECEIVERS_COLUMNS)))))); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiver.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiver.java index 8304161c9d0db..55f2712cf5e73 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiver.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiver.java @@ -27,6 +27,7 @@ import org.apache.iotdb.commons.i18n.PipeMessages; import org.apache.iotdb.commons.pipe.config.PipeConfig; import org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant; +import org.apache.iotdb.commons.pipe.receiver.runtime.PipeReceiverRuntimeRegistry; import org.apache.iotdb.commons.pipe.resource.log.PipeLogger; import org.apache.iotdb.commons.pipe.sink.payload.thrift.common.PipeTransferHandshakeConstant; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.IoTDBSinkRequestVersion; @@ -100,12 +101,21 @@ public abstract class IoTDBFileReceiver implements IoTDBReceiver { protected final AtomicBoolean shouldMarkAsPipeRequest = new AtomicBoolean(true); protected final AtomicBoolean skipIfNoPrivileges = new AtomicBoolean(false); + protected String senderClusterId = PipeReceiverRuntimeRegistry.UNKNOWN; + protected String receiverPipeName; + protected long receiverPipeCreationTime = Long.MIN_VALUE; + private final AtomicReference pipeReceiverRuntimeSessionKey = new AtomicReference<>(); + @Override public IoTDBSinkRequestVersion getVersion() { return IoTDBSinkRequestVersion.VERSION_1; } protected TPipeTransferResp handleTransferHandshakeV1(final PipeTransferHandshakeV1Req req) { + senderClusterId = PipeReceiverRuntimeRegistry.UNKNOWN; + receiverPipeName = null; + receiverPipeCreationTime = Long.MIN_VALUE; + if (!CommonDescriptor.getInstance() .getConfig() .getTimestampPrecision() @@ -340,16 +350,28 @@ protected TPipeTransferResp handleTransferHandshakeV2(final PipeTransferHandshak req.getParams() .getOrDefault(PipeTransferHandshakeConstant.HANDSHAKE_KEY_SKIP_IF, "false"))); + final String pipeNameString = + req.getParams().get(PipeTransferHandshakeConstant.HANDSHAKE_KEY_PIPE_NAME); + final String pipeCreationTimeString = + req.getParams().get(PipeTransferHandshakeConstant.HANDSHAKE_KEY_PIPE_CREATION_TIME); + // Handle the handshake request as a v1 request. // Here we construct a fake "dataNode" request to valid from v1 validation logic, though // it may not require the actual type of the v1 request. - return handleTransferHandshakeV1( - new PipeTransferHandshakeV1Req() { - @Override - protected PipeRequestType getPlanType() { - return PipeRequestType.HANDSHAKE_DATANODE_V1; - } - }.convertToTPipeTransferReq(timestampPrecision)); + final TPipeTransferResp resp = + handleTransferHandshakeV1( + new PipeTransferHandshakeV1Req() { + @Override + protected PipeRequestType getPlanType() { + return PipeRequestType.HANDSHAKE_DATANODE_V1; + } + }.convertToTPipeTransferReq(timestampPrecision)); + if (isSuccess(resp)) { + senderClusterId = clusterIdFromHandshakeRequest; + receiverPipeName = pipeNameString; + receiverPipeCreationTime = parsePipeCreationTime(pipeCreationTimeString); + } + return resp; } protected abstract String getClusterId(); @@ -873,8 +895,68 @@ protected abstract TSStatus loadFileV2( final PipeTransferFileSealReqV2 req, final List fileAbsolutePaths) throws IOException, IllegalPathException; + protected void recordPipeReceiverHandshake( + final String receiverNodeType, final int receiverNodeId, final String protocol) { + final String sessionKey = + String.format("%s-%s-%s-%s", receiverNodeType, receiverNodeId, protocol, receiverId.get()); + final String oldSessionKey = pipeReceiverRuntimeSessionKey.getAndSet(sessionKey); + if (!Objects.equals(oldSessionKey, sessionKey)) { + PipeReceiverRuntimeRegistry.getInstance().deregister(oldSessionKey); + } + PipeReceiverRuntimeRegistry.getInstance() + .registerOrUpdateSession( + sessionKey, + receiverNodeType, + receiverNodeId, + protocol, + getSenderHost(), + parseSenderPort(getSenderPort()), + username, + senderClusterId, + receiverPipeName, + receiverPipeCreationTime, + System.currentTimeMillis()); + } + + protected void recordPipeReceiverTransfer() { + PipeReceiverRuntimeRegistry.getInstance() + .markTransfer(pipeReceiverRuntimeSessionKey.get(), System.currentTimeMillis()); + } + + protected void clearPipeReceiverRuntime() { + PipeReceiverRuntimeRegistry.getInstance() + .deregister(pipeReceiverRuntimeSessionKey.getAndSet(null)); + } + + protected static boolean isSuccess(final TPipeTransferResp resp) { + return resp != null + && resp.getStatus() != null + && resp.getStatus().getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode(); + } + + private static int parseSenderPort(final String senderPort) { + try { + return Integer.parseInt(senderPort); + } catch (final Exception e) { + return -1; + } + } + + private static long parsePipeCreationTime(final String pipeCreationTime) { + if (pipeCreationTime == null) { + return Long.MIN_VALUE; + } + try { + return Long.parseLong(pipeCreationTime); + } catch (final NumberFormatException e) { + return Long.MIN_VALUE; + } + } + @Override public synchronized void handleExit() { + clearPipeReceiverRuntime(); + if (writingFileWriter != null) { try { writingFileWriter.close(); diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/runtime/PipeReceiverRuntimeRegistry.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/runtime/PipeReceiverRuntimeRegistry.java new file mode 100644 index 0000000000000..4256e3220ea33 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/runtime/PipeReceiverRuntimeRegistry.java @@ -0,0 +1,360 @@ +/* + * 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.pipe.receiver.runtime; + +import org.apache.iotdb.commons.queryengine.utils.DateTimeUtils; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.StringJoiner; +import java.util.TreeSet; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; + +public class PipeReceiverRuntimeRegistry { + + public static final String UNKNOWN = "Unknown"; + public static final String NODE_TYPE_DATA_NODE = "DataNode"; + public static final String NODE_TYPE_CONFIG_NODE = "ConfigNode"; + public static final String PROTOCOL_THRIFT = "thrift"; + public static final String PROTOCOL_AIR_GAP = "air_gap"; + public static final String PROTOCOL_WRITEBACK = "writeback"; + + private static final PipeReceiverRuntimeRegistry INSTANCE = new PipeReceiverRuntimeRegistry(); + + private final ConcurrentMap sessionInfoMap = + new ConcurrentHashMap<>(); + + private PipeReceiverRuntimeRegistry() {} + + public static PipeReceiverRuntimeRegistry getInstance() { + return INSTANCE; + } + + public void registerOrUpdateSession( + String connectionKey, + String receiverNodeType, + int receiverNodeId, + String protocol, + String senderAddress, + int senderPort, + String userName, + String senderClusterId, + String pipeName, + long pipeCreationTime, + long handshakeTime) { + if (isBlank(connectionKey)) { + return; + } + + sessionInfoMap.compute( + connectionKey, + (key, oldSession) -> { + final SessionRuntimeInfo session = + oldSession == null ? new SessionRuntimeInfo(connectionKey) : oldSession; + synchronized (session) { + session.receiverNodeType = normalize(receiverNodeType); + session.receiverNodeId = receiverNodeId; + session.protocol = normalize(protocol); + session.senderAddress = normalize(senderAddress); + session.senderPort = senderPort; + session.userName = normalize(userName); + session.senderClusterId = normalize(senderClusterId); + session.lastHandshakeTime = handshakeTime; + session.lastTransferTime = Math.max(session.lastTransferTime, handshakeTime); + session.addPipe(pipeName, pipeCreationTime); + } + return session; + }); + } + + public void markTransfer(String connectionKey, long transferTime) { + markTransfer(connectionKey, null, Long.MIN_VALUE, transferTime); + } + + public void markTransfer( + String connectionKey, String pipeName, long pipeCreationTime, long transferTime) { + if (isBlank(connectionKey)) { + return; + } + final SessionRuntimeInfo session = sessionInfoMap.get(connectionKey); + if (session == null) { + return; + } + synchronized (session) { + session.lastTransferTime = Math.max(session.lastTransferTime, transferTime); + session.addPipe(pipeName, pipeCreationTime); + } + } + + public void removePipe(String connectionKey, String pipeName, long pipeCreationTime) { + if (isBlank(connectionKey)) { + return; + } + final SessionRuntimeInfo session = sessionInfoMap.get(connectionKey); + if (session == null) { + return; + } + synchronized (session) { + session.removePipe(pipeName, pipeCreationTime); + } + } + + public void removePipeFromAllSessions(String pipeName, long pipeCreationTime) { + if (isBlank(pipeName)) { + return; + } + for (SessionRuntimeInfo session : sessionInfoMap.values()) { + synchronized (session) { + session.removePipe(pipeName, pipeCreationTime); + } + } + } + + public void deregister(String connectionKey) { + if (!isBlank(connectionKey)) { + sessionInfoMap.remove(connectionKey); + } + } + + public List snapshot() { + final Map aggregatedInfoMap = new HashMap<>(); + for (SessionRuntimeInfo session : sessionInfoMap.values()) { + synchronized (session) { + final GroupKey groupKey = + new GroupKey( + session.receiverNodeType, + session.receiverNodeId, + session.protocol, + session.senderClusterId, + session.senderAddress, + session.userName); + AggregatedRuntimeInfo aggregatedInfo = aggregatedInfoMap.get(groupKey); + if (aggregatedInfo == null) { + aggregatedInfo = new AggregatedRuntimeInfo(groupKey); + aggregatedInfoMap.put(groupKey, aggregatedInfo); + } + aggregatedInfo.senderPorts.add(session.senderPort); + aggregatedInfo.connectionCount++; + aggregatedInfo.pipeIds.addAll(session.pipeIds); + aggregatedInfo.lastHandshakeTime = + Math.max(aggregatedInfo.lastHandshakeTime, session.lastHandshakeTime); + aggregatedInfo.lastTransferTime = + Math.max(aggregatedInfo.lastTransferTime, session.lastTransferTime); + } + } + + final List aggregatedInfos = new ArrayList<>(aggregatedInfoMap.values()); + aggregatedInfos.sort(Comparator.comparing(AggregatedRuntimeInfo::getGroupKey)); + + final List snapshots = new ArrayList<>(aggregatedInfos.size()); + for (AggregatedRuntimeInfo aggregatedInfo : aggregatedInfos) { + snapshots.add(aggregatedInfo.toSnapshot()); + } + return snapshots; + } + + public void clear() { + sessionInfoMap.clear(); + } + + private static String normalize(String value) { + return isBlank(value) ? UNKNOWN : value; + } + + private static boolean isBlank(String value) { + return value == null || value.trim().isEmpty(); + } + + private static String formatPipeId(String pipeName, long pipeCreationTime) { + if (pipeCreationTime < 0) { + return pipeName + "@" + UNKNOWN; + } + return pipeName + "@" + DateTimeUtils.convertLongToDate(pipeCreationTime, "ms"); + } + + private static class SessionRuntimeInfo { + private final String connectionKey; + private String receiverNodeType = UNKNOWN; + private int receiverNodeId = -1; + private String protocol = UNKNOWN; + private String senderAddress = UNKNOWN; + private int senderPort = -1; + private String userName = UNKNOWN; + private String senderClusterId = UNKNOWN; + private long lastHandshakeTime; + private long lastTransferTime; + private final TreeSet pipeIds = new TreeSet<>(); + + private SessionRuntimeInfo(String connectionKey) { + this.connectionKey = connectionKey; + } + + private void addPipe(String pipeName, long pipeCreationTime) { + if (!isBlank(pipeName)) { + pipeIds.add(formatPipeId(pipeName, pipeCreationTime)); + } + } + + private void removePipe(String pipeName, long pipeCreationTime) { + if (!isBlank(pipeName)) { + pipeIds.remove(formatPipeId(pipeName, pipeCreationTime)); + } + } + + @Override + public int hashCode() { + return Objects.hash(connectionKey); + } + } + + private static class GroupKey implements Comparable { + private final String receiverNodeType; + private final int receiverNodeId; + private final String protocol; + private final String senderClusterId; + private final String senderAddress; + private final String userName; + + private GroupKey( + String receiverNodeType, + int receiverNodeId, + String protocol, + String senderClusterId, + String senderAddress, + String userName) { + this.receiverNodeType = receiverNodeType; + this.receiverNodeId = receiverNodeId; + this.protocol = protocol; + this.senderClusterId = senderClusterId; + this.senderAddress = senderAddress; + this.userName = userName; + } + + @Override + public int compareTo(GroupKey other) { + int result = receiverNodeType.compareTo(other.receiverNodeType); + if (result != 0) { + return result; + } + result = Integer.compare(receiverNodeId, other.receiverNodeId); + if (result != 0) { + return result; + } + result = protocol.compareTo(other.protocol); + if (result != 0) { + return result; + } + result = senderClusterId.compareTo(other.senderClusterId); + if (result != 0) { + return result; + } + result = senderAddress.compareTo(other.senderAddress); + if (result != 0) { + return result; + } + return userName.compareTo(other.userName); + } + + @Override + public boolean equals(Object object) { + if (this == object) { + return true; + } + if (!(object instanceof GroupKey)) { + return false; + } + final GroupKey groupKey = (GroupKey) object; + return receiverNodeId == groupKey.receiverNodeId + && Objects.equals(receiverNodeType, groupKey.receiverNodeType) + && Objects.equals(protocol, groupKey.protocol) + && Objects.equals(senderClusterId, groupKey.senderClusterId) + && Objects.equals(senderAddress, groupKey.senderAddress) + && Objects.equals(userName, groupKey.userName); + } + + @Override + public int hashCode() { + return Objects.hash( + receiverNodeType, receiverNodeId, protocol, senderClusterId, senderAddress, userName); + } + } + + private static class AggregatedRuntimeInfo { + private final GroupKey groupKey; + private final TreeSet senderPorts = new TreeSet<>(); + private final TreeSet pipeIds = new TreeSet<>(); + private int connectionCount; + private long lastHandshakeTime; + private long lastTransferTime; + + private AggregatedRuntimeInfo(GroupKey groupKey) { + this.groupKey = groupKey; + } + + private GroupKey getGroupKey() { + return groupKey; + } + + private PipeReceiverRuntimeSnapshot toSnapshot() { + return new PipeReceiverRuntimeSnapshot( + groupKey.receiverNodeType, + groupKey.receiverNodeId, + groupKey.protocol, + groupKey.senderAddress, + joinIntegerSet(senderPorts), + connectionCount, + pipeIds.size(), + pipeIds.isEmpty() ? UNKNOWN : joinStringSet(pipeIds), + groupKey.userName, + groupKey.senderClusterId, + lastHandshakeTime, + lastTransferTime); + } + } + + private static String joinIntegerSet(TreeSet values) { + final StringJoiner joiner = new StringJoiner(","); + boolean hasUnknown = false; + for (Integer value : values) { + if (value < 0) { + if (!hasUnknown) { + joiner.add(UNKNOWN); + hasUnknown = true; + } + continue; + } + joiner.add(String.valueOf(value)); + } + return joiner.toString(); + } + + private static String joinStringSet(TreeSet values) { + final StringJoiner joiner = new StringJoiner(";"); + for (String value : values) { + joiner.add(value); + } + return joiner.toString(); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/runtime/PipeReceiverRuntimeSnapshot.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/runtime/PipeReceiverRuntimeSnapshot.java new file mode 100644 index 0000000000000..f6effeeafd2b5 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/runtime/PipeReceiverRuntimeSnapshot.java @@ -0,0 +1,115 @@ +/* + * 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.pipe.receiver.runtime; + +public class PipeReceiverRuntimeSnapshot { + + private final String receiverNodeType; + private final int receiverNodeId; + private final String protocol; + private final String senderAddress; + private final String senderPorts; + private final int connectionCount; + private final int pipeCount; + private final String pipeIds; + private final String userName; + private final String senderClusterId; + private final long lastHandshakeTime; + private final long lastTransferTime; + + public PipeReceiverRuntimeSnapshot( + String receiverNodeType, + int receiverNodeId, + String protocol, + String senderAddress, + String senderPorts, + int connectionCount, + int pipeCount, + String pipeIds, + String userName, + String senderClusterId, + long lastHandshakeTime, + long lastTransferTime) { + this.receiverNodeType = receiverNodeType; + this.receiverNodeId = receiverNodeId; + this.protocol = protocol; + this.senderAddress = senderAddress; + this.senderPorts = senderPorts; + this.connectionCount = connectionCount; + this.pipeCount = pipeCount; + this.pipeIds = pipeIds; + this.userName = userName; + this.senderClusterId = senderClusterId; + this.lastHandshakeTime = lastHandshakeTime; + this.lastTransferTime = lastTransferTime; + } + + public String getReceiverNodeType() { + return receiverNodeType; + } + + public int getReceiverNodeId() { + return receiverNodeId; + } + + public boolean isReceiverNodeIdKnown() { + return receiverNodeId >= 0; + } + + public String getProtocol() { + return protocol; + } + + public String getSenderAddress() { + return senderAddress; + } + + public String getSenderPorts() { + return senderPorts; + } + + public int getConnectionCount() { + return connectionCount; + } + + public int getPipeCount() { + return pipeCount; + } + + public String getPipeIds() { + return pipeIds; + } + + public String getUserName() { + return userName; + } + + public String getSenderClusterId() { + return senderClusterId; + } + + public long getLastHandshakeTime() { + return lastHandshakeTime; + } + + public long getLastTransferTime() { + return lastTransferTime; + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/client/IoTDBClientManager.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/client/IoTDBClientManager.java index 1f76f5d2453f5..9d7c85d975afb 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/client/IoTDBClientManager.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/client/IoTDBClientManager.java @@ -22,12 +22,14 @@ import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.commons.audit.UserEntity; import org.apache.iotdb.commons.pipe.config.PipeConfig; +import org.apache.iotdb.commons.pipe.sink.payload.thrift.common.PipeTransferHandshakeConstant; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.net.SocketTimeoutException; import java.util.List; +import java.util.Map; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; @@ -51,6 +53,9 @@ public abstract class IoTDBClientManager { protected final boolean shouldMarkAsPipeRequest; protected final boolean skipIfNoPrivileges; + protected volatile String pipeName; + protected volatile long pipeCreationTime = Long.MIN_VALUE; + // This flag indicates whether the receiver supports mods transferring if // it is a DataNode receiver. The flag is useless for configNode receiver. protected boolean supportModsIfIsDataNodeReceiver = true; @@ -89,6 +94,21 @@ public boolean supportModsIfIsDataNodeReceiver() { return supportModsIfIsDataNodeReceiver; } + public void setPipeInfo(final String pipeName, final long pipeCreationTime) { + this.pipeName = pipeName; + this.pipeCreationTime = pipeCreationTime; + } + + protected void appendPipeInfoToHandshakeParams(final Map params) { + if (pipeName == null) { + return; + } + params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_PIPE_NAME, pipeName); + params.put( + PipeTransferHandshakeConstant.HANDSHAKE_KEY_PIPE_CREATION_TIME, + String.valueOf(pipeCreationTime)); + } + public void adjustTimeoutIfNecessary(Throwable e) { do { if (e instanceof SocketTimeoutException || e instanceof TimeoutException) { diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/client/IoTDBSyncClientManager.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/client/IoTDBSyncClientManager.java index ff62dbf477b7a..c6446f57bf357 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/client/IoTDBSyncClientManager.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/client/IoTDBSyncClientManager.java @@ -29,6 +29,7 @@ import org.apache.iotdb.commons.pipe.sink.payload.thrift.common.PipeTransferHandshakeConstant; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferHandshakeV1Req; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferHandshakeV2Req; +import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferPipeReceiverRuntimeInfoCleanupReq; import org.apache.iotdb.pipe.api.exception.PipeConnectionException; import org.apache.iotdb.rpc.TSStatusCode; import org.apache.iotdb.service.rpc.thrift.TPipeTransferResp; @@ -249,6 +250,7 @@ public void sendHandshakeReq(final Pair clientAndStatu params.put( PipeTransferHandshakeConstant.HANDSHAKE_KEY_SKIP_IF, Boolean.toString(skipIfNoPrivileges)); + appendPipeInfoToHandshakeParams(params); // Try to handshake by PipeTransferHandshakeV2Req. TPipeTransferResp resp = client.pipeTransfer(buildHandshakeV2Req(params)); @@ -327,6 +329,75 @@ public void close() { } } + public void discardReceiverRuntimeSessions() { + try { + sendPipeReceiverRuntimeInfoCleanupReq(); + } finally { + close(); + } + } + + private void sendPipeReceiverRuntimeInfoCleanupReq() { + if (pipeName == null) { + return; + } + + final PipeTransferPipeReceiverRuntimeInfoCleanupReq req; + try { + req = + PipeTransferPipeReceiverRuntimeInfoCleanupReq.toTPipeTransferReq( + pipeName, pipeCreationTime); + } catch (final IOException e) { + LOGGER.warn( + "Failed to build pipe receiver runtime info cleanup request for pipe {}@{}.", + pipeName, + pipeCreationTime, + e); + return; + } + + for (final Map.Entry> entry : + endPoint2ClientAndStatus.entrySet()) { + final Pair clientAndStatus = entry.getValue(); + if (clientAndStatus == null + || !Boolean.TRUE.equals(clientAndStatus.getRight()) + || clientAndStatus.getLeft() == null) { + continue; + } + + try { + clientAndStatus.getLeft().setTimeout(PIPE_CONFIG.getPipeSinkHandshakeTimeoutMs()); + final TPipeTransferResp resp = clientAndStatus.getLeft().pipeTransfer(req); + if (resp == null || resp.getStatus() == null) { + LOGGER.warn( + "Failed to cleanup pipe receiver runtime info for pipe {}@{} on target {}:{}, response is null.", + pipeName, + pipeCreationTime, + entry.getKey().getIp(), + entry.getKey().getPort()); + } else if (resp.getStatus().getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode() + && resp.getStatus().getCode() != TSStatusCode.PIPE_TYPE_ERROR.getStatusCode()) { + LOGGER.warn( + "Failed to cleanup pipe receiver runtime info for pipe {}@{} on target {}:{}, response status: {}.", + pipeName, + pipeCreationTime, + entry.getKey().getIp(), + entry.getKey().getPort(), + resp.getStatus()); + } + } catch (final Exception e) { + clientAndStatus.setRight(false); + LOGGER.warn( + "Failed to cleanup pipe receiver runtime info for pipe {}@{} on target {}:{}.", + pipeName, + pipeCreationTime, + entry.getKey().getIp(), + entry.getKey().getPort(), + e); + } + } + } + /////////////////////// Strategies for load balance ////////////////////////// private interface LoadBalancer { diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferHandshakeConstant.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferHandshakeConstant.java index 710fca828e938..3f084a73dc8d7 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferHandshakeConstant.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferHandshakeConstant.java @@ -29,6 +29,8 @@ public class PipeTransferHandshakeConstant { public static final String HANDSHAKE_KEY_USERNAME = "username"; public static final String HANDSHAKE_KEY_CLI_HOSTNAME = "cliHostname"; public static final String HANDSHAKE_KEY_PASSWORD = "password"; + public static final String HANDSHAKE_KEY_PIPE_NAME = "pipeName"; + public static final String HANDSHAKE_KEY_PIPE_CREATION_TIME = "pipeCreationTime"; public static final String HANDSHAKE_KEY_VALIDATE_TSFILE = "validateTsFile"; public static final String HANDSHAKE_KEY_MARK_AS_PIPE_REQUEST = "markAsPipeRequest"; public static final String HANDSHAKE_KEY_SKIP_IF = "skipIf"; diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeRequestType.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeRequestType.java index 97a9526c9f790..63c74bc7e1513 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeRequestType.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeRequestType.java @@ -61,6 +61,9 @@ public enum PipeRequestType { // Fallback Handling TRANSFER_SLICE((short) 400), + + // Lifecycle control + TRANSFER_PIPE_RECEIVER_RUNTIME_INFO_CLEANUP((short) 500), ; private final short type; diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferHandshakeV1Req.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferHandshakeV1Req.java index 54eddc453e169..4495ee1d0e330 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferHandshakeV1Req.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferHandshakeV1Req.java @@ -58,10 +58,11 @@ public final PipeTransferHandshakeV1Req convertToTPipeTransferReq(String timesta protected final PipeTransferHandshakeV1Req translateFromTPipeTransferReq( TPipeTransferReq transferReq) { - timestampPrecision = ReadWriteIOUtils.readString(transferReq.body); + timestampPrecision = ReadWriteIOUtils.readString(transferReq.body.duplicate()); version = transferReq.version; type = transferReq.type; + body = transferReq.body; return this; } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferHandshakeV2Req.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferHandshakeV2Req.java index 9a81d844e3f43..472679494729b 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferHandshakeV2Req.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferHandshakeV2Req.java @@ -64,16 +64,18 @@ protected final PipeTransferHandshakeV2Req convertToTPipeTransferReq(Map params = new HashMap<>(); - final int size = ReadWriteIOUtils.readInt(transferReq.body); + final ByteBuffer bodyBuffer = transferReq.body.duplicate(); + final int size = ReadWriteIOUtils.readInt(bodyBuffer); for (int i = 0; i < size; ++i) { - final String key = ReadWriteIOUtils.readString(transferReq.body); - final String value = ReadWriteIOUtils.readString(transferReq.body); + final String key = ReadWriteIOUtils.readString(bodyBuffer); + final String value = ReadWriteIOUtils.readString(bodyBuffer); params.put(key, value); } this.params = params; version = transferReq.version; type = transferReq.type; + body = transferReq.body; return this; } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferPipeReceiverRuntimeInfoCleanupReq.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferPipeReceiverRuntimeInfoCleanupReq.java new file mode 100644 index 0000000000000..4c261b9ee31db --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferPipeReceiverRuntimeInfoCleanupReq.java @@ -0,0 +1,103 @@ +/* + * 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.pipe.sink.payload.thrift.request; + +import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq; + +import org.apache.tsfile.utils.PublicBAOS; +import org.apache.tsfile.utils.ReadWriteIOUtils; + +import java.io.DataOutputStream; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Objects; + +public class PipeTransferPipeReceiverRuntimeInfoCleanupReq extends TPipeTransferReq { + + private String pipeName; + private long pipeCreationTime; + + private PipeTransferPipeReceiverRuntimeInfoCleanupReq() { + // Empty constructor + } + + public String getPipeName() { + return pipeName; + } + + public long getPipeCreationTime() { + return pipeCreationTime; + } + + public static PipeTransferPipeReceiverRuntimeInfoCleanupReq toTPipeTransferReq( + final String pipeName, final long pipeCreationTime) throws IOException { + final PipeTransferPipeReceiverRuntimeInfoCleanupReq req = + new PipeTransferPipeReceiverRuntimeInfoCleanupReq(); + req.pipeName = pipeName; + req.pipeCreationTime = pipeCreationTime; + req.version = IoTDBSinkRequestVersion.VERSION_1.getVersion(); + req.type = PipeRequestType.TRANSFER_PIPE_RECEIVER_RUNTIME_INFO_CLEANUP.getType(); + + try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); + final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { + ReadWriteIOUtils.write(pipeName, outputStream); + ReadWriteIOUtils.write(pipeCreationTime, outputStream); + req.body = ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size()); + } + + return req; + } + + public static PipeTransferPipeReceiverRuntimeInfoCleanupReq fromTPipeTransferReq( + final TPipeTransferReq transferReq) { + final PipeTransferPipeReceiverRuntimeInfoCleanupReq req = + new PipeTransferPipeReceiverRuntimeInfoCleanupReq(); + req.version = transferReq.version; + req.type = transferReq.type; + req.body = transferReq.body; + + final ByteBuffer bodyBuffer = transferReq.body.duplicate(); + req.pipeName = ReadWriteIOUtils.readString(bodyBuffer); + req.pipeCreationTime = ReadWriteIOUtils.readLong(bodyBuffer); + return req; + } + + @Override + public boolean equals(final Object obj) { + if (this == obj) { + return true; + } + if (obj == null || getClass() != obj.getClass()) { + return false; + } + final PipeTransferPipeReceiverRuntimeInfoCleanupReq that = + (PipeTransferPipeReceiverRuntimeInfoCleanupReq) obj; + return pipeCreationTime == that.pipeCreationTime + && version == that.version + && type == that.type + && Objects.equals(pipeName, that.pipeName) + && Objects.equals(body, that.body); + } + + @Override + public int hashCode() { + return Objects.hash(pipeName, pipeCreationTime, version, type, body); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBSink.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBSink.java index b5662aeec2ce9..7992f40880f07 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBSink.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBSink.java @@ -30,6 +30,7 @@ import org.apache.iotdb.commons.pipe.sink.compressor.PipeCompressorFactory; import org.apache.iotdb.commons.pipe.sink.limiter.GlobalRPCRateLimiter; import org.apache.iotdb.commons.pipe.sink.limiter.PipeEndPointRateLimiter; +import org.apache.iotdb.commons.pipe.sink.payload.thrift.common.PipeTransferHandshakeConstant; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferCompressedReq; import org.apache.iotdb.commons.utils.NodeUrlUtils; import org.apache.iotdb.metrics.type.Histogram; @@ -190,6 +191,8 @@ public abstract class IoTDBSink implements PipeConnector, PipeConnectorWithEvent private final AtomicLong totalCompressedSize = new AtomicLong(0); protected String attributeSortedString; protected String sinkTaskId; + protected String pipeName; + protected long creationTime = Long.MIN_VALUE; protected Timer compressionTimer; protected boolean isRealtimeFirst; @@ -396,6 +399,8 @@ public void customize( (PipeTaskSinkRuntimeEnvironment) environment; attributeSortedString = sinkEnvironment.getAttributeSortedString(); sinkTaskId = sinkEnvironment.getSinkTaskId(); + pipeName = sinkEnvironment.getPipeName(); + creationTime = sinkEnvironment.getCreationTime(); } nodeUrls.clear(); @@ -585,6 +590,10 @@ public void close() { PIPE_END_POINT_RATE_LIMITER_MAP.clear(); } + public void discardReceiverRuntimeSessions() { + // Do nothing by default. + } + public TPipeTransferReq compressIfNeeded(TPipeTransferReq req) throws IOException { // Explanation for +3: version 1 byte, type 2 bytes totalUncompressedSize.addAndGet(req.body.array().length + 3); @@ -621,6 +630,16 @@ public long getTotalUncompressedSize() { return totalUncompressedSize.get(); } + protected void appendPipeInfoToHandshakeParams(final Map params) { + if (pipeName == null) { + return; + } + params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_PIPE_NAME, pipeName); + params.put( + PipeTransferHandshakeConstant.HANDSHAKE_KEY_PIPE_CREATION_TIME, + String.valueOf(creationTime)); + } + public void rateLimitIfNeeded( final String pipeName, final long creationTime, diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBSslSyncSink.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBSslSyncSink.java index 6d3c05d3b4f21..fba0105a9c4ac 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBSslSyncSink.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBSslSyncSink.java @@ -143,6 +143,7 @@ public void customize( loadTsFileValidation, shouldMarkAsPipeRequest, skipIfNoPrivileges); + clientManager.setPipeInfo(pipeName, creationTime); } protected abstract IoTDBSyncClientManager constructClient( @@ -268,4 +269,11 @@ public void close() { super.close(); } + + @Override + public void discardReceiverRuntimeSessions() { + if (clientManager != null) { + clientManager.discardReceiverRuntimeSessions(); + } + } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/planner/plan/node/PlanNodeType.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/planner/plan/node/PlanNodeType.java index e36796fca43c3..20c3c1bf879a2 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/planner/plan/node/PlanNodeType.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/planner/plan/node/PlanNodeType.java @@ -142,6 +142,7 @@ public enum PlanNodeType { SHOW_DISK_USAGE((short) 107), TREE_COLLECT((short) 108), LOAD_TSFILE_OBJECT_PIECE((short) 109), + SHOW_RECEIVERS((short) 110), CREATE_OR_UPDATE_TABLE_DEVICE((short) 902), TABLE_DEVICE_QUERY_SCAN((short) 903), diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/schema/column/ColumnHeaderConstant.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/schema/column/ColumnHeaderConstant.java index 33c3fe451cad4..427b738045d02 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/schema/column/ColumnHeaderConstant.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/schema/column/ColumnHeaderConstant.java @@ -195,6 +195,20 @@ private ColumnHeaderConstant() { public static final String REMAINING_EVENT_COUNT = "RemainingEventCount"; public static final String ESTIMATED_REMAINING_SECONDS = "EstimatedRemainingSeconds"; + // column names for show receivers + public static final String RECEIVER_NODE_TYPE = "ReceiverNodeType"; + public static final String RECEIVER_NODE_ID = "ReceiverNodeId"; + public static final String PROTOCOL = "Protocol"; + public static final String SENDER_CLUSTER_ID = "SenderClusterId"; + public static final String SENDER_ADDRESS = "SenderAddress"; + public static final String RECEIVER_USER_NAME = "UserName"; + public static final String SENDER_PORTS = "SenderPorts"; + public static final String CONNECTION_COUNT = "ConnectionCount"; + public static final String PIPE_COUNT = "PipeCount"; + public static final String PIPE_IDS = "PipeIDs"; + public static final String LAST_HANDSHAKE_TIME = "LastHandshakeTime"; + public static final String LAST_TRANSFER_TIME = "LastTransferTime"; + // column names for select into public static final String SOURCE_DEVICE = "SourceDevice"; public static final String SOURCE_COLUMN = "SourceColumn"; @@ -271,6 +285,19 @@ private ColumnHeaderConstant() { public static final String ESTIMATED_REMAINING_SECONDS_TABLE_MODEL = "estimated_remaining_seconds"; + public static final String RECEIVER_NODE_TYPE_TABLE_MODEL = "receiver_node_type"; + public static final String RECEIVER_NODE_ID_TABLE_MODEL = "receiver_node_id"; + public static final String PROTOCOL_TABLE_MODEL = "protocol"; + public static final String SENDER_CLUSTER_ID_TABLE_MODEL = "sender_cluster_id"; + public static final String SENDER_ADDRESS_TABLE_MODEL = "sender_address"; + public static final String RECEIVER_USER_NAME_TABLE_MODEL = "user_name"; + public static final String SENDER_PORTS_TABLE_MODEL = "sender_ports"; + public static final String CONNECTION_COUNT_TABLE_MODEL = "connection_count"; + public static final String PIPE_COUNT_TABLE_MODEL = "pipe_count"; + public static final String PIPE_IDS_TABLE_MODEL = "pipe_ids"; + public static final String LAST_HANDSHAKE_TIME_TABLE_MODEL = "last_handshake_time"; + public static final String LAST_TRANSFER_TIME_TABLE_MODEL = "last_transfer_time"; + public static final String PLUGIN_NAME_TABLE_MODEL = "plugin_name"; public static final String PLUGIN_TYPE_TABLE_MODEL = "plugin_type"; public static final String CLASS_NAME_TABLE_MODEL = "class_name"; @@ -663,6 +690,21 @@ private ColumnHeaderConstant() { new ColumnHeader(CLIENT_IP_TREE_MODEL, TSDataType.STRING), new ColumnHeader(TIMEOUT, TSDataType.INT64)); + public static final List showReceiversColumnHeaders = + ImmutableList.of( + new ColumnHeader(RECEIVER_NODE_TYPE, TSDataType.TEXT), + new ColumnHeader(RECEIVER_NODE_ID, TSDataType.INT32), + new ColumnHeader(PROTOCOL, TSDataType.TEXT), + new ColumnHeader(SENDER_CLUSTER_ID, TSDataType.TEXT), + new ColumnHeader(SENDER_ADDRESS, TSDataType.TEXT), + new ColumnHeader(RECEIVER_USER_NAME, TSDataType.TEXT), + new ColumnHeader(SENDER_PORTS, TSDataType.TEXT), + new ColumnHeader(CONNECTION_COUNT, TSDataType.INT32), + new ColumnHeader(PIPE_COUNT, TSDataType.INT32), + new ColumnHeader(PIPE_IDS, TSDataType.TEXT), + new ColumnHeader(LAST_HANDSHAKE_TIME, TSDataType.TEXT), + new ColumnHeader(LAST_TRANSFER_TIME, TSDataType.TEXT)); + public static final List showDiskUsageColumnHeaders = ImmutableList.of( new ColumnHeader(DATABASE, TSDataType.TEXT), diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/schema/table/InformationSchema.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/schema/table/InformationSchema.java index 27a1dbdd9d2f7..523b29dee3468 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/schema/table/InformationSchema.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/schema/table/InformationSchema.java @@ -47,6 +47,7 @@ public class InformationSchema { public static final String COLUMNS = "columns"; public static final String REGIONS = "regions"; public static final String PIPES = "pipes"; + public static final String RECEIVERS = "receivers"; public static final String PIPE_PLUGINS = "pipe_plugins"; public static final String TOPICS = "topics"; public static final String SUBSCRIPTIONS = "subscriptions"; @@ -227,6 +228,39 @@ public class InformationSchema { ColumnHeaderConstant.ESTIMATED_REMAINING_SECONDS_TABLE_MODEL, TSDataType.DOUBLE)); schemaTables.put(PIPES, pipeTable); + final TsTable receiversTable = new TsTable(RECEIVERS); + receiversTable.addColumnSchema( + new TagColumnSchema( + ColumnHeaderConstant.RECEIVER_NODE_TYPE_TABLE_MODEL, TSDataType.STRING)); + receiversTable.addColumnSchema( + new TagColumnSchema(ColumnHeaderConstant.RECEIVER_NODE_ID_TABLE_MODEL, TSDataType.INT32)); + receiversTable.addColumnSchema( + new TagColumnSchema(ColumnHeaderConstant.PROTOCOL_TABLE_MODEL, TSDataType.STRING)); + receiversTable.addColumnSchema( + new TagColumnSchema(ColumnHeaderConstant.SENDER_CLUSTER_ID_TABLE_MODEL, TSDataType.STRING)); + receiversTable.addColumnSchema( + new TagColumnSchema(ColumnHeaderConstant.SENDER_ADDRESS_TABLE_MODEL, TSDataType.STRING)); + receiversTable.addColumnSchema( + new TagColumnSchema( + ColumnHeaderConstant.RECEIVER_USER_NAME_TABLE_MODEL, TSDataType.STRING)); + receiversTable.addColumnSchema( + new AttributeColumnSchema( + ColumnHeaderConstant.SENDER_PORTS_TABLE_MODEL, TSDataType.STRING)); + receiversTable.addColumnSchema( + new AttributeColumnSchema( + ColumnHeaderConstant.CONNECTION_COUNT_TABLE_MODEL, TSDataType.INT32)); + receiversTable.addColumnSchema( + new AttributeColumnSchema(ColumnHeaderConstant.PIPE_COUNT_TABLE_MODEL, TSDataType.INT32)); + receiversTable.addColumnSchema( + new AttributeColumnSchema(ColumnHeaderConstant.PIPE_IDS_TABLE_MODEL, TSDataType.STRING)); + receiversTable.addColumnSchema( + new AttributeColumnSchema( + ColumnHeaderConstant.LAST_HANDSHAKE_TIME_TABLE_MODEL, TSDataType.TIMESTAMP)); + receiversTable.addColumnSchema( + new AttributeColumnSchema( + ColumnHeaderConstant.LAST_TRANSFER_TIME_TABLE_MODEL, TSDataType.TIMESTAMP)); + schemaTables.put(RECEIVERS, receiversTable); + final TsTable pipePluginTable = new TsTable(PIPE_PLUGINS); pipePluginTable.addColumnSchema( new TagColumnSchema(ColumnHeaderConstant.PLUGIN_NAME_TABLE_MODEL, TSDataType.STRING)); diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/receiver/runtime/PipeReceiverRuntimeRegistryTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/receiver/runtime/PipeReceiverRuntimeRegistryTest.java new file mode 100644 index 0000000000000..73b2ae131bd74 --- /dev/null +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/receiver/runtime/PipeReceiverRuntimeRegistryTest.java @@ -0,0 +1,947 @@ +/* + * 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.pipe.receiver.runtime; + +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import java.util.List; +import java.util.concurrent.TimeUnit; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +public class PipeReceiverRuntimeRegistryTest { + + private final PipeReceiverRuntimeRegistry registry = PipeReceiverRuntimeRegistry.getInstance(); + + @Before + public void setUp() { + registry.clear(); + } + + @After + public void tearDown() { + registry.clear(); + } + + @Test + public void testSingleSenderSingleDataNodeSingleConnectionSinglePipeBasicQuery() { + registerDataSession("data-1", 1, "10.0.0.1", 9001, "root", "cluster-a", "pipe-a", 1, 100); + registry.markTransfer("data-1", 200); + + final List snapshots = registry.snapshot(); + + assertEquals(1, snapshots.size()); + final PipeReceiverRuntimeSnapshot snapshot = snapshots.get(0); + assertEquals(PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, snapshot.getReceiverNodeType()); + assertEquals(1, snapshot.getReceiverNodeId()); + assertEquals(PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, snapshot.getProtocol()); + assertEquals("10.0.0.1", snapshot.getSenderAddress()); + assertEquals("9001", snapshot.getSenderPorts()); + assertEquals(1, snapshot.getConnectionCount()); + assertEquals(1, snapshot.getPipeCount()); + assertTrue(snapshot.getPipeIds().contains("pipe-a@")); + assertEquals("root", snapshot.getUserName()); + assertEquals("cluster-a", snapshot.getSenderClusterId()); + assertEquals(100, snapshot.getLastHandshakeTime()); + assertEquals(200, snapshot.getLastTransferTime()); + } + + @Test + public void testMultipleSendersAndMultipleDataNodesClusterAggregation() { + registerDataSession("data-1", 1, "10.0.0.1", 9001, "root", "cluster-a", "pipe-a", 1, 100); + registerDataSession("data-2", 2, "10.0.0.1", 9002, "root", "cluster-a", "pipe-b", 2, 200); + registerDataSession("data-3", 2, "10.0.0.2", 9003, "root", "cluster-b", "pipe-c", 3, 300); + + final List snapshots = registry.snapshot(); + + assertEquals(3, snapshots.size()); + assertNotNull( + findSnapshot( + snapshots, + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1")); + assertNotNull( + findSnapshot( + snapshots, + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 2, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1")); + assertNotNull( + findSnapshot( + snapshots, + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 2, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.2")); + } + + @Test + public void testConfigNodeReceiversAreShownWithDataNodeResultsAndAirGapProtocol() { + registerDataSession("data-1", 1, "10.0.0.1", 9001, "root", "cluster-a", "pipe-a", 1, 100); + registry.registerOrUpdateSession( + "config-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + 0, + PipeReceiverRuntimeRegistry.PROTOCOL_AIR_GAP, + "10.0.0.2", + 9002, + "root", + "cluster-b", + "pipe-b", + 2, + 200); + + final List snapshots = registry.snapshot(); + + assertEquals(2, snapshots.size()); + final PipeReceiverRuntimeSnapshot configSnapshot = + findSnapshot( + snapshots, + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + 0, + PipeReceiverRuntimeRegistry.PROTOCOL_AIR_GAP, + "10.0.0.2"); + assertNotNull(configSnapshot); + assertEquals(PipeReceiverRuntimeRegistry.PROTOCOL_AIR_GAP, configSnapshot.getProtocol()); + assertEquals(1, configSnapshot.getConnectionCount()); + assertEquals(1, configSnapshot.getPipeCount()); + } + + @Test + public void testAggregateAndSortSnapshots() { + registry.registerOrUpdateSession( + "data-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 2, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "127.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + registry.markTransfer("data-1", 200); + registry.registerOrUpdateSession( + "data-2", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 2, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "127.0.0.1", + 9002, + "root", + "cluster-b", + "pipe-b", + 2, + 150); + registry.registerOrUpdateSession( + "config-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_AIR_GAP, + "127.0.0.2", + 9003, + "root", + PipeReceiverRuntimeRegistry.UNKNOWN, + null, + Long.MIN_VALUE, + 300); + + final List snapshots = registry.snapshot(); + + assertEquals(3, snapshots.size()); + assertEquals( + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, snapshots.get(0).getReceiverNodeType()); + assertEquals( + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, snapshots.get(1).getReceiverNodeType()); + assertEquals( + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, snapshots.get(2).getReceiverNodeType()); + + final PipeReceiverRuntimeSnapshot clusterASnapshot = snapshots.get(1); + assertEquals(1, clusterASnapshot.getConnectionCount()); + assertEquals("9001", clusterASnapshot.getSenderPorts()); + assertEquals(1, clusterASnapshot.getPipeCount()); + assertTrue(clusterASnapshot.getPipeIds().contains("pipe-a@")); + assertEquals("cluster-a", clusterASnapshot.getSenderClusterId()); + assertEquals(100, clusterASnapshot.getLastHandshakeTime()); + assertEquals(200, clusterASnapshot.getLastTransferTime()); + + final PipeReceiverRuntimeSnapshot clusterBSnapshot = snapshots.get(2); + assertEquals(1, clusterBSnapshot.getConnectionCount()); + assertEquals("9002", clusterBSnapshot.getSenderPorts()); + assertEquals(1, clusterBSnapshot.getPipeCount()); + assertTrue(clusterBSnapshot.getPipeIds().contains("pipe-b@")); + assertEquals("cluster-b", clusterBSnapshot.getSenderClusterId()); + assertEquals(150, clusterBSnapshot.getLastHandshakeTime()); + assertEquals(150, clusterBSnapshot.getLastTransferTime()); + } + + @Test + public void testDefaultSnapshotOrderingUsesAllDocumentedKeys() { + registry.registerOrUpdateSession( + "data-cluster-b", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9007, + "root", + "cluster-b", + "pipe-g", + 7, + 700); + registry.registerOrUpdateSession( + "data-user-b", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.2", + 9006, + "bob", + "cluster-a", + "pipe-f", + 6, + 600); + registry.registerOrUpdateSession( + "data-address-a", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9004, + "root", + "cluster-a", + "pipe-d", + 4, + 400); + registry.registerOrUpdateSession( + "config-node-2", + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + 2, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9002, + "root", + "cluster-a", + "pipe-b", + 2, + 200); + registry.registerOrUpdateSession( + "data-user-a", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.2", + 9005, + "alice", + "cluster-a", + "pipe-e", + 5, + 500); + registry.registerOrUpdateSession( + "config-node-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + registry.registerOrUpdateSession( + "data-air-gap", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_AIR_GAP, + "10.0.0.2", + 9003, + "root", + "cluster-a", + "pipe-c", + 3, + 300); + + final List snapshots = registry.snapshot(); + + assertEquals(7, snapshots.size()); + assertSnapshotOrderKey( + snapshots.get(0), + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "cluster-a", + "10.0.0.1", + "root"); + assertSnapshotOrderKey( + snapshots.get(1), + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + 2, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "cluster-a", + "10.0.0.1", + "root"); + assertSnapshotOrderKey( + snapshots.get(2), + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_AIR_GAP, + "cluster-a", + "10.0.0.2", + "root"); + assertSnapshotOrderKey( + snapshots.get(3), + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "cluster-a", + "10.0.0.1", + "root"); + assertSnapshotOrderKey( + snapshots.get(4), + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "cluster-a", + "10.0.0.2", + "alice"); + assertSnapshotOrderKey( + snapshots.get(5), + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "cluster-a", + "10.0.0.2", + "bob"); + assertSnapshotOrderKey( + snapshots.get(6), + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "cluster-b", + "10.0.0.1", + "root"); + } + + @Test + public void testDeregisterSession() { + registry.registerOrUpdateSession( + "data-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "127.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + + assertEquals(1, registry.snapshot().size()); + + registry.deregister("data-1"); + + assertTrue(registry.snapshot().isEmpty()); + } + + @Test + public void testPipeStopOrDropUpdatesPipeIdsAndPipeCount() { + registerDataSession("data-1", 1, "10.0.0.1", 9001, "root", "cluster-a", "pipe-a", 1, 100); + registerDataSession("data-2", 1, "10.0.0.1", 9002, "root", "cluster-a", "pipe-b", 2, 200); + + List snapshots = registry.snapshot(); + assertEquals(1, snapshots.size()); + assertEquals(2, snapshots.get(0).getConnectionCount()); + assertEquals(2, snapshots.get(0).getPipeCount()); + + registry.removePipeFromAllSessions("pipe-a", 1); + + snapshots = registry.snapshot(); + assertEquals(1, snapshots.size()); + assertEquals(2, snapshots.get(0).getConnectionCount()); + assertEquals(1, snapshots.get(0).getPipeCount()); + assertFalse(snapshots.get(0).getPipeIds().contains("pipe-a@")); + assertTrue(snapshots.get(0).getPipeIds().contains("pipe-b@")); + + registry.deregister("data-2"); + + snapshots = registry.snapshot(); + assertEquals(1, snapshots.size()); + assertEquals(1, snapshots.get(0).getConnectionCount()); + assertEquals(0, snapshots.get(0).getPipeCount()); + assertEquals(PipeReceiverRuntimeRegistry.UNKNOWN, snapshots.get(0).getPipeIds()); + } + + @Test + public void testMixedVersionPipeIdsUnknownCompatibility() { + registerDataSession( + "data-1", + 1, + "10.0.0.1", + 9001, + "root", + PipeReceiverRuntimeRegistry.UNKNOWN, + null, + Long.MIN_VALUE, + 100); + + final List snapshots = registry.snapshot(); + + assertEquals(1, snapshots.size()); + assertEquals("10.0.0.1", snapshots.get(0).getSenderAddress()); + assertEquals("9001", snapshots.get(0).getSenderPorts()); + assertEquals(1, snapshots.get(0).getConnectionCount()); + assertEquals(0, snapshots.get(0).getPipeCount()); + assertEquals(PipeReceiverRuntimeRegistry.UNKNOWN, snapshots.get(0).getPipeIds()); + } + + @Test + public void testLastTransferTimeDefaultsToLastHandshakeTimeBeforeTransfer() { + registerDataSession("data-1", 1, "10.0.0.1", 9001, "root", "cluster-a", "pipe-a", 1, 100); + + List snapshots = registry.snapshot(); + + assertEquals(1, snapshots.size()); + assertEquals(100, snapshots.get(0).getLastHandshakeTime()); + assertEquals(100, snapshots.get(0).getLastTransferTime()); + + registry.markTransfer("data-1", 50); + snapshots = registry.snapshot(); + + assertEquals(100, snapshots.get(0).getLastTransferTime()); + + registry.markTransfer("data-1", 200); + snapshots = registry.snapshot(); + + assertEquals(200, snapshots.get(0).getLastTransferTime()); + } + + @Test + public void testReceiverRestartClearsRuntimeAndAllowsReconnect() { + registerDataSession("data-1", 1, "10.0.0.1", 9001, "root", "cluster-a", "pipe-a", 1, 100); + + assertEquals(1, registry.snapshot().size()); + + registry.clear(); + + assertTrue(registry.snapshot().isEmpty()); + + registerDataSession("data-2", 1, "10.0.0.1", 9002, "root", "cluster-a", "pipe-b", 2, 200); + + final List snapshots = registry.snapshot(); + assertEquals(1, snapshots.size()); + assertEquals("9002", snapshots.get(0).getSenderPorts()); + assertTrue(snapshots.get(0).getPipeIds().contains("pipe-b@")); + } + + @Test + public void testLargeActiveSessionsSnapshotPerformanceAndTransferHotPathOverhead() { + final int activeConnectionCount = 1000; + final int pipeLifecycleOperationCount = 5000; + for (int i = 0; i < activeConnectionCount; i++) { + registerDataSession( + "data-" + i, + i % 10, + "10.0." + (i / 100) + "." + (i % 100), + 9000 + i, + "root", + "cluster-" + (i % 5), + "pipe-" + i, + i, + i); + } + + for (int i = 0; i < pipeLifecycleOperationCount; i++) { + final int connectionIndex = i % activeConnectionCount; + registerDataSession( + "data-" + connectionIndex, + connectionIndex % 10, + "10.0." + (connectionIndex / 100) + "." + (connectionIndex % 100), + 9000 + connectionIndex, + "root", + "cluster-" + (connectionIndex % 5), + "pipe-update-" + i, + i, + 10_000L + i); + } + + final long transferStartTime = System.nanoTime(); + for (int i = 0; i < pipeLifecycleOperationCount; i++) { + registry.markTransfer("data-" + (i % activeConnectionCount), 20_000L + i); + } + final long transferDurationNanos = System.nanoTime() - transferStartTime; + + final long snapshotStartTime = System.nanoTime(); + final List snapshots = registry.snapshot(); + final long snapshotDurationNanos = System.nanoTime() - snapshotStartTime; + + assertEquals( + activeConnectionCount, + snapshots.stream().mapToInt(PipeReceiverRuntimeSnapshot::getConnectionCount).sum()); + assertEquals( + 20_000L + pipeLifecycleOperationCount - 1, + snapshots.stream() + .mapToLong(PipeReceiverRuntimeSnapshot::getLastTransferTime) + .max() + .orElse(0)); + assertTrue( + "transfer updates should stay lightweight, duration ms: " + + TimeUnit.NANOSECONDS.toMillis(transferDurationNanos), + TimeUnit.NANOSECONDS.toMillis(transferDurationNanos) < 3000); + assertTrue( + "snapshot should be built from in-memory state, duration ms: " + + TimeUnit.NANOSECONDS.toMillis(snapshotDurationNanos), + TimeUnit.NANOSECONDS.toMillis(snapshotDurationNanos) < 3000); + } + + @Test + public void testRegisterOrUpdateSessionAddsPipeId() { + registry.registerOrUpdateSession( + "config-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + -1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "127.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + registry.registerOrUpdateSession( + "config-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + -1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "127.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-b", + 2, + 200); + + final List snapshots = registry.snapshot(); + + assertEquals(1, snapshots.size()); + assertEquals(2, snapshots.get(0).getPipeCount()); + assertTrue(snapshots.get(0).getPipeIds().contains("pipe-a@")); + assertTrue(snapshots.get(0).getPipeIds().contains("pipe-b@")); + } + + @Test + public void testSharedSessionCanRecordAndRemoveMultiplePipesOnTransfer() { + registerDataSession("data-1", 1, "10.0.0.1", 9001, "root", "cluster-a", "pipe-a", 1, 100); + registry.markTransfer("data-1", "pipe-b", 2, 200); + + List snapshots = registry.snapshot(); + + assertEquals(1, snapshots.size()); + assertEquals(1, snapshots.get(0).getConnectionCount()); + assertEquals(2, snapshots.get(0).getPipeCount()); + assertTrue(snapshots.get(0).getPipeIds().contains("pipe-a@")); + assertTrue(snapshots.get(0).getPipeIds().contains("pipe-b@")); + assertEquals(100, snapshots.get(0).getLastHandshakeTime()); + assertEquals(200, snapshots.get(0).getLastTransferTime()); + + registry.removePipe("data-1", "pipe-a", 1); + snapshots = registry.snapshot(); + + assertEquals(1, snapshots.size()); + assertEquals(1, snapshots.get(0).getPipeCount()); + assertFalse(snapshots.get(0).getPipeIds().contains("pipe-a@")); + assertTrue(snapshots.get(0).getPipeIds().contains("pipe-b@")); + } + + @Test + public void testRemovePipeFromAllSessions() { + registerDataSession("data-1", 1, "10.0.0.1", 9001, "root", "cluster-a", "pipe-a", 1, 100); + registerDataSession("data-2", 1, "10.0.0.1", 9002, "root", "cluster-a", "pipe-a", 1, 200); + registerDataSession("data-3", 1, "10.0.0.1", 9003, "root", "cluster-a", "pipe-b", 2, 300); + + registry.removePipeFromAllSessions("pipe-a", 1); + + final List snapshots = registry.snapshot(); + + assertEquals(1, snapshots.size()); + assertEquals(3, snapshots.get(0).getConnectionCount()); + assertEquals(1, snapshots.get(0).getPipeCount()); + assertFalse(snapshots.get(0).getPipeIds().contains("pipe-a@")); + assertTrue(snapshots.get(0).getPipeIds().contains("pipe-b@")); + } + + @Test + public void testMultipleConnectionsAndDuplicatePipeAggregationDoesNotInflatePipeCount() { + registerDataSession("data-1", 1, "10.0.0.1", 9001, "root", "cluster-a", "pipe-a", 1, 100); + registerDataSession("data-2", 1, "10.0.0.1", 9002, "root", "cluster-a", "pipe-a", 1, 200); + registerDataSession("data-3", 1, "10.0.0.1", -1, "root", "cluster-a", "pipe-b", 2, 150); + + registry.markTransfer("data-1", 300); + registry.markTransfer("data-2", 250); + registry.markTransfer("data-1", 50); + + final List snapshots = registry.snapshot(); + + assertEquals(1, snapshots.size()); + final PipeReceiverRuntimeSnapshot snapshot = snapshots.get(0); + assertEquals(3, snapshot.getConnectionCount()); + assertEquals("Unknown,9001,9002", snapshot.getSenderPorts()); + assertEquals(2, snapshot.getPipeCount()); + assertTrue(snapshot.getPipeIds().contains("pipe-a@")); + assertTrue(snapshot.getPipeIds().contains("pipe-b@")); + assertEquals(200, snapshot.getLastHandshakeTime()); + assertEquals(300, snapshot.getLastTransferTime()); + } + + @Test + public void testSameSenderAddressWithDifferentUsersAreNotAggregated() { + registerDataSession("data-user-a", 1, "10.0.0.1", 9001, "alice", "cluster-a", "pipe-a", 1, 100); + registerDataSession("data-user-b", 1, "10.0.0.1", 9002, "bob", "cluster-a", "pipe-b", 2, 200); + + final List snapshots = registry.snapshot(); + + assertEquals(2, snapshots.size()); + assertEquals("alice", snapshots.get(0).getUserName()); + assertEquals("9001", snapshots.get(0).getSenderPorts()); + assertEquals(1, snapshots.get(0).getConnectionCount()); + assertEquals(1, snapshots.get(0).getPipeCount()); + assertTrue(snapshots.get(0).getPipeIds().contains("pipe-a@")); + assertEquals("bob", snapshots.get(1).getUserName()); + assertEquals("9002", snapshots.get(1).getSenderPorts()); + assertEquals(1, snapshots.get(1).getConnectionCount()); + assertEquals(1, snapshots.get(1).getPipeCount()); + assertTrue(snapshots.get(1).getPipeIds().contains("pipe-b@")); + } + + @Test + public void testSameSenderAddressAndUserWithDifferentClusterIdsAreNotAggregated() { + registerDataSession( + "data-cluster-a", 1, "10.0.0.1", 9001, "root", "cluster-a", "pipe-a", 1, 100); + registerDataSession( + "data-cluster-b", 1, "10.0.0.1", 9002, "root", "cluster-b", "pipe-b", 2, 200); + + final List snapshots = registry.snapshot(); + + assertEquals(2, snapshots.size()); + assertEquals("cluster-a", snapshots.get(0).getSenderClusterId()); + assertEquals("9001", snapshots.get(0).getSenderPorts()); + assertEquals(1, snapshots.get(0).getConnectionCount()); + assertEquals(1, snapshots.get(0).getPipeCount()); + assertTrue(snapshots.get(0).getPipeIds().contains("pipe-a@")); + assertEquals("cluster-b", snapshots.get(1).getSenderClusterId()); + assertEquals("9002", snapshots.get(1).getSenderPorts()); + assertEquals(1, snapshots.get(1).getConnectionCount()); + assertEquals(1, snapshots.get(1).getPipeCount()); + assertTrue(snapshots.get(1).getPipeIds().contains("pipe-b@")); + } + + @Test + public void testSameSenderAddressWithDifferentProtocolsAreNotAggregated() { + registry.registerOrUpdateSession( + "data-thrift", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-thrift", + 1, + 100); + registry.registerOrUpdateSession( + "data-air-gap", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_AIR_GAP, + "10.0.0.1", + 9002, + "root", + "cluster-a", + "pipe-air-gap", + 2, + 200); + registry.registerOrUpdateSession( + "data-writeback", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_WRITEBACK, + "10.0.0.1", + 9003, + "root", + "cluster-a", + "pipe-writeback", + 3, + 300); + + final List snapshots = registry.snapshot(); + + assertEquals(3, snapshots.size()); + final PipeReceiverRuntimeSnapshot airGapSnapshot = + findSnapshot( + snapshots, + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_AIR_GAP, + "10.0.0.1"); + final PipeReceiverRuntimeSnapshot thriftSnapshot = + findSnapshot( + snapshots, + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1"); + final PipeReceiverRuntimeSnapshot writeBackSnapshot = + findSnapshot( + snapshots, + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_WRITEBACK, + "10.0.0.1"); + assertNotNull(airGapSnapshot); + assertNotNull(thriftSnapshot); + assertNotNull(writeBackSnapshot); + assertEquals("9002", airGapSnapshot.getSenderPorts()); + assertEquals(1, airGapSnapshot.getConnectionCount()); + assertTrue(airGapSnapshot.getPipeIds().contains("pipe-air-gap@")); + assertEquals("9001", thriftSnapshot.getSenderPorts()); + assertEquals(1, thriftSnapshot.getConnectionCount()); + assertTrue(thriftSnapshot.getPipeIds().contains("pipe-thrift@")); + assertEquals("9003", writeBackSnapshot.getSenderPorts()); + assertEquals(1, writeBackSnapshot.getConnectionCount()); + assertTrue(writeBackSnapshot.getPipeIds().contains("pipe-writeback@")); + } + + @Test + public void testBlankAndMalformedRuntimeFieldsFallbackToUnknown() { + registry.registerOrUpdateSession( + " ", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "10.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + assertTrue(registry.snapshot().isEmpty()); + + registry.registerOrUpdateSession( + "data-blank", " ", -1, "\t", "", -1, "\n", " ", "pipe-legacy", Long.MIN_VALUE, 0); + + final List snapshots = registry.snapshot(); + + assertEquals(1, snapshots.size()); + final PipeReceiverRuntimeSnapshot snapshot = snapshots.get(0); + assertEquals(PipeReceiverRuntimeRegistry.UNKNOWN, snapshot.getReceiverNodeType()); + assertEquals(-1, snapshot.getReceiverNodeId()); + assertFalse(snapshot.isReceiverNodeIdKnown()); + assertEquals(PipeReceiverRuntimeRegistry.UNKNOWN, snapshot.getProtocol()); + assertEquals(PipeReceiverRuntimeRegistry.UNKNOWN, snapshot.getSenderAddress()); + assertEquals(PipeReceiverRuntimeRegistry.UNKNOWN, snapshot.getSenderPorts()); + assertEquals(PipeReceiverRuntimeRegistry.UNKNOWN, snapshot.getUserName()); + assertEquals(PipeReceiverRuntimeRegistry.UNKNOWN, snapshot.getSenderClusterId()); + assertEquals(1, snapshot.getPipeCount()); + assertEquals("pipe-legacy@Unknown", snapshot.getPipeIds()); + assertEquals(0, snapshot.getLastHandshakeTime()); + assertEquals(0, snapshot.getLastTransferTime()); + } + + @Test + public void testRegisterOrUpdateSessionWithoutPipeInfoKeepsPipeId() { + registry.registerOrUpdateSession( + "config-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + -1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "127.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + registry.registerOrUpdateSession( + "config-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + -1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "127.0.0.1", + 9001, + "root", + PipeReceiverRuntimeRegistry.UNKNOWN, + null, + Long.MIN_VALUE, + 200); + + final List snapshots = registry.snapshot(); + + assertEquals(1, snapshots.size()); + assertEquals(1, snapshots.get(0).getPipeCount()); + assertTrue(snapshots.get(0).getPipeIds().contains("pipe-a@")); + } + + @Test + public void testUnknownSenderPort() { + registry.registerOrUpdateSession( + "data-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "127.0.0.1", + -1, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + + final List snapshots = registry.snapshot(); + + assertEquals(1, snapshots.size()); + assertEquals(PipeReceiverRuntimeRegistry.UNKNOWN, snapshots.get(0).getSenderPorts()); + } + + @Test + public void testDuplicateUnknownSenderPortsAreDeduplicated() { + registry.registerOrUpdateSession( + "data-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "127.0.0.1", + -1, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + registry.registerOrUpdateSession( + "data-2", + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + 1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "127.0.0.1", + -2, + "root", + "cluster-a", + "pipe-b", + 2, + 200); + + final List snapshots = registry.snapshot(); + + assertEquals(1, snapshots.size()); + assertEquals(PipeReceiverRuntimeRegistry.UNKNOWN, snapshots.get(0).getSenderPorts()); + assertEquals(2, snapshots.get(0).getConnectionCount()); + assertEquals(2, snapshots.get(0).getPipeCount()); + } + + @Test + public void testUnknownReceiverNodeId() { + registry.registerOrUpdateSession( + "config-1", + PipeReceiverRuntimeRegistry.NODE_TYPE_CONFIG_NODE, + -1, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + "127.0.0.1", + 9001, + "root", + "cluster-a", + "pipe-a", + 1, + 100); + + final List snapshots = registry.snapshot(); + + assertEquals(1, snapshots.size()); + assertEquals(-1, snapshots.get(0).getReceiverNodeId()); + assertFalse(snapshots.get(0).isReceiverNodeIdKnown()); + } + + private void registerDataSession( + String connectionKey, + int receiverNodeId, + String senderAddress, + int senderPort, + String userName, + String senderClusterId, + String pipeName, + long pipeCreationTime, + long handshakeTime) { + registry.registerOrUpdateSession( + connectionKey, + PipeReceiverRuntimeRegistry.NODE_TYPE_DATA_NODE, + receiverNodeId, + PipeReceiverRuntimeRegistry.PROTOCOL_THRIFT, + senderAddress, + senderPort, + userName, + senderClusterId, + pipeName, + pipeCreationTime, + handshakeTime); + } + + private static PipeReceiverRuntimeSnapshot findSnapshot( + List snapshots, + String receiverNodeType, + int receiverNodeId, + String protocol, + String senderAddress) { + for (PipeReceiverRuntimeSnapshot snapshot : snapshots) { + if (receiverNodeType.equals(snapshot.getReceiverNodeType()) + && receiverNodeId == snapshot.getReceiverNodeId() + && protocol.equals(snapshot.getProtocol()) + && senderAddress.equals(snapshot.getSenderAddress())) { + return snapshot; + } + } + return null; + } + + private static void assertSnapshotOrderKey( + PipeReceiverRuntimeSnapshot snapshot, + String receiverNodeType, + int receiverNodeId, + String protocol, + String senderClusterId, + String senderAddress, + String userName) { + assertEquals(receiverNodeType, snapshot.getReceiverNodeType()); + assertEquals(receiverNodeId, snapshot.getReceiverNodeId()); + assertEquals(protocol, snapshot.getProtocol()); + assertEquals(senderClusterId, snapshot.getSenderClusterId()); + assertEquals(senderAddress, snapshot.getSenderAddress()); + assertEquals(userName, snapshot.getUserName()); + } +} diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeRequestTypeTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeRequestTypeTest.java index c2e0ba949f984..9ae9e6c8db0c4 100644 --- a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeRequestTypeTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeRequestTypeTest.java @@ -53,8 +53,33 @@ public void testAllV13RequestTypesAreRecognized() { assertV13RequestType((short) 400, PipeRequestType.TRANSFER_SLICE); } + @Test + public void testLifecycleRequestTypesAreRecognized() { + assertRequestType((short) 500, PipeRequestType.TRANSFER_PIPE_RECEIVER_RUNTIME_INFO_CLEANUP); + } + + @Test + public void testPipeReceiverRuntimeInfoCleanupReqSerDe() throws Exception { + final PipeTransferPipeReceiverRuntimeInfoCleanupReq req = + PipeTransferPipeReceiverRuntimeInfoCleanupReq.toTPipeTransferReq("pipe-a", 1L); + + final PipeTransferPipeReceiverRuntimeInfoCleanupReq parsedReq = + PipeTransferPipeReceiverRuntimeInfoCleanupReq.fromTPipeTransferReq(req); + + Assert.assertEquals( + PipeRequestType.TRANSFER_PIPE_RECEIVER_RUNTIME_INFO_CLEANUP.getType(), req.type); + Assert.assertEquals("pipe-a", parsedReq.getPipeName()); + Assert.assertEquals(1L, parsedReq.getPipeCreationTime()); + Assert.assertEquals(req, parsedReq); + } + private static void assertV13RequestType( final short type, final PipeRequestType expectedRequestType) { + assertRequestType(type, expectedRequestType); + } + + private static void assertRequestType( + final short type, final PipeRequestType expectedRequestType) { Assert.assertTrue(PipeRequestType.isValidatedRequestType(type)); Assert.assertEquals(expectedRequestType, PipeRequestType.valueOf(type)); } diff --git a/iotdb-core/relational-grammar/src/main/antlr4/org/apache/iotdb/db/relational/grammar/sql/RelationalSql.g4 b/iotdb-core/relational-grammar/src/main/antlr4/org/apache/iotdb/db/relational/grammar/sql/RelationalSql.g4 index 4222754cdb9bb..a8aafb72e5834 100644 --- a/iotdb-core/relational-grammar/src/main/antlr4/org/apache/iotdb/db/relational/grammar/sql/RelationalSql.g4 +++ b/iotdb-core/relational-grammar/src/main/antlr4/org/apache/iotdb/db/relational/grammar/sql/RelationalSql.g4 @@ -102,6 +102,7 @@ statement | startPipeStatement | stopPipeStatement | showPipesStatement + | showReceiversStatement | createPipePluginStatement | dropPipePluginStatement | showPipePluginsStatement @@ -520,6 +521,10 @@ showPipesStatement : SHOW ((PIPE pipeName=identifier) | PIPES (WHERE (CONNECTOR | SINK) USED BY pipeName=identifier)?) ; +showReceiversStatement + : SHOW RECEIVERS + ; + createPipePluginStatement : CREATE PIPEPLUGIN (IF NOT EXISTS)? pluginName=identifier AS className=string uriClause ; @@ -1511,7 +1516,7 @@ nonReserved | OBJECT | OF | OFFSET | OMIT | ONE | ONLY | OPTION | ORDINALITY | OUTPUT | OVER | OVERFLOW | PARTITION | PARTITIONS | PASSING | PAST | PATH | PATTERN | PER | PERIOD | PERMUTE | PIPE | PIPEPLUGIN | PIPEPLUGINS | PIPES | PLAN | POSITION | PRECEDING | PRECISION | PRIVILEGES | PREVIOUS | PROCESSLIST | PROCESSOR | PROPERTIES | PRUNE | QUERIES | QUERY | QUOTES - | RANGE | READ | READONLY | RECONSTRUCT | REFRESH | REGION | REGIONID | REGIONS | REMOVE | RENAME | REPAIR | REPEAT | REPEATABLE | REPLACE | RESET | RESPECT | RESTRICT | RETURN | RETURNING | RETURNS | REVOKE | ROLE | ROLES | ROLLBACK | ROOT | ROW | ROWS | RPR_FIRST | RPR_LAST | RUNNING + | RANGE | READ | READONLY | RECEIVERS | RECONSTRUCT | REFRESH | REGION | REGIONID | REGIONS | REMOVE | RENAME | REPAIR | REPEAT | REPEATABLE | REPLACE | RESET | RESPECT | RESTRICT | RETURN | RETURNING | RETURNS | REVOKE | ROLE | ROLES | ROLLBACK | ROOT | ROW | ROWS | RPR_FIRST | RPR_LAST | RUNNING | SERIESSLOTID | SERVICE | SERVICES | SCALAR | SCHEMA | SCHEMAS | SECOND | SECURITY | SEEK | SERIALIZABLE | SESSION | SET | SETS | SECURITY | SHOW | SINK | SOME | SOURCE | START | STATS | STOP | SUBSCRIPTION | SUBSCRIPTIONS | SUBSET | SUBSTRING | SYSTEM | TABLES | TABLESAMPLE | TAG | TAGS | TEXT | TEXT_STRING | TIES | TIME | TIMEPARTITION | TIMER | TIMER_XL | TIMESERIES | TIMESLOTID | TIMESTAMP | TO | TOPIC | TOPICS | TRAILING | TRANSACTION | TRUNCATE | TRY_CAST | TYPE @@ -1801,6 +1806,7 @@ QUOTES: 'QUOTES'; RANGE: 'RANGE'; READ: 'READ'; READONLY: 'READONLY'; +RECEIVERS: 'RECEIVERS'; RECONSTRUCT: 'RECONSTRUCT'; RECURSIVE: 'RECURSIVE'; REFRESH: 'REFRESH';