From c3050dfea507ecc548eab96af96dea07260aac43 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Mon, 9 Nov 2026 14:10:32 +0800 Subject: [PATCH] Fix test helper JDBC cleanup and Pipe subtask stubbing race (#19160) * Stub rejected submissions before registering pipe subtasks * Fix JDBC resource cleanup in non-query test utilities * Stabilize pipe operation diagnostics IT barriers (cherry picked from commit c5f57e9305e795fbd97d6810f3234aed91a2cf39) --- .../iotdb/db/it/utils/IoTDBTestUtilsIT.java | 128 ++++++++++++++++++ .../apache/iotdb/db/it/utils/TestUtils.java | 31 +++-- .../IoTDBPipeOperationDiagnosticsIT.java | 35 ++++- 3 files changed, 183 insertions(+), 11 deletions(-) diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/utils/IoTDBTestUtilsIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/utils/IoTDBTestUtilsIT.java index 68ab36b346edb..a97b04758f8e0 100644 --- a/integration-test/src/test/java/org/apache/iotdb/db/it/utils/IoTDBTestUtilsIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/utils/IoTDBTestUtilsIT.java @@ -19,6 +19,7 @@ package org.apache.iotdb.db.it.utils; +import org.apache.iotdb.isession.SessionConfig; import org.apache.iotdb.it.env.EnvFactory; import org.apache.iotdb.it.framework.IoTDBTestRunner; import org.apache.iotdb.itbase.category.ClusterIT; @@ -34,12 +35,14 @@ import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Proxy; import java.sql.Connection; +import java.sql.SQLException; import java.sql.Statement; import java.util.Collections; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; @@ -86,6 +89,131 @@ public void testAssertDataAlwaysIncludesQueryTime() { assertEquals(1, queries.get()); } + @Test + public void testNonQueryResourceCleanupOnSuccess() { + for (boolean retry : new boolean[] {false, true}) { + for (boolean borrowedConnection : + retry ? new boolean[] {false} : new boolean[] {false, true}) { + final NonQueryResources resources = new NonQueryResources(null); + executeNonQuery(resources, retry, borrowedConnection); + assertResourcesClosed(resources, borrowedConnection); + } + } + } + + @Test + public void testNonQueryResourceCleanupOnRuntimeFailure() { + for (boolean retry : new boolean[] {false, true}) { + for (boolean borrowedConnection : + retry ? new boolean[] {false} : new boolean[] {false, true}) { + final IllegalStateException failure = + new IllegalStateException("Expected execution failure"); + final NonQueryResources resources = new NonQueryResources(failure); + assertSame( + failure, + assertThrows( + IllegalStateException.class, + () -> executeNonQuery(resources, retry, borrowedConnection))); + assertResourcesClosed(resources, borrowedConnection); + } + } + } + + @Test + public void testNonQueryResourceCleanupOnSqlFailure() { + for (boolean borrowedConnection : new boolean[] {false, true}) { + final SQLException failure = new SQLException("Expected execution failure"); + final NonQueryResources resources = new NonQueryResources(failure); + assertSame( + failure, + assertThrows( + RuntimeException.class, + () -> executeNonQuery(resources, false, borrowedConnection)) + .getCause()); + assertResourcesClosed(resources, borrowedConnection); + } + } + + private void executeNonQuery( + NonQueryResources resources, boolean retry, boolean borrowedConnection) { + if (retry) { + // The retry helper in dev/1.3 always owns its connection. + assertTrue( + TestUtils.tryExecuteNonQueriesWithRetry( + resources.env, + Collections.singletonList("flush"), + SessionConfig.DEFAULT_USER, + SessionConfig.DEFAULT_PASSWORD)); + } else { + final Connection connection = borrowedConnection ? resources.connection : null; + TestUtils.executeNonQuery(resources.env, "flush", connection); + } + } + + private void assertResourcesClosed(NonQueryResources resources, boolean borrowedConnection) { + assertEquals(1, resources.statementsClosed); + assertEquals(borrowedConnection ? 0 : 1, resources.connectionsOpened); + assertEquals(borrowedConnection ? 0 : 1, resources.connectionsClosed); + } + + private static class NonQueryResources { + private final Connection connection; + private final BaseEnv env; + private int connectionsOpened; + private int connectionsClosed; + private int statementsClosed; + + private NonQueryResources(Exception executionFailure) { + final Statement statement = + (Statement) + Proxy.newProxyInstance( + Statement.class.getClassLoader(), + new Class[] {Statement.class}, + (proxy, method, args) -> { + switch (method.getName()) { + case "execute": + if (executionFailure != null) { + throw executionFailure; + } + return false; + case "close": + statementsClosed++; + return null; + default: + throw new UnsupportedOperationException(method.getName()); + } + }); + connection = + (Connection) + Proxy.newProxyInstance( + Connection.class.getClassLoader(), + new Class[] {Connection.class}, + (proxy, method, args) -> { + switch (method.getName()) { + case "createStatement": + return statement; + case "close": + connectionsClosed++; + return null; + default: + throw new UnsupportedOperationException(method.getName()); + } + }); + env = + (BaseEnv) + Proxy.newProxyInstance( + BaseEnv.class.getClassLoader(), + new Class[] {BaseEnv.class}, + (proxy, method, args) -> { + if ("getConnection".equals(method.getName())) { + connectionsOpened++; + return connection; + } + throw new UnsupportedOperationException(method.getName()); + }); + } + } + private T delayQueries(final Class type, final T delegate, final AtomicInteger queries) { return type.cast( Proxy.newProxyInstance( diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/utils/TestUtils.java b/integration-test/src/test/java/org/apache/iotdb/db/it/utils/TestUtils.java index d375cc20e78e4..e5af90c5cb628 100644 --- a/integration-test/src/test/java/org/apache/iotdb/db/it/utils/TestUtils.java +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/utils/TestUtils.java @@ -567,7 +567,7 @@ public static void executeNonQueries( int lastIndex = 0; Connection localConnection = null; Connection connectionToUse = defaultConnection; - Statement statement; + Statement statement = null; try { // create a new connection if default is not provided or the previous is broken if (connectionToUse == null) { @@ -579,16 +579,27 @@ public static void executeNonQueries( statement.execute(sqlList.get(i)); } } catch (SQLException e) { - // the default connection should be closed by the upper level - // while the local connection should be closed here - if (connectionToUse == localConnection && localConnection != null) { - try { - localConnection.close(); - } catch (SQLException ex) { - // ignore - } - } throw new RuntimeException(e); + } finally { + closeNonQueryResources(statement, localConnection); + } + } + + private static void closeNonQueryResources(Statement statement, Connection localConnection) { + if (statement != null) { + try { + statement.close(); + } catch (Exception ignored) { + // ignore + } + } + // Caller-provided connections are never passed as localConnection. + if (localConnection != null) { + try { + localConnection.close(); + } catch (SQLException ignored) { + // ignore + } } } diff --git a/integration-test/src/test/java/org/apache/iotdb/pipe/it/single/IoTDBPipeOperationDiagnosticsIT.java b/integration-test/src/test/java/org/apache/iotdb/pipe/it/single/IoTDBPipeOperationDiagnosticsIT.java index 691e2b40b3b22..39b5c5ce58853 100644 --- a/integration-test/src/test/java/org/apache/iotdb/pipe/it/single/IoTDBPipeOperationDiagnosticsIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/pipe/it/single/IoTDBPipeOperationDiagnosticsIT.java @@ -96,12 +96,15 @@ public void testDropPipeExposesPreDeleteAndRejectsConcurrentOperations() throws new BlackholeServer( unavailableDataNode.getInternalAddress(), unavailableDataNode.getInternalPort())) { configNode.stopForcibly(); + configNode.clearLogContent(); configNode.start(); awaitLeader(); + awaitPipeLeaderChange(configNode, PIPE_NAME); try (final SyncConfigNodeIServiceClient client = getLeaderClient()) { + final int acceptedConnectionCountBeforeDrop = ignored.getAcceptedConnectionCount(); final TSStatus dropStatus = client.dropPipe(PIPE_NAME); assertDropTimeout(dropStatus, unavailableDataNodeId); - Assert.assertTrue(ignored.getAcceptedConnectionCount() > 0); + Assert.assertTrue(ignored.getAcceptedConnectionCount() > acceptedConnectionCountBeforeDrop); assertPipeState(client, PIPE_NAME, "PRE_DELETE"); assertPipeBeingDropped( @@ -141,8 +144,10 @@ public void testDropPipeProcedureRecoversWithPreDelete() throws Exception { new BlackholeServer( unavailableDataNode.getInternalAddress(), unavailableDataNode.getInternalPort())) { configNode.stopForcibly(); + configNode.clearLogContent(); configNode.start(); awaitLeader(); + awaitPipeLeaderChange(configNode, RECOVER_PIPE_NAME); try (final SyncConfigNodeIServiceClient client = getLeaderClient()) { final TSStatus dropStatus = client.dropPipe(RECOVER_PIPE_NAME); assertDropTimeout(dropStatus, unavailableDataNodeId); @@ -284,6 +289,34 @@ private static void awaitLeader() { }); } + private static void awaitPipeLeaderChange( + final ConfigNodeWrapper configNode, final String pipeName) { + // Leader readiness does not mean its asynchronous Pipe leader-change procedure has finished. + // Wait until that procedure starts pushing metadata, then use START on the already RUNNING + // pipe as an idempotent lock barrier. Otherwise DROP may time out waiting for the coordinator + // lock instead of reaching the DataNode operation that this test intends to block. + Awaitility.await() + .pollInterval(100, TimeUnit.MILLISECONDS) + .atMost(45, TimeUnit.SECONDS) + .until( + () -> + configNode.logContains( + "PipeHandleLeaderChangeProcedure: executeFromHandleOnDataNodes")); + Awaitility.await() + .pollInterval(100, TimeUnit.MILLISECONDS) + .atMost(45, TimeUnit.SECONDS) + .untilAsserted( + () -> { + try (final SyncConfigNodeIServiceClient client = getLeaderClientStatic()) { + final TSStatus status = client.startPipe(pipeName); + Assert.assertEquals( + status.getMessage(), + TSStatusCode.SUCCESS_STATUS.getStatusCode(), + status.getCode()); + } + }); + } + private static void restartDataNode(final DataNodeWrapper dataNodeWrapper) { if (!dataNodeWrapper.isAlive()) { dataNodeWrapper.start();