From c916b82ab5fbd2ad9afedbed453c36a2a9865564 Mon Sep 17 00:00:00 2001
From: trytocatch <trytocatch@163.com>
Date: Thu, 11 May 2023 19:07:45 +0800
Subject: [PATCH] clean the RpcContext in ThreadLocal after QoS invoke,
 synchronously get the result while async started (#12291)

* clean the context in ThreadLocal after qos invoke, and synchronously get the result while async started, so the invoke command could get the correct result instead of null

* add Thread.interrupt to the catch block

add Thread.interrupt to the catch block

* add Thread.interrupt to the catch block

add Thread.interrupt to the catch block
---
 .../dubbo/qos/command/impl/InvokeTelnet.java  | 22 ++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/dubbo-plugin/dubbo-qos/src/main/java/org/apache/dubbo/qos/command/impl/InvokeTelnet.java b/dubbo-plugin/dubbo-qos/src/main/java/org/apache/dubbo/qos/command/impl/InvokeTelnet.java
index 128b174972..3e7400ff05 100644
--- a/dubbo-plugin/dubbo-qos/src/main/java/org/apache/dubbo/qos/command/impl/InvokeTelnet.java
+++ b/dubbo-plugin/dubbo-qos/src/main/java/org/apache/dubbo/qos/command/impl/InvokeTelnet.java
@@ -25,6 +25,9 @@ import org.apache.dubbo.qos.api.BaseCommand;
 import org.apache.dubbo.qos.api.CommandContext;
 import org.apache.dubbo.qos.api.Cmd;
 import org.apache.dubbo.rpc.AppResponse;
+import org.apache.dubbo.rpc.AsyncContext;
+import org.apache.dubbo.rpc.AsyncContextImpl;
+import org.apache.dubbo.rpc.RpcContext;
 import org.apache.dubbo.rpc.model.FrameworkModel;
 import org.apache.dubbo.rpc.model.MethodDescriptor;
 import org.apache.dubbo.rpc.model.ProviderModel;
@@ -38,6 +41,7 @@ import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 
 import static org.apache.dubbo.common.utils.PojoUtils.realize;
 
@@ -150,9 +154,25 @@ public class InvokeTelnet implements BaseCommand {
             AppResponse result = new AppResponse();
             try {
                 Object o = invokeMethod.invoke(selectedProvider.getServiceInstance(), array);
-                result.setValue(o);
+                boolean setValueDone = false;
+                if (RpcContext.getServerAttachment().isAsyncStarted()) {
+                    AsyncContext asyncContext = RpcContext.getServerAttachment().getAsyncContext();
+                    if (asyncContext instanceof AsyncContextImpl) {
+                        CompletableFuture<Object> internalFuture = ((AsyncContextImpl) asyncContext).getInternalFuture();
+                        result.setValue(internalFuture.get());
+                        setValueDone = true;
+                    }
+                }
+                if (!setValueDone) {
+                    result.setValue(o);
+                }
             } catch (Throwable t) {
                 result.setException(t);
+                if (t instanceof InterruptedException) {
+                    Thread.currentThread().interrupt();
+                }
+            } finally {
+                RpcContext.removeContext();
             }
             long end = System.currentTimeMillis();
             buf.append("\r\nresult: ");