From 9d0859c665d9dd3322fccedbdc5ee3fccf4e0881 Mon Sep 17 00:00:00 2001 From: Chris Rice Date: Fri, 9 Oct 2026 13:28:18 -0400 Subject: [PATCH] fix(core): make Util.memoize thread-safe Util.memoize stored its cached value and "computed" flag in plain fields, so a thread racing the first call could observe the flag set without the value and get null back, and concurrent callers could each re-run the delegate. The memoizer backs function keys and anchors, the ExtensionCollection lookup maps, and relation record types, so a shared ExtensionCollection failed intermittently on first concurrent use with "NullPointerException: key" or "Duplicate key null". The flag is now volatile and initialization uses double-checked locking, so the delegate runs exactly once and every caller sees the fully computed value. The memoizer's delegate and AbstractRel's memoizer field are now final so they are safely published along with their owner. Closes #1389 --- .../io/substrait/relation/AbstractRel.java | 2 +- .../src/main/java/io/substrait/util/Util.java | 17 +++- .../test/java/io/substrait/util/UtilTest.java | 87 +++++++++++++++++++ 3 files changed, 101 insertions(+), 5 deletions(-) create mode 100644 core/src/test/java/io/substrait/util/UtilTest.java 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(); + } + } +}