Merge pull request #674 from cheenamalhotra/TCPKeepAlive2

Added cancelQueryTimeout to cancel QueryTimeout on Connection and Statement

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

@@ -6348,7 +6348,8 @@ final class TDSReaderMark {
 final class TDSReader {
     private final static Logger logger = Logger.getLogger("com.microsoft.sqlserver.jdbc.internals.TDS.Reader");
     final private String traceID;
-
+    private TimeoutTimer tcpKeepAliveTimeoutTimer;
+
     final public String toString() {
         return traceID;
     }
@@ -6389,7 +6390,12 @@ private static int nextReaderID() {
         this.tdsChannel = tdsChannel;
         this.con = con;
         this.command = command; // may be null
-        // if the logging level is not detailed than fine or more we will not have proper readerids.
+        if(null != command) {
+            //if cancelQueryTimeout is set, we should wait for the total amount of queryTimeout + cancelQueryTimeout to terminate the connection.
+            this.tcpKeepAliveTimeoutTimer = (command.getCancelQueryTimeoutSeconds() > 0 && command.getQueryTimeoutSeconds() > 0 ) ? 
+        	    (new TimeoutTimer(command.getCancelQueryTimeoutSeconds() + command.getQueryTimeoutSeconds(), null, con)) : null;
+        }
+        // if the logging level is not detailed than fine or more we will not have proper reader IDs.
         if (logger.isLoggable(Level.FINE))
             traceID = "TDSReader@" + nextReaderID() + " (" + con.toString() + ")";
         else
@@ -6487,7 +6493,12 @@ synchronized final boolean readPacket() throws SQLServerException {
                 + tdsChannel.numMsgsSent;
 
         TDSPacket newPacket = new TDSPacket(con.getTDSPacketSize());
-
+        if (null != tcpKeepAliveTimeoutTimer) {
+            if (logger.isLoggable(Level.FINEST)) {
+                logger.finest(this.toString() + ": starting timer...");
+            }
+            tcpKeepAliveTimeoutTimer.start();
+        }
         // First, read the packet header.
         for (int headerBytesRead = 0; headerBytesRead < TDS.PACKET_HEADER_SIZE;) {
             int bytesRead = tdsChannel.read(newPacket.header, headerBytesRead, TDS.PACKET_HEADER_SIZE - headerBytesRead);
@@ -6501,7 +6512,14 @@ synchronized final boolean readPacket() throws SQLServerException {
 
             headerBytesRead += bytesRead;
         }
-
+
+        // if execution was subject to timeout then stop timing
+        if (null != tcpKeepAliveTimeoutTimer) {
+            if (logger.isLoggable(Level.FINEST)) {
+                logger.finest(this.toString() + ":stopping timer...");
+            }
+            tcpKeepAliveTimeoutTimer.stop();
+        }
         // Header size is a 2 byte unsigned short integer in big-endian order.
         int packetLength = Util.readUnsignedShortBigEndian(newPacket.header, TDS.PACKET_HEADER_MESSAGE_LENGTH);
 
@@ -7109,7 +7127,8 @@ final class TimeoutTimer implements Runnable {
     private final int timeoutSeconds;
     private final TDSCommand command;
     private volatile Future<?> task;
-
+    private final SQLServerConnection con;
+
     private static final ExecutorService executor = Executors.newCachedThreadPool(new ThreadFactory() {
         private final AtomicReference<ThreadGroup> tgr = new AtomicReference<>();
         private final AtomicInteger threadNumber = new AtomicInteger(0);
@@ -7134,12 +7153,13 @@ public Thread newThread(Runnable r)
     private volatile boolean canceled = false;
 
     TimeoutTimer(int timeoutSeconds,
-            TDSCommand command) {
+            TDSCommand command,
+            SQLServerConnection con) {
         assert timeoutSeconds > 0;
-        assert null != command;
 
         this.timeoutSeconds = timeoutSeconds;
         this.command = command;
+        this.con = con;
     }
 
     final void start() {
@@ -7173,12 +7193,22 @@ public void run() {
         // If the timer wasn't canceled before it ran out of
         // time then interrupt the registered command.
         try {
-            command.interrupt(SQLServerException.getErrString("R_queryTimedOut"));
+            // If TCP Connection to server is silently dropped, exceeding the query timeout on the same connection does not throw SQLTimeoutException
+            // The application stops responding instead until SocketTimeoutException is thrown. In this case, we must manually terminate the connection.
+            if (null == command && null != con) {
+                con.terminate(SQLServerException.DRIVER_ERROR_IO_FAILED, SQLServerException.getErrString("R_connectionIsClosed"));
+            }
+            else {
+                // If the timer wasn't canceled before it ran out of
+                // time then interrupt the registered command.
+                command.interrupt(SQLServerException.getErrString("R_queryTimedOut"));
+            }
         }
         catch (SQLServerException e) {
             // Unfortunately, there's nothing we can do if we
             // fail to time out the request. There is no way
             // to report back what happened.
+            assert null != command;
             command.log(Level.FINE, "Command could not be timed out. Reason: " + e.getMessage());
         }
     }
@@ -7306,7 +7336,17 @@ protected void setProcessedResponse(boolean processedResponse) {
     // any attention ack. The command's response is read either on demand as it is processed,
     // or by detaching.
     private volatile boolean readingResponse;
+	private int queryTimeoutSeconds;
+	private int cancelQueryTimeoutSeconds;
+
+    protected int getQueryTimeoutSeconds() {
+    	return this.queryTimeoutSeconds;
+    }
 
+    protected int getCancelQueryTimeoutSeconds() {
+    	return this.cancelQueryTimeoutSeconds;
+    }
+
     final boolean readingResponse() {
         return readingResponse;
     }
@@ -7320,9 +7360,11 @@ final boolean readingResponse() {
      *            (optional) the time before which the command must complete before it is interrupted. A value of 0 means no timeout.
      */
     TDSCommand(String logContext,
-            int timeoutSeconds) {
+            int queryTimeoutSeconds, int cancelQueryTimeoutSeconds) {
         this.logContext = logContext;
-        this.timeoutTimer = (timeoutSeconds > 0) ? (new TimeoutTimer(timeoutSeconds, this)) : null;
+        this.queryTimeoutSeconds = queryTimeoutSeconds;
+        this.cancelQueryTimeoutSeconds = cancelQueryTimeoutSeconds;
+        this.timeoutTimer = (queryTimeoutSeconds > 0) ? (new TimeoutTimer(queryTimeoutSeconds, this, null)) : null;
     }
 
     /**
@@ -7770,7 +7812,7 @@ final TDSReader startResponse(boolean isAdaptive) throws SQLServerException {
  */
 abstract class UninterruptableTDSCommand extends TDSCommand {
     UninterruptableTDSCommand(String logContext) {
-        super(logContext, 0);
+        super(logContext, 0, 0);
     }
 
     final void interrupt(String reason) throws SQLServerException {

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

@@ -689,7 +689,7 @@ private void initializeDefaults() {
     private void sendBulkLoadBCP() throws SQLServerException {
         final class InsertBulk extends TDSCommand {
             InsertBulk() {
-                super("InsertBulk", 0);
+                super("InsertBulk", 0, 0);
                 int timeoutSeconds = copyOptions.getBulkCopyTimeout();
                 timeoutTimer = (timeoutSeconds > 0) ? (new BulkTimeoutTimer(timeoutSeconds, this)) : null;
             }

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

@@ -471,6 +471,18 @@ final String getResponseBuffering() {
     final int getQueryTimeoutSeconds() {
         return queryTimeoutSeconds;
     }
+    /**
+     * timeout value for canceling the query timeout
+     */
+    private int cancelQueryTimeoutSeconds;
+
+    /**
+     * Retrieves the cancelTimeout in seconds
+     * @return
+     */
+    final int getCancelQueryTimeoutSeconds() {
+        return cancelQueryTimeoutSeconds;
+    }
 
     private int socketTimeoutMilliseconds;
 
@@ -1689,6 +1701,28 @@ else if (0 == requestedPacketSize)
                 }
             }
 
+            sPropKey = SQLServerDriverIntProperty.CANCEL_QUERY_TIMEOUT.toString();
+            int cancelQueryTimeout = SQLServerDriverIntProperty.CANCEL_QUERY_TIMEOUT.getDefaultValue();
+            // use cancelQueryTimeout only if queryTimeout is set.  
+            if (activeConnectionProperties.getProperty(sPropKey) != null && activeConnectionProperties.getProperty(sPropKey).length() > 0  && queryTimeoutSeconds > defaultQueryTimeout) {
+                try {
+                    int n = Integer.parseInt(activeConnectionProperties.getProperty(sPropKey));
+                    if (n >= cancelQueryTimeout) {
+                        cancelQueryTimeoutSeconds = n;
+                    }
+                    else {
+                        MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_invalidCancelQueryTimeout"));
+                        Object[] msgArgs = {activeConnectionProperties.getProperty(sPropKey)};
+                        SQLServerException.makeFromDriverError(this, this, form.format(msgArgs), null, false);
+                    }
+                }
+                catch (NumberFormatException e) {
+                    MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_invalidCancelQueryTimeout"));
+                    Object[] msgArgs = {activeConnectionProperties.getProperty(sPropKey)};
+                    SQLServerException.makeFromDriverError(this, this, form.format(msgArgs), null, false);
+                }
+            }
+
             sPropKey = SQLServerDriverIntProperty.SERVER_PREPARED_STATEMENT_DISCARD_THRESHOLD.toString();
             if (activeConnectionProperties.getProperty(sPropKey) != null && activeConnectionProperties.getProperty(sPropKey).length() > 0) {
                 try {

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

@@ -690,6 +690,26 @@ public int getQueryTimeout() {
         return getIntProperty(connectionProps, SQLServerDriverIntProperty.QUERY_TIMEOUT.toString(),
                 SQLServerDriverIntProperty.QUERY_TIMEOUT.getDefaultValue());
     }
+
+    /**
+     * Setting the cancel timeout
+     * 
+     * @param cancelQueryTimeout
+     *            The number of seconds to wait before we wait for the query timeout to happen.
+     */
+    public void setCancelQueryTimeout(int cancelQueryTimeout) {
+        setIntProperty(connectionProps, SQLServerDriverIntProperty.CANCEL_QUERY_TIMEOUT.toString(), cancelQueryTimeout);
+    }
+
+    /**
+     * Getting the cancel timeout
+     * 
+     * @return the number of seconds to wait before we wait for the query timeout to happen.
+     */
+    public int getCancelQueryTimeout() {
+        return getIntProperty(connectionProps, SQLServerDriverIntProperty.CANCEL_QUERY_TIMEOUT.toString(),
+                SQLServerDriverIntProperty.CANCEL_QUERY_TIMEOUT.getDefaultValue());
+    }
 
     /**
      * If this configuration is false the first execution of a prepared statement will call sp_executesql and not prepare 

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

@@ -318,6 +318,7 @@ enum SQLServerDriverIntProperty {
     SOCKET_TIMEOUT                             ("socketTimeout",                           0),
     SERVER_PREPARED_STATEMENT_DISCARD_THRESHOLD("serverPreparedStatementDiscardThreshold", SQLServerConnection.DEFAULT_SERVER_PREPARED_STATEMENT_DISCARD_THRESHOLD),
     STATEMENT_POOLING_CACHE_SIZE               ("statementPoolingCacheSize",               SQLServerConnection.DEFAULT_STATEMENT_POOLING_CACHE_SIZE),
+    CANCEL_QUERY_TIMEOUT                       ("cancelQueryTimeout",                      -1),
     ;  
 
     private final String name;
@@ -428,6 +429,7 @@ public final class SQLServerDriver implements java.sql.Driver {
         new SQLServerDriverPropertyInfo(SQLServerDriverIntProperty.STATEMENT_POOLING_CACHE_SIZE.toString(),                   Integer.toString(SQLServerDriverIntProperty.STATEMENT_POOLING_CACHE_SIZE.getDefaultValue()),            false,      null),
         new SQLServerDriverPropertyInfo(SQLServerDriverStringProperty.JAAS_CONFIG_NAME.toString(),                            SQLServerDriverStringProperty.JAAS_CONFIG_NAME.getDefaultValue(),                                       false,      null),
         new SQLServerDriverPropertyInfo(SQLServerDriverStringProperty.SSL_PROTOCOL.toString(),                                SQLServerDriverStringProperty.SSL_PROTOCOL.getDefaultValue(),                                           false,      new String[] {SSLProtocol.TLS.toString(), SSLProtocol.TLS_V10.toString(), SSLProtocol.TLS_V11.toString(), SSLProtocol.TLS_V12.toString()}),
+        new SQLServerDriverPropertyInfo(SQLServerDriverIntProperty.CANCEL_QUERY_TIMEOUT.toString(),                           Integer.toString(SQLServerDriverIntProperty.CANCEL_QUERY_TIMEOUT.getDefaultValue()),                          false,      null),
     };
 
     // Properties that can only be set by using Properties.

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

@@ -476,7 +476,7 @@ private final class PrepStmtExecCmd extends TDSCommand {
 
         PrepStmtExecCmd(SQLServerPreparedStatement stmt,
                 int executeMethod) {
-            super(stmt.toString() + " executeXXX", queryTimeout);
+            super(stmt.toString() + " executeXXX", queryTimeout, cancelQueryTimeoutSeconds);
             this.stmt = stmt;
             stmt.executeMethod = executeMethod;
         }
@@ -2561,7 +2561,7 @@ private final class PrepStmtBatchExecCmd extends TDSCommand {
         long updateCounts[];
 
         PrepStmtBatchExecCmd(SQLServerPreparedStatement stmt) {
-             super(stmt.toString() + " executeBatch", queryTimeout);
+             super(stmt.toString() + " executeBatch", queryTimeout, cancelQueryTimeoutSeconds);
             this.stmt = stmt;
         }
 

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

@@ -391,5 +391,7 @@ protected Object[][] getContents() {
 				{"R_invalidDataTypeSupportForSQLVariant", "Unexpected TDS type ' '{0}' ' in SQL_VARIANT."},
 				{"R_sslProtocolPropertyDescription", "SSL protocol label from TLS, TLSv1, TLSv1.1, and TLSv1.2. The default is TLS."},
 				{"R_invalidSSLProtocol", "SSL Protocol {0} label is not valid. Only TLS, TLSv1, TLSv1.1, and TLSv1.2 are supported."},
+				{"R_cancelQueryTimeoutPropertyDescription", "The number of seconds to wait to cancel sending a query timeout."},
+				{"R_invalidCancelQueryTimeout", "The cancel timeout value {0} is not valid."},
     };
 }

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

@@ -5677,7 +5677,7 @@ final class InsertRowRPC extends TDSCommand {
             final String tableName;
 
             InsertRowRPC(String tableName) {
-                super("InsertRowRPC", 0);
+                super("InsertRowRPC", 0, 0);
                 this.tableName = tableName;
             }
 
@@ -5775,7 +5775,7 @@ public void updateRow() throws SQLServerException {
         }
         final class UpdateRowRPC extends TDSCommand {
             UpdateRowRPC() {
-                super("UpdateRowRPC", 0);
+                super("UpdateRowRPC", 0, 0);
             }
 
             final boolean doExecute() throws SQLServerException {
@@ -5853,7 +5853,7 @@ public void deleteRow() throws SQLServerException {
         }
         final class DeleteRowRPC extends TDSCommand {
             DeleteRowRPC() {
-                super("DeleteRowRPC", 0);
+                super("DeleteRowRPC", 0, 0);
             }
 
             final boolean doExecute() throws SQLServerException {
@@ -6453,7 +6453,7 @@ private final class CursorFetchCommand extends TDSCommand {
                 int fetchType,
                 int startRow,
                 int numRows) {
-            super("doServerFetch", stmt.queryTimeout);
+            super("doServerFetch", stmt.queryTimeout, stmt.cancelQueryTimeoutSeconds);
             this.serverCursorId = serverCursorId;
             this.fetchType = fetchType;
             this.startRow = startRow;

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

@@ -112,6 +112,11 @@ final boolean wasExecuted() {
      */
     int queryTimeout;
 
+    /**
+     * timeout value for canceling the query timeout
+     */
+    int cancelQueryTimeoutSeconds;
+
     /**
      * Is closeOnCompletion is enabled? If true statement will be closed when all of its dependent result sets are closed
      */
@@ -576,6 +581,7 @@ else if (ResultSet.TYPE_SCROLL_SENSITIVE == nType) {
         setResponseBuffering(connection.getResponseBuffering());
 
         setDefaultQueryTimeout();
+        setDefaultQueryCancelTimeout();
 
         if (stmtlogger.isLoggable(java.util.logging.Level.FINER)) {
             stmtlogger.finer("Properties for " + toString() + ":" + " Result type:" + appResultSetType + " (" + resultSetType + ")" + " Concurrency:"
@@ -588,7 +594,14 @@ else if (ResultSet.TYPE_SCROLL_SENSITIVE == nType) {
         }
     }
 
-    // add query timeout to statement
+    private void setDefaultQueryCancelTimeout() {
+        int cancelQueryTimeoutSeconds = this.connection.getCancelQueryTimeoutSeconds();
+        if (cancelQueryTimeoutSeconds > 0) {
+            this.cancelQueryTimeoutSeconds = cancelQueryTimeoutSeconds;
+        }
+	}
+
+	// add query timeout to statement
     private void setDefaultQueryTimeout() {
         int queryTimeoutSeconds = this.connection.getQueryTimeoutSeconds();
         if (queryTimeoutSeconds > 0) {
@@ -752,7 +765,7 @@ private final class StmtExecCmd extends TDSCommand {
                 String sql,
                 int executeMethod,
                 int autoGeneratedKeys) {
-            super(stmt.toString() + " executeXXX", stmt.queryTimeout);
+            super(stmt.toString() + " executeXXX", stmt.queryTimeout, stmt.cancelQueryTimeoutSeconds);
             this.stmt = stmt;
             this.sql = sql;
             this.executeMethod = executeMethod;
@@ -883,7 +896,7 @@ private final class StmtBatchExecCmd extends TDSCommand {
         final SQLServerStatement stmt;
 
         StmtBatchExecCmd(SQLServerStatement stmt) {
-            super(stmt.toString() + " executeBatch", stmt.queryTimeout);
+            super(stmt.toString() + " executeBatch", stmt.queryTimeout, stmt.cancelQueryTimeoutSeconds);
             this.stmt = stmt;
         }
 
@@ -1133,6 +1146,25 @@ public final void setLargeMaxRows(long max) throws SQLServerException {
         loggerExternal.exiting(getClassNameLogging(), "setQueryTimeout");
     }
 
+    /* L0 */ public final int getCancelQueryTimeout() throws SQLServerException {
+        loggerExternal.entering(getClassNameLogging(), "getCancelQueryTimeout");
+        checkClosed();
+        loggerExternal.exiting(getClassNameLogging(), "getCancelQueryTimeout", cancelQueryTimeoutSeconds);
+        return cancelQueryTimeoutSeconds;
+    }
+
+    /* L0 */ public final void setCancelQueryTimeout(int seconds) throws SQLServerException {
+        loggerExternal.entering(getClassNameLogging(), "setCancelQueryTimeout", seconds);
+        checkClosed();
+        if (seconds < 0) {
+            MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_invalidCancelQueryTimeout"));
+            Object[] msgArgs = {seconds};
+            SQLServerException.makeFromDriverError(connection, this, form.format(msgArgs), null, true);
+        }
+        cancelQueryTimeoutSeconds = seconds;
+        loggerExternal.exiting(getClassNameLogging(), "setCancelQueryTimeout");
+    }
+
     public final void cancel() throws SQLServerException {
         loggerExternal.entering(getClassNameLogging(), "cancel");
         checkClosed();