Add CompletableFuture and CompletionStage accessors to ReactiveResponseConsumer

httpcore5-reactive/src/main/java/org/apache/hc/core5/reactive/ReactiveResponseConsumer.java

@@ -31,6 +31,8 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
 import java.util.concurrent.Future;
 
 import org.apache.hc.core5.annotation.Contract;
@@ -62,6 +64,8 @@ public final class ReactiveResponseConsumer implements AsyncResponseConsumer<Voi
     private final ReactiveDataConsumer reactiveDataConsumer = new ReactiveDataConsumer();
     private final List<Header> trailers = Collections.synchronizedList(new ArrayList<>());
     private final BasicFuture<Message<HttpResponse, Publisher<ByteBuffer>>> responseFuture;
+    private final CompletableFuture<Message<HttpResponse, Publisher<ByteBuffer>>> responseCompletableFuture;
+    private final CompletableFuture<Void> responseCompletionFuture;
 
     private volatile BasicFuture<Void> responseCompletion;
     private volatile HttpResponse informationResponse;
@@ -72,6 +76,8 @@ public final class ReactiveResponseConsumer implements AsyncResponseConsumer<Voi
      */
     public ReactiveResponseConsumer() {
         this.responseFuture = new BasicFuture<>(null);
+        this.responseCompletableFuture = new CompletableFuture<>();
+        this.responseCompletionFuture = new CompletableFuture<>();
     }
 
     /**
@@ -82,12 +88,54 @@ public ReactiveResponseConsumer() {
      */
     public ReactiveResponseConsumer(final FutureCallback<Message<HttpResponse, Publisher<ByteBuffer>>> responseCallback) {
         this.responseFuture = new BasicFuture<>(Args.notNull(responseCallback, "responseCallback"));
+        this.responseCompletableFuture = new CompletableFuture<>();
+        this.responseCompletionFuture = new CompletableFuture<>();
     }
 
     public Future<Message<HttpResponse, Publisher<ByteBuffer>>> getResponseFuture() {
         return responseFuture;
     }
 
