fix(core): make Util.memoize thread-safe

core/src/main/java/io/substrait/relation/AbstractRel.java

@@ -10,7 +10,7 @@
  */
 public abstract class AbstractRel implements Rel {
 
-  private Supplier<Type.Struct> recordType =
+  private final Supplier<Type.Struct> recordType =
       Util.memoize(
           () -> {
             Type.Struct s = deriveRecordType();

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.
    *
+   * <p>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 <T> the supplied value type
    * @return a memoizing supplier delegating to {@code supplier}
@@ -19,9 +23,10 @@ public static <T> Supplier<T> memoize(Supplier<T> supplier) {
 
   private static class Memoizer<T> implements Supplier<T> {
 
-    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<T> delegate;
+    private final Supplier<T> delegate;
 
     public Memoizer(Supplier<T> delegate) {
       this.delegate = delegate;
@@ -30,8 +35,12 @@ public Memoizer(Supplier<T> delegate) {
     @Override
     public T get() {
       if (!retrieved) {
-        value = delegate.get();
-        retrieved = true;
+        synchronized (this) {
+          if (!retrieved) {
+            value = delegate.get();
+            retrieved = true;
+          }
+        }
       }
       return value;
     }

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<Object> 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<Object> 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<Object> 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<Future<Object>> 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<Object> result : results) {
+        assertSame(first, result.get(10, TimeUnit.SECONDS));
+      }
+      assertEquals(1, calls.get());
+    } finally {
+      executor.shutdownNow();
+    }
+  }
+}