fix Context loss when async invoke issue#13666

dubbo-rpc/dubbo-rpc-api/src/main/java/org/apache/dubbo/rpc/AsyncRpcResult.java

@@ -21,6 +21,7 @@
 import org.apache.dubbo.common.logger.LoggerFactory;
 import org.apache.dubbo.common.threadpool.ThreadlessExecutor;
 import org.apache.dubbo.common.utils.SystemPropertyConfigUtils;
+import org.apache.dubbo.rpc.RpcContext.RestoreContext;
 import org.apache.dubbo.rpc.model.ConsumerMethodModel;
 import org.apache.dubbo.rpc.protocol.dubbo.FutureAdapter;
 
@@ -81,6 +82,16 @@ public AsyncRpcResult(CompletableFuture<AppResponse> future, Invocation invocati
                 && !future.isDone()) {
             async = true;
             this.storedContext = RpcContext.clearAndStoreContext();
+
+            // this is to fix https://github.com/apache/dubbo/issues/13666
+            this.responseFuture = this.responseFuture.thenApply((appResponse) -> {
+                RestoreContext tempStoredContext = RpcContext.clearAndStoreContext();
+                this.storedContext.restore();
+                appResponse.addObjectAttachments(
+                        RpcContext.getServerResponseContext().getObjectAttachments());
+                tempStoredContext.restore();
+                return appResponse;
+            });
         } else {
             async = false;
         }

dubbo-rpc/dubbo-rpc-api/src/main/java/org/apache/dubbo/rpc/RpcContextAttachment.java

