From 116263e98d956a05abd8755cab651111d4f2f14f Mon Sep 17 00:00:00 2001
From: Zeus <40241629+XDanwar@users.noreply.github.com>
Date: Wed, 14 Dec 2022 15:40:21 +0800
Subject: [PATCH] =?UTF-8?q?=E3=80=90fix=E3=80=91forking=20cluster=20(#1106?=
 =?UTF-8?q?3)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* 【fix】forking cluster

* 【fix】fix code style

* 【fix】unit test unwrap exception

* 【fix】unit test unwrap CompletionException

* 【fix】code review

Co-authored-by: kanji.ainiwaer <kanji.ainiwaer@trip.com>
---
 .../support/ForkingClusterInvoker.java        | 27 ++++++++++---------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/dubbo-cluster/src/main/java/org/apache/dubbo/rpc/cluster/support/ForkingClusterInvoker.java b/dubbo-cluster/src/main/java/org/apache/dubbo/rpc/cluster/support/ForkingClusterInvoker.java
index eaf8f69353..a1334aab07 100644
--- a/dubbo-cluster/src/main/java/org/apache/dubbo/rpc/cluster/support/ForkingClusterInvoker.java
+++ b/dubbo-cluster/src/main/java/org/apache/dubbo/rpc/cluster/support/ForkingClusterInvoker.java
@@ -30,6 +30,8 @@ import org.apache.dubbo.rpc.cluster.LoadBalance;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
@@ -84,27 +86,28 @@ public class ForkingClusterInvoker<T> extends AbstractClusterInvoker<T> {
             RpcContext.getServiceContext().setInvokers((List) selected);
             final AtomicInteger count = new AtomicInteger();
             final BlockingQueue<Object> ref = new LinkedBlockingQueue<>(1);
-            for (final Invoker<T> invoker : selected) {
+            selected.forEach(invoker -> {
                 URL consumerUrl = RpcContext.getServiceContext().getConsumerUrl();
-                executor.execute(() -> {
-                    try {
-                        if (ref.size() > 0) {
-                            return;
-                        }
-                        Result result = invokeWithContextAsync(invoker, invocation, consumerUrl);
-                        ref.offer(result);
-                    } catch (Throwable e) {
+                CompletableFuture.<Object>supplyAsync(() -> {
+                    if (ref.size() > 0) {
+                        return null;
+                    }
+                    return invokeWithContextAsync(invoker, invocation, consumerUrl);
+                }, executor).whenComplete((v, t) -> {
+                    if (t == null) {
+                        ref.offer(v);
+                    } else {
                         int value = count.incrementAndGet();
                         if (value >= selected.size()) {
-                            ref.offer(e);
+                            ref.offer(t);
                         }
                     }
                 });
-            }
+            });
             try {
                 Object ret = ref.poll(timeout, TimeUnit.MILLISECONDS);
                 if (ret instanceof Throwable) {
-                    Throwable e = (Throwable) ret;
+                    Throwable e = ret instanceof CompletionException ? ((CompletionException) ret).getCause() : (Throwable) ret;
                     throw new RpcException(e instanceof RpcException ? ((RpcException) e).getCode() : RpcException.UNKNOWN_EXCEPTION,
                         "Failed to forking invoke provider " + selected + ", but no luck to perform the invocation. " +
                             "Last error is: " + e.getMessage(), e.getCause() != null ? e.getCause() : e);