WFLY-3810 Graceful shutdown support for mod_cluster

wildfly · Nov 10, 2014 · 99dfd12 · 99dfd12
1 parent 5e47a06
commit 99dfd12
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 5 deletions.
diff --git a/.../undertow/src/main/java/org/wildfly/mod_cluster/undertow/UndertowEventHandlerAdapter.java b/.../undertow/src/main/java/org/wildfly/mod_cluster/undertow/UndertowEventHandlerAdapter.java
@@ -25,11 +25,16 @@
 import io.undertow.servlet.api.Deployment;
 
 import java.security.AccessController;
+import java.util.HashSet;
+import java.util.Set;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 
+import org.jboss.as.server.suspend.ServerActivity;
+import org.jboss.as.server.suspend.ServerActivityCallback;
+import org.jboss.as.server.suspend.SuspendController;
 import org.jboss.logging.Logger;
 import org.jboss.modcluster.container.Connector;
 import org.jboss.modcluster.container.ContainerEventHandler;
@@ -52,22 +57,25 @@
  *
  * @author Paul Ferraro
  */
-public class UndertowEventHandlerAdapter implements UndertowEventListener, Service<Void>, Runnable {
+public class UndertowEventHandlerAdapter implements UndertowEventListener, Service<Void>, Runnable, ServerActivity {
     private static final Logger log = Logger.getLogger(UndertowEventHandlerAdapter.class);
 
     @SuppressWarnings("rawtypes")
     private final Value<ListenerService> listener;
     private final Value<UndertowService> service;
     private final Value<ContainerEventHandler> eventHandler;
+    private final Value<SuspendController> suspendController;
+    private final Set<Context> contexts = new HashSet<>();
     private volatile ScheduledExecutorService executor;
     private volatile Server server;
     private volatile Connector connector;
     private int statusInterval;
 
-    public UndertowEventHandlerAdapter(Value<ContainerEventHandler> eventHandler, Value<UndertowService> service, @SuppressWarnings("rawtypes") Value<ListenerService> listener, int statusInterval) {
+    public UndertowEventHandlerAdapter(Value<ContainerEventHandler> eventHandler, Value<UndertowService> service, @SuppressWarnings("rawtypes") Value<ListenerService> listener, Value<SuspendController> suspendController, int statusInterval) {
         this.eventHandler = eventHandler;
         this.service = service;
         this.listener = listener;
+        this.suspendController = suspendController;
         this.statusInterval = statusInterval;
     }
 
@@ -95,10 +103,12 @@ public void start(StartContext context) {
         ThreadFactory factory = new JBossThreadFactory(group, Boolean.FALSE, null, "%G - %t", null, null, AccessController.doPrivileged(GetAccessControlContextAction.getInstance()));
         this.executor = Executors.newScheduledThreadPool(1, factory);
         this.executor.scheduleWithFixedDelay(this, 0, statusInterval, TimeUnit.SECONDS);
+        suspendController.getValue().registerActivity(this);
     }
 
     @Override
     public void stop(StopContext context) {
+        suspendController.getValue().unRegisterActivity(this);
         this.service.getValue().unregisterListener(this);
 
         this.executor.shutdown();
@@ -112,21 +122,23 @@ private Context createContext(Deployment deployment, Host host) {
     }
 
     @Override
-    public void onDeploymentStart(Deployment deployment, Host host) {
+    public synchronized void onDeploymentStart(Deployment deployment, Host host) {
         Context context = this.createContext(deployment, host);
         this.eventHandler.getValue().add(context);
 
         // TODO break into onDeploymentAdd once implemented in Undertow
         this.eventHandler.getValue().start(context);
+        contexts.add(context);
     }
 
     @Override
-    public void onDeploymentStop(Deployment deployment, Host host) {
+    public synchronized void onDeploymentStop(Deployment deployment, Host host) {
         Context context = this.createContext(deployment, host);
         this.eventHandler.getValue().stop(context);
 
         // TODO break into onDeploymentRemove once implemented in Undertow
         this.eventHandler.getValue().remove(context);
+        contexts.remove(context);
     }
 
     @Override
@@ -164,4 +176,27 @@ public void run() {
             log.error(e.getMessage(), e);
         }
     }
+
+    @Override
+    public synchronized void preSuspend(ServerActivityCallback listener) {
+        try {
+            for (Context context : contexts) {
+                this.eventHandler.getValue().stop(context);
+            }
+        } finally {
+            listener.done();
+        }
+    }
+
+    @Override
+    public void suspended(ServerActivityCallback listener) {
+        listener.done();
+    }
+
+    @Override
+    public void resume() {
+        for (Context context : contexts) {
+            this.eventHandler.getValue().start(context);
+        }
+    }
 }
diff --git a/...ow/src/main/java/org/wildfly/mod_cluster/undertow/UndertowEventHandlerAdapterBuilder.java b/...ow/src/main/java/org/wildfly/mod_cluster/undertow/UndertowEventHandlerAdapterBuilder.java
@@ -23,6 +23,7 @@
 package org.wildfly.mod_cluster.undertow;
 
 import org.jboss.as.clustering.msc.AsynchronousService;
+import org.jboss.as.server.suspend.SuspendController;
 import org.jboss.modcluster.container.ContainerEventHandler;
 import org.jboss.msc.service.ServiceBuilder;
 import org.jboss.msc.service.ServiceName;
@@ -40,12 +41,14 @@ public class UndertowEventHandlerAdapterBuilder implements ContainerEventHandler
     public ServiceBuilder<?> build(ServiceTarget target, String connector, int statusInterval) {
         InjectedValue<ContainerEventHandler> eventHandler = new InjectedValue<>();
         InjectedValue<UndertowService> undertowService = new InjectedValue<>();
+        InjectedValue<SuspendController> suspendController = new InjectedValue<>();
         @SuppressWarnings("rawtypes")
         InjectedValue<ListenerService> listener = new InjectedValue<>();
-        return AsynchronousService.addService(target, SERVICE_NAME, new UndertowEventHandlerAdapter(eventHandler, undertowService, listener, statusInterval))
+        return AsynchronousService.addService(target, SERVICE_NAME, new UndertowEventHandlerAdapter(eventHandler, undertowService, listener, suspendController, statusInterval))
                 .addDependency(ContainerEventHandlerService.SERVICE_NAME, ContainerEventHandler.class, eventHandler)
                 .addDependency(UndertowService.UNDERTOW, UndertowService.class, undertowService)
                 .addDependency(UndertowService.listenerName(connector), ListenerService.class, listener)
+                .addDependency(SuspendController.SERVICE_NAME, SuspendController.class, suspendController)
         ;
     }
 }