package com.vladmihalcea.book.hpjp.hibernate.statistics;

import org.hibernate.stat.internal.ConcurrentStatisticsImpl;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicLong;

/**
 * <code>TransactionStatistics</code> - Transaction Statistics
 *
 * @author Vlad Mihalcea
 */
public class TransactionStatistics extends ConcurrentStatisticsImpl {

    private ConcurrentMap<Long, Long> transactionStartNanos =
        new ConcurrentHashMap<>();

    private ConcurrentMap<Long, AtomicLong> connectionCounter =
        new ConcurrentHashMap<>();

    private StatisticsReport report = new StatisticsReport();

    @Override public void connect() {
        long threadId = Thread.currentThread().getId();
        AtomicLong counter = connectionCounter.get(threadId);
        if(counter == null) {
            counter = new AtomicLong();
            connectionCounter.put(threadId, counter);
        }
        counter.incrementAndGet();
        transactionStartNanos.putIfAbsent(threadId, System.nanoTime());
        super.connect();
    }

    @Override public void endTransaction(boolean success) {
        long threadId = Thread.currentThread().getId();
        Long startNanos = transactionStartNanos.remove(threadId);
        if (startNanos != null)
            report.transactionTime(System.nanoTime() - startNanos);
        AtomicLong connectionCounter = this.connectionCounter.remove(threadId);
        if (connectionCounter != null)
            report.connectionsCount(connectionCounter.longValue());
        report.generate();
        super.endTransaction(success);
    }
}

@@ -117,98 +117,8 @@
 l;%0A%0A
-import java.util.concurrent.ConcurrentHashMap;%0Aimport java.util.concurrent.ConcurrentMap;%0A
 impo
@@ -347,439 +347,381 @@
 ate 
-ConcurrentMap%3CLong, Long%3E transactionStartNanos =%0A        new ConcurrentHashMap%3C%3E();%0A%0A    private ConcurrentMap%3CLong, AtomicLong%3E connectionCounter =%0A        new ConcurrentHashMap%3C%3E();%0A%0A    private StatisticsReport report = new StatisticsReport
+static final ThreadLocal%3CAtomicLong%3E transactionStartNanos = new ThreadLocal%3CAtomicLong%3E() %7B%0A        @Override%0A        protected AtomicLong initialValue() %7B%0A            return new AtomicLong
 ();%0A
-%0A
     
-@Override public void connect() %7B%0A        long threadId = Thread.currentThread().getId();%0A        AtomicLong counter = connectionCounter.get(threadId);%0A        if(counter == null
+    %7D%0A    %7D;%0A%0A    private static final ThreadLocal%3CAtomicLong%3E connectionCounter = new ThreadLocal%3CAtomicLong%3E() %7B%0A        @Override%0A        protected AtomicLong initialValue(
 ) %7B%0A
@@ -736,17 +736,14 @@
     
-counter =
+return
  new
@@ -769,79 +769,150 @@
     
+%7D%0A
     
-connectionCounter.put(threadId, counter
+%7D;%0A%0A    private StatisticsReport report = new StatisticsReport(
 );%0A
+%0A
     
-    %7D%0A        counter
+@Override public void connect() %7B%0A        connectionCounter.get()
 .inc
@@ -961,30 +961,28 @@
 nos.
-putIfAbsent(threadId, 
+get().compareAndSet(
 Syst
@@ -994,16 +994,19 @@
 noTime()
+, 0
 );%0A     
@@ -1104,241 +1104,108 @@
     
-long threadId = Thread.currentThread().getId();%0A        Long startNanos = transactionStartNanos.remove(threadId);%0A        if (startNanos != null)%0A            report.transactionTime(System.nanoTime() - startNanos);%0A        AtomicLong 
+report.transactionTime(System.nanoTime() - transactionStartNanos.get().longValue());%0A        report.
 conn
@@ -1214,23 +1214,15 @@
 tion
+s
 Count
-er = this.
+(
 conn
@@ -1239,131 +1239,93 @@
 ter.
-remove(threadId);%0A        if (connectionCounter != null)%0A            report.connectionsCount(connectionCounter.longValu
+get().longValue());%0A        report.generate();%0A        transactionStartNanos.remov
 e()
-)
 ;%0A  
@@ -1326,38 +1326,47 @@
 ();%0A        
-report.generat
+connectionCounter.remov
 e();%0A       

Switch to using TL