Replace timeout handling to use SharedTimer

Replaces the timeout handling for basic and bulk TDS commands to use a new SharedTimer
class. SharedTimer provides a static method for fetching an existing static object or
creating one on demand. Usage is tracked through reference counting and callers are
required to call removeRef() when they will no longer be using the SharedTimer. If the
SharedTimer does not have any more references then its internal ScheduledThreadPoolExecutor
will be shutdown.

The SharedTimer is cached at the Connection level so that repeated invocations do not
create new timers. Connections only create timers on first use so if no actions involve
a timeout then no timer is fetched or created. If a Connection does create a timer then
it will be released when the Connection closed.

Properly written JDBC applications that always close their Connection objects when
they are finished using them should not have any extra threads running after they are
all closed. Applications that do not use query timeouts will not have any extra threads
created as they are only done on demand.  Applications that use timeouts and use
a JDBC connection pool will have a single shared object across all JDBC connections as
long as there are some open connections in the pool with timeouts enabled.

Interrupt actions to handle a timeout are executed in their own thread. A handler thread is
created when the timeout occurs with the thread name matching the connection id of the
client connection that created the timeout. If the timeout is canceled prior to the interrupt
action being executed, say because the command finished, then no handler thread is created.

Note that the sharing of the timers happens across all Connections, not just Connections
with the same JDBC URL and properties.

Author: sehrope

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

@@ -0,0 +1,29 @@
+package com.microsoft.sqlserver.jdbc;
+
+import java.util.UUID;
+import java.util.concurrent.atomic.AtomicLong;
+
+
+abstract class AbstractTimeoutTask implements SqlServerTimerTask {
+    private static final AtomicLong COUNTER = new AtomicLong(0);
+    private final UUID connectionId;
+
+    AbstractTimeoutTask(UUID connectionId) {
+        this.connectionId = connectionId;
+    }
+
+    /**
+     * Action to interrupt the command implemented by concrete subclasses.
+     */
+    protected abstract void interrupt();
+
+    @Override
+    public final void run() {
+        // Create a new thread to run the interrupt to ensure that blocking operations performed
+        // by the interrupt do not hang the primary timer thread.
+        String name = "mssql-timeout-task-" + COUNTER.incrementAndGet() + "-" + connectionId;
+        Thread thread = new Thread(this::interrupt, name);
+        thread.setDaemon(true);
+        thread.start();
+    }
+}

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

@@ -52,6 +52,7 @@
 import java.util.Set;
 import java.util.SimpleTimeZone;
 import java.util.TimeZone;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.SynchronousQueue;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
@@ -3016,6 +3017,10 @@ void setDataLoggable(boolean value) {
         dataIsLoggable = value;
     }
 
+    SharedTimer getSharedTimer() {
+        return con.getSharedTimer();
+    }
+
     private TDSCommand command = null;
 
     // TDS message type (Query, RPC, DTC, etc.) sent at the beginning
