HADOOP-19235. IPC client uses CompletableFuture to support asynchronous operations.

hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/Client.java

@@ -96,8 +96,8 @@ public class Client implements AutoCloseable {
   private static final ThreadLocal<Integer> retryCount = new ThreadLocal<Integer>();
   private static final ThreadLocal<Object> EXTERNAL_CALL_HANDLER
       = new ThreadLocal<>();
-  private static final ThreadLocal<AsyncGet<? extends Writable, IOException>>
-      ASYNC_RPC_RESPONSE = new ThreadLocal<>();
+  private static final ThreadLocal<CompletableFuture<Writable>> ASYNC_RPC_RESPONSE
+      = new ThreadLocal<>();
   private static final ThreadLocal<Boolean> asynchronousMode =
       new ThreadLocal<Boolean>() {
         @Override
@@ -110,7 +110,50 @@ protected Boolean initialValue() {
   @Unstable
   public static <T extends Writable> AsyncGet<T, IOException>
       getAsyncRpcResponse() {
-    return (AsyncGet<T, IOException>) ASYNC_RPC_RESPONSE.get();
+    CompletableFuture<Writable> responseFuture = ASYNC_RPC_RESPONSE.get();
+    return new AsyncGet<T, IOException>() {
+      @Override
+      public T get(long timeout, TimeUnit unit)
+          throws IOException, TimeoutException, InterruptedException {
+        try {
+          if (unit == null || timeout < 0) {
+            return (T) responseFuture.get();
+          }
+          return (T) responseFuture.get(timeout, unit);
+        } catch (ExecutionException e) {
+          Throwable cause = e.getCause();
+          if (cause == null) {
+            throw new IOException(e);
+          }
+          if (cause instanceof IOException) {
+            throw (IOException) cause;
+          } else {
+            throw new IOException(cause);
+          }
+        }
+      }
+
+      @Override
+      public boolean isDone() {
+        return responseFuture.isDone();
+      }
+    };
+  }
+
+  /**
+   * Retrieves the current response future from the thread-local storage.
+   *
+   * @return A {@link CompletableFuture} of type T that represents the
+   *         asynchronous operation. If no response future is present in
+   *         the thread-local storage, this method returns {@code null}.
+   * @param <T> The type of the value completed by the returned
+   *            {@link CompletableFuture}. It must be a subclass of
+   *            {@link Writable}.
+   * @see CompletableFuture
+   * @see Writable
+   */
+  public static <T extends Writable> CompletableFuture<T> getResponseFuture() {
+    return (CompletableFuture<T>) ASYNC_RPC_RESPONSE.get();
   }
 
   /**
@@ -283,6 +326,7 @@ static class Call {
     boolean done;               // true when call is done
     private final Object externalHandler;
     private AlignmentContext alignmentContext;
+    private CompletableFuture<Object> completableFuture;
 
     private Call(RPC.RpcKind rpcKind, Writable param) {
       this.rpcKind = rpcKind;
@@ -304,6 +348,9 @@ private Call(RPC.RpcKind rpcKind, Writable param) {
       }
 
       this.externalHandler = EXTERNAL_CALL_HANDLER.get();
+      if (Client.isAsynchronousMode()) {
+        completableFuture = new CompletableFuture<>();
+      }
     }
 
     @Override
@@ -322,6 +369,9 @@ protected synchronized void callComplete() {
           externalHandler.notify();
         }
       }
+      if (completableFuture != null) {
+        completableFuture.complete(this);
+      }
     }
 
     /**
@@ -1495,36 +1545,16 @@ Writable call(RPC.RpcKind rpcKind, Writable rpcRequest,
     }
 
     if (isAsynchronousMode()) {
-      final AsyncGet<Writable, IOException> asyncGet
-          = new AsyncGet<Writable, IOException>() {
-        @Override
-        public Writable get(long timeout, TimeUnit unit)
-            throws IOException, TimeoutException{
-          boolean done = true;
-          try {
-            final Writable w = getRpcResponse(call, connection, timeout, unit);
-            if (w == null) {
-              done = false;
-              throw new TimeoutException(call + " timed out "
-                  + timeout + " " + unit);
-            }
-            return w;
-          } finally {
-            if (done) {
-              releaseAsyncCall();
-            }
-          }
-        }
-
-        @Override
-        public boolean isDone() {
-          synchronized (call) {
-            return call.done;
-          }
+      CompletableFuture<Writable> result = call.completableFuture.thenApply(o -> {
+        try {
+          return getRpcResponse(call, connection, -1, null);
+        } catch (IOException e) {
+          throw new CompletionException(e);
+        } finally {
+          releaseAsyncCall();
         }
-      };
-
-      ASYNC_RPC_RESPONSE.set(asyncGet);
+      });
+      ASYNC_RPC_RESPONSE.set(result);
       return null;
     } else {
       return getRpcResponse(call, connection, -1, null);

hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/ipc/TestAsyncIPC.java

@@ -28,6 +28,7 @@
 import org.apache.hadoop.ipc.protobuf.RpcHeaderProtos.RpcResponseHeaderProto;
 import org.apache.hadoop.net.NetUtils;
 import org.apache.hadoop.util.StringUtils;
+import org.apache.hadoop.util.Time;
 import org.apache.hadoop.util.concurrent.AsyncGetFuture;
 import org.junit.Assert;
 import org.junit.Before;
@@ -38,13 +39,16 @@
 import java.io.IOException;
 import java.net.InetSocketAddress;
 import java.util.*;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 public class TestAsyncIPC {
 
@@ -137,6 +141,60 @@ void assertReturnValues(long timeout, TimeUnit unit)
     }
   }
 
+  /**
+   * For testing the asynchronous calls of the RPC client
+   * implemented with CompletableFuture.
+   */
+  static class AsyncCompletableFutureCaller extends Thread {
+    private final Client client;
+    private final InetSocketAddress server;
+    private final int count;
+    private final List<CompletableFuture<Writable>> completableFutures;
+    private final List<Long> expectedValues;
+
+    AsyncCompletableFutureCaller(Client client, InetSocketAddress server, int count) {
+      this.client = client;
+      this.server = server;
+      this.count = count;
+      this.completableFutures = new ArrayList<>(count);
+      this.expectedValues = new ArrayList<>(count);
+      setName("Async CompletableFuture Caller");
+    }
+
+    @Override
+    public void run() {
+      // Set the RPC client to use asynchronous mode.
+      Client.setAsynchronousMode(true);
+      long startTime = Time.monotonicNow();
+      try {
+        for (int i = 0; i < count; i++) {
+          final long param = TestIPC.RANDOM.nextLong();
+          TestIPC.call(client, param, server, conf);
+          expectedValues.add(param);
+          completableFutures.add(Client.getResponseFuture());
+        }
+        // Since the run method is asynchronous,
+        // it does not need to wait for a response after sending a request,
+        // so the time taken by the run method is less than count * 100
+        // (where 100 is the time taken by the server to process a request).
+        long cost = Time.monotonicNow() - startTime;
+        assertTrue(cost < count * 100L);
+        LOG.info("[{}] run cost {}ms", Thread.currentThread().getName(), cost);
+      } catch (Exception e) {
+        fail();
+      }
+    }
+
+    public void assertReturnValues()
+        throws InterruptedException, ExecutionException {
+      for (int i = 0; i < count; i++) {
+        LongWritable value = (LongWritable) completableFutures.get(i).get();
+        Assert.assertEquals("call" + i + " failed.",
+            expectedValues.get(i).longValue(), value.get());
+      }
+    }
+  }
+
   static class AsyncLimitlCaller extends Thread {
     private Client client;
     private InetSocketAddress server;
@@ -538,4 +596,37 @@ public void run() {
       assertEquals(startID + i, callIds.get(i).intValue());
     }
   }
+
+  @Test(timeout = 60000)
+  public void testAsyncCallWithCompletableFuture() throws IOException,
+      InterruptedException, ExecutionException {
+    // Override client to store the call id
+    final Client client = new Client(LongWritable.class, conf);
+
+    // Construct an RPC server, which includes a handler thread.
+    final TestServer server = new TestIPC.TestServer(1, false, conf);
+    server.callListener = () -> {
+      try {
+        // The server requires at least 100 milliseconds to process a request.
+        Thread.sleep(100);
+      } catch (InterruptedException e) {
+        throw new RuntimeException(e);
+      }
+    };
+
+    try {
+      InetSocketAddress addr = NetUtils.getConnectAddress(server);
+      server.start();
+      // Send 10 asynchronous requests.
+      final AsyncCompletableFutureCaller caller =
+          new AsyncCompletableFutureCaller(client, addr, 10);
+      caller.start();
+      caller.join();
+      // Check if the values returned by the asynchronous call meet the expected values.
+      caller.assertReturnValues();
+    } finally {
+      client.stop();
+      server.stop();
+    }
+  }
 }