Instrumentation support for dataloader - updated ScheduledDataLoaderRegistry

Author: bbakerman

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

@@ -19,28 +19,45 @@
 /**
  * This allows data loaders to be registered together into a single place, so
  * they can be dispatched as one.  It also allows you to retrieve data loaders by
- * name from a central place
+ * name from a central place.
+ * <p>
+ * Notes on {@link DataLoaderInstrumentation} : A {@link DataLoaderRegistry} can have an instrumentation
+ * associated with it.  As each {@link DataLoader} is added to the registry, the {@link DataLoaderInstrumentation}
+ * of the registry is applied to that {@link DataLoader}.
+ * <p>
+ * The {@link DataLoader} is changed and hence the object in the registry is not the
+ * same one as was originally registered.  So you MUST get access to the {@link DataLoader} via {@link DataLoaderRegistry#getDataLoader(String)} methods
+ * and not use the original {@link DataLoader} object.
+ * <p>
+ * If the {@link DataLoader} has no {@link DataLoaderInstrumentation} then the registry one is added to it.  If it does have one already
+ * then a {@link ChainedDataLoaderInstrumentation} is created with the registry {@link DataLoaderInstrumentation} in it first and then any other
+ * {@link DataLoaderInstrumentation}s added after that.
  */
 @PublicApi
 public class DataLoaderRegistry {
-    protected final Map<String, DataLoader<?, ?>> dataLoaders = new ConcurrentHashMap<>();
+    protected final Map<String, DataLoader<?, ?>> dataLoaders;
     protected final DataLoaderInstrumentation instrumentation;
 
 
     public DataLoaderRegistry() {
-        instrumentation = null;
+        this(new ConcurrentHashMap<>(),null);
     }
 
     private DataLoaderRegistry(Builder builder) {
-        instrument(builder.instrumentation, builder.dataLoaders);
-        this.instrumentation = builder.instrumentation;
+        this(builder.dataLoaders,builder.instrumentation);
     }
 
-    private void instrument(DataLoaderInstrumentation registryInstrumentation, Map<String, DataLoader<?, ?>> incomingDataLoaders) {
-        this.dataLoaders.putAll(incomingDataLoaders);
+    protected DataLoaderRegistry(Map<String, DataLoader<?, ?>> dataLoaders,  DataLoaderInstrumentation instrumentation ) {
+        this.dataLoaders = instrumentDLs(dataLoaders, instrumentation);
+        this.instrumentation = instrumentation;
+    }
+
+    private Map<String, DataLoader<?, ?>> instrumentDLs(Map<String, DataLoader<?, ?>> incomingDataLoaders, DataLoaderInstrumentation registryInstrumentation) {
+        Map<String, DataLoader<?, ?>> dataLoaders = new ConcurrentHashMap<>(incomingDataLoaders);
         if (registryInstrumentation != null) {
-            this.dataLoaders.replaceAll((k, existingDL) -> instrumentDL(registryInstrumentation, existingDL));
+            dataLoaders.replaceAll((k, existingDL) -> instrumentDL(registryInstrumentation, existingDL));
         }
+        return dataLoaders;
     }
 
     /**

src/main/java/org/dataloader/registries/ScheduledDataLoaderRegistry.java

@@ -3,6 +3,7 @@
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderRegistry;
 import org.dataloader.annotations.ExperimentalApi;
+import org.dataloader.instrumentation.DataLoaderInstrumentation;
 
 import java.time.Duration;
 import java.util.LinkedHashMap;
@@ -64,8 +65,7 @@ public class ScheduledDataLoaderRegistry extends DataLoaderRegistry implements A
     private volatile boolean closed;
 
     private ScheduledDataLoaderRegistry(Builder builder) {
-        super();
-        this.dataLoaders.putAll(builder.dataLoaders);
+        super(builder.dataLoaders, builder.instrumentation);
         this.scheduledExecutorService = builder.scheduledExecutorService;
         this.defaultExecutorUsed = builder.defaultExecutorUsed;
         this.schedule = builder.schedule;
@@ -271,6 +271,8 @@ public static class Builder {
         private boolean defaultExecutorUsed = false;
         private Duration schedule = Duration.ofMillis(10);
         private boolean tickerMode = false;
+        private DataLoaderInstrumentation instrumentation;
+
 
         /**
          * If you provide a {@link ScheduledExecutorService} then it will NOT be shutdown when
@@ -363,6 +365,11 @@ public Builder tickerMode(boolean tickerMode) {
             return this;
         }
 
+        public Builder instrumentation(DataLoaderInstrumentation instrumentation) {
+            this.instrumentation = instrumentation;
+            return this;
+        }
+
         /**
          * @return the newly built {@link ScheduledDataLoaderRegistry}
          */

src/test/java/org/dataloader/instrumentation/ChainedDataLoaderInstrumentationTest.java

@@ -18,7 +18,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.equalTo;
 
-class ChainedDataLoaderInstrumentationTest {
+public class ChainedDataLoaderInstrumentationTest {
 
     CapturingInstrumentation capturingA;
     CapturingInstrumentation capturingB;

src/test/java/org/dataloader/instrumentation/DataLoaderInstrumentationTest.java

@@ -22,7 +22,7 @@
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.is;
 
-class DataLoaderInstrumentationTest {
+public class DataLoaderInstrumentationTest {
 
     BatchLoader<String, String> snoozingBatchLoader = keys -> CompletableFuture.supplyAsync(() -> {
         TestKit.snooze(100);

src/test/java/org/dataloader/instrumentation/DataLoaderRegistryInstrumentationTest.java

@@ -4,6 +4,7 @@
 import org.dataloader.DataLoaderRegistry;
 import org.dataloader.fixtures.TestKit;
 import org.dataloader.fixtures.parameterized.TestDataLoaderFactory;
+import org.dataloader.registries.ScheduledDataLoaderRegistry;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -18,7 +19,7 @@
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.instanceOf;
 
-class DataLoaderRegistryInstrumentationTest {
+public class DataLoaderRegistryInstrumentationTest {
     DataLoader<String, String> dlX;
     DataLoader<String, String> dlY;
     DataLoader<String, String> dlZ;
@@ -177,6 +178,29 @@ void chainedInstrumentationsWillBeCombined() {
         assertThat(instrumentations, equalTo(List.of(instrA, instrB)));
     }
 
+    @SuppressWarnings("resource")
+    @Test
+    void canInstrumentScheduledRegistryViaBuilder() {
+
+        assertThat(dlX.getOptions().getInstrumentation(), equalTo(DataLoaderInstrumentationHelper.NOOP_INSTRUMENTATION));
+
+        ScheduledDataLoaderRegistry registry = ScheduledDataLoaderRegistry.newScheduledRegistry()
+                .instrumentation(chainedInstrA)
+                .register("X", dlX)
+                .register("Y", dlY)
+                .register("Z", dlZ)
+                .build();
+
+        assertThat(registry.getInstrumentation(), equalTo(chainedInstrA));
+
+        for (String key : List.of("X", "Y", "Z")) {
+            DataLoaderInstrumentation instrumentation = registry.getDataLoader(key).getOptions().getInstrumentation();
+            assertThat(instrumentation, instanceOf(ChainedDataLoaderInstrumentation.class));
+            List<DataLoaderInstrumentation> instrumentations = ((ChainedDataLoaderInstrumentation) instrumentation).getInstrumentations();
+            assertThat(instrumentations, equalTo(List.of(instrA)));
+        }
+    }
+
     @ParameterizedTest
     @MethodSource("org.dataloader.fixtures.parameterized.TestDataLoaderFactories#get")
     public void endToEndIntegrationTest(TestDataLoaderFactory factory) {
@@ -200,6 +224,5 @@ public void endToEndIntegrationTest(TestDataLoaderFactory factory) {
         assertThat(instrA.methods, equalTo(List.of("A_beginDispatch",
                 "A_beginBatchLoader", "A_beginBatchLoader_onDispatched", "A_beginBatchLoader_onCompleted",
                 "A_beginDispatch_onDispatched", "A_beginDispatch_onCompleted")));
-
     }
 }