POC- Orchestration of DLs

src/main/java/org/dataloader/orchestration/Orchestrator.java

@@ -0,0 +1,104 @@
+package org.dataloader.orchestration;
+
+import org.dataloader.DataLoader;
+import org.dataloader.impl.Assertions;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Function;
+
+public class Orchestrator<K, V> {
+
+    private final Tracker tracker;
+    private final DataLoader<K, V> startingDL;
+    private final List<Step<?, ?>> steps = new ArrayList<>();
+
+    /**
+     * This will create a new {@link Orchestrator} that can allow multiple calls to multiple data-loader's
+     * to be orchestrated so they all run optimally.
+     *
+     * @param dataLoader the data loader to start with
+     * @param <K>        the key type
+     * @param <V>        the value type
+     * @return a new {@link Orchestrator}
+     */
+    public static <K, V> Orchestrator<K, V> orchestrate(DataLoader<K, V> dataLoader) {
+        return new Orchestrator<>(new Tracker(), dataLoader);
+    }
+
+    public Tracker getTracker() {
+        return tracker;
+    }
+
+    private Orchestrator(Tracker tracker, DataLoader<K, V> dataLoader) {
+        this.tracker = tracker;
+        this.startingDL = dataLoader;
+    }
+
+
+    public Step<K, V> load(K key) {
+        return load(key, null);
+    }
+
+    public Step<K, V> load(K key, Object keyContext) {
+        return Step.loadImpl(this, castAs(startingDL), key, keyContext);
+    }
+
+    static <T> T castAs(Object o) {
+        //noinspection unchecked
+        return (T) o;
+    }
+
+
+    <KT, VT> void record(Step<KT, VT> step) {
+        steps.add(step);
+        tracker.incrementStepCount();
+    }
+
+    /**
+     * This is the callback point saying to start the DataLoader loading process.
+     * <p>
+     * The type of object returned here depends on the value type of the last Step.  We cant be truly generic
+     * here and must be case.
+     *
+     * @param <VT> the value type
+     * @return the final composed value
+     */
+    <VT> CompletableFuture<VT> execute() {
+        Assertions.assertState(!steps.isEmpty(), () -> "How can the steps to run be empty??");
+        int index = 0;
+        Step<?, ?> firstStep = steps.get(index);
+
+        CompletableFuture<Object> currentCF = castAs(firstStep.codeToRun().apply(null)); // first load uses variable capture
+        whenComplete(index, firstStep, currentCF);
+
+        for (index++; index < steps.size(); index++) {
+            Step<?, ?> nextStep = steps.get(index);
+            Function<Object, CompletableFuture<?>> codeToRun = castAs(nextStep.codeToRun());
+            CompletableFuture<Object> nextCF = currentCF.thenCompose(value -> castAs(codeToRun.apply(value)));
+            currentCF = nextCF;
+
+            // side effect when this step is complete
+            whenComplete(index, nextStep, nextCF);
+        }
+        return castAs(currentCF);
+
+    }
+
+    private void whenComplete(int index, Step<?, ?> step, CompletableFuture<Object> cf) {
+        cf.whenComplete((v, throwable) -> {
+            getTracker().loadCallComplete(step.dataLoader());
+            // replace with instrumentation code
+            if (throwable != null) {
+                // TODO - should we be cancelling future steps here - no need for dispatch tracking if they will never run
+                System.out.println("A throwable has been thrown on step  " + index + ": " + throwable.getMessage());
+                throwable.printStackTrace(System.out);
+            } else {
+                System.out.println("step " + index + " returned : " + v);
+            }
+        });
+    }
+
+
+}

src/main/java/org/dataloader/orchestration/Step.java

