fix test

Author: KeeProMise

hadoop-hdfs-project/hadoop-hdfs-rbf/src/main/java/org/apache/hadoop/hdfs/protocolPB/AsyncRpcProtocolPBUtil.java

@@ -18,28 +18,28 @@
 
 package org.apache.hadoop.hdfs.protocolPB;
 
+import org.apache.hadoop.hdfs.server.federation.router.ThreadLocalContext;
 import org.apache.hadoop.hdfs.server.federation.router.async.ApplyFunction;
 import org.apache.hadoop.io.Writable;
-import org.apache.hadoop.ipc.CallerContext;
 import org.apache.hadoop.ipc.Client;
 import org.apache.hadoop.ipc.ProtobufRpcEngine2;
-import org.apache.hadoop.ipc.Server;
 import org.apache.hadoop.ipc.internal.ShadedProtobufHelper;
 import org.apache.hadoop.util.concurrent.AsyncGet;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
 
 import static org.apache.hadoop.hdfs.server.federation.router.async.Async.warpCompletionException;
-import static org.apache.hadoop.hdfs.server.federation.router.async.AsyncUtil.asyncApply;
 import static org.apache.hadoop.hdfs.server.federation.router.async.AsyncUtil.asyncCompleteWith;
 import static org.apache.hadoop.hdfs.server.federation.router.async.AsyncUtil.asyncReturn;
 import static org.apache.hadoop.ipc.internal.ShadedProtobufHelper.ipc;
 
 public final class AsyncRpcProtocolPBUtil {
   public static final Logger LOG = LoggerFactory.getLogger(AsyncRpcProtocolPBUtil.class);
+  private static Executor worker;
 
   private AsyncRpcProtocolPBUtil() {}
 
@@ -50,20 +50,24 @@ public static <T, R> R asyncIpcClient(
     AsyncGet<T, Exception> asyncReqMessage =
         (AsyncGet<T, Exception>) ProtobufRpcEngine2.getAsyncReturnMessage();
     CompletableFuture<Writable> responseFuture = Client.getResponseFuture();
-    // transfer originCall & callerContext to worker threads of executor.
-    final Server.Call originCall = Server.getCurCall().get();
-    final CallerContext originContext = CallerContext.getCurrent();
-    asyncCompleteWith(responseFuture);
-    asyncApply(o -> {
+    // transfer thread local context to worker threads of executor.
+    ThreadLocalContext threadLocalContext = new ThreadLocalContext();
+    asyncCompleteWith(responseFuture.handleAsync((result, e) -> {
+      threadLocalContext.transfer();
+      if (e != null) {
+        throw warpCompletionException(e);
+      }
       try {
-        Server.getCurCall().set(originCall);
-        CallerContext.setCurrent(originContext);
         T res = asyncReqMessage.get(-1, null);
         return response.apply(res);
-      } catch (Exception e) {
-        throw warpCompletionException(e);
+      } catch (Exception ex) {
+        throw warpCompletionException(ex);
       }
-    });
+    }, worker));
     return asyncReturn(clazz);
   }
+
+  public static void setWorker(Executor worker) {
+    AsyncRpcProtocolPBUtil.worker = worker;
+  }
 }

hadoop-hdfs-project/hadoop-hdfs-rbf/src/main/java/org/apache/hadoop/hdfs/server/federation/router/RouterRpcServer.java

@@ -67,6 +67,7 @@
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.hdfs.HAUtil;
 import org.apache.hadoop.hdfs.protocol.UnresolvedPathException;
+import org.apache.hadoop.hdfs.protocolPB.AsyncRpcProtocolPBUtil;
 import org.apache.hadoop.thirdparty.com.google.common.cache.CacheBuilder;
 import org.apache.hadoop.thirdparty.com.google.common.cache.CacheLoader;
 import org.apache.hadoop.thirdparty.com.google.common.cache.LoadingCache;
@@ -461,18 +462,17 @@ protected void initAsyncThreadPool() {
         DFS_ROUTER_RPC_ASYNC_HANDLER_COUNT_DEFAULT);
     int asyncResponderCount = conf.getInt(DFS_ROUTER_RPC_ASYNC_RESPONDER_COUNT,
         DFS_ROUTER_RPC_ASYNC_RESPONDER_COUNT_DEFAULT);