@@ -6236,7 +6241,7 @@ 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 TimeoutCommand<TDSCommand> timeoutCommand;
+    private ScheduledFuture<?> timeout;
 
     final public String toString() {
         return traceID;
@@ -6390,9 +6395,8 @@ synchronized final boolean readPacket() throws SQLServerException {
             // terminate the connection.
             if ((command.getCancelQueryTimeoutSeconds() > 0 && command.getQueryTimeoutSeconds() > 0)) {
                 // if a timeout is configured with this object, add it to the timeout poller
-                int timeout = command.getCancelQueryTimeoutSeconds() + command.getQueryTimeoutSeconds();
-                this.timeoutCommand = new TdsTimeoutCommand(timeout, this.command, this.con);
-                TimeoutPoller.getTimeoutPoller().addTimeoutCommand(this.timeoutCommand);
+                int seconds = command.getCancelQueryTimeoutSeconds() + command.getQueryTimeoutSeconds();
+                this.timeout = con.getSharedTimer().schedule(new TdsCommandTimeoutTask(command, con), seconds);
             }
         }
         // First, read the packet header.
@@ -6413,8 +6417,9 @@ synchronized final boolean readPacket() throws SQLServerException {
         }
 
         // if execution was subject to timeout then stop timing
-        if (this.timeoutCommand != null) {
-            TimeoutPoller.getTimeoutPoller().remove(this.timeoutCommand);
+        if (this.timeout != null) {
+            this.timeout.cancel(false);
+            this.timeout = null;
         }
         // Header size is a 2 byte unsigned short integer in big-endian order.
         int packetLength = Util.readUnsignedShortBigEndian(newPacket.header, TDS.PACKET_HEADER_MESSAGE_LENGTH);
@@ -7003,42 +7008,6 @@ final void trySetSensitivityClassification(SensitivityClassification sensitivity
 }
 
 
-/**
- * The tds default implementation of a timeout command
- */
-class TdsTimeoutCommand extends TimeoutCommand<TDSCommand> {
-    public TdsTimeoutCommand(int timeout, TDSCommand command, SQLServerConnection sqlServerConnection) {
-        super(timeout, command, sqlServerConnection);
-    }
-
-    public void interrupt() {
-        TDSCommand command = getCommand();
-        SQLServerConnection sqlServerConnection = getSqlServerConnection();
-        try {
-            // 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 != sqlServerConnection) {
-                sqlServerConnection.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());
-        }
-    }
-}
-
 /**
  * TDSCommand encapsulates an interruptable TDS conversation.
  *
@@ -7160,7 +7129,7 @@ protected void setProcessedResponse(boolean processedResponse) {
     private volatile boolean readingResponse;
     private int queryTimeoutSeconds;
     private int cancelQueryTimeoutSeconds;
-    private TdsTimeoutCommand timeoutCommand;
+    private ScheduledFuture<?> timeout;
 
     protected int getQueryTimeoutSeconds() {
         return this.queryTimeoutSeconds;
@@ -7576,8 +7545,8 @@ final TDSReader startResponse(boolean isAdaptive) throws SQLServerException {
         // If command execution is subject to timeout then start timing until
         // the server returns the first response packet.
         if (queryTimeoutSeconds > 0) {
-            this.timeoutCommand = new TdsTimeoutCommand(queryTimeoutSeconds, this, null);
-            TimeoutPoller.getTimeoutPoller().addTimeoutCommand(this.timeoutCommand);
+            SQLServerConnection conn = tdsReader != null ? tdsReader.getConnection() : null;
+            this.timeout = tdsWriter.getSharedTimer().schedule(new TdsCommandTimeoutTask(this, conn), queryTimeoutSeconds);
         }
 
         if (logger.isLoggable(Level.FINEST))
@@ -7600,8 +7569,9 @@ final TDSReader startResponse(boolean isAdaptive) throws SQLServerException {
         } finally {
             // If command execution was subject to timeout then stop timing as soon
             // as the server returns the first response packet or errors out.
-            if (this.timeoutCommand != null) {
-                TimeoutPoller.getTimeoutPoller().remove(this.timeoutCommand);
+            if (this.timeout != null) {
+                this.timeout.cancel(false);
+                this.timeout = null;
             }
         }
 

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

@@ -42,6 +42,7 @@
 import java.util.SimpleTimeZone;
 import java.util.TimeZone;
 import java.util.UUID;
+import java.util.concurrent.ScheduledFuture;
 import java.util.logging.Level;
 
 import javax.sql.RowSet;
@@ -246,31 +247,7 @@ class BulkColumnMetaData {
      */
     private int srcColumnCount;
 
-    /**
-     * Timeout for the bulk copy command
-     */
-    private final class BulkTimeoutCommand extends TimeoutCommand<TDSCommand> {
-        public BulkTimeoutCommand(int timeout, TDSCommand command, SQLServerConnection sqlServerConnection) {
-            super(timeout, command, sqlServerConnection);
-        }
-
-        @Override
-        public void interrupt() {
-            TDSCommand command = getCommand();
-            // If the timer wasn't canceled before it ran out of
-            // time then interrupt the registered command.
-            try {
-                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.
-                command.log(Level.FINE, "Command could not be timed out. Reason: " + e.getMessage());
-            }
-        }
-    }
-
-    private BulkTimeoutCommand timeoutCommand;
+    private ScheduledFuture<?> timeout;
 
     /**
      * The maximum temporal precision we can send when using varchar(precision) in bulkcommand, to send a
@@ -646,16 +623,13 @@ private void sendBulkLoadBCP() throws SQLServerException {
         final class InsertBulk extends TDSCommand {
             InsertBulk() {
                 super("InsertBulk", 0, 0);
-                int timeoutSeconds = copyOptions.getBulkCopyTimeout();
-                timeoutCommand = timeoutSeconds > 0 ? new BulkTimeoutCommand(timeoutSeconds, this, null) : null;
             }
 
             final boolean doExecute() throws SQLServerException {
-                if (null != timeoutCommand) {
-                    if (logger.isLoggable(Level.FINEST))
-                        logger.finest(this.toString() + ": Starting bulk timer...");
-
-                    TimeoutPoller.getTimeoutPoller().addTimeoutCommand(timeoutCommand);
+                int timeoutSeconds = copyOptions.getBulkCopyTimeout();
+                if (timeoutSeconds > 0) {
+                    connection.checkClosed();
+                    timeout = connection.getSharedTimer().schedule(new TdsBulkCommandTimeoutTask(this, connection), timeoutSeconds);
                 }
 
                 // doInsertBulk inserts the rows in one batch. It returns true if there are more rows in
@@ -671,21 +645,27 @@ final boolean doExecute() throws SQLServerException {
                     }
 
                     // Check whether it is a timeout exception.
-                    if (rootCause instanceof SQLException) {
-                        checkForTimeoutException((SQLException) rootCause, timeoutCommand);
+                    if (rootCause instanceof SQLException && timeout != null && timeout.isDone()) {
+                        SQLException sqlEx = (SQLException) rootCause;
+                        if (sqlEx.getSQLState() != null
+                                && sqlEx.getSQLState().equals(SQLState.STATEMENT_CANCELED.getSQLStateCode())) {
+                            // If SQLServerBulkCopy is managing the transaction, a rollback is needed.
+                            if (copyOptions.isUseInternalTransaction()) {
+                                connection.rollback();
+                            }
+                            throw new SQLServerException(SQLServerException.getErrString("R_queryTimedOut"),
+                                    SQLState.STATEMENT_CANCELED, DriverError.NOT_SET, sqlEx);
+                        }
                     }
 
                     // It is not a timeout exception. Re-throw.
                     throw topLevelException;
                 }
 
-                if (null != timeoutCommand) {
-                    if (logger.isLoggable(Level.FINEST))
-                        logger.finest(this.toString() + ": Stopping bulk timer...");
-
-                    TimeoutPoller.getTimeoutPoller().remove(timeoutCommand);
+                if (timeout != null) {
+                    timeout.cancel(true);
+                    timeout = null;
                 }
-
                 return true;
             }
         }
@@ -1145,22 +1125,6 @@ private void writeColumnMetaData(TDSWriter tdsWriter) throws SQLServerException
         }
     }
 
-    /**
-     * Helper method that throws a timeout exception if the cause of the exception was that the query was cancelled
-     */
-    private void checkForTimeoutException(SQLException e, BulkTimeoutCommand timeoutCommand) throws SQLServerException {
-        if ((null != e.getSQLState()) && (e.getSQLState().equals(SQLState.STATEMENT_CANCELED.getSQLStateCode()))
-                && timeoutCommand.canTimeout()) {
-            // If SQLServerBulkCopy is managing the transaction, a rollback is needed.
-            if (copyOptions.isUseInternalTransaction()) {
-                connection.rollback();
-            }
-
-            throw new SQLServerException(SQLServerException.getErrString("R_queryTimedOut"),
-                    SQLState.STATEMENT_CANCELED, DriverError.NOT_SET, e);
-        }
-    }
-
     /**
      * Validates whether the source JDBC types are compatible with the destination table data types. We need to do this
      * only once for the whole bulk copy session.

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

@@ -141,6 +141,18 @@ public class SQLServerConnection implements ISQLServerConnection, java.io.Serial
 
     private Boolean isAzureDW = null;
 
+    private SharedTimer sharedTimer;
+
+    SharedTimer getSharedTimer() {
+        if (state == State.Closed) {
+            throw new IllegalStateException("Connection is closed");
+        }
+        if (sharedTimer == null) {
+            this.sharedTimer = SharedTimer.getTimer();
+        }
+        return this.sharedTimer;
+    }
+
     static class CityHash128Key implements java.io.Serializable {
 
         /**
@@ -3174,6 +3186,11 @@ public void close() throws SQLServerException {
         // with the connection.
         setState(State.Closed);
 
+        if (sharedTimer != null) {
+            sharedTimer.removeRef();
+            sharedTimer = null;
+        }
+
         // Close the TDS channel. When the channel is closed, the server automatically
         // rolls back any pending transactions and closes associated resources like
         // prepared handles.

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

@@ -0,0 +1,64 @@
+/*
+ * Microsoft JDBC Driver for SQL Server Copyright(c) Microsoft Corporation All rights reserved. This program is made
+ * available under the terms of the MIT License. See the LICENSE file in the project root for more information.
+ */
+package com.microsoft.sqlserver.jdbc;
+
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+
+
+class SharedTimer {
+    private static final AtomicLong CORE_THREAD_COUNTER = new AtomicLong();
+    private static SharedTimer INSTANCE;
+
+    private final long id = CORE_THREAD_COUNTER.getAndIncrement();
+    private int refCount = 0;
+    private ScheduledThreadPoolExecutor executor;
+
+    private SharedTimer() {
+        executor = new ScheduledThreadPoolExecutor(1, new ThreadFactory() {
+            @Override
+            public Thread newThread(Runnable task) {
+                return new Thread(task, "mssql-jdbc-shared-timer-core-" + id);
+            }
+        });
+        executor.setRemoveOnCancelPolicy(true);
+    }
+
+    public synchronized void removeRef() {
+        if (refCount <= 0) {
+            throw new IllegalStateException("removeRef() called more than actual references");
+        }
+        refCount -= 1;
+        if (refCount == 0) {
+            // Removed last reference so perform cleanup
+            executor.shutdownNow();
+            executor = null;
+            INSTANCE = null;
+        }
+    }
+
+    public static synchronized SharedTimer getTimer() {
+        if (INSTANCE == null) {
+            // No shared object exists so create a new one
+            INSTANCE = new SharedTimer();
+        }
+        INSTANCE.refCount += 1;
+        return INSTANCE;
+    }
+
+    public ScheduledFuture<?> schedule(SqlServerTimerTask task, long delaySeconds) {
+        return schedule(task, delaySeconds, TimeUnit.SECONDS);
+    }
+
+    public ScheduledFuture<?> schedule(SqlServerTimerTask task, long delay, TimeUnit unit) {
+        if (executor == null) {
+            throw new IllegalStateException("Cannot schedule tasks after shutdown");
+        }
+        return executor.schedule(task, delay, unit);
+    }
+}

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

@@ -0,0 +1,3 @@
+package com.microsoft.sqlserver.jdbc;
+
+public interface SqlServerTimerTask extends Runnable {}