Skip to content

Commit

Permalink
Fixed Idle Connection recovery so that unprocessedResponseCount isn't…
Browse files Browse the repository at this point in the history
… over decremented (#1989)

* Updated onDone with reason/comment for unecessary decrement for future reference

* Added more ICR tests

* Added missing brackets

* Test table names needed to be unique

* Drop tables within test if exist

Co-authored-by: lilgreenbird <v-susanh@microsoft.com>
  • Loading branch information
tkyc and lilgreenbird authored Dec 13, 2022
1 parent 135a1e3 commit fd16cb8
Show file tree
Hide file tree
Showing 6 changed files with 424 additions and 9 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
15 changes: 12 additions & 3 deletions src/main/java/com/microsoft/sqlserver/jdbc/SQLServerResultSet.java
Original file line number Diff line number Diff line change
Expand Up @@ -390,16 +390,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;
}
}
Expand Down Expand Up @@ -5389,6 +5394,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();
Expand All @@ -5397,7 +5407,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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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++) {
Expand Down
Loading

0 comments on commit fd16cb8

Please sign in to comment.