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

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

@@ -17,26 +17,30 @@ 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);
+    static void clear() {
+        if (null != activityIdTlsMap) {
+            activityIdTlsMap.clear();
         }
     }
+
+    static void cleanupActivityId() {
+        // remove ActivityIds that belongs to this thread or no longer have an associated thread.
+        activityIdTlsMap.entrySet().removeIf(e ->
+                e.getValue() == null ||
+                e.getValue().getThread() == null ||
+                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
@@ -56,6 +60,10 @@ static void setCurrentActivityIdSentFlag() {
         activityId.setSentFlag();
     }
 
+    static Map<Long, ActivityId> getActivityIdTlsMap() {
+        return activityIdTlsMap;
+    }
+
     /*
      * Prevent instantiation.
      */
@@ -65,14 +73,20 @@ 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

@@ -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;
@@ -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

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

@@ -156,21 +156,27 @@ static String getErrString(String errCode) {
         super(errText, errState, errNum);
         initCause(cause);
         logException(null, errText, true);
-        ActivityCorrelator.setCurrentActivityIdSentFlag(); // set the activityid flag so that we don't send the current
-                                                           // ActivityId later.
+        if (Util.IsActivityTraceOn()) {
+            // set the activityid flag so that we don't send the current  ActivityId later.
+            ActivityCorrelator.setCurrentActivityIdSentFlag(); 
+        }
     }
 
     SQLServerException(String errText, Throwable cause) {
         super(errText);
         initCause(cause);
         logException(null, errText, true);
-        ActivityCorrelator.setCurrentActivityIdSentFlag();
+        if (Util.IsActivityTraceOn()) {
+            ActivityCorrelator.setCurrentActivityIdSentFlag(); 
+        }
     }
 
     SQLServerException(Object obj, String errText, String errState, int errNum, boolean bStack) {
         super(errText, errState, errNum);
         logException(obj, errText, bStack);
-        ActivityCorrelator.setCurrentActivityIdSentFlag();
+        if (Util.IsActivityTraceOn()) {
+            ActivityCorrelator.setCurrentActivityIdSentFlag(); 
+        }
     }
 
     /**

src/test/java/com/microsoft/sqlserver/jdbc/ActivityIDTest.java

@@ -0,0 +1,171 @@
+package com.microsoft.sqlserver.jdbc;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.fail;
+
+import java.io.ByteArrayInputStream;
+import java.io.InputStream;
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.logging.LogManager;
+import java.util.logging.Logger;
+
+import javax.sql.PooledConnection;
+
+import org.junit.ComparisonFailure;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.platform.runner.JUnitPlatform;
+import org.junit.runner.RunWith;
+
+import com.microsoft.sqlserver.testframework.AbstractTest;
+import com.zaxxer.hikari.HikariConfig;
+import com.zaxxer.hikari.HikariDataSource;
+
+@RunWith(JUnitPlatform.class)
+public class ActivityIDTest extends AbstractTest {
+
+    static final Logger logger = Logger.getLogger("ActivityIDTest");
+
+    @Test
+    public void testActivityID() throws Exception {
+        int numExecution = 20;
+        ExecutorService es = Executors.newFixedThreadPool(numExecution);
+        CountDownLatch latch = new CountDownLatch(numExecution);
+        for (int i = 0; i < numExecution; i++) {
+            es.execute(new Runnable() {
+                public void run() {
+                    try (Connection con = getConnection(); Statement stmt = con.createStatement()) {
+                        stmt.execute("SELECT @@VERSION AS 'SQL Server Version'");
+                    } catch (SQLException e) {
+                        fail(e.toString());
+                    }
+                    latch.countDown();
+                }
+            });
+        }
+        latch.await();
+        es.shutdown();
+        assertEquals(0, ActivityCorrelator.getActivityIdTlsMap().size());
+    }
+
+    @Test
+    public void testActivityIDPooled() throws Exception {
+        int poolsize = 10;
+        int numPooledExecution = 200;
+
+        HikariConfig config = new HikariConfig();
+        config.setJdbcUrl(connectionString);
+        config.setMaximumPoolSize(poolsize);
+        ExecutorService es = Executors.newFixedThreadPool(poolsize);
+        CountDownLatch latchPoolOuterThread = new CountDownLatch(1);
+        Thread t = new Thread(new Runnable() {
+            CountDownLatch latchPool = new CountDownLatch(numPooledExecution);
+            public void run() {
+                HikariDataSource ds = new HikariDataSource(config);
+                for (int i = 0; i < numPooledExecution; i++) {
+                    es.execute(new Runnable() {
+                        public void run() {
+                            try {
+                                try (Connection con = ds.getConnection(); Statement stmt = con.createStatement()) {
+                                    stmt.execute("SELECT @@VERSION AS 'SQL Server Version'");
+                                }
+                            } catch (SQLException e) {
+                                fail(e.toString());
+                            }
+                            latchPool.countDown();
+                        }
+                    });
+                }
+                try {
+                    latchPool.await();
+                } catch (InterruptedException e) {
+                    fail(e.toString());
+                } finally {
+                    if (null != ds) {
+                        es.shutdown();
+                        ds.close();
+                    }
+                }
+                latchPoolOuterThread.countDown();
+            }
+        });
+        t.run();
+        latchPoolOuterThread.await();
+
+        try {
+            try (Connection con = getConnection(); Statement stmt = con.createStatement()) {
+                stmt.execute("SELECT @@VERSION AS 'SQL Server Version'");
+            }
+        } catch (SQLException e) {
+            fail(e.toString());
+        }
+
+        try {
+            assertEquals(0, ActivityCorrelator.getActivityIdTlsMap().size());
+        } catch (ComparisonFailure e) {
+            Set<Thread> threadSet = Thread.getAllStackTraces().keySet();
+            logger.fine("List of threads alive:");
+            for (Thread thread: threadSet) {
+                logger.fine(thread.toString());
+            }
+            logger.fine("List of entries in the ActivityID map:");
+            logger.fine(ActivityCorrelator.getActivityIdTlsMap().toString());
+            throw new Exception(e);
+        }
+    }
+
+    @Test
+    public void testActivityIDPooledConnection() throws Exception {
+        int poolsize = 10;
+        int numPooledExecution = 200;
+
+        PooledConnection pooledCon = ((SQLServerConnectionPoolDataSource) dsPool).getPooledConnection();
+        ExecutorService es = Executors.newFixedThreadPool(poolsize);
+        try {
+            CountDownLatch latchPool = new CountDownLatch(numPooledExecution);
+            es.execute(new Runnable() {
+                public void run() {
+                    for (int i = 0; i < numPooledExecution; i++) {
+                        try (Connection con = pooledCon.getConnection(); Statement stmt = con.createStatement()) {
+                            stmt.execute("SELECT @@VERSION AS 'SQL Server Version'");
+                        } catch (SQLException e) {
+                            fail(e.toString());
+                        }
+                        latchPool.countDown();
+                    }
+                }
+            });
+            latchPool.await();
+        } finally {
+            es.shutdown();
+            pooledCon.close();
+        }
+        assertEquals(0, ActivityCorrelator.getActivityIdTlsMap().size());
+    }
+
+    @AfterAll
+    public static void teardown() throws Exception {
+        String ActivityIDTraceOff = Util.ActivityIdTraceProperty + "=off";
+        try (InputStream is = new ByteArrayInputStream(ActivityIDTraceOff.getBytes());) {
+            LogManager lm = LogManager.getLogManager();
+            lm.readConfiguration(is);
+        }
+    }
+
+    @BeforeAll
+    public static void testSetup() throws Exception {
+        String ActivityIDTraceOn = Util.ActivityIdTraceProperty + "=on";
+        try (InputStream is = new ByteArrayInputStream(ActivityIDTraceOn.getBytes());) {
+            LogManager lm = LogManager.getLogManager();
+            lm.readConfiguration(is);
+        }
+    }
+}
+