[Dubbo-1876] Enhancements for the new async way of Dubbo

dubbo-common/src/main/java/org/apache/dubbo/common/Constants.java

@@ -19,7 +19,6 @@
 
 import java.util.concurrent.ExecutorService;
 import java.util.regex.Pattern;
-
 /**
  * Constants
  */
@@ -281,7 +280,8 @@ public class Constants {
 
     public static final String ASYNC_KEY = "async";
 
-    public static final String FUTURE_KEY = "async_future";
+    public static final String FUTURE_GENERATED_KEY = "future_generated";
+    public static final String FUTURE_RETURNTYPE_KEY = "future_returntype";
 
     public static final String ASYNC_SUFFIX = "Async";
 

dubbo-common/src/main/java/org/apache/dubbo/common/config/AsyncFor.java

@@ -39,3 +39,4 @@
     Class<?> value();
 
 }
+

dubbo-compatible/src/main/java/com/alibaba/dubbo/rpc/Result.java

@@ -63,6 +63,16 @@ public Map<String, String> getAttachments() {
             return delegate.getAttachments();
         }
 
+        @Override
+        public void addAttachments(Map<String, String> map) {
+            delegate.addAttachments(map);
+        }
+
+        @Override
+        public void setAttachments(Map<String, String> map) {
+            delegate.setAttachments(map);
+        }
+
         @Override
         public String getAttachment(String key) {
             return delegate.getAttachment(key);
@@ -72,5 +82,10 @@ public String getAttachment(String key) {
         public String getAttachment(String key, String defaultValue) {
             return delegate.getAttachment(key, defaultValue);
         }
+
+        @Override
+        public void setAttachment(String key, String value) {
+            delegate.setAttachment(key, value);
+        }
     }
 }

dubbo-demo/dubbo-demo-consumer/src/main/java/org/apache/dubbo/demo/consumer/Consumer.java

@@ -17,7 +17,6 @@
 package org.apache.dubbo.demo.consumer;
 
 import org.apache.dubbo.demo.DemoService;
-
 import org.springframework.context.support.ClassPathXmlApplicationContext;
 
 public class Consumer {

dubbo-remoting/dubbo-remoting-api/src/main/java/org/apache/dubbo/remoting/exchange/support/ExchangeHandlerAdapter.java

@@ -33,4 +33,4 @@ public CompletableFuture<Object> reply(ExchangeChannel channel, Object msg) thro
         return null;
     }
 
-}
+}

dubbo-remoting/dubbo-remoting-api/src/main/java/org/apache/dubbo/remoting/exchange/support/header/HeaderExchangeHandler.java

@@ -36,6 +36,7 @@
 import java.net.InetSocketAddress;
 import java.util.concurrent.CompletableFuture;
 
+
 /**
  * ExchangeReceiver
  */
@@ -102,7 +103,7 @@ void handleRequest(ExchangeChannel channel, Request req) throws RemotingExceptio
                 channel.send(res);
                 return;
             }
