[ISSUE #3267] Refactor RequestContext (#3268)

mxsm · web-flow · commit 94562e99a8de · 2023-03-04T13:42:49.000+08:00
* [ISSUE #3267] Refactor RequestContext * polish code
diff --git a/eventmesh-sdk-java/src/main/java/org/apache/eventmesh/client/tcp/common/RequestContext.java b/eventmesh-sdk-java/src/main/java/org/apache/eventmesh/client/tcp/common/RequestContext.java
@@ -19,22 +19,23 @@
 
 import org.apache.eventmesh.common.protocol.tcp.Package;
 
-import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 import lombok.extern.slf4j.Slf4j;
 
 @Slf4j
 public class RequestContext {
 
-    private transient Object key;
-    private transient Package request;
-    private transient Package response;
-    private transient CountDownLatch latch;
+    private Object key;
+    private Package request;
+    private final CompletableFuture<Package> future = new CompletableFuture<>();
 
-    public RequestContext(final Object key, final Package request, final CountDownLatch latch) {
+    public RequestContext(final Object key, final Package request) {
         this.key = key;
         this.request = request;
-        this.latch = latch;
     }
 
     public Object getKey() {
@@ -53,33 +54,28 @@ public void setRequest(final Package request) {
         this.request = request;
     }
 
-    public Package getResponse() {
-        return response;
+    public CompletableFuture<Package> future() {
+        return this.future;
     }
 
-    public void setResponse(final Package response) {
-        this.response = response;
+    public Package getResponse(long timeout, TimeUnit timeUnit) throws ExecutionException, InterruptedException, TimeoutException {
+        return this.future.get(timeout, timeUnit);
     }
 
-    public CountDownLatch getLatch() {
-        return latch;
-    }
-
-    public void setLatch(final CountDownLatch latch) {
-        this.latch = latch;
+    public Package getResponse(long timeout) throws ExecutionException, InterruptedException, TimeoutException {
+        return this.future.get(timeout, TimeUnit.MILLISECONDS);
     }
 
     public void finish(final Package msg) {
-        this.response = msg;
-        latch.countDown();
+        this.future.complete(msg);
     }
 
-    public static RequestContext context(final Object key, final Package request, final CountDownLatch latch) throws Exception {
-        final RequestContext c = new RequestContext(key, request, latch);
+    public static RequestContext context(final Object key, final Package request) throws Exception {
+        final RequestContext context = new RequestContext(key, request);
         if (log.isInfoEnabled()) {
             log.info("_RequestContext|create|key={}", key);
         }
-        return c;
+        return context;
     }
 
 
diff --git a/eventmesh-sdk-java/src/main/java/org/apache/eventmesh/client/tcp/common/TcpClient.java b/eventmesh-sdk-java/src/main/java/org/apache/eventmesh/client/tcp/common/TcpClient.java
@@ -28,12 +28,14 @@
 import java.io.Closeable;
 import java.net.InetSocketAddress;
 import java.util.Random;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.function.Supplier;
 
 import io.netty.bootstrap.Bootstrap;
 import io.netty.buffer.PooledByteBufAllocator;
@@ -176,20 +178,23 @@ protected void send(Package msg) throws Exception {
 
     protected Package io(Package msg, long timeout) throws Exception {
         Object key = RequestContext.key(msg);
-        CountDownLatch latch = new CountDownLatch(1);
-        RequestContext c = RequestContext.context(key, msg, latch);
-        if (!contexts.contains(c)) {
-            contexts.put(key, c);
+        RequestContext context = RequestContext.context(key, msg);
+        if (!contexts.contains(context)) {
+            contexts.put(key, context);
         } else {
             if (log.isInfoEnabled()) {
                 log.info("duplicate key : {}", key);
             }
         }
         send(msg);
-        if (!c.getLatch().await(timeout, TimeUnit.MILLISECONDS)) {
-            throw new TimeoutException("operation timeout, context.key=" + c.getKey());
-        }
-        return c.getResponse();
+        Supplier<Package> supplier = () -> {
+            try {
+                return context.getResponse(timeout);
+            } catch (ExecutionException | InterruptedException | TimeoutException exception) {
+                throw new RuntimeException(exception);
+            }
+        };
+        return CompletableFuture.supplyAsync(supplier).get();
     }
 
     // todo: remove hello
