Better coalescing supplier

atlasdb-commons/src/main/java/com/palantir/common/concurrent/CoalescingSupplier.java

@@ -17,8 +17,7 @@
 
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Supplier;
 
 import com.google.common.base.Throwables;
@@ -29,52 +28,62 @@
  * requested; requests will not receive results for computations that started prior to the request.
  */
 public class CoalescingSupplier<T> implements Supplier<T> {
-
     private final Supplier<T> delegate;
-    private volatile CompletableFuture<T> nextResult = new CompletableFuture<T>();
-    private final Lock fairLock = new ReentrantLock(true);
+    private volatile Round round = new Round();
 
     public CoalescingSupplier(Supplier<T> delegate) {
         this.delegate = delegate;
     }
 
     @Override
     public T get() {
-        CompletableFuture<T> future = nextResult;
-
-        completeOrWaitForCompletion(future);
-
-        return getResult(future);
+        Round present = round;
+        if (present.isFirstToArrive()) {
+            present.execute();
+            return present.getResult();
+        }
+        Round next = present.awaitDone();
+        if (next.isFirstToArrive()) {
+            next.execute();
+        }
+        return next.getResult();
     }
 
-    private void completeOrWaitForCompletion(CompletableFuture<T> future) {
-        fairLock.lock();
-        try {
-            resetAndCompleteIfNotCompleted(future);
-        } finally {
-            fairLock.unlock();
+    private final class Round {
+        private final AtomicBoolean hasStarted = new AtomicBoolean(false);
+        private final CompletableFuture<T> future = new CompletableFuture<>();
+        private volatile Round next;
+
+        boolean isFirstToArrive() {
+            // adding the get benchmarks as faster, expected because compareAndSet forces an exclusive cache line
+            return !hasStarted.get() && hasStarted.compareAndSet(false, true);
         }
-    }
 
-    private void resetAndCompleteIfNotCompleted(CompletableFuture<T> future) {
-        if (future.isDone()) {
-            return;
+        Round awaitDone() {
+            try {
+                future.join();
+            } catch (CompletionException e) {
+                // ignore
+            }
+            return next;
         }
 
-        nextResult = new CompletableFuture<T>();
-        try {
-            future.complete(delegate.get());
-        } catch (Throwable t) {
-            future.completeExceptionally(t);
+        void execute() {
+            next = new Round();
+            try {
+                future.complete(delegate.get());
+            } catch (Throwable t) {
+                future.completeExceptionally(t);
+            }
+            round = next;
         }
-    }
 
-    private T getResult(CompletableFuture<T> future) {
-        try {
-            return future.getNow(null);
-        } catch (CompletionException ex) {
-            throw Throwables.propagate(ex.getCause());
+        T getResult() {
+            try {
+                return future.join();
+            } catch (CompletionException e) {
+                throw Throwables.propagate(e.getCause());
+            }
         }
     }
-
 }

changelog/@unreleased/pr-4508.v2.yml

@@ -0,0 +1,5 @@
+type: improvement
+improvement:
+  description: More predictable coalescing supplier
+  links:
+  - https://github.com/palantir/atlasdb/pull/4508