11.2.2 backport of 1989 and 1982 - Fixed Idle Connection recovery so …

…that unprocessedResponseCount isn't over decremented and clear cache on recovery (#2004)
microsoft · Dec 14, 2022 · eeef9c8 · eeef9c8
1 parent 25a4bd6
commit eeef9c8
Show file tree

Hide file tree

Showing 7 changed files with 433 additions and 10 deletions.
diff --git a/pom.xml b/pom.xml
@@ -74,7 +74,7 @@
 
 		<!-- JUnit Test Dependencies -->
 		<junit.platform.version>[1.3.2, 1.9.0]</junit.platform.version>
-		<junit.jupiter.version>5.5.2</junit.jupiter.version>
+		<junit.jupiter.version>5.8.2</junit.jupiter.version>
 		<hikaricp.version>3.4.2</hikaricp.version>
 		<dbcp2.version>2.7.0</dbcp2.version>
 		<slf4j.nop.version>1.7.30</slf4j.nop.version>

diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerCallableStatement.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerCallableStatement.java
@@ -255,7 +255,10 @@ boolean onDone(TDSReader tdsReader) throws SQLServerException {
                 // Consume the done token and decide what to do with it...
                 StreamDone doneToken = new StreamDone();
                 doneToken.setFromTDS(tdsReader);
-                connection.getSessionRecovery().decrementUnprocessedResponseCount();
+
+                if (doneToken.isFinal()) {
+                    connection.getSessionRecovery().decrementUnprocessedResponseCount();
+                }
 
                 // If this is a non-final batch-terminating DONE token,
                 // then stop parsing the response now and set up for

diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerConnection.java
@@ -3882,6 +3882,10 @@ boolean executeCommand(TDSCommand newCommand) throws SQLServerException {
                                         false);
                             }
                             try {
+                                if (null != preparedStatementHandleCache) {
+                                    preparedStatementHandleCache.clear();
+                                }
+
                                 sessionRecovery.reconnect(newCommand);
                             } catch (InterruptedException e) {
                                 // re-interrupt thread

diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResultSet.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResultSet.java
@@ -381,10 +381,16 @@ boolean onDone(TDSReader tdsReader) throws SQLServerException {
                 // following the column metadata indicates an empty result set.
                 rowCount = 0;
 
-                short status = tdsReader.peekStatusFlag();
-                stmt.connection.getSessionRecovery().decrementUnprocessedResponseCount();
+                // decrementUnprocessedResponseCount() outside the "if" is not necessary here. It will over decrement if added.
 
+                short status = tdsReader.peekStatusFlag();
                 if ((status & TDS.DONE_ERROR) != 0 || (status & TDS.DONE_SRVERROR) != 0) {
+                    StreamDone doneToken = new StreamDone();
+                    doneToken.setFromTDS(tdsReader);
+                    if (doneToken.isFinal()) {
+                        stmt.connection.getSessionRecovery().decrementUnprocessedResponseCount();
+                    }
+                    SQLServerError databaseError = this.getDatabaseError();
                     MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_serverError"));
                     Object[] msgArgs = {status};
                     SQLServerException.makeFromDriverError(stmt.connection, stmt, form.format(msgArgs), null, false);
@@ -5376,7 +5382,18 @@ boolean onDone(TDSReader tdsReader) throws SQLServerException {
 
                 StreamDone doneToken = new StreamDone();
                 doneToken.setFromTDS(tdsReader);
-                stmt.connection.getSessionRecovery().decrementUnprocessedResponseCount();
+
+                if (doneToken.isFinal()) {
+                    stmt.connection.getSessionRecovery().decrementUnprocessedResponseCount();
+                }
+
+                if (doneToken.isFinal() && doneToken.isError()) {
+                    short status = tdsReader.peekStatusFlag();
+                    SQLServerError databaseError = getDatabaseError();
+                    MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_serverError"));
+                    Object[] msgArgs = {status, (databaseError != null) ? databaseError.getErrorMessage() : ""};
+                    SQLServerException.makeFromDriverError(stmt.connection, stmt, form.format(msgArgs), null, false);
+                }
 
                 // Done with all the rows in this fetch buffer and done with parsing
                 // unless it's a server cursor, in which case there is a RETSTAT and

diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerStatement.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerStatement.java
@@ -1480,7 +1480,10 @@ boolean onDone(TDSReader tdsReader) throws SQLServerException {
                 // Handling DONE/DONEPROC/DONEINPROC tokens is a little tricky...
                 StreamDone doneToken = new StreamDone();
                 doneToken.setFromTDS(tdsReader);
-                connection.getSessionRecovery().decrementUnprocessedResponseCount();
+
+                if (doneToken.isFinal()) {
+                    connection.getSessionRecovery().decrementUnprocessedResponseCount();
+                }
 
                 // If the done token has the attention ack bit set, then record
                 // it as the attention ack DONE token. We may or may not throw

diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/BasicConnectionTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/resiliency/BasicConnectionTest.java
@@ -14,17 +14,20 @@
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.UUID;
 
 import javax.sql.PooledConnection;
 
-import org.junit.jupiter.api.BeforeAll;
-import org.junit.jupiter.api.Tag;
-import org.junit.jupiter.api.Test;
-
 import com.microsoft.sqlserver.jdbc.RandomUtil;
+import com.microsoft.sqlserver.jdbc.SQLServerConnection;
 import com.microsoft.sqlserver.jdbc.SQLServerConnectionPoolDataSource;
+import com.microsoft.sqlserver.jdbc.SQLServerPreparedStatement;
 import com.microsoft.sqlserver.jdbc.TestResource;
 import com.microsoft.sqlserver.jdbc.TestUtils;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+
 import com.microsoft.sqlserver.testframework.AbstractTest;
 import com.microsoft.sqlserver.testframework.Constants;
 
@@ -266,6 +269,68 @@ public void testPooledConnectionLang() throws SQLException {
         }
     }
 
+    @Test
+    public void testPreparedStatementCacheShouldBeCleared() throws SQLException {
+        try (SQLServerConnection con = (SQLServerConnection) ResiliencyUtils.getConnection(connectionString)) {
+            int cacheSize = 2;
+            String query = String.format("/*testPreparedStatementCacheShouldBeCleared_%s*/SELECT 1; -- ",
+                    UUID.randomUUID().toString());
+            int discardedStatementCount = 1;
+
+            // enable caching
+            con.setDisableStatementPooling(false);
+            con.setStatementPoolingCacheSize(cacheSize);
+            con.setServerPreparedStatementDiscardThreshold(discardedStatementCount);
+
+            // add new statements to fill cache
+            for (int i = 0; i < cacheSize; ++i) {
+                try (SQLServerPreparedStatement pstmt = (SQLServerPreparedStatement) con
+                        .prepareStatement(query + i)) {
+                    pstmt.execute();
+                    pstmt.execute();
+                }
+            }
+
+            // nothing should be discarded yet
+            assertEquals(0, con.getDiscardedServerPreparedStatementCount());
+
+            ResiliencyUtils.killConnection(con, connectionString, 1);
+
+            // add 1 more - if cache was not cleared this would cause it to be discarded
+            try (SQLServerPreparedStatement pstmt = (SQLServerPreparedStatement) con.prepareStatement(query)) {
+                pstmt.execute();
+                pstmt.execute();
+            }
+            assertEquals(0, con.getDiscardedServerPreparedStatementCount());
+        }
+    }
+
+    @Test
+    public void testUnprocessedResponseCountSuccessfulIdleConnectionRecovery() throws SQLException {
+        try (SQLServerConnection con = (SQLServerConnection) ResiliencyUtils.getConnection(connectionString)) {
+            int queriesToSend = 5;
+            String query = String.format("/*testUnprocessedResponseCountSuccessfulIdleConnectionRecovery_%s*/SELECT 1; -- ",
+                    UUID.randomUUID());
+
+            for (int i = 0; i < queriesToSend; ++i) {
+                try (SQLServerPreparedStatement pstmt = (SQLServerPreparedStatement) con
+                        .prepareStatement(query + i)) {
+                    pstmt.executeQuery();
+                    pstmt.executeQuery();
+                }
+            }
+
+            // Kill the connection. If the unprocessedResponseCount is negative, test will fail.
+            ResiliencyUtils.killConnection(con, connectionString, 1);
+
+            // Should successfully recover.
+            try (SQLServerPreparedStatement pstmt = (SQLServerPreparedStatement) con.prepareStatement(query)) {
+                pstmt.executeQuery();
+                pstmt.executeQuery();
+            }
+        }
+    }
+
     private void basicReconnect(String connectionString) throws SQLException {
         // Ensure reconnects can happen multiple times over the same connection and subsequent connections
         for (int i1 = 0; i1 < 2; i1++) {