+    /**
+     * Returns a {@link CompletableFuture} that completes when the response head and body {@link Publisher}
+     * are available.
+     *
+     * @since 5.5
+     */
+    public CompletableFuture<Message<HttpResponse, Publisher<ByteBuffer>>> getResponseCompletableFuture() {
+        return responseCompletableFuture;
+    }
+
+    /**
+     * Returns a {@link CompletableFuture} that completes when the response exchange is complete
+     * (end-of-stream reached and trailers processed, if any).
+     *
+     * @since 5.5
+     */
+    public CompletableFuture<Void> getResponseCompletionFuture() {
+        return responseCompletionFuture;
+    }
+
+    /**
+     * Returns a {@link CompletionStage} that completes when the response head and body {@link Publisher}
+     * are available.
+     *
+     * @since 5.5
+     */
+    public CompletionStage<Message<HttpResponse, Publisher<ByteBuffer>>> getResponseStage() {
+        return responseCompletableFuture;
+    }
+
+    /**
+     * Returns a {@link CompletionStage} that completes when the response exchange is complete
+     * (end-of-stream reached and trailers processed, if any).
+     *
+     * @since 5.5
+     */
+    public CompletionStage<Void> getResponseCompletionStage() {
+        return responseCompletionFuture;
+    }
+
     /**
      * Returns the intermediate (1xx) HTTP response if one was received.
      *
@@ -117,14 +165,18 @@ public List<Header> getTrailers() {
 
     @Override
     public void consumeResponse(
-        final HttpResponse response,
-        final EntityDetails entityDetails,
-        final HttpContext httpContext,
-        final FutureCallback<Void> resultCallback
+            final HttpResponse response,
+            final EntityDetails entityDetails,
+            final HttpContext httpContext,
+            final FutureCallback<Void> resultCallback
     ) {
         this.entityDetails = entityDetails;
         this.responseCompletion = new BasicFuture<>(resultCallback);
-        this.responseFuture.completed(new Message<>(response, reactiveDataConsumer));
+
+        final Message<HttpResponse, Publisher<ByteBuffer>> message = new Message<>(response, reactiveDataConsumer);
+        this.responseFuture.completed(message);
+        this.responseCompletableFuture.complete(message);
+
         if (entityDetails == null) {
             streamEnd(null);
         }
@@ -138,9 +190,15 @@ public void informationResponse(final HttpResponse response, final HttpContext h
     @Override
     public void failed(final Exception cause) {
         reactiveDataConsumer.failed(cause);
+
+        // Complete stage/futures regardless of whether consumeResponse() has been invoked yet.
         responseFuture.failed(cause);
-        if (responseCompletion != null) {
-            responseCompletion.failed(cause);
+        responseCompletableFuture.completeExceptionally(cause);
+        responseCompletionFuture.completeExceptionally(cause);
+
+        final BasicFuture<Void> completion = responseCompletion;
+        if (completion != null) {
+            completion.failed(cause);
         }
     }
 
@@ -160,15 +218,32 @@ public void streamEnd(final List<? extends Header> trailers) {
             this.trailers.addAll(trailers);
         }
         reactiveDataConsumer.streamEnd(trailers);
-        responseCompletion.completed(null);
+
+        // Complete CF before BasicFuture.completed(...) (it may trigger releaseResources()).
+        responseCompletionFuture.complete(null);
+
+        final BasicFuture<Void> completion = responseCompletion;
+        if (completion != null) {
+            completion.completed(null);
+        }
     }
 
     @Override
     public void releaseResources() {
         reactiveDataConsumer.releaseResources();
+
         responseFuture.cancel();
-        if (responseCompletion != null) {
-            responseCompletion.cancel();
+
+        if (!responseCompletableFuture.isDone()) {
+            responseCompletableFuture.cancel(true);
+        }
+        if (!responseCompletionFuture.isDone()) {
+            responseCompletionFuture.cancel(true);
+        }
+
+        final BasicFuture<Void> completion = responseCompletion;
+        if (completion != null) {
+            completion.cancel();
         }
     }
 }

httpcore5-reactive/src/test/java/org/apache/hc/core5/reactive/examples/ReactiveClientCompletionStageExample.java

@@ -0,0 +1,237 @@
+/*
+ * ====================================================================
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ * ====================================================================
+ *
+ * This software consists of voluntary contributions made by many
+ * individuals on behalf of the Apache Software Foundation.  For more
+ * information on the Apache Software Foundation, please see
+ * <http://www.apache.org/>.
+ *
+ */
+package org.apache.hc.core5.reactive.examples;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+import java.net.InetAddress;
+import java.net.URI;
+import java.nio.ByteBuffer;
+import java.util.Random;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import io.reactivex.rxjava3.core.Flowable;
+import io.reactivex.rxjava3.core.Observable;
+import org.apache.hc.core5.http.ContentType;
+import org.apache.hc.core5.http.Header;
+import org.apache.hc.core5.http.HttpConnection;
+import org.apache.hc.core5.http.HttpRequest;
+import org.apache.hc.core5.http.HttpResponse;
+import org.apache.hc.core5.http.impl.Http1StreamListener;
+import org.apache.hc.core5.http.impl.bootstrap.AsyncRequesterBootstrap;
+import org.apache.hc.core5.http.impl.bootstrap.HttpAsyncRequester;
+import org.apache.hc.core5.http.message.RequestLine;
+import org.apache.hc.core5.http.message.StatusLine;
+import org.apache.hc.core5.http.nio.AsyncRequestProducer;
+import org.apache.hc.core5.http.nio.support.AsyncRequestBuilder;
+import org.apache.hc.core5.io.CloseMode;
+import org.apache.hc.core5.reactive.ReactiveEntityProducer;
+import org.apache.hc.core5.reactive.ReactiveResponseConsumer;
+import org.apache.hc.core5.reactor.IOReactorConfig;
+import org.apache.hc.core5.util.Timeout;
+import org.reactivestreams.Publisher;
+
+/**
+ * Client demo using CompletionStage accessors on ReactiveResponseConsumer (Java 8).
+ */
+public final class ReactiveClientCompletionStageExample {
+
+    private ReactiveClientCompletionStageExample() {
+    }
+
+    private static <T> CompletionStage<T> withTimeout(
+            final CompletableFuture<T> future,
+            final ScheduledExecutorService scheduler,
+            final long timeout,
+            final TimeUnit unit) {
+
+        final CompletableFuture<T> timeoutFuture = new CompletableFuture<>();
+        final java.util.concurrent.ScheduledFuture<?> task = scheduler.schedule(
+                () -> timeoutFuture.completeExceptionally(new TimeoutException("Timeout after " + timeout + " " + unit)),
+                timeout, unit);
+
+        final CompletableFuture<T> combined = future.applyToEither(timeoutFuture, t -> t);
+        combined.whenComplete((v, ex) -> task.cancel(false));
+        return combined;
+    }
+
+    private static CompletableFuture<String> readBodyAsString(final Publisher<ByteBuffer> publisher) {
+        final CompletableFuture<String> bodyFuture = new CompletableFuture<>();
+
+        Observable.fromPublisher(publisher)
+                .map(buf -> {
+                    final byte[] bytes = new byte[buf.remaining()];
+                    buf.get(bytes);
+                    return new String(bytes, UTF_8);
+                })
+                .reduce(new StringBuilder(), StringBuilder::append)
+                .map(StringBuilder::toString)
+                .subscribe(
+                        bodyFuture::complete,
+                        bodyFuture::completeExceptionally);
+
+        return bodyFuture;
+    }
+
+    private static boolean isLoopbackHost(final String host) {
+        if (host == null) {
+            return false;
+        }
+        return "localhost".equalsIgnoreCase(host) || "127.0.0.1".equals(host) || "::1".equals(host);
+    }
+
+    private static URI normalizeAuthority(final URI uri) {
+        if (!isLoopbackHost(uri.getHost())) {
+            return uri;
+        }
+        try {
+            final InetAddress local = InetAddress.getLocalHost();
+            final String canonical = local.getCanonicalHostName();
+            if (canonical != null && !canonical.isEmpty() && !isLoopbackHost(canonical)) {
+                return new URI(uri.getScheme(), null, canonical, uri.getPort(), uri.getPath(), uri.getQuery(), uri.getFragment());
+            }
+            final String addr = local.getHostAddress();
+            if (addr != null && !addr.isEmpty() && !isLoopbackHost(addr)) {
+                return new URI(uri.getScheme(), null, addr, uri.getPort(), uri.getPath(), uri.getQuery(), uri.getFragment());
+            }
+        } catch (final Exception ignore) {
+            // fall back
+        }
+        return uri;
+    }
+
+    public static void main(final String[] args) throws Exception {
+        String endpoint = "http://localhost:8080/echo";
+        if (args.length >= 1) {
+            endpoint = args[0];
+        }
+
+        final HttpAsyncRequester requester = AsyncRequesterBootstrap.bootstrap()
+                .setIOReactorConfig(IOReactorConfig.custom().setSoTimeout(5, TimeUnit.SECONDS).build())
+                .setStreamListener(new Http1StreamListener() {
+
+                    @Override
+                    public void onRequestHead(final HttpConnection connection, final HttpRequest request) {
+                        System.out.println(connection.getRemoteAddress() + " " + new RequestLine(request));
+                    }
+
+                    @Override
+                    public void onResponseHead(final HttpConnection connection, final HttpResponse response) {
+                        System.out.println(connection.getRemoteAddress() + " " + new StatusLine(response));
+                    }
+
+                    @Override
+                    public void onExchangeComplete(final HttpConnection connection, final boolean keepAlive) {
+                        if (keepAlive) {
+                            System.out.println(connection.getRemoteAddress()
+                                    + " exchange completed (connection kept alive)");
+                        } else {
+                            System.out.println(connection.getRemoteAddress()
+                                    + " exchange completed (connection closed)");
+                        }
+                    }
+
+                })
+                .create();
+
+        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
+            System.out.println("HTTP requester shutting down");
+            requester.close(CloseMode.GRACEFUL);
+        }));
+
+        requester.start();
+
+        final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(r -> {
+            final Thread t = new Thread(r, "timeout-scheduler");
+            t.setDaemon(true);
+            return t;
+        });
+
+        try {
+            final URI uri = normalizeAuthority(new URI(endpoint));
+
+            final Random random = new Random();
+            final Flowable<ByteBuffer> requestBody = Flowable.range(1, 100)
+                    .map(i -> ByteBuffer.wrap((i + ":" + random.nextDouble() + "\n").getBytes(UTF_8)));
+
+            final AsyncRequestProducer requestProducer = AsyncRequestBuilder.post(uri)
+                    .setEntity(new ReactiveEntityProducer(requestBody, -1, ContentType.TEXT_PLAIN, null))
+                    .build();
+
+            final ReactiveResponseConsumer consumer = new ReactiveResponseConsumer();
+            requester.execute(requestProducer, consumer, Timeout.ofSeconds(30), null);
+
+            final CompletableFuture<Void> printedAndDrained = new CompletableFuture<>();
+
+            consumer.getResponseStage().whenComplete((msg, ex) -> {
+                if (ex != null) {
+                    printedAndDrained.completeExceptionally(ex);
+                    return;
+                }
+                try {
+                    System.out.println(msg.getHead());
+                    for (final Header h : msg.getHead().getHeaders()) {
+                        System.out.println(h);
+                    }
+                    System.out.println();
+
+                    readBodyAsString(msg.getBody()).whenComplete((body, ex2) -> {
+                        if (ex2 != null) {
+                            printedAndDrained.completeExceptionally(ex2);
+                        } else {
+                            if (body != null && !body.isEmpty()) {
+                                System.out.print(body);
+                                if (!body.endsWith("\n")) {
+                                    System.out.println();
+                                }
+                            }
+                            printedAndDrained.complete(null);
+                        }
+                    });
+                } catch (final RuntimeException e) {
+                    printedAndDrained.completeExceptionally(e);
+                }
+            });
+
+            final CompletableFuture<Void> exchangeDone = consumer.getResponseCompletionStage().toCompletableFuture();
+
+            final CompletableFuture<Void> both = CompletableFuture.allOf(printedAndDrained, exchangeDone);
+            withTimeout(both, scheduler, 60, TimeUnit.SECONDS).toCompletableFuture().get();
+
+            System.out.println("Shutting down I/O reactor");
+            requester.initiateShutdown();
+
+        } finally {
+            scheduler.shutdownNow();
+        }
+    }
+}