POC- Orchestration of DLs - Added a tracker that can observe via a callback

Author: bbakerman

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

@@ -2,6 +2,9 @@
 
 import org.dataloader.DataLoader;
 import org.dataloader.impl.Assertions;
+import org.dataloader.orchestration.executors.ImmediateExecutor;
+import org.dataloader.orchestration.observation.Tracker;
+import org.dataloader.orchestration.observation.TrackingObserver;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -25,19 +28,14 @@ public class Orchestrator<K, V> {
      * @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, ImmediateExecutor.INSTANCE);
+    public static <K, V> Builder<K, V> orchestrate(DataLoader<K, V> dataLoader) {
+        return new Builder<>(dataLoader);
     }
 
-    // TODO - make this a builder
-    public static <K, V> Orchestrator<K, V> orchestrate(DataLoader<K, V> dataLoader, Executor executor) {
-        return new Orchestrator<>(new Tracker(), dataLoader, executor);
-    }
-
-    private Orchestrator(Tracker tracker, DataLoader<K, V> dataLoader, Executor executor) {
-        this.tracker = tracker;
-        this.startingDL = dataLoader;
-        this.executor = executor;
+    public Orchestrator(Builder<K, V> builder) {
+        this.tracker = new Tracker(builder.trackingObserver);
+        this.executor = builder.executor;
+        this.startingDL = builder.dataLoader;
     }
 
     public Tracker getTracker() {
@@ -79,6 +77,10 @@ <KT, VT> void record(Step<KT, VT> step) {
      */
     <VT> CompletableFuture<VT> execute() {
         Assertions.assertState(!steps.isEmpty(), () -> "How can the steps to run be empty??");
+
+        // tell the tracker we are under way
+        getTracker().startingExecution();
+
         int index = 0;
         Step<?, ?> firstStep = steps.get(index);
 
@@ -94,23 +96,47 @@ <VT> CompletableFuture<VT> execute() {
             // side effect when this step is complete
             whenComplete(index, nextStep, nextCF);
         }
+
         return castAs(currentCF);
 
     }
 
-    private void whenComplete(int index, Step<?, ?> step, CompletableFuture<Object> cf) {
+    private void whenComplete(int stepIndex, 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());
+                System.out.println("A throwable has been thrown on step  " + stepIndex + ": " + throwable.getMessage());
                 throwable.printStackTrace(System.out);
             } else {
-                System.out.println("step " + index + " returned : " + v);
+                System.out.println("step " + stepIndex + " returned : " + v);
             }
+            getTracker().loadCallComplete(stepIndex, step.dataLoader(), throwable);
         });
     }
 
 
+    public static class Builder<K, V> {
+        private Executor executor = ImmediateExecutor.INSTANCE;
+        private DataLoader<K, V> dataLoader;
+        private TrackingObserver trackingObserver;
+
+        Builder(DataLoader<K, V> dataLoader) {
+            this.dataLoader = dataLoader;
+        }
+
+        public Builder<K, V> executor(Executor executor) {
+            this.executor = executor;
+            return this;
+        }
+
+        public Builder<K, V> observer(TrackingObserver trackingObserver) {
+            this.trackingObserver = trackingObserver;
+            return this;
+        }
+
+        public Orchestrator<K, V> build() {
+            return new Orchestrator<>(this);
+        }
+    }
+
 }

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

@@ -1,6 +1,8 @@
 package org.dataloader.orchestration;
 
 import org.dataloader.DataLoader;
