Ability to atomically compute a data loader if absent

Author: Antoine Boyer

src/main/java/org/dataloader/DataLoaderRegistry.java

@@ -1,13 +1,13 @@
 package org.dataloader;
 
-import org.dataloader.stats.Statistics;
-
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
+import org.dataloader.stats.Statistics;
 
 /**
  * This allows data loaders to be registered together into a single place so
@@ -30,6 +30,26 @@ public DataLoaderRegistry register(String key, DataLoader<?, ?> dataLoader) {
         return this;
     }
 
+    /**
+     * Computes a data loader if absent or return it if it was
+     * already registered at that key.
+     *
+     * Note: The entire method invocation is performed atomically,
+     * so the function is applied at most once per key.
+     *
+     * @param key the key of the data loader
+     * @param mappingFunction the function to compute a data loader
+     * @param <K> the type of keys
+     * @param <V> the type of values
+     *
+     * @return a data loader
+     */
+    @SuppressWarnings("unchecked")
+    public <K, V> DataLoader<K, V> computeIfAbsent(final String key,
+                                                   final Function<String, DataLoader<?, ?>> mappingFunction) {
+        return (DataLoader<K, V>) dataLoaders.computeIfAbsent(key, mappingFunction);
+    }
+
     /**
      * This will combine all the current data loaders in this registry and all the data loaders from the specified registry
      * and return a new combined registry

src/test/java/org/dataloader/DataLoaderRegistryTest.java

@@ -1,10 +1,9 @@
 package org.dataloader;
 
+import java.util.concurrent.CompletableFuture;
 import org.dataloader.stats.Statistics;
 import org.junit.Test;
 
-import java.util.concurrent.CompletableFuture;
-
 import static java.util.Arrays.asList;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasItems;
@@ -101,4 +100,35 @@ public void stats_can_be_collected() throws Exception {
         assertThat(statistics.getLoadErrorCount(), equalTo(0L));
         assertThat(statistics.getBatchLoadExceptionCount(), equalTo(0L));
     }
-}
+
+    @Test
+    public void computeIfAbsent_creates_a_data_loader_if_there_was_no_value_at_key() {
+
+        DataLoaderRegistry registry = new DataLoaderRegistry();
+
+        DataLoader<Object, Object> dlA = new DataLoader<>(identityBatchLoader);
+        DataLoader<Object, Object> registered = registry.computeIfAbsent("a", (key) -> dlA);
+
+        assertThat(registered, equalTo(dlA));
+        assertThat(registry.getKeys(), hasItems("a"));
+        assertThat(registry.getDataLoaders(), hasItems(dlA));
+    }
+
+    @Test
+    public void computeIfAbsent_returns_an_existing_data_loader_if_there_was_a_value_at_key() {
+
+        DataLoaderRegistry registry = new DataLoaderRegistry();
+
+        DataLoader<Object, Object> dlA = new DataLoader<>(identityBatchLoader);
+        registry.computeIfAbsent("a", (key) -> dlA);
+
+        // register again at same key
+        DataLoader<Object, Object> dlA2 = new DataLoader<>(identityBatchLoader);
+        DataLoader<Object, Object> registered = registry.computeIfAbsent("a", (key) -> dlA2);
+
+        assertThat(registered, equalTo(dlA));
+        assertThat(registry.getKeys(), hasItems("a"));
+        assertThat(registry.getDataLoaders(), hasItems(dlA));
+    }
+
+}