diff --git a/core/src/main/java/io/substrait/relation/AbstractRel.java b/core/src/main/java/io/substrait/relation/AbstractRel.java index 8794951a3..f80081e57 100644 --- a/core/src/main/java/io/substrait/relation/AbstractRel.java +++ b/core/src/main/java/io/substrait/relation/AbstractRel.java @@ -10,7 +10,7 @@ */ public abstract class AbstractRel implements Rel { - private Supplier recordType = + private final Supplier recordType = Util.memoize( () -> { Type.Struct s = deriveRecordType(); diff --git a/core/src/main/java/io/substrait/util/Util.java b/core/src/main/java/io/substrait/util/Util.java index 32df0835c..f204c2487 100644 --- a/core/src/main/java/io/substrait/util/Util.java +++ b/core/src/main/java/io/substrait/util/Util.java @@ -9,6 +9,10 @@ public class Util { * Wraps a supplier so that its value is computed at most once, on first access, and cached for * subsequent calls. * + *

The returned supplier is thread-safe: once {@code supplier} returns, its value is cached and + * every caller observes it, and concurrent first calls wait rather than invoking {@code supplier} + * again. If {@code supplier} throws, the exception propagates and the next call retries. + * * @param supplier the supplier to memoize * @param the supplied value type * @return a memoizing supplier delegating to {@code supplier} @@ -19,9 +23,10 @@ public static Supplier memoize(Supplier supplier) { private static class Memoizer implements Supplier { - private boolean retrieved; + // Written after value; a reader that sees true via this volatile read also sees value. + private volatile boolean retrieved; private T value; - private Supplier delegate; + private final Supplier delegate; public Memoizer(Supplier delegate) { this.delegate = delegate; @@ -30,8 +35,12 @@ public Memoizer(Supplier delegate) { @Override public T get() { if (!retrieved) { - value = delegate.get(); - retrieved = true; + synchronized (this) { + if (!retrieved) { + value = delegate.get(); + retrieved = true; + } + } } return value; } diff --git a/core/src/test/java/io/substrait/util/UtilTest.java b/core/src/test/java/io/substrait/util/UtilTest.java new file mode 100644 index 000000000..5db9806f9 --- /dev/null +++ b/core/src/test/java/io/substrait/util/UtilTest.java @@ -0,0 +1,87 @@ +package io.substrait.util; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Supplier; +import org.junit.jupiter.api.Test; + +class UtilTest { + + @Test + void memoizeComputesOnce() { + AtomicInteger calls = new AtomicInteger(); + Supplier memoized = Util.memoize(() -> calls.incrementAndGet()); + + assertEquals(1, memoized.get()); + assertEquals(1, memoized.get()); + assertEquals(1, calls.get()); + } + + @Test + void memoizeCachesNull() { + AtomicInteger calls = new AtomicInteger(); + Supplier memoized = + Util.memoize( + () -> { + calls.incrementAndGet(); + return null; + }); + + memoized.get(); + memoized.get(); + assertEquals(1, calls.get()); + } + + @Test + void memoizeComputesOnceUnderConcurrentFirstAccess() throws Exception { + int threads = 8; + AtomicInteger calls = new AtomicInteger(); + // Holds the delegate open until every thread has entered it, so an unsynchronized memoizer + // lets all of them through; a synchronized one admits a single caller, whose wait expires. + CountDownLatch entered = new CountDownLatch(threads); + Supplier memoized = + Util.memoize( + () -> { + calls.incrementAndGet(); + entered.countDown(); + try { + entered.await(200, TimeUnit.MILLISECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return new Object(); + }); + + CyclicBarrier start = new CyclicBarrier(threads); + ExecutorService executor = Executors.newFixedThreadPool(threads); + try { + List> results = new ArrayList<>(); + for (int i = 0; i < threads; i++) { + results.add( + executor.submit( + () -> { + start.await(); + return memoized.get(); + })); + } + + Object first = results.get(0).get(10, TimeUnit.SECONDS); + for (Future result : results) { + assertSame(first, result.get(10, TimeUnit.SECONDS)); + } + assertEquals(1, calls.get()); + } finally { + executor.shutdownNow(); + } + } +}