Merge pull request #3109 from bstansberry/WFCORE-3597

[WFCORE-3597] Allow reload handler to signal when to send the early r…

controller/src/main/java/org/jboss/as/controller/AbstractControllerService.java

@@ -60,6 +60,7 @@
 import org.jboss.msc.service.ServiceContainer;
 import org.jboss.msc.service.ServiceController;
 import org.jboss.msc.service.ServiceTarget;
+import org.jboss.msc.service.StabilityMonitor;
 import org.jboss.msc.service.StartContext;
 import org.jboss.msc.service.StartException;
 import org.jboss.msc.service.StopContext;
@@ -161,6 +162,7 @@ private static int getBootStackSize() {
     private final InjectedValue<ControllerInstabilityListener> injectedInstabilityListener = new InjectedValue<ControllerInstabilityListener>();
     private final ExpressionResolver expressionResolver;
     private volatile ModelControllerImpl controller;
+    private volatile StabilityMonitor stabilityMonitor;
     private ConfigurationPersister configurationPersister;
     private final ManagedAuditLogger auditLogger;
     private final BootErrorCollector bootErrorCollector;
@@ -330,7 +332,7 @@ public void start(final StartContext context) throws StartException {
         ManagementResourceRegistration rootResourceRegistration = ManagementResourceRegistration.Factory.forProcessType(processType).createRegistration(rootResourceDefinition, authorizerConfig, capabilityRegistry);
         final ModelControllerImpl controller = new ModelControllerImpl(container, target,
                 rootResourceRegistration,
-                new ContainerStateMonitor(container),
+                new ContainerStateMonitor(container, getStabilityMonitor()),
                 configurationPersister, processType, runningModeControl, prepareStep,
                 processState, executorService, expressionResolver, authorizer, securityIdentitySupplier, auditLogger, notificationSupport,
                 bootErrorCollector, createExtraValidationStepHandler(), capabilityRegistry, getPartialModelIndicator(),
@@ -550,11 +552,19 @@ protected PartialModelIndicator getPartialModelIndicator() {
         return PartialModelIndicator.DEFAULT;
     }
 
+    protected final StabilityMonitor getStabilityMonitor() {
+        if (stabilityMonitor == null) {
+            stabilityMonitor = new StabilityMonitor();
+        }
+        return stabilityMonitor;
+    }
+
     public void stop(final StopContext context) {
         capabilityRegistry.clear();
         capabilityRegistry.publish();
         ServiceNameFactory.clearCache();
         controller = null;
+        stabilityMonitor = null;
         processState.setStopping();
         Runnable r = new Runnable() {
             @Override

controller/src/main/java/org/jboss/as/controller/AbstractOperationContext.java

@@ -826,7 +826,7 @@ public void commit() {
                     public void rollback() {
                         ref.set(ResultAction.ROLLBACK);
                     }
-                }, primaryResponse);
+                }, primaryResponse, this);
             }
             resultAction = ref.get();
 

controller/src/main/java/org/jboss/as/controller/ContainerStateMonitor.java

@@ -48,15 +48,16 @@
 final class ContainerStateMonitor {
 
     private final ServiceRegistry serviceRegistry;
-    private final StabilityMonitor monitor = new StabilityMonitor();
-    final Set<ServiceController<?>> failed = new HashSet<ServiceController<?>>();
-    final Set<ServiceController<?>> problems = new HashSet<ServiceController<?>>();
+    private final StabilityMonitor monitor;
+    private final Set<ServiceController<?>> failed = new HashSet<ServiceController<?>>();
+    private final Set<ServiceController<?>> problems = new HashSet<ServiceController<?>>();
 
     private Set<ServiceName> previousMissingDepSet = new HashSet<ServiceName>();
     private Set<ServiceController<?>> previousFailedSet = new HashSet<>();
 
-    ContainerStateMonitor(final ServiceRegistry registry) {
+    ContainerStateMonitor(final ServiceRegistry registry, final StabilityMonitor stabilityMonitor) {
         serviceRegistry = registry;
+        monitor = stabilityMonitor;
     }
 
     /**

controller/src/main/java/org/jboss/as/controller/ModelController.java

@@ -115,6 +115,22 @@ public interface ModelController {
      */
     interface OperationTransactionControl {
 
+        /**
+         * Notify that an operation is complete and may be committed or rolled back.
+         *
+         * <p><strong>It is the responsibility of the user of this {@code OperationTransactionControl} to ensure that
+         * {@link OperationTransaction#commit()} or {@link OperationTransaction#rollback()} is eventually called on
+         * the provided {@code transaction}.
+         * </strong></p>
+         *
+         * @param transaction the transaction to control the fate of the operation. Cannot be {@code null}
+         * @param result the result. Cannot be {@code null}
+         * @param context the {@code OperationContext} for the operation that is prepared
+         */
+        default void operationPrepared(OperationTransaction transaction, ModelNode result, OperationContext context) {
+            operationPrepared(transaction, result);
+        }
+
         /**
          * Notify that an operation is complete and may be committed or rolled back.
          *

controller/src/main/java/org/jboss/as/controller/ModelControllerImpl.java

@@ -358,6 +358,11 @@ protected OperationResponse internalExecute(final ModelNode operation, final Ope
         // Report the correct operation response, otherwise the preparedResult would only contain
         // the result of the last active step in a composite operation
         final OperationTransactionControl originalResultTxControl = control == null ? null : new OperationTransactionControl() {
+            @Override
+            public void operationPrepared(OperationTransaction transaction, ModelNode result, OperationContext context) {
+                control.operationPrepared(transaction, responseNode, context);
+            }
+
             @Override
             public void operationPrepared(OperationTransaction transaction, ModelNode result) {
                 control.operationPrepared(transaction, responseNode);

controller/src/main/java/org/jboss/as/controller/ProxyStepHandler.java

@@ -154,6 +154,16 @@ public void close() throws IOException {
 
                     @Override
                     public void operationPrepared(ModelController.OperationTransaction transaction, ModelNode response) {
+                        proxyControl.operationPrepared(transaction, transformResponse(response));
+                    }
+
+                    @Override
+                    public void operationPrepared(ModelController.OperationTransaction transaction, ModelNode response,
+                                                  OperationContext context1) {
+                        proxyControl.operationPrepared(transaction, transformResponse(response), context1);
+                    }
+
+                    private ModelNode transformResponse(ModelNode response) {
                         final ModelNode transformed;
                         // Check if we have to reject the operation
                         if(result.rejectOperation(response)) {
@@ -164,7 +174,7 @@ public void operationPrepared(ModelController.OperationTransaction transaction,
                         } else {
                             transformed = response;
                         }
-                        proxyControl.operationPrepared(transaction, transformed);
+                        return transformed;
                     }
                 };
                 proxyController.execute(transformedOperation, messageHandler, transformingProxyControl,

controller/src/main/java/org/jboss/as/controller/operations/common/ProcessReloadHandler.java

@@ -22,11 +22,6 @@
 
 package org.jboss.as.controller.operations.common;
 
-import static org.jboss.as.controller.AbstractControllerService.EXECUTOR_CAPABILITY;
-
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Future;
 import org.jboss.as.controller.AttributeDefinition;
 import org.jboss.as.controller.ControlledProcessState;
 import org.jboss.as.controller.OperationContext;
@@ -35,7 +30,7 @@
 import org.jboss.as.controller.RunningModeControl;
 import org.jboss.as.controller.SimpleAttributeDefinitionBuilder;
 import org.jboss.as.controller.descriptions.ModelDescriptionConstants;
-import org.jboss.as.controller.logging.ControllerLogger;
+import org.jboss.as.controller.remote.EarlyResponseSendListener;
 import org.jboss.dmr.ModelNode;
 import org.jboss.dmr.ModelType;
 import org.jboss.msc.service.AbstractServiceListener;
@@ -78,37 +73,49 @@ public void execute(OperationContext context, ModelNode operation) throws Operat
             public void execute(final OperationContext context, final ModelNode operation) throws OperationFailedException {
                 final ReloadContext<T> reloadContext = initializeReloadContext(context, operation);
                 final ServiceController<?> service = context.getServiceRegistry(true).getRequiredService(rootService);
-                final ExecutorService executor = (ExecutorService) context.getServiceRegistry(false).getRequiredService(EXECUTOR_CAPABILITY.getCapabilityServiceName()).getValue();
                 context.completeStep(new OperationContext.ResultHandler() {
                     @Override
                     public void handleResult(OperationContext.ResultAction resultAction, OperationContext context, ModelNode operation) {
-                        if(resultAction == OperationContext.ResultAction.KEEP) {
-                            service.addListener(new AbstractServiceListener<Object>() {
-                                @Override
-                                public void listenerAdded(final ServiceController<?> controller) {
-                                    Future<?> stopping = executor.submit(() -> {
+                        final EarlyResponseSendListener sendListener = context.getAttachment(EarlyResponseSendListener.ATTACHMENT_KEY);
+                        try {
+                            if (resultAction == OperationContext.ResultAction.KEEP) {
+                                service.addListener(new AbstractServiceListener<Object>() {
+                                    @Override
+                                    public void listenerAdded(final ServiceController<?> controller) {
                                         reloadContext.reloadInitiated(runningModeControl);
                                         processState.setStopping();
-                                        controller.setMode(ServiceController.Mode.NEVER);
-                                    });
-                                    try {
-                                        stopping.get();
-                                    } catch (InterruptedException ex) {
-                                        Thread.currentThread().interrupt();
-                                    } catch (ExecutionException ex) {
-                                        ControllerLogger.ROOT_LOGGER.errorStoppingServer(ex);
+                                        try {
+                                            // If we were interrupted during setStopping (i.e. while calling process state listeners)
+                                            // we want to clear that so we don't disrupt the reload of MSC services.
+                                            // Once we set STOPPING state we proceed. And we don't want this thread
+                                            // to have interrupted status as that will just mess up checking for
+                                            // container stability
+                                            Thread.interrupted();
+                                            // Now that we're in STOPPING state we can send the response to the caller
+                                            if (sendListener != null) {
+                                                sendListener.sendEarlyResponse(resultAction);
+                                            }
+                                        } finally {
+                                            // If we set the process state to STOPPING, we must stop
+                                            controller.setMode(ServiceController.Mode.NEVER);
+                                        }
                                     }
-                                }
 
-                                @Override
-                                public void transition(final ServiceController<? extends Object> controller, final ServiceController.Transition transition) {
-                                    if (transition == ServiceController.Transition.STOPPING_to_DOWN) {
-                                        controller.removeListener(this);
-                                        reloadContext.doReload(runningModeControl);
-                                        controller.setMode(ServiceController.Mode.ACTIVE);
+                                    @Override
+                                    public void transition(final ServiceController<? extends Object> controller, final ServiceController.Transition transition) {
+                                        if (transition == ServiceController.Transition.STOPPING_to_DOWN) {
+                                            controller.removeListener(this);
+                                            reloadContext.doReload(runningModeControl);
+                                            controller.setMode(ServiceController.Mode.ACTIVE);
+                                        }
                                     }
-                                }
-                            });
+                                });
+                            }
+                        } finally {
+                            if (sendListener != null) {
+                                // even if we called this in the try block, it's ok to call it again.
+                                sendListener.sendEarlyResponse(resultAction);
+                            }
                         }
                     }
                 });

controller/src/main/java/org/jboss/as/controller/remote/EarlyResponseSendListener.java

@@ -0,0 +1,47 @@
+/*
+Copyright 2018 Red Hat, 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 org.jboss.as.controller.remote;
+
+import org.jboss.as.controller.OperationContext;
+
+/**
+ * Callback that can be provided to operation step handlers for operations like 'reload' and 'shutdown'
+ * where the response needs to be sent to the caller before the operation completes. The handler will
+ * invoke the callback when the operation has reached the point where it is safe to send the response.
+ * <p>
+ * No callback should be attached before the operation is committed.
+ *
+ * @author Brian Stansberry
+ */
+public interface EarlyResponseSendListener {
+    /**
+     * Key under which a listener would be
+     * {@link OperationContext#attach(OperationContext.AttachmentKey, Object) attached to an operation context}
+     * if notification that it's safe
+     */
+    OperationContext.AttachmentKey<EarlyResponseSendListener> ATTACHMENT_KEY = OperationContext.AttachmentKey.create(EarlyResponseSendListener.class);
+
+    /**
+     * Informs the management kernel that it is ok to send an early response to the operation.
+     * <strong>Note:</strong> It is valid for this method to be invoked multiple times for the same
+     * listener. It is the responsibility of the listener implementation to ensure that only one
+     * response is sent to the caller.
+     *
+     * @param resultAction the result of the operation for which an early response is being sent
+     */
+    void sendEarlyResponse(OperationContext.ResultAction resultAction);
+}

controller/src/main/java/org/jboss/as/controller/remote/ModelControllerClientOperationHandler.java

@@ -48,8 +48,8 @@
 import java.net.InetSocketAddress;
 import java.security.PrivilegedActionException;
 import java.security.PrivilegedExceptionAction;
-import java.util.Collections;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
@@ -58,6 +58,7 @@
 
 import org.jboss.as.controller.AccessAuditContext;
 import org.jboss.as.controller.ModelController;
+import org.jboss.as.controller.OperationContext;
 import org.jboss.as.controller.PathAddress;
 import org.jboss.as.controller.client.OperationResponse;
 import org.jboss.as.controller.client.impl.ModelControllerProtocol;
@@ -200,9 +201,31 @@ private void doExecute(final ModelNode operation, final int attachmentsLength,
             // Send the prepared response for :reload operations
             final boolean sendPreparedOperation = sendPreparedResponse(operation);
             final ModelController.OperationTransactionControl transactionControl = sendPreparedOperation ? new ModelController.OperationTransactionControl() {
+
                 @Override
                 public void operationPrepared(ModelController.OperationTransaction transaction, ModelNode preparedResult) {
+                    // Shouldn't be called but if it is, send the result immediately
+                    operationPrepared(transaction, preparedResult, null);
+                }
+
+                @Override
+                public void operationPrepared(ModelController.OperationTransaction transaction, final ModelNode preparedResult, OperationContext context) {
                     transaction.commit();
+                    if (context == null || !RELOAD.equals(operation.get(OP).asString())) { // TODO deal with shutdown as well,
+                                                                                           // the handlers for which have some
+                                                                                           // subtleties that need thought
+                        sendResponse(preparedResult);
+                    } else {
+                        context.attach(EarlyResponseSendListener.ATTACHMENT_KEY, new EarlyResponseSendListener() {
+                            @Override
+                            public void sendEarlyResponse(OperationContext.ResultAction resultAction) {
+                                sendResponse(preparedResult);
+                            }
+                        });
+                    }
+                }
+
+                private void sendResponse(ModelNode preparedResult) {
                     // Fix prepared result
                     preparedResult.get(OUTCOME).set(SUCCESS);
                     preparedResult.get(RESULT);