Exceptions of async operations are not wrapped in CompletionEx anymore

Author: SgtSilvio

src/main/java/com/hivemq/client/internal/mqtt/MqttAsyncClient.java

@@ -46,16 +46,54 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
-import java.util.function.Function;
 
 /**
  * @author Silvio Giebl
  */
 public class MqttAsyncClient implements Mqtt5AsyncClient {
 
-    private static final @NotNull Function<Mqtt5SubAck, Mqtt5SubAck> SUBACK_HANDLER = MqttBlockingClient::handleSubAck;
-    private static final @NotNull Function<Mqtt5UnsubAck, Mqtt5UnsubAck> UNSUBACK_HANDLER =
-            MqttBlockingClient::handleUnsubAck;
+    private static @NotNull CompletableFuture<@NotNull Mqtt5SubAck> handleSubAck(
+            final @NotNull CompletableFuture<@NotNull Mqtt5SubAck> future, final @NotNull MqttSubscribe subscribe) {
+
+        if (subscribe.getSubscriptions().size() == 1) {
+            return future;
+        }
+        final CompletableFuture<Mqtt5SubAck> mappedFuture = new CompletableFuture<>();
+        future.whenComplete((subAck, throwable) -> {
+            if (throwable != null) {
+                mappedFuture.completeExceptionally(throwable);
+            } else {
+                try {
+                    mappedFuture.complete(MqttBlockingClient.handleSubAck(subAck));
+                } catch (final Throwable t) {
+                    mappedFuture.completeExceptionally(t);
+                }
+            }
+        });
+        return mappedFuture;
+    }
+
+    private static @NotNull CompletableFuture<@NotNull Mqtt5UnsubAck> handleUnsubAck(
+            final @NotNull CompletableFuture<@NotNull Mqtt5UnsubAck> future,
+            final @NotNull MqttUnsubscribe unsubscribe) {
+
+        if (unsubscribe.getTopicFilters().size() == 1) {
+            return future;
+        }
+        final CompletableFuture<Mqtt5UnsubAck> mappedFuture = new CompletableFuture<>();
+        future.whenComplete((unsubAck, throwable) -> {
+            if (throwable != null) {
+                mappedFuture.completeExceptionally(throwable);
+            } else {
+                try {
+                    mappedFuture.complete(MqttBlockingClient.handleUnsubAck(unsubAck));
+                } catch (final Throwable t) {
+                    mappedFuture.completeExceptionally(t);
+                }
+            }
+        });
+        return mappedFuture;
+    }
 
     private final @NotNull MqttRxClient delegate;
 
@@ -74,7 +112,7 @@ public class MqttAsyncClient implements Mqtt5AsyncClient {
     public @NotNull CompletableFuture<@NotNull Mqtt5SubAck> subscribe(final @Nullable Mqtt5Subscribe subscribe) {
         final MqttSubscribe mqttSubscribe = MqttChecks.subscribe(subscribe);
 
-        return RxFutureConverter.toFuture(delegate.subscribe(mqttSubscribe)).thenApply(SUBACK_HANDLER);
+        return handleSubAck(RxFutureConverter.toFuture(delegate.subscribe(mqttSubscribe)), mqttSubscribe);
     }
 
     @Override
@@ -84,9 +122,9 @@ public class MqttAsyncClient implements Mqtt5AsyncClient {
         final MqttSubscribe mqttSubscribe = MqttChecks.subscribe(subscribe);
         Checks.notNull(callback, "Callback");
 
-        return delegate.subscribeStream(mqttSubscribe)
-                .subscribeSingleFuture(new CallbackSubscriber(callback))
-                .thenApply(SUBACK_HANDLER);
+        return handleSubAck(
+                delegate.subscribeStream(mqttSubscribe).subscribeSingleFuture(new CallbackSubscriber(callback)),
+                mqttSubscribe);
     }
 
     @Override
@@ -98,10 +136,10 @@ public class MqttAsyncClient implements Mqtt5AsyncClient {
         Checks.notNull(callback, "Callback");
         Checks.notNull(executor, "Executor");
 
-        return delegate.subscribeStreamUnsafe(mqttSubscribe)
-                .observeOnBoth(Schedulers.from(executor), true)
-                .subscribeSingleFuture(new CallbackSubscriber(callback))
-                .thenApply(SUBACK_HANDLER);
+        return handleSubAck(
+                delegate.subscribeStreamUnsafe(mqttSubscribe)
+                .observeOnBoth(Schedulers.from(executor), true).subscribeSingleFuture(new CallbackSubscriber(callback)),
+                mqttSubscribe);
     }
 
     @Override
@@ -134,7 +172,7 @@ public void publishes(
 
         final MqttUnsubscribe mqttUnsubscribe = MqttChecks.unsubscribe(unsubscribe);
 
-        return RxFutureConverter.toFuture(delegate.unsubscribe(mqttUnsubscribe)).thenApply(UNSUBACK_HANDLER);
+        return handleUnsubAck(RxFutureConverter.toFuture(delegate.unsubscribe(mqttUnsubscribe)), mqttUnsubscribe);
     }
 
     @Override

src/main/java/com/hivemq/client/internal/mqtt/mqtt3/Mqtt3AsyncClientView.java

@@ -41,61 +41,33 @@
 import com.hivemq.client.mqtt.mqtt3.message.subscribe.Mqtt3Subscribe;
 import com.hivemq.client.mqtt.mqtt3.message.subscribe.suback.Mqtt3SubAck;
 import com.hivemq.client.mqtt.mqtt3.message.unsubscribe.Mqtt3Unsubscribe;
-import com.hivemq.client.mqtt.mqtt5.message.connect.connack.Mqtt5ConnAck;
 import com.hivemq.client.mqtt.mqtt5.message.publish.Mqtt5Publish;
-import com.hivemq.client.mqtt.mqtt5.message.publish.Mqtt5PublishResult;
 import com.hivemq.client.mqtt.mqtt5.message.subscribe.suback.Mqtt5SubAck;
-import com.hivemq.client.mqtt.mqtt5.message.unsubscribe.unsuback.Mqtt5UnsubAck;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.CompletionException;
 import java.util.concurrent.Executor;
-import java.util.function.BiFunction;
 import java.util.function.Consumer;
-import java.util.function.Function;
 
 /**
  * @author Silvio Giebl
  */
 public class Mqtt3AsyncClientView implements Mqtt3AsyncClient {
 
-    private static final @NotNull BiFunction<Mqtt5ConnAck, Throwable, Mqtt3ConnAck> CONNACK_MAPPER =
-            (connAck, throwable) -> {
-                if (throwable != null) {
-                    throw new CompletionException(Mqtt3ExceptionFactory.map(throwable));
-                }
-                return Mqtt3ConnAckView.of(connAck);
-            };
-
-    private static final @NotNull BiFunction<Mqtt5SubAck, Throwable, Mqtt3SubAck> SUBACK_MAPPER =
-            (subAck, throwable) -> {
-                if (throwable != null) {
-                    throw new CompletionException(Mqtt3ExceptionFactory.map(throwable));
-                }
-                return Mqtt3SubAckView.of(subAck);
-            };
-
-    private static final @NotNull BiFunction<Mqtt5UnsubAck, Throwable, Void> UNSUBACK_MAPPER =
-            (unsubAck, throwable) -> {
-                if (throwable != null) {
-                    throw new CompletionException(Mqtt3ExceptionFactory.map(throwable));
-                }
-                return null;
-            };
-
-    private static final @NotNull BiFunction<Mqtt5PublishResult, Throwable, Mqtt3Publish> PUBLISH_RESULT_MAPPER =
-            (publishResult, throwable) -> {
-                if (throwable != null) {
-                    throw new CompletionException(Mqtt3ExceptionFactory.map(throwable));
-                }
-                return Mqtt3PublishView.of(publishResult.getPublish());
-            };
+    private static @NotNull CompletableFuture<@NotNull Mqtt3SubAck> handleSubAck(
+            final @NotNull CompletableFuture<@NotNull Mqtt5SubAck> future) {
 
-    private static final @NotNull Function<Throwable, Void> DISCONNECT_MAPPER = throwable -> {
-        throw new CompletionException(Mqtt3ExceptionFactory.map(throwable));
-    };
+        final CompletableFuture<Mqtt3SubAck> mappedFuture = new CompletableFuture<>();
+        future.whenComplete((subAck, throwable) -> {
+            if (throwable != null) {
+                mappedFuture.completeExceptionally(Mqtt3ExceptionFactory.map(throwable));
+            } else {
+                mappedFuture.complete(Mqtt3SubAckView.of(subAck));
+            }
+        });
+        return mappedFuture;
+    }
 
     private static @NotNull Consumer<Mqtt5Publish> callbackView(final @NotNull Consumer<Mqtt3Publish> callback) {
         return publish -> callback.accept(Mqtt3PublishView.of(publish));
@@ -113,14 +85,22 @@ public class Mqtt3AsyncClientView implements Mqtt3AsyncClient {
     public @NotNull CompletableFuture<@NotNull Mqtt3ConnAck> connect(final @Nullable Mqtt3Connect connect) {
         final MqttConnect mqttConnect = MqttChecks.connect(connect);
 
-        return delegate.connect(mqttConnect).handle(CONNACK_MAPPER);
+        final CompletableFuture<Mqtt3ConnAck> future = new CompletableFuture<>();
+        delegate.connect(mqttConnect).whenComplete((connAck, throwable) -> {
+            if (throwable != null) {
+                future.completeExceptionally(Mqtt3ExceptionFactory.map(throwable));
+            } else {
+                future.complete(Mqtt3ConnAckView.of(connAck));
+            }
+        });
+        return future;
     }
 
     @Override
     public @NotNull CompletableFuture<@NotNull Mqtt3SubAck> subscribe(final @Nullable Mqtt3Subscribe subscribe) {
         final MqttSubscribe mqttSubscribe = MqttChecks.subscribe(subscribe);
 
-        return delegate.subscribe(mqttSubscribe).handle(SUBACK_MAPPER);
+        return handleSubAck(delegate.subscribe(mqttSubscribe));
     }
 
     @Override
@@ -130,7 +110,7 @@ public class Mqtt3AsyncClientView implements Mqtt3AsyncClient {
         final MqttSubscribe mqttSubscribe = MqttChecks.subscribe(subscribe);
         Checks.notNull(callback, "Callback");
 
-        return delegate.subscribe(mqttSubscribe, callbackView(callback)).handle(SUBACK_MAPPER);
+        return handleSubAck(delegate.subscribe(mqttSubscribe, callbackView(callback)));
     }
 
     @Override
@@ -142,7 +122,7 @@ public class Mqtt3AsyncClientView implements Mqtt3AsyncClient {
         Checks.notNull(callback, "Callback");
         Checks.notNull(executor, "Executor");
 
-        return delegate.subscribe(mqttSubscribe, callbackView(callback), executor).handle(SUBACK_MAPPER);
+        return handleSubAck(delegate.subscribe(mqttSubscribe, callbackView(callback), executor));
     }
 
     @Override
@@ -171,19 +151,43 @@ public void publishes(
     public @NotNull CompletableFuture<Void> unsubscribe(final @Nullable Mqtt3Unsubscribe unsubscribe) {
         final MqttUnsubscribe mqttUnsubscribe = MqttChecks.unsubscribe(unsubscribe);
 
-        return delegate.unsubscribe(mqttUnsubscribe).handle(UNSUBACK_MAPPER);
+        final CompletableFuture<Void> future = new CompletableFuture<>();
+        delegate.unsubscribe(mqttUnsubscribe).whenComplete((unsubAck, throwable) -> {
+            if (throwable != null) {
+                future.completeExceptionally(Mqtt3ExceptionFactory.map(throwable));
+            } else {
+                future.complete(null);
+            }
+        });
+        return future;
     }
 
     @Override
     public @NotNull CompletableFuture<@NotNull Mqtt3Publish> publish(final @Nullable Mqtt3Publish publish) {
         final MqttPublish mqttPublish = MqttChecks.publish(publish);
 
-        return delegate.publish(mqttPublish).handle(PUBLISH_RESULT_MAPPER);
+        final CompletableFuture<Mqtt3Publish> future = new CompletableFuture<>();
+        delegate.publish(mqttPublish).whenComplete((publishResult, throwable) -> {
+            if (throwable != null) {
+                future.completeExceptionally(Mqtt3ExceptionFactory.map(throwable));
+            } else {
+                future.complete(Mqtt3PublishView.of(publishResult.getPublish()));
+            }
+        });
+        return future;
     }
 
     @Override
     public @NotNull CompletableFuture<Void> disconnect() {
-        return delegate.disconnect(Mqtt3DisconnectView.DELEGATE).exceptionally(DISCONNECT_MAPPER);
+        final CompletableFuture<Void> future = new CompletableFuture<>();
+        delegate.disconnect(Mqtt3DisconnectView.DELEGATE).whenComplete((ignored, throwable) -> {
+            if (throwable != null) {
+                future.completeExceptionally(Mqtt3ExceptionFactory.map(throwable));
+            } else {
+                future.complete(null);
+            }
+        });
+        return future;
     }
 
     @Override