add shutdown handler #1796

- remove shutdown handler from undertow server

Author: jknack

docs/asciidoc/handlers.adoc

@@ -8,6 +8,8 @@ include::handlers/cors.adoc[]
 
 include::handlers/csrf.adoc[]
 
+include::handlers/graceful-shutdown.adoc[]
+
 include::handlers/head.adoc[]
 
 include::handlers/rate-limit.adoc[]

docs/asciidoc/handlers/graceful-shutdown.adoc

@@ -0,0 +1,42 @@
+=== GracefulShutdown
+
+The javadoc:GracefulShutdown[] extension waits for existing requests to finish.
+
+.Example
+[source, java, role = "primary"]
+----
+import io.jooby.Jooby;
+import io.jooby.GracefulShutdown;
+...
+{
+  
+  install(new GracefulShutdown()); <1>
+
+  // other routes go here
+}
+----
+
+.Kotlin
+[source, kotlin, role = "secondary"]
+----
+import io.jooby.Jooby
+import io.jooby.GracefulShutdown
+...
+{
+  install(GracefulShutdown())      <1>
+  
+  // other routes go here
+}
+----
+
+<1> Install GracefulShutdown.
+
+Incoming request are resolved as `Service Unavailable(503)`. Optionally you can specify a max
+amount of time to wait before shutdown:
+
+    install(new GracefulShutdown(Duration.ofMinutes(1)));
+
+[INFO]
+====
+This extension must be installed at very beginning of your route pipeline.
+====

jooby/src/main/java/io/jooby/GracefulShutdown.java

@@ -0,0 +1,47 @@
+/**
+ * Jooby https://jooby.io
+ * Apache License Version 2.0 https://jooby.io/LICENSE.txt
+ * Copyright 2014 Edgar Espina
+ */
+package io.jooby;
+
+import io.jooby.internal.GracefulShutdownHandler;
+
+import javax.annotation.Nonnull;
+import java.time.Duration;
+
+/**
+ * Install a handler that at application shutdown time:
+ *
+ * - Waits for existing requests to finished with an optional timeout
+ * - Incoming requests are resolved as Service Unavailable(503)
+ *
+ * NOTE: This extension must be installed at very beginning of your route pipeline.
+ *
+ * @author edgar
+ */
+public class GracefulShutdown implements Extension {
+  private Duration await;
+
+  /**
+   * Creates a new shutdown handler and waits for existing requests to finish or for specified
+   * amount of time.
+   *
+   * @param await Max time to wait for handlers to complete.
+   */
+  public GracefulShutdown(@Nonnull Duration await) {
+    this.await = await;
+  }
+
+  /**
+   * Creates a new shutdown handler and waits for existing request to finish.
+   */
+  public GracefulShutdown() {
+  }
+
+  @Override public void install(@Nonnull Jooby application) throws Exception {
+    GracefulShutdownHandler handler = new GracefulShutdownHandler(await);
+    application.decorator(handler);
+    application.onStop(handler::shutdown);
+  }
+}

jooby/src/main/java/io/jooby/internal/GracefulShutdownHandler.java