-    synchronized (RouterRpcServer.class) {
-      if (asyncRouterHandler == null) {
-        LOG.info("init router async handler count: {}", asyncHandlerCount);
-        asyncRouterHandler = Executors.newFixedThreadPool(
-            asyncHandlerCount, new AsyncThreadFactory("router async handler "));
-      }
-      if (asyncRouterResponder == null) {
-        LOG.info("init router async responder count: {}", asyncResponderCount);
-        asyncRouterResponder = Executors.newFixedThreadPool(
-            asyncResponderCount, new AsyncThreadFactory("router async responder "));
-      }
+    if (asyncRouterHandler == null) {
+      LOG.info("init router async handler count: {}", asyncHandlerCount);
+      asyncRouterHandler = Executors.newFixedThreadPool(
+          asyncHandlerCount, new AsyncThreadFactory("router async handler "));
+    }
+    if (asyncRouterResponder == null) {
+      LOG.info("init router async responder count: {}", asyncResponderCount);
+      asyncRouterResponder = Executors.newFixedThreadPool(
+          asyncResponderCount, new AsyncThreadFactory("router async responder "));
     }
+    AsyncRpcProtocolPBUtil.setWorker(asyncRouterResponder);
   }
 
   /**

hadoop-hdfs-project/hadoop-hdfs-rbf/src/main/java/org/apache/hadoop/hdfs/server/federation/router/async/AsyncUtil.java

@@ -80,11 +80,14 @@ public static <R> R asyncReturn(Class<R> clazz) {
     if (clazz == null) {
       return null;
     }
-    if (clazz.equals(Boolean.class)) {
+    if (clazz.equals(Boolean.class)
+        | clazz.equals(boolean.class)) {
       return (R) BOOLEAN_RESULT;
-    } else if (clazz.equals(Long.class)) {
+    } else if (clazz.equals(Long.class)
+        | clazz.equals(long.class)) {
       return (R) LONG_RESULT;
-    } else if (clazz.equals(Integer.class)) {
+    } else if (clazz.equals(Integer.class)
+        | clazz.equals(int.class)) {
       return (R) INT_RESULT;
     }
     return (R) NULL_RESULT;

hadoop-hdfs-project/hadoop-hdfs-rbf/src/test/java/org/apache/hadoop/hdfs/protocolPB/TestAsyncRpcProtocolPBUtil.java

@@ -40,6 +40,7 @@
 
 import java.io.IOException;
 import java.net.InetSocketAddress;
+import java.util.concurrent.ForkJoinPool;
 
 import static org.apache.hadoop.hdfs.server.federation.router.async.AsyncUtil.syncReturn;
 import static org.junit.Assert.assertEquals;
@@ -53,6 +54,7 @@ public class TestAsyncRpcProtocolPBUtil {
 
   @Before
   public void setUp() throws IOException {
+    AsyncRpcProtocolPBUtil.setWorker(ForkJoinPool.commonPool());
     Configuration conf = new Configuration();
     RPC.setProtocolEngine(conf, TestRpcBase.TestRpcService.class,
         ProtobufRpcEngine2.class);

hadoop-hdfs-project/hadoop-hdfs-rbf/src/test/java/org/apache/hadoop/hdfs/protocolPB/TestRouterClientSideTranslatorPB.java

@@ -52,6 +52,7 @@
 import java.net.InetSocketAddress;
 import java.util.EnumSet;
 import java.util.List;
+import java.util.concurrent.ForkJoinPool;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import static org.apache.hadoop.crypto.CryptoProtocolVersion.ENCRYPTION_ZONES;
@@ -81,6 +82,7 @@ public class TestRouterClientSideTranslatorPB {
 
   @BeforeClass
   public static void setUp() throws Exception {
+    AsyncRpcProtocolPBUtil.setWorker(ForkJoinPool.commonPool());
     conf = new HdfsConfiguration();
     cluster = (new MiniDFSCluster.Builder(conf))
         .numDataNodes(1).build();