@@ -0,0 +1,67 @@
+package org.dataloader.orchestration;
+
+import org.dataloader.DataLoader;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Function;
+
+import static org.dataloader.orchestration.Orchestrator.castAs;
+
+public class Step<K, V> {
+    private final Orchestrator<?, ?> orchestrator;
+    private final DataLoader<Object, Object> dl;
+    private final Function<K, CompletableFuture<V>> codeToRun;
+
+    Step(Orchestrator<?, ?> orchestrator, DataLoader<?, ?> dataLoader, Function<K, CompletableFuture<V>> codeToRun) {
+        this.orchestrator = orchestrator;
+        this.dl = castAs(dataLoader);
+        this.codeToRun = codeToRun;
+    }
+
+    DataLoader<Object, Object> dataLoader() {
+        return dl;
+    }
+
+    public Function<K, CompletableFuture<V>> codeToRun() {
+        return codeToRun;
+    }
+
+    public <KT, VT> With<KT, VT> with(DataLoader<KT, VT> dataLoader) {
+        return new With<>(orchestrator, dataLoader);
+    }
+
+    public Step<K, V> load(K key, Object keyContext) {
+        return loadImpl(orchestrator, dl, key, keyContext);
+    }
+
+    public Step<V, V> thenLoad(Function<V, K> codeToRun) {
+        return thenLoadImpl(orchestrator, dl, codeToRun);
+    }
+
+    static <K, V> Step<K, V> loadImpl(Orchestrator<?, ?> orchestrator, DataLoader<Object, Object> dl, K key, Object keyContext) {
+        Function<K, CompletableFuture<V>> codeToRun = k -> {
+            CompletableFuture<V> cf = castAs(dl.load(key, keyContext));
+            orchestrator.getTracker().loadCall(dl);
+            return cf;
+        };
+        Step<K, V> step = new Step<>(orchestrator, dl, codeToRun);
+        orchestrator.record(step);
+        return step;
+    }
+
+    static <K, V> Step<V, V> thenLoadImpl(Orchestrator<?, ?> orchestrator, DataLoader<Object, Object> dl, Function<V, K> codeToRun) {
+        Function<V, CompletableFuture<V>> actualCodeToRun = v -> {
+            K key = codeToRun.apply(v);
+            CompletableFuture<V> cf = castAs(dl.load(key));
+            orchestrator.getTracker().loadCall(dl);
+            return cf;
+        };
+        Step<V, V> step = new Step<>(orchestrator, dl, actualCodeToRun);
+        orchestrator.record(step);
+        return step;
+    }
+
+    public CompletableFuture<V> toCompletableFuture() {
+        return orchestrator.execute();
+    }
+}

src/main/java/org/dataloader/orchestration/Tracker.java

@@ -0,0 +1,57 @@
+package org.dataloader.orchestration;
+
+import org.dataloader.DataLoader;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * This needs HEAPS more work - heaps more - I am not sure if its just counts of call backs or what.
+ * <p>
+ * This is just POC stuff for now
+ */
+public class Tracker {
+    private final AtomicInteger stepCount = new AtomicInteger();
+    private final Map<DataLoader<?,?>, AtomicInteger> counters = new HashMap<>();
+
+    public int getOutstandingLoadCount(DataLoader<?,?> dl) {
+        synchronized (this) {
+            return getDLCounter(dl).intValue();
+        }
+    }
+
+    public int getOutstandingLoadCount() {
+        int count = 0;
+        synchronized (this) {
+            for (AtomicInteger atomicInteger : counters.values()) {
+                count += atomicInteger.get();
+            }
+        }
+        return count;
+    }
+
+    public int getStepCount() {
+        return stepCount.get();
+    }
+
+    void incrementStepCount() {
+        this.stepCount.incrementAndGet();
+    }
+
+    void loadCall(DataLoader<?,?> dl) {
+        synchronized (this) {
+            getDLCounter(dl).incrementAndGet();
+        }
+    }
+
+    void loadCallComplete(DataLoader<?,?> dl) {
+        synchronized (this) {
+            getDLCounter(dl).decrementAndGet();
+        }
+    }
+
+    private AtomicInteger getDLCounter(DataLoader<?, ?> dl) {
+        return counters.computeIfAbsent(dl, key -> new AtomicInteger());
+    }
+}

src/main/java/org/dataloader/orchestration/With.java

@@ -0,0 +1,36 @@
+package org.dataloader.orchestration;
+
+import org.dataloader.DataLoader;
+
+import java.util.function.Function;
+
+import static org.dataloader.orchestration.Orchestrator.castAs;
+import static org.dataloader.orchestration.Step.loadImpl;
+
+/**
+ * A transitional step that allows a new step to be started with a new data loader in play
+ *
+ * @param <K> the key type
+ * @param <V> the value type
+ */
+public class With<K, V> {
+    private final Orchestrator<?, ?> orchestrator;
+    private final DataLoader<K, V> dl;
+
+    public With(Orchestrator<?, ?> orchestrator, DataLoader<K, V> dl) {
+        this.orchestrator = orchestrator;
+        this.dl = dl;
+    }
+
+    public Step<K, V> load(K key) {
+        return load(key, null);
+    }
+
+    public Step<K, V> load(K key, Object keyContext) {
+        return loadImpl(orchestrator, castAs(dl), key,keyContext);
+    }
+
+    public Step<V, V> thenLoad(Function<V, K> codeToRun) {
+        return Step.thenLoadImpl(orchestrator, castAs(dl), codeToRun);
+    }
+}

src/test/java/org/dataloader/fixtures/TestKit.java

@@ -11,8 +11,8 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
-import java.util.LinkedHashSet;
 import java.util.HashMap;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -60,6 +60,26 @@ public static <K, V> BatchLoader<K, V> keysAsValues(List<List<K>> loadCalls) {
         };
     }
 
