diff --git a/integration-test/src/main/java/org/apache/iotdb/it/env/MultiEnvFactory.java b/integration-test/src/main/java/org/apache/iotdb/it/env/MultiEnvFactory.java index 5832f1c485bc..de09307fb887 100644 --- a/integration-test/src/main/java/org/apache/iotdb/it/env/MultiEnvFactory.java +++ b/integration-test/src/main/java/org/apache/iotdb/it/env/MultiEnvFactory.java @@ -33,11 +33,20 @@ public class MultiEnvFactory { private static final List envList = new ArrayList<>(); private static final Logger logger = IoTDBTestLogger.logger; private static String currentMethodName; + private static String currentTestClassName; private MultiEnvFactory() { // Empty constructor } + static { + Runtime.getRuntime() + .addShutdownHook( + new Thread( + MultiEnvFactory::cleanupEnvs, + MultiEnvFactory.class.getSimpleName() + "-ShutdownHook")); + } + public static void setTestMethodName(final String testMethodName) { currentMethodName = testMethodName; envList.forEach(baseEnv -> baseEnv.setTestMethodName(testMethodName)); @@ -49,8 +58,15 @@ public static BaseEnv getEnv(final int index) throws IndexOutOfBoundsException { } /** Create several environments according to the specific number. */ - public static void createEnv(final int num) { - // Not judge EnvType for individual test convenience + public static synchronized void createEnv(final int num) { + final String requestedTestClassName = detectCurrentTestClassName(); + if (requestedTestClassName.equals(currentTestClassName) && envList.size() == num) { + return; + } + + cleanupEnvs(); + + currentTestClassName = requestedTestClassName; final long startTime = System.currentTimeMillis(); for (int i = 0; i < num; ++i) { try { @@ -62,4 +78,30 @@ public static void createEnv(final int num) { } } } + + public static synchronized void cleanupEnvs() { + for (final BaseEnv baseEnv : envList) { + try { + baseEnv.cleanClusterEnvironment(); + } catch (final Exception e) { + logger.warn("Cleanup env error", e); + } + } + envList.clear(); + currentTestClassName = null; + } + + private static String detectCurrentTestClassName() { + final StackTraceElement[] stack = Thread.currentThread().getStackTrace(); + for (final StackTraceElement stackTraceElement : stack) { + final String className = stackTraceElement.getClassName(); + if (className.endsWith("IT")) { + final String simpleClassName = className.substring(className.lastIndexOf('.') + 1); + if (!simpleClassName.startsWith("Abstract")) { + return className; + } + } + } + return "UNKNOWN-IT"; + } } 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 0e7eefb3a418..1d3aaa9b9add 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 @@ -106,6 +106,7 @@ public abstract class AbstractEnv implements BaseEnv { protected int index = 0; protected long startTime; protected int retryCount = 30; + protected boolean clusterRunning = false; private IClientManager clientManager; private List configNodeKillPoints = new ArrayList<>(); private List dataNodeKillPoints = new ArrayList<>(); @@ -177,6 +178,14 @@ protected void initEnvironment( final int dataNodesNum, final int retryCount, final boolean addAINode) { + if (clusterRunning) { + logger.info( + "Reuse running cluster for test class {} and method {}", + getTestClassName(), + testMethodName); + return; + } + this.retryCount = retryCount; this.configNodeWrapperList = new ArrayList<>(); this.dataNodeWrapperList = new ArrayList<>(); @@ -264,6 +273,7 @@ protected void initEnvironment( } checkClusterStatusWithoutUnknown(); + clusterRunning = true; } private ConfigNodeWrapper newConfigNode() { @@ -609,9 +619,16 @@ public void cleanClusterEnvironment() { } if (clientManager != null) { clientManager.close(); + clientManager = null; } + configNodeWrapperList = Collections.emptyList(); + dataNodeWrapperList = Collections.emptyList(); + aiNodeWrapperList = Collections.emptyList(); + configNodeKillPoints = new ArrayList<>(); + dataNodeKillPoints = new ArrayList<>(); testMethodName = null; clusterConfig = new MppClusterConfig(); + clusterRunning = false; } @Override diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/PipeEnvReuseManager.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/PipeEnvReuseManager.java new file mode 100644 index 000000000000..48b4d09c1322 --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/PipeEnvReuseManager.java @@ -0,0 +1,209 @@ +/* + * 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; + +import org.apache.iotdb.it.framework.IoTDBTestLogger; +import org.apache.iotdb.itbase.env.BaseEnv; + +import org.slf4j.Logger; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashSet; +import java.util.List; +import java.util.Locale; +import java.util.Objects; +import java.util.Set; +import java.util.stream.Collectors; + +public final class PipeEnvReuseManager { + + private static final Logger LOGGER = IoTDBTestLogger.logger; + private static final Set BUILTIN_USERS = + new HashSet<>( + Arrays.asList( + "root", "sys_admin", "security_admin", "audit_admin", "__internal_auditor")); + private static final Set BUILTIN_TABLE_DATABASES = + new HashSet<>(Arrays.asList("information_schema")); + + private PipeEnvReuseManager() { + // Utility class + } + + public static void prepareForNextTest(final BaseEnv... envs) { + final List distinctEnvs = + Arrays.stream(envs).filter(Objects::nonNull).distinct().collect(Collectors.toList()); + try { + for (final BaseEnv env : distinctEnvs) { + resetEnvironment(env); + } + } catch (final Exception e) { + LOGGER.warn("Failed to logically reset pipe IT environment, fallback to full cleanup", e); + for (final BaseEnv env : distinctEnvs) { + try { + env.cleanClusterEnvironment(); + } catch (final Exception cleanupException) { + LOGGER.warn("Failed to fully cleanup pipe IT environment", cleanupException); + } + } + } + } + + private static void resetEnvironment(final BaseEnv env) throws SQLException { + dropAllPipes(env); + dropAllCustomPipePlugins(env); + clearTreeDatabases(env); + dropAllTemplates(env); + clearUsers(env); + clearRoles(env); + clearTableDatabases(env); + } + + private static void dropAllPipes(final BaseEnv env) throws SQLException { + try (final Connection connection = env.getConnection(); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery("show pipes")) { + final List pipeNames = new ArrayList<>(); + while (resultSet.next()) { + final String pipeName = resultSet.getString(1); + if (pipeName != null && !pipeName.startsWith("__consensus")) { + pipeNames.add(pipeName); + } + } + for (final String pipeName : pipeNames) { + statement.execute("drop pipe " + quoteTreeIdentifier(pipeName)); + } + } + } + + private static void dropAllCustomPipePlugins(final BaseEnv env) throws SQLException { + try (final Connection connection = env.getConnection(); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery("show pipeplugins")) { + final List pluginNames = new ArrayList<>(); + while (resultSet.next()) { + final String pluginName = resultSet.getString(1); + final String pluginType = resultSet.getString(2); + if (pluginName != null && !"Builtin".equalsIgnoreCase(pluginType)) { + pluginNames.add(pluginName); + } + } + for (final String pluginName : pluginNames) { + statement.execute("drop pipePlugin " + quoteTreeIdentifier(pluginName)); + } + } + } + + private static void clearTreeDatabases(final BaseEnv env) throws SQLException { + try (final Connection connection = env.getConnection(); + final Statement statement = connection.createStatement()) { + try { + statement.execute("delete database root.**"); + } catch (final SQLException ignored) { + try { + statement.execute("drop database root.**"); + } catch (final SQLException ignoredAgain) { + // Ignore when there is no tree database to drop. + } + } + } + } + + private static void dropAllTemplates(final BaseEnv env) throws SQLException { + try (final Connection connection = env.getConnection(); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery("show device templates")) { + final List templateNames = new ArrayList<>(); + while (resultSet.next()) { + final String templateName = resultSet.getString(1); + if (templateName != null) { + templateNames.add(templateName); + } + } + for (final String templateName : templateNames) { + statement.execute("drop device template " + quoteTreeIdentifier(templateName)); + } + } + } + + private static void clearUsers(final BaseEnv env) throws SQLException { + try (final Connection connection = env.getConnection(); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery("list user")) { + final List users = new ArrayList<>(); + while (resultSet.next()) { + final String user = resultSet.getString(1); + if (user != null && !BUILTIN_USERS.contains(user.toLowerCase(Locale.ROOT))) { + users.add(user); + } + } + for (final String user : users) { + statement.execute("drop user " + quoteTreeIdentifier(user)); + } + } + } + + private static void clearRoles(final BaseEnv env) throws SQLException { + try (final Connection connection = env.getConnection(); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery("list role")) { + final List roles = new ArrayList<>(); + while (resultSet.next()) { + final String role = resultSet.getString(1); + if (role != null) { + roles.add(role); + } + } + for (final String role : roles) { + statement.execute("drop role " + quoteTreeIdentifier(role)); + } + } + } + + private static void clearTableDatabases(final BaseEnv env) throws SQLException { + try (final Connection connection = env.getConnection(BaseEnv.TABLE_SQL_DIALECT); + final Statement statement = connection.createStatement(); + final ResultSet resultSet = statement.executeQuery("show databases")) { + final List databases = new ArrayList<>(); + while (resultSet.next()) { + final String database = resultSet.getString(1); + if (database != null + && !BUILTIN_TABLE_DATABASES.contains(database.toLowerCase(Locale.ROOT))) { + databases.add(database); + } + } + for (final String database : databases) { + statement.execute("drop database if exists " + quoteTableIdentifier(database)); + } + } + } + + private static String quoteTreeIdentifier(final String identifier) { + return "`" + identifier.replace("`", "``") + "`"; + } + + private static String quoteTableIdentifier(final String identifier) { + return "\"" + identifier.replace("\"", "\"\"") + "\""; + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/manual/AbstractPipeTableModelDualManualIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/manual/AbstractPipeTableModelDualManualIT.java index 3b3fae80902f..08a97d9e19c6 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/manual/AbstractPipeTableModelDualManualIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/manual/AbstractPipeTableModelDualManualIT.java @@ -22,6 +22,7 @@ import org.apache.iotdb.consensus.ConsensusFactory; import org.apache.iotdb.it.env.MultiEnvFactory; import org.apache.iotdb.itbase.env.BaseEnv; +import org.apache.iotdb.pipe.it.PipeEnvReuseManager; import org.junit.After; import org.junit.Before; @@ -72,7 +73,6 @@ protected void setupConfig() { @After public final void tearDown() { - senderEnv.cleanClusterEnvironment(); - receiverEnv.cleanClusterEnvironment(); + PipeEnvReuseManager.prepareForNextTest(senderEnv, receiverEnv); } } diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/auto/AbstractPipeDualTreeModelAutoIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/auto/AbstractPipeDualTreeModelAutoIT.java index a7fae02f6d16..f6c00d82f983 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/auto/AbstractPipeDualTreeModelAutoIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/auto/AbstractPipeDualTreeModelAutoIT.java @@ -25,6 +25,7 @@ import org.apache.iotdb.it.env.MultiEnvFactory; import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper; import org.apache.iotdb.itbase.env.BaseEnv; +import org.apache.iotdb.pipe.it.PipeEnvReuseManager; import org.awaitility.Awaitility; import org.junit.After; @@ -79,8 +80,7 @@ protected void setupConfig() { @After public final void tearDown() { - senderEnv.cleanClusterEnvironment(); - receiverEnv.cleanClusterEnvironment(); + PipeEnvReuseManager.prepareForNextTest(senderEnv, receiverEnv); } protected static void awaitUntilFlush(BaseEnv env) { diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/manual/AbstractPipeDualTreeModelManualIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/manual/AbstractPipeDualTreeModelManualIT.java index 11f70d944b63..c09238ea07d0 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/manual/AbstractPipeDualTreeModelManualIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/treemodel/manual/AbstractPipeDualTreeModelManualIT.java @@ -25,6 +25,7 @@ import org.apache.iotdb.it.env.MultiEnvFactory; import org.apache.iotdb.it.env.cluster.node.DataNodeWrapper; import org.apache.iotdb.itbase.env.BaseEnv; +import org.apache.iotdb.pipe.it.PipeEnvReuseManager; import org.awaitility.Awaitility; import org.junit.After; @@ -78,8 +79,7 @@ protected void setupConfig() { @After public final void tearDown() { - senderEnv.cleanClusterEnvironment(); - receiverEnv.cleanClusterEnvironment(); + PipeEnvReuseManager.prepareForNextTest(senderEnv, receiverEnv); } protected static void awaitUntilFlush(BaseEnv env) { diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/single/AbstractPipeSingleIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/single/AbstractPipeSingleIT.java index 3ade13c7209d..efaa063ba2aa 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/single/AbstractPipeSingleIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/single/AbstractPipeSingleIT.java @@ -21,6 +21,7 @@ import org.apache.iotdb.it.env.MultiEnvFactory; import org.apache.iotdb.itbase.env.BaseEnv; +import org.apache.iotdb.pipe.it.PipeEnvReuseManager; import org.junit.After; import org.junit.Before; @@ -44,6 +45,6 @@ public void setUp() { @After public final void tearDown() { - env.cleanClusterEnvironment(); + PipeEnvReuseManager.prepareForNextTest(env); } } diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/AbstractPipeTripleManualIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/AbstractPipeTripleManualIT.java index f4e63e1d2f85..5c8c662b150a 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/AbstractPipeTripleManualIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/triple/AbstractPipeTripleManualIT.java @@ -22,6 +22,7 @@ import org.apache.iotdb.consensus.ConsensusFactory; import org.apache.iotdb.it.env.MultiEnvFactory; import org.apache.iotdb.itbase.env.BaseEnv; +import org.apache.iotdb.pipe.it.PipeEnvReuseManager; import org.junit.After; import org.junit.Before; @@ -81,8 +82,6 @@ protected void setupConfig() { @After public final void tearDown() { - env1.cleanClusterEnvironment(); - env2.cleanClusterEnvironment(); - env3.cleanClusterEnvironment(); + PipeEnvReuseManager.prepareForNextTest(env1, env2, env3); } }