/*
 * 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.UUID;

/**
 * ActivityCorrelator provides the APIs to access the ActivityId in TLS
 */
final class ActivityCorrelator {

    private static ThreadLocal<ActivityId> ActivityIdTls = new ThreadLocal<ActivityId>() {
        protected ActivityId initialValue() {
            return new ActivityId();
        }
    };

    // Get the current ActivityId in TLS
    static ActivityId getCurrent() {
        // get the value in TLS, not reference
        return ActivityIdTls.get();
    }

    // Increment the Sequence number of the ActivityId in TLS
    // and return the ActivityId with new Sequence number
    static ActivityId getNext() {
        // We need to call get() method on ThreadLocal to get
        // the current value of ActivityId stored in TLS,
        // then increment the sequence number.

        // Get the current ActivityId in TLS
        ActivityId activityId = getCurrent();

        // Increment the Sequence number
        activityId.Increment();

        return activityId;
    }

    static void setCurrentActivityIdSentFlag() {
        ActivityId activityId = getCurrent();
        activityId.setSentFlag();
    }

}

class ActivityId {
    private final UUID Id;
    private long Sequence;
    private boolean isSentToServer;

    ActivityId() {
        Id = UUID.randomUUID();
        Sequence = 0;
        isSentToServer = false;
    }

    UUID getId() {
        return Id;
    }

    long getSequence() {
        return Sequence;
    }

    void Increment() {
        if (Sequence < 0xffffffffl) // to get to 32-bit unsigned
        {
            ++Sequence;
        }
        else {
            Sequence = 0;
        }

        isSentToServer = false;
    }

    void setSentFlag() {
        isSentToServer = true;
    }

    boolean IsSentToServer() {
        return isSentToServer;
    }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder();
        sb.append(Id.toString());
        sb.append("-");
        sb.append(Sequence);
        return sb.toString();
    }
}

@@ -280,16 +280,42 @@
 .jdbc;%0A%0A
+import java.util.HashMap;%0A
 import j
@@ -329,16 +329,16 @@
 l.UUID;%0A
-
 %0A/**%0A * 
@@ -463,28 +463,30 @@
  static 
-ThreadLocal%3C
+HashMap%3CLong, 
 Activity
@@ -506,27 +506,80 @@
 dTls
+Map
  = new 
-ThreadLocal%3C
+HashMap%3CLong, ActivityId%3E();%0A    %0A    static void checkAndInit
 Acti
@@ -584,17 +584,16 @@
 tivityId
-%3E
 () %7B%0A   
@@ -601,64 +601,347 @@
     
-protected ActivityId initialValue() %7B%0A            return
+long uniqueThreadId = Thread.currentThread().getId();%0A        %0A        //Since the Id for each thread is unique, this assures that the below code is run only once per *thread*,%0A        //the first time this constructor is called.%0A        if (!ActivityIdTlsMap.containsKey(uniqueThreadId)) %7B%0A            ActivityIdTlsMap.put(uniqueThreadId,
  new
@@ -949,24 +949,25 @@
 ActivityId()
+)
 ;%0A        %7D%0A
@@ -971,17 +971,16 @@
  %7D%0A    %7D
-;
 %0A%0A    //
@@ -1043,32 +1043,75 @@
  getCurrent() %7B%0A
+        checkAndInitActivityId();%0A        %0A
         // get t
@@ -1141,16 +1141,87 @@
 ference%0A
+        long uniqueThreadId = Thread.currentThread().getId();%0A        %0A
         
@@ -1244,13 +1244,30 @@
 dTls
+Map
 .get(
+uniqueThreadId
 );%0A 
@@ -1422,24 +1422,58 @@
 getNext() %7B%0A
+        checkAndInitActivityId();%0A
         // W
@@ -1869,32 +1869,75 @@
 yIdSentFlag() %7B%0A
+        checkAndInitActivityId();%0A        %0A
         Activity
@@ -2000,23 +2000,22 @@
 Flag();%0A
+
     %7D%0A
-%0A
 %7D%0A%0Aclass

Fix classloader leak issue