+    public static BatchLoader<String, String> upperCaseBatchLoader() {
+        return keys -> CompletableFuture.completedFuture(keys.stream().map(String::toUpperCase).collect(toList()));
+    }
+
+    public static BatchLoader<String, String> lowerCaseBatchLoader() {
+        return keys -> CompletableFuture.completedFuture(keys.stream().map(String::toLowerCase).collect(toList()));
+    }
+
+    public static BatchLoader<String, String> reverseBatchLoader() {
+        return keys -> CompletableFuture.completedFuture(keys.stream().map(TestKit::reverse).collect(toList()));
+    }
+
+    public static String reverse(String s) {
+        StringBuilder sb = new StringBuilder();
+        for (int i = s.length() - 1; i >= 0; i--) {
+            sb.append(s.charAt(i));
+        }
+        return sb.toString();
+    }
+
     public static <K, V> DataLoader<K, V> idLoader() {
         return idLoader(null, new ArrayList<>());
     }
@@ -104,7 +124,7 @@ public static <T> Set<T> asSet(Collection<T> elements) {
 
     public static boolean areAllDone(CompletableFuture<?>... cfs) {
         for (CompletableFuture<?> cf : cfs) {
-            if (! cf.isDone()) {
+            if (!cf.isDone()) {
                 return false;
             }
         }

src/test/java/org/dataloader/orchestration/OrchestratorTest.java

@@ -0,0 +1,77 @@
+package org.dataloader.orchestration;
+
+import org.dataloader.DataLoader;
+import org.dataloader.DataLoaderOptions;
+import org.dataloader.DataLoaderRegistry;
+import org.junit.jupiter.api.Test;
+
+import java.util.concurrent.CompletableFuture;
+
+import static org.awaitility.Awaitility.await;
+import static org.dataloader.DataLoaderFactory.newDataLoader;
+import static org.dataloader.fixtures.TestKit.lowerCaseBatchLoader;
+import static org.dataloader.fixtures.TestKit.reverseBatchLoader;
+import static org.dataloader.fixtures.TestKit.upperCaseBatchLoader;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+
+class OrchestratorTest {
+
+    DataLoaderOptions cachingAndBatchingOptions = DataLoaderOptions.newOptions().setBatchingEnabled(true).setCachingEnabled(true);
+
+    DataLoader<String, String> dlUpper = newDataLoader(upperCaseBatchLoader(), cachingAndBatchingOptions);
+    DataLoader<String, String> dlLower = newDataLoader(lowerCaseBatchLoader(), cachingAndBatchingOptions);
+    DataLoader<String, String> dlReverse = newDataLoader(reverseBatchLoader(), cachingAndBatchingOptions);
+
+    @Test
+    void canOrchestrate() {
+
+        DataLoaderRegistry registry = DataLoaderRegistry.newRegistry()
+                .register("upper", dlUpper)
+                .register("lower", dlLower)
+                .register("reverse", dlReverse)
+                .build();
+
+        Orchestrator<String, String> orchestrator = Orchestrator.orchestrate(dlUpper);
+        Step<String, String> step1 = orchestrator.load("aBc", null);
+        With<String, String> with1 = step1.with(dlLower);
+        Step<String, String> step2 = with1.thenLoad(key -> key);
+        With<String, String> with2 = step2.with(dlReverse);
+        Step<String, String> step3 = with2.thenLoad(key -> key);
+        CompletableFuture<String> cf = step3.toCompletableFuture();
+
+        // because all the dls are dispatched in "perfect order" here they all end up dispatching
+        // at JUST the right time.  A change in order would be different
+        registry.dispatchAll();
+
+        await().until(cf::isDone);
+
+        assertThat(cf.join(), equalTo("cba"));
+    }
+
+    @Test
+    void canOrchestrateWhenNotInPerfectOrder() {
+
+        DataLoaderRegistry registry = DataLoaderRegistry.newRegistry()
+                .register("reverse", dlReverse)
+                .register("lower", dlLower)
+                .register("upper", dlUpper)
+                .build();
+
+        Orchestrator<String, String> orchestrator = Orchestrator.orchestrate(dlUpper);
+        CompletableFuture<String> cf = orchestrator.load("aBc", null)
+                .with(dlLower).thenLoad(key1 -> key1)
+                .with(dlReverse).thenLoad(key -> key)
+                .toCompletableFuture();
+
+        registry.dispatchAll();
+
+        assertThat(cf.isDone(), equalTo(false));
+
+        assertThat(orchestrator.getTracker().getOutstandingLoadCount(),equalTo(2));
+
+        await().until(cf::isDone);
+
+        assertThat(cf.join(), equalTo("cba"));
+    }
+}