Performance | Improve performance of driver by continuously cleaning up ActivityIDs stored in internal Map

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

@@ -18,25 +18,20 @@ final class ActivityCorrelator {
     private static Map<Long, ActivityId> activityIdTlsMap = new ConcurrentHashMap<>();
 
     static void cleanupActivityId() {
-        // remove the ActivityId that belongs to this thread.
-        long uniqueThreadId = Thread.currentThread().getId();
-
-        if (activityIdTlsMap.containsKey(uniqueThreadId)) {
-            activityIdTlsMap.remove(uniqueThreadId);
-        }
+        // remove ActivityIds that belongs to this thread or no longer have an associated thread.
+        activityIdTlsMap.entrySet().removeIf(e -> null == e.getValue() || null == e.getValue().getThread()
+                || e.getValue().getThread() == Thread.currentThread() || !e.getValue().getThread().isAlive());
     }
 
     // Get the current ActivityId in TLS
     static ActivityId getCurrent() {
         // get the value in TLS, not reference
-        long uniqueThreadId = Thread.currentThread().getId();
-
-        // Since the Id for each thread is unique, this assures that the below if statement is run only once per thread.
-        if (!activityIdTlsMap.containsKey(uniqueThreadId)) {
-            activityIdTlsMap.put(uniqueThreadId, new ActivityId());
+        Thread thread = Thread.currentThread();
+        if (!activityIdTlsMap.containsKey(thread.getId())) {
+            activityIdTlsMap.put(thread.getId(), new ActivityId(thread));
         }
 
-        return activityIdTlsMap.get(uniqueThreadId);
+        return activityIdTlsMap.get(thread.getId());
     }
 
     // Increment the Sequence number of the ActivityId in TLS
@@ -55,7 +50,11 @@ static void setCurrentActivityIdSentFlag() {
         ActivityId activityId = getCurrent();
         activityId.setSentFlag();
     }
-
+
+    static Map<Long, ActivityId> getActivityIdTlsMap() {
+        return activityIdTlsMap;
+    }
+
     /*
      * Prevent instantiation.
      */
@@ -65,15 +64,21 @@ private ActivityCorrelator() {}
 
 class ActivityId {
     private final UUID id;
+    private final Thread thread;
     private long sequence;
     private boolean isSentToServer;
 
-    ActivityId() {
+    ActivityId(Thread thread) {
         id = UUID.randomUUID();
+        this.thread = thread;
         sequence = 0;
         isSentToServer = false;
     }
 
+    Thread getThread() {
+        return thread;
+    }
+
     UUID getId() {
         return id;
     }

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

@@ -3115,8 +3115,8 @@ void writeMessageHeader() throws SQLServerException {
             boolean includeTraceHeader = false;
             int totalHeaderLength = TDS.MESSAGE_HEADER_LENGTH;
             if (TDS.PKT_QUERY == tdsMessageType || TDS.PKT_RPC == tdsMessageType) {
-                if (con.isDenaliOrLater() && !ActivityCorrelator.getCurrent().isSentToServer()
-                        && Util.IsActivityTraceOn()) {
+                if (con.isDenaliOrLater() && Util.isActivityTraceOn()
+                        && !ActivityCorrelator.getCurrent().isSentToServer()) {
                     includeTraceHeader = true;
                     totalHeaderLength += TDS.TRACE_HEADER_LENGTH;
                 }

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

@@ -888,7 +888,7 @@ final boolean isSessionUnAvailable() {
     final void setMaxFieldSize(int limit) throws SQLServerException {
         // assert limit >= 0;
         if (maxFieldSize != limit) {
-            if (loggerExternal.isLoggable(Level.FINER) && Util.IsActivityTraceOn()) {
+            if (loggerExternal.isLoggable(Level.FINER) && Util.isActivityTraceOn()) {
                 loggerExternal.finer(toString() + " ActivityId: " + ActivityCorrelator.getNext().toString());
             }
             // If no limit on field size, set text size to max (2147483647), NOT default (0 --> 4K)
@@ -918,7 +918,7 @@ final void initResettableValues() {
     final void setMaxRows(int limit) throws SQLServerException {
         // assert limit >= 0;
         if (maxRows != limit) {
-            if (loggerExternal.isLoggable(Level.FINER) && Util.IsActivityTraceOn()) {
+            if (loggerExternal.isLoggable(Level.FINER) && Util.isActivityTraceOn()) {
                 loggerExternal.finer(toString() + " ActivityId: " + ActivityCorrelator.getNext().toString());
             }
             connectionCommand("SET ROWCOUNT " + limit, "setMaxRows");
@@ -2541,8 +2541,6 @@ void Prelogin(String serverName, int portNumber) throws SQLServerException {
         final byte[] preloginResponse = new byte[TDS.INITIAL_PACKET_SIZE];
         String preloginErrorLogString = " Prelogin error: host " + serverName + " port " + portNumber;
 
-        ActivityId activityId = ActivityCorrelator.getNext();
-        final byte[] actIdByteArray = Util.asGuidByteArray(activityId.getId());
         final byte[] conIdByteArray = Util.asGuidByteArray(clientConnectionId);
 
         int offset;
@@ -2558,18 +2556,23 @@ void Prelogin(String serverName, int portNumber) throws SQLServerException {
         System.arraycopy(conIdByteArray, 0, preloginRequest, offset, conIdByteArray.length);
         offset += conIdByteArray.length;
 
-        // copy ActivityId
-        System.arraycopy(actIdByteArray, 0, preloginRequest, offset, actIdByteArray.length);
-        offset += actIdByteArray.length;
-
-        long seqNum = activityId.getSequence();
-        Util.writeInt((int) seqNum, preloginRequest, offset);
-        offset += 4;
+        if (Util.isActivityTraceOn()) {
+            ActivityId activityId = ActivityCorrelator.getNext();
+            final byte[] actIdByteArray = Util.asGuidByteArray(activityId.getId());
+            System.arraycopy(actIdByteArray, 0, preloginRequest, offset, actIdByteArray.length);
+            offset += actIdByteArray.length;
+            long seqNum = activityId.getSequence();
+            Util.writeInt((int) seqNum, preloginRequest, offset);
+            offset += 4;
+
+            if (connectionlogger.isLoggable(Level.FINER)) {
+                connectionlogger.finer(toString() + " ActivityId " + activityId.toString());
+            }
+        }
 
         if (connectionlogger.isLoggable(Level.FINER)) {
             connectionlogger.finer(
                     toString() + " Requesting encryption level:" + TDS.getEncryptionLevel(requestedEncryptionLevel));
-            connectionlogger.finer(toString() + " ActivityId " + activityId.toString());
         }
 
         // Write the entire prelogin request
@@ -2585,7 +2588,9 @@ void Prelogin(String serverName, int portNumber) throws SQLServerException {
             throw e;
         }
 
-        ActivityCorrelator.setCurrentActivityIdSentFlag(); // indicate current ActivityId is sent
+        if (Util.isActivityTraceOn()) {
+            ActivityCorrelator.setCurrentActivityIdSentFlag(); // indicate current ActivityId is sent
+        }
 
         // Read the entire prelogin response
         int responseLength = preloginResponse.length;
@@ -3097,7 +3102,7 @@ public String nativeSQL(String sql) throws SQLServerException {
     public void setAutoCommit(boolean newAutoCommitMode) throws SQLServerException {
         if (loggerExternal.isLoggable(Level.FINER)) {
             loggerExternal.entering(getClassNameLogging(), "setAutoCommit", newAutoCommitMode);
-            if (Util.IsActivityTraceOn())
+            if (Util.isActivityTraceOn())
                 loggerExternal.finer(toString() + " ActivityId: " + ActivityCorrelator.getNext().toString());
         }
         String commitPendingTransaction = "";
@@ -3139,7 +3144,7 @@ final byte[] getTransactionDescriptor() {
     @Override
     public void commit() throws SQLServerException {
         loggerExternal.entering(getClassNameLogging(), "commit");
-        if (loggerExternal.isLoggable(Level.FINER) && Util.IsActivityTraceOn()) {
+        if (loggerExternal.isLoggable(Level.FINER) && Util.isActivityTraceOn()) {
             loggerExternal.finer(toString() + " ActivityId: " + ActivityCorrelator.getNext().toString());
         }
 
@@ -3152,7 +3157,7 @@ public void commit() throws SQLServerException {
     @Override
     public void rollback() throws SQLServerException {
         loggerExternal.entering(getClassNameLogging(), "rollback");
-        if (loggerExternal.isLoggable(Level.FINER) && Util.IsActivityTraceOn()) {
+        if (loggerExternal.isLoggable(Level.FINER) && Util.isActivityTraceOn()) {
             loggerExternal.finer(toString() + " ActivityId: " + ActivityCorrelator.getNext().toString());
         }
         checkClosed();
@@ -3242,7 +3247,9 @@ private void clearConnectionResources() {
         // Clean-up queue etc. related to batching of prepared statement discard actions (sp_unprepare).
         cleanupPreparedStatementDiscardActions();
 
-        ActivityCorrelator.cleanupActivityId();
+        if (Util.isActivityTraceOn()) {
+            ActivityCorrelator.cleanupActivityId();
+        }
     }
 
     // This function is used by the proxy for notifying the pool manager that this connection proxy is closed
@@ -3261,12 +3268,13 @@ final void poolCloseEventNotify() throws SQLServerException {
                 connectionCommand("IF @@TRANCOUNT > 0 ROLLBACK TRAN" /* +close connection */, "close connection");
             }
             notifyPooledConnection(null);
-            ActivityCorrelator.cleanupActivityId();
+            if (Util.isActivityTraceOn()) {
+                ActivityCorrelator.cleanupActivityId();
+            }
             if (connectionlogger.isLoggable(Level.FINER)) {
                 connectionlogger.finer(toString() + " Connection closed and returned to connection pool");
             }
         }
-
     }
 
     @Override
@@ -3308,7 +3316,7 @@ public boolean isReadOnly() throws SQLServerException {
     @Override
     public void setCatalog(String catalog) throws SQLServerException {
         loggerExternal.entering(getClassNameLogging(), "setCatalog", catalog);
-        if (loggerExternal.isLoggable(Level.FINER) && Util.IsActivityTraceOn()) {
+        if (loggerExternal.isLoggable(Level.FINER) && Util.isActivityTraceOn()) {
             loggerExternal.finer(toString() + " ActivityId: " + ActivityCorrelator.getNext().toString());
         }
         checkClosed();
@@ -3335,7 +3343,7 @@ String getSCatalog() throws SQLServerException {
     public void setTransactionIsolation(int level) throws SQLServerException {
         if (loggerExternal.isLoggable(Level.FINER)) {
             loggerExternal.entering(getClassNameLogging(), "setTransactionIsolation", level);
-            if (Util.IsActivityTraceOn()) {
+            if (Util.isActivityTraceOn()) {
                 loggerExternal.finer(toString() + " ActivityId: " + ActivityCorrelator.getNext().toString());
             }
         }
@@ -5275,7 +5283,7 @@ final private Savepoint setNamedSavepoint(String sName) throws SQLServerExceptio
     @Override
     public Savepoint setSavepoint(String sName) throws SQLServerException {
         loggerExternal.entering(getClassNameLogging(), "setSavepoint", sName);
-        if (loggerExternal.isLoggable(Level.FINER) && Util.IsActivityTraceOn()) {
+        if (loggerExternal.isLoggable(Level.FINER) && Util.isActivityTraceOn()) {
             loggerExternal.finer(toString() + " ActivityId: " + ActivityCorrelator.getNext().toString());
         }
         checkClosed();
@@ -5287,7 +5295,7 @@ public Savepoint setSavepoint(String sName) throws SQLServerException {
     @Override
     public Savepoint setSavepoint() throws SQLServerException {
         loggerExternal.entering(getClassNameLogging(), "setSavepoint");
-        if (loggerExternal.isLoggable(Level.FINER) && Util.IsActivityTraceOn()) {
+        if (loggerExternal.isLoggable(Level.FINER) && Util.isActivityTraceOn()) {
             loggerExternal.finer(toString() + " ActivityId: " + ActivityCorrelator.getNext().toString());
         }
         checkClosed();
@@ -5299,7 +5307,7 @@ public Savepoint setSavepoint() throws SQLServerException {
     @Override
     public void rollback(Savepoint s) throws SQLServerException {
         loggerExternal.entering(getClassNameLogging(), "rollback", s);
-        if (loggerExternal.isLoggable(Level.FINER) && Util.IsActivityTraceOn()) {
+        if (loggerExternal.isLoggable(Level.FINER) && Util.isActivityTraceOn()) {
             loggerExternal.finer(toString() + " ActivityId: " + ActivityCorrelator.getNext().toString());
         }
         checkClosed();
@@ -5324,7 +5332,7 @@ public int getHoldability() throws SQLServerException {
     public void setHoldability(int holdability) throws SQLServerException {
         loggerExternal.entering(getClassNameLogging(), "setHoldability", holdability);
 
-        if (loggerExternal.isLoggable(Level.FINER) && Util.IsActivityTraceOn()) {
+        if (loggerExternal.isLoggable(Level.FINER) && Util.isActivityTraceOn()) {
             loggerExternal.finer(toString() + " ActivityId: " + ActivityCorrelator.getNext().toString());
         }
         checkValidHoldability(holdability);