@@ -245,9 +245,10 @@ public Object get(String key) {
      * @return a copy of RpcContextAttachment with deep copied attachments
      */
     public RpcContextAttachment copyOf(boolean needCopy) {
-        if (!isValid()) {
-            return null;
-        }
+        // to fix https://github.com/apache/dubbo/issues/13666
+        //        if (!isValid()) {
+        //            return null;
+        //        }
 
         if (needCopy) {
             RpcContextAttachment copy = new RpcContextAttachment();

dubbo-test/dubbo-test-common/src/main/java/org/apache/dubbo/test/common/api/AsyncContextService.java

@@ -0,0 +1,23 @@
+/*
+ * 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.
+ */
+package org.apache.dubbo.test.common.api;
+
+import java.util.concurrent.CompletableFuture;
+
+public interface AsyncContextService {
+    CompletableFuture<String> getServerSideContext();
+}

dubbo-test/dubbo-test-common/src/main/java/org/apache/dubbo/test/common/impl/AsyncContextServiceImpl.java

@@ -0,0 +1,40 @@
+/*
+ * 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.
+ */
+package org.apache.dubbo.test.common.impl;
+
+import org.apache.dubbo.rpc.RpcContext;
+import org.apache.dubbo.rpc.RpcContextAttachment;
+import org.apache.dubbo.test.common.api.AsyncContextService;
+
+import java.util.concurrent.CompletableFuture;
+
+public class AsyncContextServiceImpl implements AsyncContextService {
+
+    @Override
+    public CompletableFuture<String> getServerSideContext() {
+        RpcContextAttachment serverResponseContext = RpcContext.getServerResponseContext();
+        return CompletableFuture.supplyAsync(() -> {
+            try {
+                Thread.sleep(1000);
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+            serverResponseContext.setAttachment("theToken", "Hello, Dubbo");
+            return "finish";
+        });
+    }
+}

dubbo-test/dubbo-test-spring/src/main/java/org/apache/dubbo/test/spring/SpringRpcContextTest.java

@@ -0,0 +1,88 @@
+/*
+ * 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.
+ */
+package org.apache.dubbo.test.spring;
+
+import org.apache.dubbo.config.bootstrap.DubboBootstrap;
+import org.apache.dubbo.rpc.RpcContext;
+import org.apache.dubbo.test.common.SysProps;
+import org.apache.dubbo.test.common.api.AsyncContextService;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.springframework.context.support.ClassPathXmlApplicationContext;
+
+import static org.apache.dubbo.common.constants.CommonConstants.SHUTDOWN_WAIT_KEY;
+
+public class SpringRpcContextTest {
+    private static ClassPathXmlApplicationContext providerContext;
+
+    @BeforeAll
+    public static void beforeAll() {
+        DubboBootstrap.reset();
+    }
+
+    @AfterAll
+    public static void afterAll() {
+        DubboBootstrap.reset();
+        providerContext.close();
+    }
+
+    private void startProvider() {
+        providerContext = new ClassPathXmlApplicationContext("/spring/dubbo-demo-provider.xml");
+    }
+
+    @Test
+    public void test() {
+        SysProps.setProperty(SHUTDOWN_WAIT_KEY, "10000");
+        startProvider();
+        ClassPathXmlApplicationContext applicationContext = null;
+        try {
+            applicationContext = new ClassPathXmlApplicationContext("/spring/dubbo-demo.xml");
+            AsyncContextService asyncContextService =
+                    applicationContext.getBean("asyncContextService", AsyncContextService.class);
+
+            AtomicBoolean contextCorrect = new AtomicBoolean(true);
+
+            for (int i = 0; i < 5; i++) {
+                CompletableFuture<String> serverSideContext = asyncContextService.getServerSideContext();
+                serverSideContext.whenComplete((s, throwable) -> {
+                    if (!"Hello, Dubbo"
+                            .equals(RpcContext.getClientResponseContext().getAttachment("theToken"))) {
+                        contextCorrect.set(false);
+                    }
+                    RpcContext.removeContext();
+                });
+                serverSideContext.get();
+            }
+            Assertions.assertTrue(contextCorrect.get());
+
+        } catch (ExecutionException | InterruptedException e) {
+            throw new RuntimeException(e);
+        } finally {
+            SysProps.clear();
+            if (applicationContext != null) {
+                applicationContext.close();
+            }
+        }
+    }
+}

dubbo-test/dubbo-test-spring/src/main/resources/spring/dubbo-demo-provider.xml

@@ -33,10 +33,12 @@
     <bean id="demoServiceImpl" class="org.apache.dubbo.test.common.impl.DemoServiceImpl"/>
     <bean id="greetingServiceImpl" class="org.apache.dubbo.test.common.impl.GreetingServiceImpl"/>
     <bean id="restDemoServiceImpl" class="org.apache.dubbo.test.common.impl.RestDemoServiceImpl"/>
+    <bean id="asyncContextServiceImpl" class="org.apache.dubbo.test.common.impl.AsyncContextServiceImpl"/>
 
     <dubbo:service interface="org.apache.dubbo.test.common.api.DemoService" timeout="3000" ref="demoServiceImpl" protocol="dubbo"/>
     <dubbo:service version="1.0.0" group="greeting" timeout="5000" interface="org.apache.dubbo.test.common.api.GreetingService"
                    ref="greetingServiceImpl" protocol="dubbo"/>
     <dubbo:service version="1.0.0" timeout="5000" interface="org.apache.dubbo.test.common.api.RestDemoService"
                    ref="restDemoServiceImpl" protocol="tri"/>
+    <dubbo:service version="1.0.0" interface="org.apache.dubbo.test.common.api.AsyncContextService" timeout="3000" ref="asyncContextServiceImpl" protocol="dubbo"/>
 </beans>

dubbo-test/dubbo-test-spring/src/main/resources/spring/dubbo-demo.xml

@@ -40,4 +40,6 @@
 
     <dubbo:reference  id="restDemoService" version="1.0.0" protocol="tri"
                       interface="org.apache.dubbo.test.common.api.RestDemoService"/>
+
+    <dubbo:reference  id="asyncContextService" version="1.0.0" interface="org.apache.dubbo.test.common.api.AsyncContextService"/>
 </beans>