Shutsdown executor if its was auto added by us

Author: bbakerman

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

@@ -58,6 +58,8 @@ public class ScheduledDataLoaderRegistry extends DataLoaderRegistry implements A
     private final Map<DataLoader<?, ?>, DispatchPredicate> dataLoaderPredicates = new ConcurrentHashMap<>();
     private final DispatchPredicate dispatchPredicate;
     private final ScheduledExecutorService scheduledExecutorService;
+    private final boolean defaultExecutorUsed;
+
     private final Duration schedule;
     private final boolean tickerMode;
     private volatile boolean closed;
@@ -66,6 +68,7 @@ private ScheduledDataLoaderRegistry(Builder builder) {
         super();
         this.dataLoaders.putAll(builder.dataLoaders);
         this.scheduledExecutorService = builder.scheduledExecutorService;
+        this.defaultExecutorUsed = builder.defaultExecutorUsed;
         this.schedule = builder.schedule;
         this.tickerMode = builder.tickerMode;
         this.closed = false;
@@ -79,6 +82,16 @@ private ScheduledDataLoaderRegistry(Builder builder) {
     @Override
     public void close() {
         closed = true;
+        if (defaultExecutorUsed) {
+            scheduledExecutorService.shutdown();
+        }
+    }
+
+    /**
+     * @return executor being used by this registry
+     */
+    public ScheduledExecutorService getScheduledExecutorService() {
+        return scheduledExecutorService;
     }
 
     /**
@@ -258,9 +271,18 @@ public static class Builder {
         private final Map<DataLoader<?, ?>, DispatchPredicate> dataLoaderPredicates = new LinkedHashMap<>();
         private DispatchPredicate dispatchPredicate = DispatchPredicate.DISPATCH_ALWAYS;
         private ScheduledExecutorService scheduledExecutorService;
+        private boolean defaultExecutorUsed = false;
         private Duration schedule = Duration.ofMillis(10);
         private boolean tickerMode = false;
 
+        /**
+         * If you provide a {@link ScheduledExecutorService} then it will NOT be shutdown when
+         * {@link ScheduledDataLoaderRegistry#close()} is called.  This is left to the code that made this setup code
+         *
+         * @param executorService the executor service to run the ticker on
+         *
+         * @return this builder for a fluent pattern
+         */
         public Builder scheduledExecutorService(ScheduledExecutorService executorService) {
             this.scheduledExecutorService = nonNull(executorService);
             return this;
@@ -350,6 +372,7 @@ public Builder tickerMode(boolean tickerMode) {
         public ScheduledDataLoaderRegistry build() {
             if (scheduledExecutorService == null) {
                 scheduledExecutorService = Executors.newSingleThreadScheduledExecutor();
+                defaultExecutorUsed = true;
             }
             return new ScheduledDataLoaderRegistry(this);
         }

src/test/java/org/dataloader/registries/ScheduledDataLoaderRegistryTest.java

@@ -12,6 +12,7 @@
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -285,7 +286,7 @@ public void test_can_tick_after_first_dispatch_for_chain_data_loaders() {
         assertThat(registry.isTickerMode(), equalTo(true));
 
         int count = registry.dispatchAllWithCount();
-        assertThat(count,equalTo(1));
+        assertThat(count, equalTo(1));
 
         await().atMost(TWO_SECONDS).untilAtomic(done, is(true));
 
@@ -314,7 +315,7 @@ public void test_chain_data_loaders_will_hang_if_not_in_ticker_mode() {
         assertThat(registry.isTickerMode(), equalTo(false));
 
         int count = registry.dispatchAllWithCount();
-        assertThat(count,equalTo(1));
+        assertThat(count, equalTo(1));
 
         try {
             await().atMost(TWO_SECONDS).untilAtomic(done, is(true));
@@ -323,4 +324,25 @@ public void test_chain_data_loaders_will_hang_if_not_in_ticker_mode() {
         }
         registry.close();
     }
+
+    public void test_executors_are_shutdown() {
+        ScheduledDataLoaderRegistry registry = ScheduledDataLoaderRegistry.newScheduledRegistry().build();
+
+        ScheduledExecutorService executorService = registry.getScheduledExecutorService();
+        assertThat(executorService.isShutdown(), equalTo(false));
+        registry.close();
+        assertThat(executorService.isShutdown(), equalTo(true));
+
+        executorService = Executors.newSingleThreadScheduledExecutor();
+        registry = ScheduledDataLoaderRegistry.newScheduledRegistry()
+                .scheduledExecutorService(executorService).build();
+
+        executorService = registry.getScheduledExecutorService();
+        assertThat(executorService.isShutdown(), equalTo(false));
+        registry.close();
+        // if they provide the executor, we don't close it down
+        assertThat(executorService.isShutdown(), equalTo(false));
+
+
+    }
 }