From 167325501d6c5bcf611a763e20469c0c4b9848ec Mon Sep 17 00:00:00 2001
From: Isaac Levy <isaac.levy@dealertrack.com>
Date: Tue, 20 Sep 2016 15:04:27 -0400
Subject: [PATCH] Example CAS for safer counter collection

---
 .../src/main/java/kamon/jsr166/LongAdder.java     | 15 +++++++++++++++
 .../src/main/java/kamon/jsr166/Striped64.java     | 11 +++++++++++
 .../scala/kamon/metric/instrument/Counter.scala   |  2 +-
 3 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/kamon-core/src/main/java/kamon/jsr166/LongAdder.java b/kamon-core/src/main/java/kamon/jsr166/LongAdder.java
index 18f917592..7e47ae63b 100644
--- a/kamon-core/src/main/java/kamon/jsr166/LongAdder.java
+++ b/kamon-core/src/main/java/kamon/jsr166/LongAdder.java
@@ -151,6 +151,21 @@ public long sumThenReset() {
         return sum;
     }
 
+    public long sumAndReset() {
+        long sum = getAndSetBase(0L);
+        Cell[] as = cells;
+        if (as != null) {
+            int n = as.length;
+            for (int i = 0; i < n; ++i) {
+                Cell a = as[i];
+                if (a != null) {
+                    sum += a.getAndSet(0L);
+                }
+            }
+        }
+        return sum;
+    }
+
     /**
      * Returns the String representation of the {@link #sum}.
      * @return the String representation of the {@link #sum}
diff --git a/kamon-core/src/main/java/kamon/jsr166/Striped64.java b/kamon-core/src/main/java/kamon/jsr166/Striped64.java
index 4840c8d70..579e26f59 100644
--- a/kamon-core/src/main/java/kamon/jsr166/Striped64.java
+++ b/kamon-core/src/main/java/kamon/jsr166/Striped64.java
@@ -105,6 +105,10 @@ final boolean cas(long cmp, long val) {
             return UNSAFE.compareAndSwapLong(this, valueOffset, cmp, val);
         }
 
+        final long getAndSet(long val) {
+            return UNSAFE.getAndSetLong(this, valueOffset, val);
+        }
+
         // Unsafe mechanics
         private static final sun.misc.Unsafe UNSAFE;
         private static final long valueOffset;
@@ -181,6 +185,13 @@ final boolean casBase(long cmp, long val) {
         return UNSAFE.compareAndSwapLong(this, baseOffset, cmp, val);
     }
 
+    /**
+     * CASes the base field.
+     */
+    final long getAndSetBase(long val) {
+        return UNSAFE.getAndSetLong(this, baseOffset, val);
+    }
+
     /**
      * CASes the busy field from 0 to 1 to acquire lock.
      */
diff --git a/kamon-core/src/main/scala/kamon/metric/instrument/Counter.scala b/kamon-core/src/main/scala/kamon/metric/instrument/Counter.scala
index 349a12bde..b7ab60de9 100644
--- a/kamon-core/src/main/scala/kamon/metric/instrument/Counter.scala
+++ b/kamon-core/src/main/scala/kamon/metric/instrument/Counter.scala
@@ -48,7 +48,7 @@ class LongAdderCounter extends Counter {
     counter.add(times)
   }
 
-  def collect(context: CollectionContext): Counter.Snapshot = CounterSnapshot(counter.sumThenReset())
+  def collect(context: CollectionContext): Counter.Snapshot = CounterSnapshot(counter.sumAndReset())
 
   def cleanup: Unit = {}
 }