@@ -0,0 +1,137 @@
+/**
+ * Jooby https://jooby.io
+ * Apache License Version 2.0 https://jooby.io/LICENSE.txt
+ * Copyright 2014 Edgar Espina
+ */
+package io.jooby.internal;
+
+import io.jooby.Route;
+import io.jooby.StatusCode;
+
+import javax.annotation.Nonnull;
+import java.time.Duration;
+import java.util.concurrent.atomic.AtomicLongFieldUpdater;
+import java.util.function.LongUnaryOperator;
+
+public class GracefulShutdownHandler implements Route.Decorator {
+  private static final long SHUTDOWN_MASK = 1L << 63;
+  private static final long ACTIVE_COUNT_MASK = (1L << 63) - 1;
+
+  private static final LongUnaryOperator incrementActive = current -> {
+    long incrementedActiveCount = activeCount(current) + 1;
+    return incrementedActiveCount | (current & ~ACTIVE_COUNT_MASK);
+  };
+
+  private static final LongUnaryOperator incrementActiveAndShutdown =
+      incrementActive.andThen(current -> current | SHUTDOWN_MASK);
+
+  private static final LongUnaryOperator decrementActive = current -> {
+    long decrementedActiveCount = activeCount(current) - 1;
+    return decrementedActiveCount | (current & ~ACTIVE_COUNT_MASK);
+  };
+
+  private final Object lock = new Object();
+
+  private final Duration await;
+
+  private volatile long state = 0;
+  private static final AtomicLongFieldUpdater<GracefulShutdownHandler> stateUpdater =
+      AtomicLongFieldUpdater.newUpdater(GracefulShutdownHandler.class, "state");
+
+  public GracefulShutdownHandler(Duration await) {
+    this.await = await;
+  }
+
+  @Nonnull @Override public Route.Handler apply(@Nonnull Route.Handler next) {
+    return ctx -> {
+      long snapshot = stateUpdater.updateAndGet(this, incrementActive);
+      if (isShutdown(snapshot)) {
+        decrementRequests();
+        return ctx.setResponseCode(StatusCode.SERVICE_UNAVAILABLE)
+            .send(StatusCode.SERVICE_UNAVAILABLE);
+      } else {
+        ctx.onComplete(context -> {
+          decrementRequests();
+        });
+        return next.apply(ctx);
+      }
+    };
+  }
+
+  private static boolean isShutdown(long state) {
+    return (state & SHUTDOWN_MASK) != 0;
+  }
+
+  private static long activeCount(long state) {
+    return state & ACTIVE_COUNT_MASK;
+  }
+
+  public void shutdown() throws InterruptedException {
+    //the request count is never zero when shutdown is set to true
+    stateUpdater.updateAndGet(this, incrementActiveAndShutdown);
+    decrementRequests();
+
+    if (await == null) {
+      awaitShutdown();
+    } else {
+      awaitShutdown(await.toMillis());
+    }
+  }
+
+  public void start() {
+    synchronized (lock) {
+      stateUpdater.updateAndGet(this, current -> current & ACTIVE_COUNT_MASK);
+    }
+  }
+
+  private void shutdownComplete() {
+    synchronized (lock) {
+      lock.notifyAll();
+    }
+  }
+
+  /**
+   * Waits for the handler to shutdown.
+   */
+  private void awaitShutdown() throws InterruptedException {
+    synchronized (lock) {
+      if (!isShutdown(stateUpdater.get(this))) {
+        return;
+      }
+      while (activeCount(stateUpdater.get(this)) > 0) {
+        lock.wait();
+      }
+    }
+  }
+
+  /**
+   * Waits a set length of time for the handler to shut down
+   *
+   * @param millis The length of time
+   * @return <code>true</code> If the handler successfully shut down
+   */
+  private boolean awaitShutdown(long millis) throws InterruptedException {
+    synchronized (lock) {
+      if (!isShutdown(stateUpdater.get(this))) {
+        return false;
+      }
+      long end = System.currentTimeMillis() + millis;
+      while (activeCount(stateUpdater.get(this)) != 0) {
+        long left = end - System.currentTimeMillis();
+        if (left <= 0) {
+          return false;
+        }
+        lock.wait(left);
+      }
+      return true;
+    }
+  }
+
+  private void decrementRequests() {
+    long snapshot = stateUpdater.updateAndGet(this, decrementActive);
+    // Shutdown has completed when the activeCount portion is zero, and shutdown is set.
+    if (snapshot == SHUTDOWN_MASK) {
+      shutdownComplete();
+    }
+  }
+}

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

@@ -13,7 +13,6 @@
 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.ContentEncodingRepository;
 import io.undertow.server.handlers.encoding.DeflateEncodingProvider;
 import io.undertow.server.handlers.encoding.EncodingHandler;
@@ -27,7 +26,6 @@
 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>.
@@ -45,8 +43,6 @@ public class Utow extends Server.Base {
 
   private Undertow server;
 
-  private GracefulShutdownHandler shutdown;
-
   private List<Jooby> applications = new ArrayList<>();
 
   private ServerOptions options = new ServerOptions()
@@ -80,8 +76,6 @@ public class Utow extends Server.Base {
             .addEncodingHandler("deflate", new DeflateEncodingProvider(compressionLevel), _10));
       }
 
-      shutdown = new GracefulShutdownHandler(handler);
-
       Undertow.Builder builder = Undertow.builder()
           .addHttpListener(options.getPort(), options.getHost())
           .setBufferSize(options.getBufferSize())
@@ -98,7 +92,7 @@ public class Utow extends Server.Base {
           .setIoThreads(options.getIoThreads())
           .setWorkerOption(Options.WORKER_NAME, "worker")
           .setWorkerThreads(options.getWorkerThreads())
-          .setHandler(shutdown);
+          .setHandler(handler);
 
       SSLContext sslContext = options.getSSLContext(application.getEnvironment().getClassLoader());
       if (sslContext != null) {
@@ -123,11 +117,9 @@ 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));
+      fireStop(applications);
+      applications = null;
     } catch (Exception x) {
       throw SneakyThrows.propagate(x);
     } finally {