shutdown hook for UI server (#9005)

* shutdown hook for UI server

Signed-off-by: Tamás Fenyvesi <tamas.fenyvesi@doknet.hu>

* logging exception

Signed-off-by: Tamás Fenyvesi <tamas.fenyvesi@doknet.hu>

deeplearning4j/deeplearning4j-ui-parent/deeplearning4j-vertx/src/main/java/org/deeplearning4j/ui/VertxUIServer.java

@@ -75,7 +75,8 @@ public class VertxUIServer extends AbstractVerticle implements UIServer {
     private static Function<String, StatsStorage> statsStorageProvider;
 
     private static Integer instancePort;
-    private static Thread autoStopThread;
+    @Getter
+    private static Thread shutdownHook;
 
     /**
      * Get (and, initialize if necessary) the UI server. This synchronous function will wait until the server started.
@@ -188,19 +189,17 @@ public class VertxUIServer extends AbstractVerticle implements UIServer {
         Vertx vertx = Vertx.vertx();
         vertx.deployVerticle(VertxUIServer.class.getName(), promise);
 
-        Thread currentThread = Thread.currentThread();
+        VertxUIServer.shutdownHook = new Thread(() -> {
-        VertxUIServer.autoStopThread = new Thread(() -> {
+            if (VertxUIServer.instance != null && !VertxUIServer.instance.isStopped()) {
-            try {
+                log.info("Deeplearning4j UI server is auto-stopping in shutdown hook.");
-                currentThread.join();
+                try {
-                if (VertxUIServer.instance != null && !VertxUIServer.instance.isStopped()) {
-                    log.info("Deeplearning4j UI server is auto-stopping after thread (name: {}) died.",
-                            currentThread.getName());
                     instance.stop();
+                } catch (InterruptedException e) {
+                    log.error("Interrupted stopping of Deeplearning4j UI server in shutdown hook.", e);
                 }
-            } catch (InterruptedException e) {
-                log.error("Deeplearning4j UI server auto-stop thread was interrupted.", e);
             }
         });
+        Runtime.getRuntime().addShutdownHook(shutdownHook);
     }
 
 
@@ -385,7 +384,6 @@ public class VertxUIServer extends AbstractVerticle implements UIServer {
                                 + server.actualPort(), result.cause()));
                     }
                 });
-        VertxUIServer.autoStopThread.start();
     }
 
     private List<String> extractArgsFromRoute(String path, RoutingContext rc) {

deeplearning4j/deeplearning4j-ui-parent/deeplearning4j-vertx/src/main/java/org/deeplearning4j/ui/api/UIServer.java

@@ -168,4 +168,14 @@ public interface UIServer {
      * @param stopCallback callback {@link Promise} to notify on completion
      */
     void stopAsync(Promise<Void> stopCallback);
+
+    /**
+     * Get shutdown hook of UI server, that will stop the server when the Runtime is stopped.
+     * You may de-register this shutdown hook with {@link Runtime#removeShutdownHook(Thread)},
+     * and add your own hook with {@link Runtime#addShutdownHook(Thread)}
+     * @return shutdown hook
+     */
+    static Thread getShutdownHook() {
+        return VertxUIServer.getShutdownHook();
+    };
 }

deeplearning4j/deeplearning4j-ui-parent/deeplearning4j-vertx/src/test/java/org/deeplearning4j/ui/TestVertxUI.java

@@ -346,12 +346,16 @@ public class TestVertxUI extends BaseDL4JTest {
     }
 
     @Test
-    public void testUIAutoStopOnThreadExit() throws InterruptedException {
+    public void testUIShutdownHook() throws InterruptedException {
-        AtomicReference<UIServer> uiServer = new AtomicReference<>();
+        UIServer uIServer = UIServer.getInstance();
-        Thread thread = new Thread(() -> uiServer.set(UIServer.getInstance()));
+        Thread shutdownHook = UIServer.getShutdownHook();
-        thread.start();
+        shutdownHook.start();
-        thread.join();
+        shutdownHook.join();
-        Thread.sleep(1_000);
+        /*
-        assertTrue(uiServer.get().isStopped());
+         * running the shutdown hook thread before the Runtime is terminated
+         * enables us to check if the UI server has been shut down or not
+         */
+        assertTrue(uIServer.isStopped());
+        log.info("Deeplearning4j UI server stopped in shutdown hook.");
     }
 }