From df3c7b3c1d5900090e4c5a8afe91295022b01c48 Mon Sep 17 00:00:00 2001
From: David Engel <v-davidengel@microsoft.com>
Date: Tue, 28 Nov 2023 13:11:57 -0800
Subject: [PATCH] Fix ActivityCorrelator behavior (#2254)

---
 .../sqlserver/jdbc/ActivityCorrelator.java    | 59 ++++++++---------
 .../sqlserver/jdbc/ActivityIDTest.java        | 63 +++++++++++++++++++
 2 files changed, 94 insertions(+), 28 deletions(-)
 create mode 100644 src/test/java/com/microsoft/sqlserver/jdbc/ActivityIDTest.java

diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/ActivityCorrelator.java b/src/main/java/com/microsoft/sqlserver/jdbc/ActivityCorrelator.java
index fe195632f..974ac3c47 100644
--- a/src/main/java/com/microsoft/sqlserver/jdbc/ActivityCorrelator.java
+++ b/src/main/java/com/microsoft/sqlserver/jdbc/ActivityCorrelator.java
@@ -6,8 +6,6 @@
 package com.microsoft.sqlserver.jdbc;
 
 import java.util.UUID;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
 
 
 /**
@@ -15,32 +13,21 @@
  */
 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,10 +40,14 @@ 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() {
@@ -64,24 +55,36 @@ UUID getId() {
     }
 
     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();
     }
 }
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/ActivityIDTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/ActivityIDTest.java
new file mode 100644
index 000000000..87919a7d5
--- /dev/null
+++ b/src/test/java/com/microsoft/sqlserver/jdbc/ActivityIDTest.java
@@ -0,0 +1,63 @@
+/*
+ * 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);
+
+        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();
+            UUID uuid = id.getId();
+            assertFalse("UUID should be unique across threads.", usedIds.contains(id.getId()));
+            assertEquals(1L, id.getSequence(), "First sequence should be 1.");
+            id = ActivityCorrelator.getNext();
+            assertEquals(uuid, id.getId(), "UUID should remain the same for the same thread.");
+            assertEquals(2L, id.getSequence(), "Second sequence should be 2.");
+            id = ActivityCorrelator.getNext();
+            assertEquals(3L, id.getSequence(), "Third sequence should be 3.");
+            assertEquals(uuid, id.getId(), "UUID should remain the same for the same thread.");
+        }
+    }
+}