[WFLY-1511] don't try to reconnect on shutdown

host-controller/src/main/java/org/jboss/as/host/controller/RemoteDomainConnection.java

@@ -22,6 +22,8 @@
 
 package org.jboss.as.host.controller;
 
+import javax.net.ssl.SSLContext;
+import javax.security.auth.callback.CallbackHandler;
 import java.io.DataInput;
 import java.io.IOException;
 import java.net.URI;
@@ -34,9 +36,6 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
-import javax.net.ssl.SSLContext;
-import javax.security.auth.callback.CallbackHandler;
-
 import org.jboss.as.controller.descriptions.ModelDescriptionConstants;
 import org.jboss.as.domain.controller.SlaveRegistrationException;
 import org.jboss.as.domain.management.CallbackHandlerFactory;
@@ -48,7 +47,6 @@
 import org.jboss.as.protocol.ProtocolConnectionConfiguration;
 import org.jboss.as.protocol.ProtocolConnectionManager;
 import org.jboss.as.protocol.ProtocolConnectionUtils;
-import org.jboss.as.protocol.ProtocolMessages;
 import org.jboss.as.protocol.StreamUtils;
 import org.jboss.as.protocol.mgmt.AbstractManagementRequest;
 import org.jboss.as.protocol.mgmt.ActiveOperation;
