From 56a76a9b1de2c5b2cf03fbc97bc13955b5916946 Mon Sep 17 00:00:00 2001 From: Terry Chow Date: Mon, 15 Aug 2022 16:17:28 -0700 Subject: [PATCH] Fixed query cancellation bug that intermittently occurs in batch queries --- .../microsoft/sqlserver/jdbc/IOBuffer.java | 2 +- .../jdbc/SQLServerPreparedStatement.java | 14 ++++ .../sqlserver/jdbc/TestResource.java | 1 + .../unit/statement/BatchExecutionTest.java | 69 +++++++++++++++++++ 4 files changed, 85 insertions(+), 1 deletion(-) diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java b/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java index b8b3a459ae..c49989ccf4 100644 --- a/src/main/java/com/microsoft/sqlserver/jdbc/IOBuffer.java +++ b/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() { + public boolean wasInterrupted() { return wasInterrupted; } diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java index e91c466f9c..6865659f7f 100644 --- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java +++ b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java @@ -2868,6 +2868,20 @@ 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 request will be sent regardless of query cancellation. This will + // cause a TDS token error in the post processing when we close 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; diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java b/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java index da6773116c..5ccc77b775 100644 --- a/src/test/java/com/microsoft/sqlserver/jdbc/TestResource.java +++ b/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"}}; } diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/BatchExecutionTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/BatchExecutionTest.java index b74d99cb5b..5756e646f2 100644 --- a/src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/BatchExecutionTest.java +++ b/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();