Undertow: add Gracefull shutdown handler

Author: jknack

modules/jooby-utow/src/main/java/io/jooby/utow/Utow.java

@@ -13,6 +13,7 @@
 import io.undertow.Undertow;
 import io.undertow.UndertowOptions;
 import io.undertow.server.HttpHandler;
+import io.undertow.server.handlers.GracefulShutdownHandler;
 import io.undertow.server.handlers.encoding.EncodingHandler;
 import org.xnio.Options;
 
@@ -23,6 +24,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Web server implementation using <a href="http://undertow.io/">Undertow</a>.
@@ -36,6 +38,8 @@ public class Utow extends Server.Base {
 
   private Undertow server;
 
+  private GracefulShutdownHandler shutdown;
+
   private List<Jooby> applications = new ArrayList<>();
 
   private ServerOptions options = new ServerOptions()
@@ -66,6 +70,8 @@ public class Utow extends Server.Base {
         handler = new EncodingHandler.Builder().build(null).wrap(handler);
       }
 
+      shutdown = new GracefulShutdownHandler(handler);
+
       Undertow.Builder builder = Undertow.builder()
           .addHttpListener(options.getPort(), options.getHost())
           .setBufferSize(options.getBufferSize())
@@ -82,7 +88,7 @@ public class Utow extends Server.Base {
           .setIoThreads(options.getIoThreads())
           .setWorkerOption(Options.WORKER_NAME, "worker")
           .setWorkerThreads(options.getWorkerThreads())
-          .setHandler(handler);
+          .setHandler(shutdown);
 
       SSLContext sslContext = options.getSSLContext(application.getEnvironment().getClassLoader());
       if (sslContext != null) {
@@ -109,14 +115,25 @@ public class Utow extends Server.Base {
   @Nonnull @Override public synchronized Server stop() {
     fireStop(applications);
     applications = null;
+    try {
+      shutdown.shutdown();
+      shutdown.awaitShutdown(TimeUnit.SECONDS.toMillis(1));
+    } catch (Exception x) {
+      throw SneakyThrows.propagate(x);
+    } finally {
+      shutdownServer();
+    }
+    return this;
+  }
+
+  private void shutdownServer() {
     if (server != null) {
       try {
         server.stop();
       } finally {
         server = null;
       }
     }
-    return this;
   }
 
 }