Fix lazy-singleton double init and publish registries as immutable snapshots (#16)

java/src/main/java/com/amazon/demanddriventrafficevaluator/factory/ExperimentManagerFactory.java

@@ -26,9 +26,9 @@ public class ExperimentManagerFactory {
 
     private static final ExperimentManagerFactory INSTANCE = new ExperimentManagerFactory();
 
-    private ExperimentManager experimentManager;
-    private ConfigurationProvider<ExperimentConfiguration> experimentConfigurationProvider;
-    private TreatmentAllocator treatmentAllocator;
+    private volatile ExperimentManager experimentManager;
+    private volatile ConfigurationProvider<ExperimentConfiguration> experimentConfigurationProvider;
+    private volatile TreatmentAllocator treatmentAllocator;
 
     public static ExperimentManagerFactory getInstance() {
         return INSTANCE;
@@ -39,6 +39,9 @@ public ExperimentManager provideExperimentManager() {
             return experimentManager;
         }
         synchronized (this) {
+            if (experimentManager != null) {
+                return experimentManager;
+            }
             ConfigurationProvider<ExperimentConfiguration> provider = provideExperimentConfigurationProvider();
             TreatmentAllocator allocator = provideTreatmentAllocator();
             ExperimentHandler handler = new ExperimentHandler();
@@ -52,6 +55,9 @@ public ConfigurationProvider<ExperimentConfiguration> provideExperimentConfigura
             return experimentConfigurationProvider;
         }
         synchronized (this) {
+            if (experimentConfigurationProvider != null) {
+                return experimentConfigurationProvider;
+            }
             LocalCacheRegistry localCacheRegistry = DefaultLocalCacheRegistryFactory.getInstance().getDefaultLocalCacheRegistrySingleton();
             Dao<String, ExperimentConfiguration> cacheDao = new LocalCacheDao<String, ExperimentConfiguration>(localCacheRegistry);
             experimentConfigurationProvider = new ExperimentConfigurationProvider(cacheDao);
@@ -64,7 +70,9 @@ public TreatmentAllocator provideTreatmentAllocator() {
             return treatmentAllocator;
         }
         synchronized (this) {
-            treatmentAllocator = new TreatmentAllocatorOnRandom();
+            if (treatmentAllocator == null) {
+                treatmentAllocator = new TreatmentAllocatorOnRandom();
+            }
             return treatmentAllocator;
         }
     }

java/src/main/java/com/amazon/demanddriventrafficevaluator/factory/ExtractorRegistryFactory.java

@@ -18,7 +18,7 @@
 public class ExtractorRegistryFactory {
 
     private static final ExtractorRegistryFactory INSTANCE = new ExtractorRegistryFactory();
-    private ExtractorRegistry extractorRegistry;
+    private volatile ExtractorRegistry extractorRegistry;
 
     private ExtractorRegistryFactory() {
     }
@@ -32,7 +32,9 @@ public ExtractorRegistry getSingleton() {
             return extractorRegistry;
         }
         synchronized (this) {
-            extractorRegistry = new ExtractorRegistry(new HashMap<>());
+            if (extractorRegistry == null) {
+                extractorRegistry = new ExtractorRegistry(new HashMap<>());
+            }
             return extractorRegistry;
         }
     }

java/src/main/java/com/amazon/demanddriventrafficevaluator/factory/LocalCacheRegistryFactory.java

@@ -24,13 +24,16 @@
  */
 public abstract class LocalCacheRegistryFactory {
 
-    LocalCacheRegistry defaultLocalCacheRegistry;
+    volatile LocalCacheRegistry defaultLocalCacheRegistry;
 
     public LocalCacheRegistry getDefaultLocalCacheRegistrySingleton() {
         if (defaultLocalCacheRegistry != null) {
             return defaultLocalCacheRegistry;
         }
         synchronized (this) {
+            if (defaultLocalCacheRegistry != null) {
+                return defaultLocalCacheRegistry;
+            }
             boolean useDefaultLocalCacheBuilderConfig = true;
             ConcurrentMap<String, ImmutablePair<Integer, Cache>> cacheMap = new ConcurrentHashMap<>();
             Map<String, LocalCacheBuilder> localCacheBuilderMap = getLocalCacheBuilderMap();

java/src/main/java/com/amazon/demanddriventrafficevaluator/factory/TransformerRegistryFactory.java

@@ -18,7 +18,7 @@
 public class TransformerRegistryFactory {
 
     private static final TransformerRegistryFactory INSTANCE = new TransformerRegistryFactory();
-    private TransformerRegistry transformerRegistry;
+    private volatile TransformerRegistry transformerRegistry;
 
     private TransformerRegistryFactory() {
     }
@@ -32,7 +32,9 @@ public TransformerRegistry getSingleton() {
             return transformerRegistry;
         }
         synchronized (this) {
-            transformerRegistry = new TransformerRegistry(new HashMap<>());
+            if (transformerRegistry == null) {
+                transformerRegistry = new TransformerRegistry(new HashMap<>());
+            }
             return transformerRegistry;
         }
     }

java/src/main/java/com/amazon/demanddriventrafficevaluator/modelfeature/Registry.java

@@ -3,6 +3,8 @@
 
 package com.amazon.demanddriventrafficevaluator.modelfeature;
 
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 
 /**
@@ -12,15 +14,33 @@
  * objects of a specific type. It uses a Map to store objects, with string keys representing
  * the names of the registered objects.
  * </p>
+ * <p>
+ * Thread safety: the registry is filled by a setup task on a scheduler thread while request
+ * threads read it. The records are held as an immutable snapshot in a volatile field. Each
+ * {@code register} call is synchronized, copies the current snapshot, adds the new entry and
+ * publishes the copy, so a reader always sees a complete map and never a map being changed.
+ * Reads take no lock: one volatile read and one lookup in the snapshot (plus a second lookup
+ * only when the name is missing). Registration is expected to happen a handful of times at
+ * startup, so the copy on each write is cheap.
+ * </p>
  *
  * @param <T> The type of objects stored in this registry.
  */
 public abstract class Registry<T> {
 
-    private final Map<String, T> records;
+    private volatile Map<String, T> records;
 
+    /**
+     * Creates a registry holding a snapshot of the given records.
+     * <p>
+     * The map is copied; later changes to it are not seen by the registry.
+     * </p>
+     *
+     * @param records The initial records.
+     * @throws NullPointerException if {@code records} is null.
+     */
     protected Registry(Map<String, T> records) {
-        this.records = records;
+        this.records = snapshot(new HashMap<>(records));
     }
 
     /**
@@ -43,16 +63,17 @@ protected Registry(Map<String, T> records) {
      * @throws IllegalArgumentException if an object with the same name is already registered
      *                                  or if the object cannot be instantiated.
      */
-    public void register(String name, Class<? extends T> object) {
+    public synchronized void register(String name, Class<? extends T> object) {
         if (records.containsKey(name)) {
             throw new IllegalArgumentException(getRegistryType() + " with name " + name + " already registered");
-        } else {
-            try {
-                records.put(name, object.getDeclaredConstructor().newInstance());
-            } catch (Exception e) {
-                throw new IllegalArgumentException("Failed to register " + getRegistryType() + " with " + name, e);
-            }
         }
+        T instance;
+        try {
+            instance = object.getDeclaredConstructor().newInstance();
+        } catch (Exception e) {
+            throw new IllegalArgumentException("Failed to register " + getRegistryType() + " with " + name, e);
+        }
+        publish(name, instance);
     }
 
     /**
@@ -66,12 +87,11 @@ public void register(String name, Class<? extends T> object) {
      * @param object The object to register.
      * @throws IllegalArgumentException if an object with the same name is already registered.
      */
-    public void register(String name, T object) {
+    public synchronized void register(String name, T object) {
         if (records.containsKey(name)) {
             throw new IllegalArgumentException(getRegistryType() + " with name " + name + " already registered");
-        } else {
-            records.put(name, object);
         }
+        publish(name, object);
     }
 
     /**
@@ -82,20 +102,40 @@ public void register(String name, T object) {
      * @throws IllegalArgumentException if no object is registered with the given name.
      */
     public T get(String name) {
-        if (records.containsKey(name)) {
-            return records.get(name);
-        } else {
+        Map<String, T> current = records;
+        T object = current.get(name);
+        // containsKey only on a miss, so a name registered with a null value still returns null
+        if (object == null && !current.containsKey(name)) {
             throw new IllegalArgumentException(getRegistryType() + " with name " + name + " not registered");
         }
+        return object;
     }
 
     /**
      * Returns a map of all registered objects.
+     * <p>
+     * The map is an unmodifiable snapshot of the registry at the time of the call; objects
+     * registered later are not added to it, so call this method again to see them. The
+     * registered objects themselves are not copied.
+     * </p>
      *
      * @return An unmodifiable map of all objects in the registry, where keys are object names
      * and values are the registered objects.
      */
     public Map<String, T> getRecords() {
         return records;
     }
+
+    // Callers hold the lock: copy, add, then publish the copy through the volatile field.
+    private void publish(String name, T object) {
+        Map<String, T> next = new HashMap<>(records);
+        next.put(name, object);
+        records = snapshot(next);
+    }
+
+    // Collections.unmodifiableMap over a private HashMap, not Map.copyOf: Map.copyOf rejects null
+    // keys and values (HashMap and this class accepted them) and randomizes iteration order per JVM.
+    private static <T> Map<String, T> snapshot(Map<String, T> map) {
+        return Collections.unmodifiableMap(map);
+    }
 }

java/src/test/java/com/amazon/demanddriventrafficevaluator/factory/LazySingletonConcurrencyTest.java

@@ -0,0 +1,107 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package com.amazon.demanddriventrafficevaluator.factory;
+
+import org.junit.jupiter.api.Test;
+
+import java.lang.reflect.Constructor;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+
+/**
+ * Concurrent first-call tests for the lazy singletons (issue #16): every caller
+ * racing on a fresh factory must receive the same instance.
+ */
+class LazySingletonConcurrencyTest {
+
+    private static final int THREADS = 16;
+    private static final int ROUNDS = 50;
+
+    @Test
+    void extractorRegistryFactory_concurrentFirstCalls_returnSameInstance() throws Exception {
+        assertSameInstanceUnderRace(() -> newViaPrivateConstructor(ExtractorRegistryFactory.class),
+                ExtractorRegistryFactory::getSingleton);
+    }
+
+    @Test
+    void transformerRegistryFactory_concurrentFirstCalls_returnSameInstance() throws Exception {
+        assertSameInstanceUnderRace(() -> newViaPrivateConstructor(TransformerRegistryFactory.class),
+                TransformerRegistryFactory::getSingleton);
+    }
+
+    @Test
+    void localCacheRegistryFactory_concurrentFirstCalls_returnSameInstance() throws Exception {
+        assertSameInstanceUnderRace(DefaultLocalCacheRegistryFactory::new,
+                LocalCacheRegistryFactory::getDefaultLocalCacheRegistrySingleton);
+    }
+
+    @Test
+    void experimentManagerFactory_concurrentFirstCalls_returnSameExperimentManager() throws Exception {
+        assertSameInstanceUnderRace(ExperimentManagerFactory::new,
+                ExperimentManagerFactory::provideExperimentManager);
+    }
+
+    @Test
+    void experimentManagerFactory_concurrentFirstCalls_returnSameConfigurationProvider() throws Exception {
+        assertSameInstanceUnderRace(ExperimentManagerFactory::new,
+                ExperimentManagerFactory::provideExperimentConfigurationProvider);
+    }
+
+    @Test
+    void experimentManagerFactory_concurrentFirstCalls_returnSameTreatmentAllocator() throws Exception {
+        assertSameInstanceUnderRace(ExperimentManagerFactory::new,
+                ExperimentManagerFactory::provideTreatmentAllocator);
+    }
+
+    private static <F> F newViaPrivateConstructor(Class<F> type) {
+        try {
+            Constructor<F> constructor = type.getDeclaredConstructor();
+            constructor.setAccessible(true);
+            return constructor.newInstance();
+        } catch (ReflectiveOperationException e) {
+            throw new IllegalStateException(e);
+        }
+    }
+
+    private static <F, T> void assertSameInstanceUnderRace(Supplier<F> freshFactory, Function<F, T> getter)
+            throws Exception {
+        ExecutorService executor = Executors.newFixedThreadPool(THREADS);
+        try {
+            for (int round = 0; round < ROUNDS; round++) {
+                F factory = freshFactory.get();
+                CountDownLatch ready = new CountDownLatch(THREADS);
+                CountDownLatch start = new CountDownLatch(1);
+                List<Future<T>> futures = new ArrayList<>();
+                for (int i = 0; i < THREADS; i++) {
+                    futures.add(executor.submit(() -> {
+                        ready.countDown();
+                        start.await();
+                        return getter.apply(factory);
+                    }));
+                }
+                ready.await();
+                start.countDown();
+                T first = futures.get(0).get(10, TimeUnit.SECONDS);
+                assertNotNull(first);
+                for (Future<T> future : futures) {
+                    assertSame(first, future.get(10, TimeUnit.SECONDS),
+                            "concurrent first callers received different instances in round " + round);
+                }
+            }
+        } finally {
+            executor.shutdownNow();
+            executor.awaitTermination(10, TimeUnit.SECONDS);
+        }
+    }
+}