Add a Guice shutdown listener (#56)

Author: BrandonArp

src/main/java/com/arpnetworking/metrics/mad/AppShutdown.java

@@ -0,0 +1,63 @@
+/**
+ * Copyright 2017 Inscope Metrics, Inc
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.arpnetworking.metrics.mad;
+
+import com.arpnetworking.steno.Logger;
+import com.arpnetworking.steno.LoggerFactory;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ExecutionException;
+import java.util.function.Supplier;
+
+/**
+ * Assists Guice-instantiated classes in shutting down.
+ *
+ * @author Brandon Arp (brandon dot arp at inscopemetrics dot com)
+ */
+public class AppShutdown implements LifecycleRegistration {
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    public void registerShutdown(final Supplier<CompletionStage<Void>> callback) {
+        _shutdownCallbacks.offer(callback);
+    }
+
+    /**
+     * Executes all of the registered shutdown hooks and waits for them to complete.
+     */
+    public void shutdown() {
+        final CompletableFuture<?>[] completionStages = _shutdownCallbacks.stream()
+                .map(Supplier::get)
+                .map(CompletionStage::toCompletableFuture)
+                .toArray(CompletableFuture[]::new);
+
+        try {
+            CompletableFuture.allOf(completionStages).get();
+        } catch (final InterruptedException | ExecutionException e) {
+            LOGGER.error()
+                    .setMessage("Exception in application shutdown")
+                    .setThrowable(e)
+                    .log();
+        }
+    }
+
+    private final ConcurrentLinkedQueue<Supplier<CompletionStage<Void>>> _shutdownCallbacks = new ConcurrentLinkedQueue<>();
+
+    private static final Logger LOGGER = LoggerFactory.getLogger(AppShutdown.class);
+}

src/main/java/com/arpnetworking/metrics/mad/LifecycleRegistration.java

@@ -0,0 +1,34 @@
+/**
+ * Copyright 2017 Inscope Metrics, Inc
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.arpnetworking.metrics.mad;
+
+import java.util.concurrent.CompletionStage;
+import java.util.function.Supplier;
+
+/**
+ * Provides a way to register for lifecycle events, ensuring that Guice-instantiated classes can be notified.
+ *
+ * @author Brandon Arp (brandon dot arp at inscopemetrics dot com)
+ */
+public interface LifecycleRegistration {
+    /**
+     * Registers a method to be called when the application is shutting down.  The shutdown
+     * will wait for all CompletionStages to complete before shutting down.
+     *
+     * @param callback The callback to register
+     */
+    void registerShutdown(Supplier<CompletionStage<Void>> callback);
+}

src/main/java/com/arpnetworking/metrics/mad/Main.java

@@ -273,8 +273,11 @@ private Injector launchGuice(final ActorSystem actorSystem) {
                                         .build()))
                 .build();
 
+        final AppShutdown shutdown = new AppShutdown();
+        _guiceAppShutdown = shutdown;
+
         // Instantiate Guice
-        return Guice.createInjector(new MainModule(actorSystem, metricsFactory));
+        return Guice.createInjector(new MainModule(actorSystem, metricsFactory, shutdown));
     }
 
     private ActorSystem launchAkka() {
@@ -319,6 +322,9 @@ private void shutdownAkka() {
 
     private void shutdownGuice() {
         LOGGER.info().setMessage("Stopping guice").log();
+        if (_guiceAppShutdown != null) {
+            _guiceAppShutdown.shutdown();
+        }
     }
 
     private static Builder<? extends JsonNodeSource> getFileSourceBuilder(final File configurationFile) {
@@ -338,6 +344,7 @@ private static Builder<? extends JsonNodeSource> getFileSourceBuilder(final File
     private volatile ScheduledExecutorService _jvmMetricsCollector;
 
     private volatile ActorSystem _actorSystem;
+    private volatile AppShutdown _guiceAppShutdown;
 
     private static final Long INITIAL_DELAY_IN_MILLIS = 0L;
     private static final TimeUnit TIME_UNIT = TimeUnit.MILLISECONDS;
@@ -480,18 +487,21 @@ public void run() {
 
     private static final class MainModule extends AbstractModule {
 
-        MainModule(final ActorSystem actorSystem, final MetricsFactory metricsFactory) {
+        MainModule(final ActorSystem actorSystem, final MetricsFactory metricsFactory, final AppShutdown shutdown) {
             this._actorSystem = actorSystem;
             this._metricsFactory = metricsFactory;
+            this._shutdown = shutdown;
         }
 
         @Override
         public void configure() {
             bind(ActorSystem.class).toInstance(_actorSystem);
             bind(MetricsFactory.class).toInstance(_metricsFactory);
+            bind(LifecycleRegistration.class).toInstance(_shutdown);
         }
 
         private final ActorSystem _actorSystem;
         private final MetricsFactory _metricsFactory;
+        private final AppShutdown _shutdown;
     }
 }