Fix ActivityCorrelator behavior

Author: David-Engel

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

@@ -6,41 +6,28 @@
 package com.microsoft.sqlserver.jdbc;
 
 import java.util.UUID;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
 
 
 /**
  * ActivityCorrelator provides the APIs to access the ActivityId in TLS
  */
 final class ActivityCorrelator {
 
-    private static ActivityId activityId;
-    private static Lock lockObject = new ReentrantLock();
-
-    // Get the current ActivityId in TLS
-    static ActivityId getCurrent() {
-        if (activityId == null) {
-            lockObject.lock();
-            if (activityId == null) {
-                activityId = new ActivityId();
-            }
-            lockObject.unlock();
+    private static ThreadLocal<ActivityId> t_ActivityId = new ThreadLocal<ActivityId>() {
+        @Override
+        protected ActivityId initialValue() {
+            return new ActivityId();
         }
+    };
 
-        return activityId;
+    static ActivityId getCurrent() {
+        return t_ActivityId.get();
     }
 
     // Increment the Sequence number of the ActivityId in TLS
     // and return the ActivityId with new Sequence number
     static ActivityId getNext() {
-        // Get the current ActivityId in TLS
-        ActivityId activityId = getCurrent();
-
-        // Increment the Sequence number
-        activityId.increment();
-
-        return activityId;
+        return getCurrent().getIncrement();
     }
 
     /*
@@ -53,35 +40,51 @@ private ActivityCorrelator() {}
 class ActivityId {
     private final UUID id;
     private long sequence;
+    // Cache the string since it gets frequently referenced.
+    private String cachedToString;
 
     ActivityId() {
         id = UUID.randomUUID();
-        sequence = 1;
+        // getNext() is called during prelogin and will be the logical "first call" after
+        // instantiation, incrementing this to >= 1 before any activity logs are written.
+        sequence = 0;
     }
 
     UUID getId() {
         return id;
     }
 
     long getSequence() {
+        // Edge case: A new thread re-uses an existing connection. Ensure sequence > 0.
+        if (sequence == 0L) {
+            ++sequence;
+        }
+
         return sequence;
     }
 
-    void increment() {
+    ActivityId getIncrement() {
+        cachedToString = null;
         if (sequence < 0xffffffffl) // to get to 32-bit unsigned
         {
             ++sequence;
         } else {
             sequence = 0;
         }
+
+        return this;
     }
 
     @Override
     public String toString() {
-        StringBuilder sb = new StringBuilder();
-        sb.append(id.toString());
-        sb.append("-");
-        sb.append(sequence);
-        return sb.toString();
+        if (cachedToString == null) {
+            StringBuilder sb = new StringBuilder(38);
+            sb.append(id.toString());
+            sb.append("-");
+            sb.append(getSequence());
+            cachedToString = sb.toString();
+        }
+
+        return cachedToString.toString();
     }
 }

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

@@ -0,0 +1,67 @@
+/*
+ * 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 static org.junit.Assert.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.UUID;
+
+import org.junit.jupiter.api.Test;
+import org.junit.platform.runner.JUnitPlatform;
+import org.junit.runner.RunWith;
+
+import com.microsoft.sqlserver.testframework.AbstractTest;
+
+
+@RunWith(JUnitPlatform.class)
+public class ActivityIDTest extends AbstractTest {
+
+    @Test
+    public void testActivityID() throws Exception {
+        int numThreads = 20;
+        List<UUID> usedIds = new ArrayList<UUID>(numThreads);
+        ActivityId id = ActivityCorrelator.getCurrent();
+        usedIds.add(id.getId());
+        assertEquals(1L, id.getSequence());
+        ActivityCorrelator.getNext();
+        assertEquals(2L, id.getSequence());
+        ActivityCorrelator.getNext();
+        assertEquals(3L, id.getSequence());
+
+        for (int i = 0; i < numThreads; i++) {
+            MyThread t = new MyThread(usedIds);
+            t.start();
+            t.join();
+            usedIds.add(t.getUUID());
+        }
+    }
+
+    public class MyThread extends Thread {
+
+        public MyThread(List<UUID> usedIdsIn) {
+            usedIds = usedIdsIn;
+        }
+
+        private List<UUID> usedIds;
+        private ActivityId id;
+
+        public UUID getUUID() {
+            return id.getId();
+        }
+
+        public void run() {
+            id = ActivityCorrelator.getCurrent();
+            assertFalse(usedIds.contains(id.getId()));
+            assertEquals(1L, id.getSequence());
+            id = ActivityCorrelator.getNext();
+            assertEquals(2L, id.getSequence());
+            id = ActivityCorrelator.getNext();
+            assertEquals(3L, id.getSequence());
+        }
+    }
+}