fix(uat): bypass send DISCONNECT on dead connection in SDK client (#376)

aws-greengrass · Jul 24, 2023 · 6d366ab · 6d366ab
1 parent bb149d0
commit 6d366ab
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 14 deletions.
diff --git a/...rc/main/java/com/aws/greengrass/testing/mqtt311/client/sdkmqtt/Mqtt311ConnectionImpl.java b/...rc/main/java/com/aws/greengrass/testing/mqtt311/client/sdkmqtt/Mqtt311ConnectionImpl.java
@@ -211,12 +211,17 @@ public ConnectResult start(long timeout, int connectionId) throws MqttException
     public void disconnect(long timeout, int reasonCode, List<Mqtt5Properties> userProperties) throws MqttException {
 
         checkUserProperties(userProperties);
-        if (!isClosing.getAndSet(true)) {
-            CompletableFuture<Void> disconnnectFuture = connection.disconnect();
+
+        if (isClosing.compareAndSet(false, true)) {
             try {
                 final long deadline = System.nanoTime() + timeout * 1_000_000_000;
 
-                disconnnectFuture.get(timeout, TimeUnit.SECONDS);
+                if (isConnected.compareAndSet(true, false)) {
+                    CompletableFuture<Void> disconnnectFuture = connection.disconnect();
+                    disconnnectFuture.get(timeout, TimeUnit.SECONDS);
+                } else {
+                    logger.atWarn().log("DISCONNECT was not sent on the dead connection");
+                }
 
                 long remaining = deadline - System.nanoTime();
                 if (remaining < MIN_SHUTDOWN_NS) {

diff --git a/...sdk/src/main/java/com/aws/greengrass/testing/mqtt5/client/sdkmqtt/MqttConnectionImpl.java b/...sdk/src/main/java/com/aws/greengrass/testing/mqtt5/client/sdkmqtt/MqttConnectionImpl.java
@@ -277,23 +277,28 @@ public ConnectResult start(long timeout, int connectionId) throws MqttException
         }
     }
 
-    @SuppressWarnings({"PMD.UseTryWithResources", "PMD.AvoidCatchingGenericException"})
+    @SuppressWarnings({"PMD.UseTryWithResources", "PMD.AvoidCatchingGenericException", "PMD.CognitiveComplexity"})
     @Override
     public void disconnect(long timeout, int reasonCode, List<Mqtt5Properties> userProperties) throws MqttException {
 
         if (isClosing.compareAndSet(false, true)) {
-            final DisconnectPacket.DisconnectReasonCode disconnectReason
-                    = DisconnectPacket.DisconnectReasonCode.getEnumValueFromInteger(reasonCode);
-            DisconnectPacket.DisconnectPacketBuilder builder = new DisconnectPacket.DisconnectPacketBuilder()
-                    .withReasonCode(disconnectReason);
+            if (isConnected.compareAndSet(true, false)) {
+                final DisconnectPacket.DisconnectReasonCode disconnectReason
+                        = DisconnectPacket.DisconnectReasonCode.getEnumValueFromInteger(reasonCode);
+                DisconnectPacket.DisconnectPacketBuilder builder = new DisconnectPacket.DisconnectPacketBuilder()
+                        .withReasonCode(disconnectReason);
+
+                if (userProperties != null && !userProperties.isEmpty()) {
+                    builder.withUserProperties(convertToUserProperties(userProperties, "DISCONNECT"));
+                }
 
-            if (userProperties != null && !userProperties.isEmpty()) {
-                builder.withUserProperties(convertToUserProperties(userProperties, "DISCONNECT"));
+                client.stop(builder.build());
+            } else {
+                logger.atWarn().log("DISCONNECT was not sent on the dead connection");
             }
 
-            client.stop(builder.build());
-
             try {
+                client.close();
                 final long deadline = System.nanoTime() + timeout * 1_000_000_000;
                 lifecycleEvents.stoppedFuture.get(timeout, TimeUnit.SECONDS);
 
@@ -312,8 +317,6 @@ public void disconnect(long timeout, int reasonCode, List<Mqtt5Properties> userP
             } catch (Exception ex) {
                 logger.atError().withThrowable(ex).log("Failed during disconnecting from MQTT broker");
                 throw new MqttException("Could not disconnect", ex);
-            } finally {
-                client.close();
             }
         }
     }

diff --git a/...est/java/com/aws/greengrass/testing/mqtt311/client/sdkmqtt/Mqtt311ConnectionImplTest.java b/...est/java/com/aws/greengrass/testing/mqtt311/client/sdkmqtt/Mqtt311ConnectionImplTest.java
@@ -135,6 +135,9 @@ void GIVEN_disconnect_successful_WHEN_disconnect_THEN_client_methods_are_called(
         when(connection.disconnect()).thenReturn(connectFuture);
         connectFuture.complete(null);
 
+        // move to connected state
+        mqttConnectionImpl.connectionEvents.onConnectionResumed(true);
+
         // WHEN
         mqttConnectionImpl.disconnect(timeoutSeconds, reasonCode, null);
 
@@ -152,6 +155,9 @@ void GIVEN_timedout_WHEN_disconnect_THEN_exception() throws MqttException {
         final CompletableFuture<Void> connectFuture = new CompletableFuture<>();
         when(connection.disconnect()).thenReturn(connectFuture);
 
+        // move to connected state
+        mqttConnectionImpl.connectionEvents.onConnectionResumed(true);
+
         // WHEN, THEN
         assertThrows(MqttException.class, () -> {
             mqttConnectionImpl.disconnect(timeoutSeconds, reasonCode, null);

diff --git a/...src/test/java/com/aws/greengrass/testing/mqtt5/client/sdkmqtt/MqttConnectionImplTest.java b/...src/test/java/com/aws/greengrass/testing/mqtt5/client/sdkmqtt/MqttConnectionImplTest.java
@@ -254,6 +254,9 @@ void GIVEN_disconnect_successful_WHEN_disconnect_THEN_client_methods_are_called(
 
         mqttConnectionImpl.lifecycleEvents.onStopped(client, onStoppedReturn);
 
+        // move to connected state
+        mqttConnectionImpl.lifecycleEvents.onConnectionSuccess(client, null);
+
         // WHEN
         mqttConnectionImpl.disconnect(timeoutSeconds, reasonCode, null);
 
@@ -273,6 +276,9 @@ void GIVEN_timedout_WHEN_disconnect_THEN_exception() throws MqttException {
         final long timeoutSeconds = SHORT_TIMEOUT_SEC;
         final int reasonCode = 4;
 
+        // move to connected state
+        mqttConnectionImpl.lifecycleEvents.onConnectionSuccess(client, null);
+
         // WHEN
         assertThrows(MqttException.class, () -> {
             mqttConnectionImpl.disconnect(timeoutSeconds, reasonCode, null);