+import org.dataloader.orchestration.executors.ObservingExecutor;
+import org.dataloader.orchestration.observation.Tracker;
 
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
@@ -46,9 +48,11 @@ public Step<V, V> thenLoadAsync(Function<V, K> codeToRun) {
     }
 
     static <K, V> Step<K, V> loadImpl(Orchestrator<?, ?> orchestrator, DataLoader<Object, Object> dl, K key, Object keyContext) {
+        Tracker tracker = orchestrator.getTracker();
+        int stepIndex = tracker.getStepCount();
         Function<K, CompletableFuture<V>> codeToRun = k -> {
             CompletableFuture<V> cf = castAs(dl.load(key, keyContext));
-            orchestrator.getTracker().loadCall(dl);
+            orchestrator.getTracker().loadCall(stepIndex, dl);
             return cf;
         };
         Step<K, V> step = new Step<>(orchestrator, dl, codeToRun);
@@ -70,19 +74,21 @@ static <K, V> Step<V, V> thenLoadImpl(Orchestrator<?, ?> orchestrator, DataLoade
     }
 
     private static <K, V> Function<V, CompletableFuture<V>> mkSyncLoadLambda(DataLoader<Object, Object> dl, Function<V, K> codeToRun, Tracker tracker) {
+        int stepIndex = tracker.getStepCount();
         return v -> {
             K key = codeToRun.apply(v);
             CompletableFuture<V> cf = castAs(dl.load(key));
-            tracker.loadCall(dl);
+            tracker.loadCall(stepIndex, dl);
             return cf;
         };
     }
 
     private static <K, V> Function<V, CompletableFuture<V>> mkAsyncLoadLambda(Orchestrator<?, ?> orchestrator, DataLoader<Object, Object> dl, Function<V, K> codeToRun, Tracker tracker) {
+        int stepIndex = tracker.getStepCount();
         return v -> {
             Executor executor = orchestrator.getExecutor();
             Consumer<String> callback = atSomePointWeNeedMoreStateButUsingStringForNowToMakeItCompile -> {
-                tracker.loadCall(dl);
+                tracker.loadCall(stepIndex, dl);
             };
             ObservingExecutor<String> observingExecutor = new ObservingExecutor<>(executor, "state", callback);
             Supplier<CompletableFuture<V>> dataLoaderCall = () -> {

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

@@ -0,0 +1,15 @@
+package org.dataloader.orchestration.executors;
+
+import org.dataloader.annotations.Internal;
+
+import java.util.concurrent.Executor;
+
+@Internal
+public class ImmediateExecutor implements Executor {
+    public static final ImmediateExecutor INSTANCE = new ImmediateExecutor();
+
+    @Override
+    public void execute(Runnable command) {
+        command.run();
+    }
+}

src/main/java/org/dataloader/orchestration/executors/ImmediateExecutor.java

@@ -1,9 +1,12 @@
-package org.dataloader.orchestration;
+package org.dataloader.orchestration.executors;
+
+import org.dataloader.annotations.Internal;
 
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
 
-class ObservingExecutor<T> implements Executor {
+@Internal
+public class ObservingExecutor<T> implements Executor {
 
     private final Executor delegate;
     private final T state;

src/main/java/org/dataloader/orchestration/ObservingExecutor.java renamed to src/main/java/org/dataloader/orchestration/executors/ObservingExecutor.java

@@ -1,4 +1,4 @@
-package org.dataloader.orchestration;
+package org.dataloader.orchestration.observation;
 
 import org.dataloader.DataLoader;
 
@@ -14,6 +14,11 @@
 public class Tracker {
     private final AtomicInteger stepCount = new AtomicInteger();
     private final Map<DataLoader<?,?>, AtomicInteger> counters = new HashMap<>();
+    private final TrackingObserver trackingObserver;
+
+    public Tracker(TrackingObserver trackingObserver) {
+        this.trackingObserver = trackingObserver;
+    }
 
     public int getOutstandingLoadCount(DataLoader<?,?> dl) {
         synchronized (this) {
@@ -35,19 +40,25 @@ public int getStepCount() {
         return stepCount.get();
     }
 
-    void incrementStepCount() {
+    public void incrementStepCount() {
         this.stepCount.incrementAndGet();
     }
 
-    void loadCall(DataLoader<?,?> dl) {
+    public void startingExecution() {
+        trackingObserver.onStart(this);
+    }
+
+    public void loadCall(int stepIndex, DataLoader<?,?> dl) {
         synchronized (this) {
             getDLCounter(dl).incrementAndGet();
+            trackingObserver.onLoad(this, stepIndex, dl);
         }
     }
 
-    void loadCallComplete(DataLoader<?,?> dl) {
+    public void loadCallComplete(int stepIndex, DataLoader<?,?> dl, Throwable throwable) {
         synchronized (this) {
             getDLCounter(dl).decrementAndGet();
+            trackingObserver.onLoadComplete(this,stepIndex,dl, throwable);
         }
     }
 

src/main/java/org/dataloader/orchestration/Tracker.java renamed to src/main/java/org/dataloader/orchestration/observation/Tracker.java

@@ -0,0 +1,18 @@
+package org.dataloader.orchestration.observation;
+
+import org.dataloader.DataLoader;
+import org.dataloader.annotations.PublicSpi;
+
+/**
+ * This callback is invoked when the {@link org.dataloader.orchestration.Orchestrator} starts execution and then
+ * as each {@link DataLoader} is invoked and then again when it completes
+ */
+@PublicSpi
+public interface TrackingObserver {
+    void onStart(Tracker tracker);
+
+    void onLoad(Tracker tracker, int stepIndex, DataLoader<?, ?> dl);
+
+    // TODO - should this have an exception should it fail in CF terms ???
+    void onLoadComplete(Tracker tracker, int stepIndex, DataLoader<?, ?> dl, Throwable throwable);
+}

src/main/java/org/dataloader/orchestration/observation/TrackingObserver.java

@@ -4,13 +4,13 @@
 import org.dataloader.DataLoaderOptions;
 import org.dataloader.DataLoaderRegistry;
 import org.dataloader.fixtures.TestKit;
+import org.dataloader.orchestration.observation.Tracker;
+import org.dataloader.orchestration.observation.TrackingObserver;
 import org.junit.jupiter.api.Test;
 
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.Executor;
 import java.util.concurrent.ForkJoinPool;
-import java.util.concurrent.Semaphore;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.awaitility.Awaitility.await;
 import static org.dataloader.DataLoaderFactory.newDataLoader;
@@ -40,7 +40,7 @@ void canOrchestrate() {
                 .register("alternateCase", dlAlternateCase)
                 .build();
 
-        Orchestrator<String, String> orchestrator = Orchestrator.orchestrate(dlUpper);
+        Orchestrator<String, String> orchestrator = Orchestrator.orchestrate(dlUpper).build();
         Step<String, String> step1 = orchestrator.load("aBc", null);
         With<String, String> with1 = step1.with(dlLower);
         Step<String, String> step2 = with1.thenLoad(key -> key);
@@ -68,7 +68,7 @@ void canOrchestrateWhenNotInPerfectOrder() {
                 .build();
 
         ForkJoinPool forkJoinPool = ForkJoinPool.commonPool();
-        Orchestrator<String, String> orchestrator = Orchestrator.orchestrate(dlUpper, forkJoinPool);
+        Orchestrator<String, String> orchestrator = Orchestrator.orchestrate(dlUpper).executor(forkJoinPool).build();
         CompletableFuture<String> cf = orchestrator.load("aBc", null)
                 .with(dlLower).thenLoad(key1 -> key1)
                 .with(dlReverse).thenLoad(key -> key)
@@ -85,4 +85,51 @@ void canOrchestrateWhenNotInPerfectOrder() {
 
         assertThat(cf.join(), equalTo("cBa"));
     }
+
+    @Test
+    void can_observe_orchestration_happening() {
+
+        DataLoaderRegistry registry = DataLoaderRegistry.newRegistry()
+                .register("upper", dlUpper)
+                .register("lower", dlLower)
+                .register("reverse", dlReverse)
+                .register("alternateCase", dlAlternateCase)
+                .build();
+
+        AtomicInteger stepCount = new AtomicInteger();
+        TrackingObserver observer = new TrackingObserver() {
+            @Override
+            public void onStart(Tracker tracker) {
+                System.out.println("starting - step count : " + tracker.getStepCount());
+                stepCount.set(tracker.getStepCount());
+            }
+
+            @Override
+            public void onLoad(Tracker tracker, int stepIndex, DataLoader<?, ?> dl) {
+                System.out.println("onLoad : " + stepIndex);
+            }
+
+            @Override
+            public void onLoadComplete(Tracker tracker, int stepIndex, DataLoader<?, ?> dl, Throwable throwable) {
+                System.out.println("onLoadComplete : " + stepIndex);
+            }
+        };
+
+        Orchestrator<String, String> orchestrator = Orchestrator.orchestrate(dlUpper).observer(observer).build();
+        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"));
+        assertThat(stepCount.get(), equalTo(3));
+    }
 }