diff --git a/src/main/java/org/dataloader/registries/ScheduledDataLoaderRegistry.java b/src/main/java/org/dataloader/registries/ScheduledDataLoaderRegistry.java index 2eb9843..7e54fab 100644 --- a/src/main/java/org/dataloader/registries/ScheduledDataLoaderRegistry.java +++ b/src/main/java/org/dataloader/registries/ScheduledDataLoaderRegistry.java @@ -58,6 +58,7 @@ public class ScheduledDataLoaderRegistry extends DataLoaderRegistry implements A private final Map, 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 +67,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 +81,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 +270,18 @@ public static class Builder { private final Map, 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 +371,7 @@ public Builder tickerMode(boolean tickerMode) { public ScheduledDataLoaderRegistry build() { if (scheduledExecutorService == null) { scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(); + defaultExecutorUsed = true; } return new ScheduledDataLoaderRegistry(this); } diff --git a/src/test/java/org/dataloader/registries/ScheduledDataLoaderRegistryTest.java b/src/test/java/org/dataloader/registries/ScheduledDataLoaderRegistryTest.java index 5e0cd9a..146c186 100644 --- a/src/test/java/org/dataloader/registries/ScheduledDataLoaderRegistryTest.java +++ b/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)); + + + } } \ No newline at end of file