Fixed Idle Connection recovery so that unprocessedResponseCount isn't over decremented

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

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

@@ -3996,6 +3996,10 @@ boolean executeCommand(TDSCommand newCommand) throws SQLServerException {
                                         false);
                             }
                             try {
+                                if (null != preparedStatementHandleCache) {
+                                    preparedStatementHandleCache.clear();
+                                }
+
                                 sessionRecovery.reconnect(newCommand);
                             } catch (InterruptedException e) {
                                 // re-interrupt thread

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

@@ -385,16 +385,21 @@ boolean onDone(TDSReader tdsReader) throws SQLServerException {
                 // following the column metadata indicates an empty result set.
                 rowCount = 0;
 
+                // 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, (databaseError != null) ? databaseError.getErrorMessage() : ""};
                     SQLServerException.makeFromDriverError(stmt.connection, stmt, form.format(msgArgs), null, false);
                 }
 
-                stmt.connection.getSessionRecovery().decrementUnprocessedResponseCount();
-
                 return false;
             }
         }
@@ -5382,6 +5387,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 +5400,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

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

src/test/java/com/microsoft/sqlserver/jdbc/resiliency/BasicConnectionTest.java

@@ -14,18 +14,21 @@
 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.RandomUtil;
+import com.microsoft.sqlserver.jdbc.SQLServerConnection;
+import com.microsoft.sqlserver.jdbc.SQLServerConnectionPoolDataSource;
 import com.microsoft.sqlserver.jdbc.SQLServerPooledConnection;
+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.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 +304,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());
+
+            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++) {