-            future.whenCompleteAsync((result, t) -> {
+            future.whenComplete((result, t) -> {
                 try {
                     if (t == null) {
                         res.setStatus(Response.OK);

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

@@ -0,0 +1,35 @@
+/*
+ * 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.rpc;
+
+/**
+ *
+ */
+public abstract class AbstractPostProcessFilter implements PostProcessFilter {
+    @Override
+    public Result postProcessResult(Result result, Invoker<?> invoker, Invocation invocation) {
+        if (result instanceof AsyncRpcResult) {
+            AsyncRpcResult asyncResult = (AsyncRpcResult) result;
+            asyncResult.thenApplyWithContext(r -> doPostProcess(r, invoker, invocation));
+            return asyncResult;
+        } else {
+            return doPostProcess(result, invoker, invocation);
+        }
+    }
+
+    protected abstract Result doPostProcess(Result result, Invoker<?> invoker, Invocation invocation);
+}

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

@@ -0,0 +1,71 @@
+/*
+ * 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.rpc;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ *
+ */
+public abstract class AbstractResult implements Result {
+    protected Map<String, String> attachments = new HashMap<String, String>();
+
+    protected Object result;
+
+    protected Throwable exception;
+
+    @Override
+    public Map<String, String> getAttachments() {
+        return attachments;
+    }
+
+    @Override
+    public void setAttachments(Map<String, String> map) {
+        this.attachments = map == null ? new HashMap<String, String>() : map;
+    }
+
+    @Override
+    public void addAttachments(Map<String, String> map) {
+        if (map == null) {
+            return;
+        }
+        if (this.attachments == null) {
+            this.attachments = new HashMap<String, String>();
+        }
+        this.attachments.putAll(map);
+    }
+
+    @Override
+    public String getAttachment(String key) {
+        return attachments.get(key);
+    }
+
+    @Override
+    public String getAttachment(String key, String defaultValue) {
+        String result = attachments.get(key);
+        if (result == null || result.length() == 0) {
+            result = defaultValue;
+        }
+        return result;
+    }
+
+    public void setAttachment(String key, String value) {
+        attachments.put(key, value);
+    }
+
+}

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

@@ -16,6 +16,8 @@
  */
 package org.apache.dubbo.rpc;
 
+import java.util.concurrent.CompletableFuture;
+
 /**
  * AsyncContext works like {@see javax.servlet.AsyncContext} in the Servlet 3.0.
  * An AsyncContext is stated by a call to {@link RpcContext#startAsync()}.
@@ -25,7 +27,7 @@
  */
 public interface AsyncContext {
 
-    void addListener(Runnable run);
+    CompletableFuture getInternalFuture();
 
     /**
      * write value and complete the async context.
@@ -45,7 +47,9 @@ public interface AsyncContext {
     boolean stop();
 
     /**
-     * change the context state to stop
+     * change the context state to start
      */
     void start();
+
+    void signalContextSwitch();
 }

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

@@ -26,30 +26,30 @@ public class AsyncContextImpl implements AsyncContext {
     private static final Logger logger = LoggerFactory.getLogger(AsyncContextImpl.class);
 
     private final AtomicBoolean started = new AtomicBoolean(false);
+    private final AtomicBoolean stoped = new AtomicBoolean(false);
 
     private CompletableFuture<Object> future;
 
+    private RpcContext storedContext;
+    private RpcContext storedServerContext;
+
     public AsyncContextImpl() {
     }
 
     public AsyncContextImpl(CompletableFuture<Object> future) {
         this.future = future;
-    }
-
-    @Override
-    public void addListener(Runnable run) {
-
+        this.storedContext = RpcContext.getContext();
+        this.storedServerContext = RpcContext.getServerContext();
     }
 
     @Override
     public void write(Object value) {
-        if (stop()) {
+        if (isAsyncStarted() && stop()) {
             if (value instanceof Throwable) {
-                // TODO check exception type like ExceptionFilter do.
                 Throwable bizExe = (Throwable) value;
-                future.complete(new RpcResult(bizExe));
+                future.completeExceptionally(bizExe);
             } else {
-                future.complete(new RpcResult(value));
+                future.complete(value);
             }
         } else {
             throw new IllegalStateException("The async response has probably been wrote back by another thread, or the asyncContext has been closed.");
@@ -63,11 +63,22 @@ public boolean isAsyncStarted() {
 
     @Override
     public boolean stop() {
-        return started.compareAndSet(true, false);
+        return stoped.compareAndSet(false, true);
     }
 
     @Override
     public void start() {
         this.started.set(true);
     }
+
+    public void signalContextSwitch() {
+        RpcContext.restoreContext(storedContext);
+        RpcContext.restoreServerContext(storedServerContext);
+        // Restore any other contexts in here if necessary.
+    }
+
+    @Override
+    public CompletableFuture getInternalFuture() {
+        return future;
+    }
 }