Fix test helper JDBC cleanup and Pipe subtask stubbing race

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.TableClusterIT;
@@ -34,6 +35,7 @@
 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;
@@ -42,6 +44,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
@@ -173,6 +176,129 @@ private <T> T delayQueries(final Class<T> type, final T delegate, final AtomicIn
             }));
   }
 
+  @Test
+  public void testNonQueryResourceCleanupOnSuccess() {
+    for (boolean retry : new boolean[] {false, true}) {
+      for (boolean borrowedConnection : 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 : 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) {
+    final Connection connection = borrowedConnection ? resources.connection : null;
+    if (retry) {
+      assertTrue(
+          TestUtils.tryExecuteNonQueriesWithRetry(
+              resources.env,
+              Collections.singletonList("flush"),
+              SessionConfig.DEFAULT_USER,
+              SessionConfig.DEFAULT_PASSWORD,
+              connection));
+    } else {
+      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 void assertObservationPeriod(final Runnable assertion) {
     final long start = System.nanoTime();
     assertion.run();

integration-test/src/test/java/org/apache/iotdb/db/it/utils/TestUtils.java

@@ -1312,7 +1312,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) {
@@ -1333,16 +1333,9 @@ 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);
     }
   }
 
@@ -1413,26 +1406,9 @@ public static boolean tryExecuteNonQueriesWithRetry(
         }
         return true;
       } catch (SQLException e) {
-        if (statement != null) {
-          try {
-            statement.close();
-          } catch (Exception ex) {
-            // ignore
-          } finally {
-            statement = null;
-          }
-        }
-        // 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
-          } finally {
-            localConnection = null;
-          }
-        }
+        closeNonQueryResources(statement, localConnection);
+        statement = null;
+        localConnection = null;
         connectionToUse = null;
 
         if (retryCountLeft > 0) {
@@ -1444,11 +1420,31 @@ public static boolean tryExecuteNonQueriesWithRetry(
           e.printStackTrace();
           return false;
         }
+      } finally {
+        closeNonQueryResources(statement, localConnection);
       }
     }
     return false;
   }
 
+  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
+      }
+    }
+  }
+
   public static void executeNonQueryOnSpecifiedDataNodeWithRetry(
       BaseEnv env, DataNodeWrapper wrapper, String sql) {
     for (int retryCountLeft = 10; retryCountLeft >= 0; retryCountLeft--) {

integration-test/src/test/java/org/apache/iotdb/pipe/it/single/IoTDBPipeOperationDiagnosticsIT.java

@@ -24,6 +24,7 @@
 import org.apache.iotdb.commons.client.sync.SyncConfigNodeIServiceClient;
 import org.apache.iotdb.commons.cluster.NodeStatus;
 import org.apache.iotdb.confignode.i18n.ConfigNodeMessages;
+import org.apache.iotdb.confignode.i18n.ProcedureMessages;
 import org.apache.iotdb.confignode.rpc.thrift.TAlterPipeReq;
 import org.apache.iotdb.confignode.rpc.thrift.TCreatePipeReq;
 import org.apache.iotdb.confignode.rpc.thrift.TShowPipeReq;
@@ -98,12 +99,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(
@@ -152,8 +156,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);
@@ -313,6 +319,35 @@ private static void restartDataNode(final DataNodeWrapper dataNodeWrapper) {
             Collections.singletonList(NodeStatus.Running));
   }
 
+  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(
+                    ProcedureMessages
+                        .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 final class BlackholeServer implements AutoCloseable {
     private final ServerSocket serverSocket;
     private final CopyOnWriteArrayList<Socket> acceptedSockets = new CopyOnWriteArrayList<>();

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/PipeSubtaskExecutorTest.java

@@ -124,8 +124,9 @@ public void testStop() {
    */
   @Test
   public void testStartWithRejectedSubmission() {
-    executor.register(subtask);
+    // Processor registration starts workers that also invoke this spy. Stub it before sharing it.
     Mockito.doThrow(new RejectedExecutionException()).when(subtask).submitSelf();
+    executor.register(subtask);
 
     try {
       executor.start(subtask.getTaskID());