Added CDL so main thread waits for subscribed messages to complete be… (#114)

* Added CDL so main thread waits for subscribed messages to complete before closing connection

Co-authored-by: Kevin Weber <[email protected]>

Author: webermaster

samples/BasicPubSub/src/main/java/pubsub/PubSub.java

@@ -22,8 +22,10 @@
 import software.amazon.awssdk.iot.iotjobs.model.RejectedError;
 
 import java.io.UnsupportedEncodingException;
+import java.nio.charset.StandardCharsets;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 
 class PubSub {
@@ -291,13 +293,12 @@ public void onConnectionResumed(boolean sessionPresent) {
                     throw new RuntimeException("Exception occurred during connect", ex);
                 }
 
+                CountDownLatch countDownLatch = new CountDownLatch(messagesToPublish);
+
                 CompletableFuture<Integer> subscribed = connection.subscribe(topic, QualityOfService.AT_LEAST_ONCE, (message) -> {
-                    try {
-                        String payload = new String(message.getPayload(), "UTF-8");
-                        System.out.println("MESSAGE: " + payload);
-                    } catch (UnsupportedEncodingException ex) {
-                        System.out.println("Unable to decode payload: " + ex.getMessage());
-                    }
+                    String payload = new String(message.getPayload(), StandardCharsets.UTF_8);
+                    System.out.println("MESSAGE: " + payload);
+                    countDownLatch.countDown();
                 });
 
                 subscribed.get();
@@ -308,6 +309,8 @@ public void onConnectionResumed(boolean sessionPresent) {
                     published.get();
                     Thread.sleep(1000);
                 }
+                
+                countDownLatch.await();
 
                 CompletableFuture<Void> disconnected = connection.disconnect();
                 disconnected.get();

samples/RawPubSub/src/main/java/rawpubsub/RawPubSub.java

@@ -16,10 +16,12 @@
 import software.amazon.awssdk.iot.iotjobs.model.RejectedError;
 
 import java.io.UnsupportedEncodingException;
+import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import java.util.List;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 
 class RawPubSub {
@@ -216,13 +218,12 @@ public void onConnectionResumed(boolean sessionPresent) {
                         throw new RuntimeException("Exception occurred during connect", ex);
                     }
 
+                    CountDownLatch countDownLatch = new CountDownLatch(messagesToPublish);
+
                     CompletableFuture<Integer> subscribed = connection.subscribe(topic, QualityOfService.AT_LEAST_ONCE, (message) -> {
-                        try {
-                            String payload = new String(message.getPayload(), "UTF-8");
-                            System.out.println("MESSAGE: " + payload);
-                        } catch (UnsupportedEncodingException ex) {
-                            System.out.println("Unable to decode payload: " + ex.getMessage());
-                        }
+                        String payload = new String(message.getPayload(), StandardCharsets.UTF_8);
+                        System.out.println("MESSAGE: " + payload);
+                        countDownLatch.countDown();
                     });
 
                     subscribed.get();
@@ -233,6 +234,8 @@ public void onConnectionResumed(boolean sessionPresent) {
                         published.get();
                         Thread.sleep(1000);
                     }
+                    
+                    countDownLatch.await();
 
                     CompletableFuture<Void> disconnected = connection.disconnect();
                     disconnected.get();