@@ -109,7 +107,6 @@ class RemoteDomainConnection extends FutureManagementChannel {
     private final ManagementPongRequestHandler pongHandler = new ManagementPongRequestHandler();
     private final List<DiscoveryOption> discoveryOptions;
     private URI uri;
-    private volatile boolean closing = false;
 
     RemoteDomainConnection(final String localHostName, final ModelNode localHostInfo,
                            final ProtocolChannelClient.Configuration configuration, final SecurityRealm realm,
@@ -151,22 +148,7 @@ protected ManagementChannelHandler getChannelHandler() {
 
     @Override
     public Channel getChannel() throws IOException {
-        Channel result;
-        if (!closing) {
-            // Normal case
-            result = awaitChannel();
-        } else {
-            // TODO WFLY-1511 this 'closing' stuff is just a quick and dirty fix; do it better
-            // We're closing so we don't need to wait for a channel if there isn't one
-            // If there isn't one the master will either be shut down itself or will detect the close
-            // of the existing channel
-            result = super.getChannel();
-            if (result == null) {
-                // Better than returning null with a subsequent NPE
-                throw ProtocolMessages.MESSAGES.channelClosed();
-            }
-        }
-        return result;
+        return awaitChannel();
     }
 
     /**
@@ -180,24 +162,20 @@ protected void setUri(URI uri) {
 
     @Override
     public void close() throws IOException {
-        synchronized (this) {
-            closing = true;
-            try {
-                if(isConnected()) {
-                    try {
-                        channelHandler.executeRequest(new UnregisterModelControllerRequest(), null).getResult().await();
-                    } catch (InterruptedException e) {
-                        Thread.currentThread().interrupt();
-                    }
-                }
-            } finally {
+        try {
+            if(prepareClose() && isConnected()) {
                 try {
-                    connectionManager.shutdown();
-                } finally {
-                    super.close();
-                    closing = false;
+                    channelHandler.executeRequest(new UnregisterModelControllerRequest(), null).getResult().await();
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
                 }
             }
+        } finally {
+            try {
+                super.close();
+            } finally {
+                connectionManager.shutdown();
+            }
         }
     }
 

host-controller/src/main/java/org/jboss/as/host/controller/RemoteDomainConnectionService.java

@@ -173,8 +173,8 @@ public class RemoteDomainConnectionService implements MasterDomainControllerClie
     private ExecutorService executor;
     private ScheduledExecutorService scheduledExecutorService;
 
-    private RemoteDomainConnection connection;
     private ManagementChannelHandler handler;
+    private volatile RemoteDomainConnection connection;
 
     private RemoteDomainConnectionService(final ModelController controller, final ExtensionRegistry extensionRegistry,
                                           final LocalHostControllerInfo localHostControllerInfo, final ProductConfig productConfig,

protocol/src/main/java/org/jboss/as/protocol/ProtocolConnectionManager.java

@@ -35,9 +35,9 @@
  */
 public class ProtocolConnectionManager {
 
-    private Connection connection;
     private ConnectTask connectTask;
-    private boolean shutdown;
+    private volatile boolean shutdown;
+    private volatile Connection connection;
 
     protected ProtocolConnectionManager(final ConnectTask initial) {
         if(initial == null) {

protocol/src/main/java/org/jboss/as/protocol/mgmt/FutureManagementChannel.java

@@ -43,12 +43,19 @@ public abstract class FutureManagementChannel extends ManagementClientChannelStr
 
     private final Object lock = new Object();
     private volatile Channel channel;
-    private volatile boolean closed;
+    private volatile State state = State.OPEN;
+
+    static enum State {
+        OPEN,
+        CLOSING,
+        CLOSED,
+        ;
+    }
 
     @Override
     public Channel getChannel() throws IOException {
         final Channel channel = this.channel;
-        if(channel == null && closed) {
+        if(channel == null && state != State.OPEN) {
             throw ProtocolMessages.MESSAGES.channelClosed();
         }
         return channel;
@@ -57,10 +64,10 @@ public Channel getChannel() throws IOException {
     @Override
     public void close() throws IOException {
         synchronized (lock) {
-            if(closed) {
+            if(state == State.CLOSED) {
                 return;
             }
-            closed = true;
+            state = State.CLOSED;
             StreamUtils.safeClose(channel);
             lock.notifyAll();
         }
@@ -72,7 +79,7 @@ public void close() throws IOException {
      * @return {@code true} if the connection is open, {@code false} otherwise
      */
     protected boolean isConnected() {
-        return channel != null && !closed;
+        return channel != null && state != State.CLOSED;
     }
 
     /**
@@ -88,13 +95,16 @@ protected Channel awaitChannel() throws IOException {
         }
         synchronized (lock) {
             for(;;) {
-                if(closed) {
+                if(state == State.CLOSED) {
                     throw ProtocolMessages.MESSAGES.channelClosed();
                 }
                 channel = this.channel;
                 if(channel != null) {
                     return channel;
                 }
+                if(state == State.CLOSING) {
+                    throw ProtocolMessages.MESSAGES.channelClosed();
+                }
                 try {
                     lock.wait();
                 } catch (InterruptedException e) {
@@ -104,6 +114,24 @@ protected Channel awaitChannel() throws IOException {
         }
     }
 
+    /**
+     * Signal that we are about to close the channel. This will not have any affect on the underlying channel, however
+     * prevent setting a new channel.
+     *
+     * @return whether the closing state was set successfully
+     */
+    protected boolean prepareClose() {
+        synchronized (lock) {
+            final State state = this.state;
+            if (state == State.OPEN) {
+                this.state = State.CLOSING;
+                lock.notifyAll();
+                return true;
+            }
+        }
+        return false;
+    }
+
     /**
      * Open a channel.
      *
@@ -129,7 +157,7 @@ protected boolean setChannel(final Channel newChannel) {
             return false;
         }
         synchronized (lock) {
-            if(closed || channel != null) {
+            if(state != State.OPEN || channel != null) {
                 return false;
             }
             this.channel = newChannel;
@@ -195,9 +223,9 @@ public void connectionOpened(final Connection connection) throws IOException {
         @Override
         public void close() throws IOException {
             try {
-                connectionManager.shutdown();
-            } finally {
                 super.close();
+            } finally {
+                connectionManager.shutdown();
             }
         }
 

server/src/main/java/org/jboss/as/server/mgmt/domain/HostControllerConnection.java

@@ -209,9 +209,9 @@ public void connectionOpened(final Connection connection) throws IOException {
     @Override
     public void close() throws IOException {
         try {
-            connectionManager.shutdown();
-        } finally {
             super.close();
+        } finally {
+            connectionManager.shutdown();
         }
     }
 

testsuite/domain/src/test/java/org/jboss/as/test/integration/respawn/TestControllerClient.java

@@ -83,9 +83,9 @@ public void connectionOpened(final Connection connection) throws IOException {
         @Override
         public void close() throws IOException {
             try {
-                connectionManager.shutdown();
-            } finally {
                 super.close();
+            } finally {
+                connectionManager.shutdown();
             }
         }
     }