Fixed ICR recovery so that unprocessedResponseCount isn't over decrem…

…ented
microsoft · Dec 5, 2022 · 5fca34a · 5fca34a
1 parent eb59ba7
commit 5fca34a
Show file tree

Hide file tree

Showing 5 changed files with 80 additions and 10 deletions.
diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerCallableStatement.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerCallableStatement.java
@@ -258,7 +258,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
@@ -3996,6 +3996,9 @@ 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
@@ -393,8 +393,6 @@ boolean onDone(TDSReader tdsReader) throws SQLServerException {
                     SQLServerException.makeFromDriverError(stmt.connection, stmt, form.format(msgArgs), null, false);
                 }
 
-                stmt.connection.getSessionRecovery().decrementUnprocessedResponseCount();
-
                 return false;
             }
         }
@@ -5382,6 +5380,11 @@ boolean onDone(TDSReader tdsReader) throws SQLServerException {
 
                 StreamDone doneToken = new StreamDone();
                 doneToken.setFromTDS(tdsReader);
+
+                if (doneToken.isFinal()) {
+                    stmt.connection.getSessionRecovery().decrementUnprocessedResponseCount();
+                }
+
                 if (doneToken.isFinal() && doneToken.isError()) {
                     short status = tdsReader.peekStatusFlag();
                     SQLServerError databaseError = getDatabaseError();
@@ -5390,7 +5393,6 @@ boolean onDone(TDSReader tdsReader) throws SQLServerException {
                     SQLServerException.makeFromDriverError(stmt.connection, stmt, form.format(msgArgs), null, false);
                 }
 
-                stmt.connection.getSessionRecovery().decrementUnprocessedResponseCount();
 
                 // 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
@@ -1494,7 +1494,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,18 +14,15 @@
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.UUID;
 
 import javax.sql.PooledConnection;
 
-import com.microsoft.sqlserver.jdbc.SQLServerPooledConnection;
+import com.microsoft.sqlserver.jdbc.*;
 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.SQLServerConnectionPoolDataSource;
-import com.microsoft.sqlserver.jdbc.TestResource;
-import com.microsoft.sqlserver.jdbc.TestUtils;
 import com.microsoft.sqlserver.testframework.AbstractTest;
 import com.microsoft.sqlserver.testframework.Constants;
 
@@ -301,6 +298,68 @@ public void testDSPooledConnectionAccessTokenCallbackIdleConnectionResiliency()
         }
     }
 
+    @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().toString());
+
+            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++) {