Fixed query cancellation bug that intermittently occurs in batch queries

src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java

@@ -7500,7 +7500,7 @@ protected void setInterruptsEnabled(boolean interruptsEnabled) {
     // Flag set to indicate that an interrupt has happened.
     private volatile boolean wasInterrupted = false;
 
-    private boolean wasInterrupted() {
+    boolean wasInterrupted() {
         return wasInterrupted;
     }
 

src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java

@@ -2868,6 +2868,21 @@ final void doExecutePreparedStatementBatch(PrepStmtBatchExecCmd batchCommand) th
             for (int attempt = 1; attempt <= 2; ++attempt) {
                 try {
 
+                    // If the command was interrupted, that means the TDS.PKT_CANCEL_REQ was sent to the server.
+                    // Since the cancellation request was sent, stop processing the batch query and process the
+                    // cancellation request and then return.
+                    //
+                    // Otherwise, if we do continue processing the batch query, in the case where a query requires
+                    // prepexec/sp_prepare, the TDS request for prepexec/sp_prepare will be sent regardless of
+                    // query cancellation. This will cause a TDS token error in the post processing when we
+                    // close the query.
+                    if (batchCommand.wasInterrupted()) {
+                        ensureExecuteResultsReader(batchCommand.startResponse(getIsResponseBufferingAdaptive()));
+                        startResults();
+                        getNextResult(true);
+                        return;
+                    }
+
                     // Re-use handle if available, requires parameter definitions which are not available until here.
                     if (reuseCachedHandle(hasNewTypeDefinitions, 1 < attempt)) {
                         hasNewTypeDefinitions = false;

src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java

@@ -201,5 +201,6 @@ protected Object[][] getContents() {
             {"R_objectNullOrEmpty", "The {0} is null or empty."},
             {"R_cekDecryptionFailed", "Failed to decrypt a column encryption key using key store provider: {0}."},
             {"R_connectTimedOut", "connect timed out"},
+            {"R_queryCancelled", "The query was canceled."},
             {"R_sessionKilled", "Cannot continue the execution because the session is in the kill state"}};
 }

src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/BatchExecutionTest.java

@@ -5,6 +5,7 @@
 package com.microsoft.sqlserver.jdbc.unit.statement;
 
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.fail;
 
 import java.lang.reflect.Field;
@@ -69,6 +70,51 @@ public void testBatchSpPrepare() throws Exception {
         testExecuteBatch1UseBulkCopyAPI();
     }
 
+    @Test
+    public void testBatchStatementCancellation() throws Exception {
+        try (Connection connection = PrepUtil.getConnection(connectionString)) {
+            connection.setAutoCommit(false);
+
+            try (PreparedStatement statement = connection.prepareStatement(
+                    "if object_id('test_table') is not null drop table test_table")) {
+                statement.execute();
+            }
+            connection.commit();
+
+            try (PreparedStatement statement = connection.prepareStatement(
+                    "create table test_table (column_name bit)")) {
+                statement.execute();
+            }
+            connection.commit();
+
+            for (long delayInMilliseconds : new long[] { 1, 2, 4, 8, 16, 32, 64, 128 }) {
+                for (int numberOfCommands : new int[] { 1, 2, 4, 8, 16, 32, 64 }) {
+                    int parameterCount = 512;
+
+                    try (PreparedStatement statement = connection.prepareStatement(
+                            "insert into test_table values (?)" + repeat(",(?)", parameterCount - 1))) {
+
+                        for (int i = 0; i < numberOfCommands; i++) {
+                            for (int j = 0; j < parameterCount; j++) {
+                                statement.setBoolean(j + 1, true);
+                            }
+                            statement.addBatch();
+                        }
+
+                        Thread cancelThread = cancelAsync(statement, delayInMilliseconds);
+                        try {
+                            statement.executeBatch();
+                        } catch (SQLException e) {
+                            assertEquals(TestResource.getResource("R_queryCancelled"), e.getMessage());
+                        }
+                        cancelThread.join();
+                    }
+                    connection.commit();
+                }
+            }
+        }
+    }
+
     /**
      * Get a PreparedStatement object and call the addBatch() method with 3 SQL statements and call the executeBatch()
      * method and it should return array of Integer values of length 3
@@ -240,6 +286,29 @@ private void modifyConnectionForBulkCopyAPI(SQLServerConnection con) throws Exce
         con.setUseBulkCopyForBatchInsert(true);
     }
 
+    private static String repeat(String string, int count) {
+        StringBuilder sb = new StringBuilder();
+        for (int i = 0; i < count; i++) {
+            sb.append(string);
+        }
+        return sb.toString();
+    }
+
+    private static Thread cancelAsync(Statement statement, long delayInMilliseconds) {
+        Thread thread = new Thread(() -> {
+            try {
+                Thread.sleep(delayInMilliseconds);
+                statement.cancel();
+            } catch (SQLException | InterruptedException e) {
+                // does not/must not happen
+                e.printStackTrace();
+                throw new IllegalStateException(e);
+            }
+        });
+        thread.start();
+        return thread;
+    }
+
     @BeforeAll
     public static void testSetup() throws TestAbortedException, Exception {
         setConnection();