Better coalescing supplier

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

@@ -17,8 +17,8 @@
 
 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.concurrent.atomic.AtomicIntegerFieldUpdater;
 import java.util.function.Supplier;
 
 import com.google.common.base.Throwables;
@@ -29,52 +29,61 @@
  * 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 nextResult = 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 = nextResult;
+        if (present.isFirstToArrive()) {
+            present.execute();
+            return present.getResult();
+        }
+        awaitDone(present.future);
+        Round next = present.next;
+        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() {
+            return !hasStarted.get() && hasStarted.compareAndSet(false, true);
         }
-    }
 
-    private void resetAndCompleteIfNotCompleted(CompletableFuture<T> future) {
-        if (future.isDone()) {
-            return;
+        void execute() {
+            next = new Round();
+            nextResult = next;
+            try {
+                future.complete(delegate.get());
+            } catch (Throwable t) {
+                future.completeExceptionally(t);
+            }
         }
 
-        nextResult = new CompletableFuture<T>();
-        try {
-            future.complete(delegate.get());
-        } catch (Throwable t) {
-            future.completeExceptionally(t);
+        T getResult() {
+            try {
+                return future.join();
+            } catch (CompletionException e) {
+                throw Throwables.propagate(e.getCause());
+            }
         }
     }
 
-    private T getResult(CompletableFuture<T> future) {
+    private static void awaitDone(CompletableFuture<?> future) {
         try {
-            return future.getNow(null);
-        } catch (CompletionException ex) {
-            throw Throwables.propagate(ex.getCause());
+            future.join();
+        } catch (CompletionException e) {
+            // ignore
         }
     }
-
 }

atlasdb-perf/src/main/java/com/palantir/atlasdb/performance/benchmarks/CoalescingSupplierTests.java

@@ -0,0 +1,52 @@
+/*
+ * (c) Copyright 2020 Palantir Technologies Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.palantir.atlasdb.performance.benchmarks;
+
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
+
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Threads;
+import org.openjdk.jmh.annotations.Warmup;
+
+import com.palantir.common.concurrent.CoalescingSupplier;
+
+@State(Scope.Benchmark)
+public class CoalescingSupplierTests {
+    private final Supplier<String> supplier = new CoalescingSupplier<>(() -> {
+        try {
+            Thread.sleep(2);
+            return "result";
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            throw new RuntimeException(e);
+        }
+    });
+
+    @Benchmark
+    @Threads(16)
+    @Warmup(time = 1, timeUnit = TimeUnit.SECONDS, iterations = 4)
+    @Measurement(time = 1, timeUnit = TimeUnit.SECONDS, iterations = 10)
+    @Fork(1)
+    public String benchmark() {
+        return supplier.get();
+    }
+}

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

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