Add:RouterCacheAdmin supports asynchronous rpc

Author: Archie-wang

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

@@ -18,52 +18,91 @@
 
 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;
 
+/**
+ * <p>This utility class encapsulates the logic required to initiate asynchronous RPCs,
+ * handle responses, and propagate exceptions. It works in conjunction with
+ * {@link ProtobufRpcEngine2} and {@link Client} to facilitate the asynchronous
+ * nature of the operations.
+ *
+ * @see ProtobufRpcEngine2
+ * @see Client
+ * @see CompletableFuture
+ */
 public final class AsyncRpcProtocolPBUtil {
   public static final Logger LOG = LoggerFactory.getLogger(AsyncRpcProtocolPBUtil.class);
+  /** The executor used for handling responses asynchronously. */
+  private static Executor worker;
 
   private AsyncRpcProtocolPBUtil() {}
 
+  /**
+   * Asynchronously invokes an RPC call and applies a response transformation function
+   * to the result. This method is generic and can be used to handle any type of
+   * RPC call.
+   *
+   * <p>The method uses the {@link ShadedProtobufHelper.IpcCall} to prepare the RPC call
+   * and the {@link ApplyFunction} to process the response. It also handles exceptions
+   * that may occur during the RPC call and wraps them in a user-friendly manner.
+   *
+   * @param call The IPC call encapsulating the RPC request.
+   * @param response The function to apply to the response of the RPC call.
+   * @param clazz The class object representing the type {@code R} of the response.
+   * @param <T> Type of the call's result.
+   * @param <R> Type of method return.
+   * @return An object of type {@code R} that is the result of applying the response
+   *         function to the RPC call result.
+   * @throws IOException If an I/O error occurs during the asynchronous RPC call.
+   */
   public static <T, R> R asyncIpcClient(
       ShadedProtobufHelper.IpcCall<T> call, ApplyFunction<T, R> response,
       Class<R> clazz) throws IOException {
     ipc(call);
     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);
   }
+
+  /**
+   * Sets the executor used for handling responses asynchronously within
+   * the utility class.
+   *
+   * @param worker The executor to be used for handling responses asynchronously.
+   */
+  public static void setWorker(Executor worker) {
+    AsyncRpcProtocolPBUtil.worker = worker;
+  }
 }

hadoop-hdfs-project/hadoop-hdfs-rbf/src/main/java/org/apache/hadoop/hdfs/server/federation/metrics/FederationRPCPerformanceMonitor.java

@@ -52,9 +52,9 @@ public class FederationRPCPerformanceMonitor implements RouterRpcMonitor {
 
 
   /** Time for an operation to be received in the Router. */
-  private static final ThreadLocal<Long> START_TIME = new ThreadLocal<>();
+  private static final ThreadLocal<Long> START_TIME = ThreadLocal.withInitial(() -> -1L);
   /** Time for an operation to be sent to the Namenode. */
-  private static final ThreadLocal<Long> PROXY_TIME = new ThreadLocal<>();
+  private static final ThreadLocal<Long> PROXY_TIME = ThreadLocal.withInitial(() -> -1L);
 
   /** Configuration for the performance monitor. */
   private Configuration conf;
@@ -141,6 +141,14 @@ public void startOp() {
     START_TIME.set(monotonicNow());
   }
 
+  public static long getStartOpTime() {
+    return START_TIME.get();
+  }
+
+  public static void setStartOpTime(long startOpTime) {
+    START_TIME.set(startOpTime);
+  }
+
   @Override
   public long proxyOp() {
     PROXY_TIME.set(monotonicNow());
@@ -151,6 +159,14 @@ public long proxyOp() {
     return Thread.currentThread().getId();
   }
 
+  public static long getProxyOpTime() {
+    return PROXY_TIME.get();
+  }
+
+  public static void setProxyOpTime(long proxyOpTime) {
+    PROXY_TIME.set(proxyOpTime);
+  }
+
   @Override
   public void proxyOpComplete(boolean success, String nsId,
       FederationNamenodeServiceState state) {

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

@@ -72,6 +72,15 @@ public class RBFConfigKeys extends CommonConfigurationKeysPublic {
   public static final String DFS_ROUTER_RPC_ENABLE =
       FEDERATION_ROUTER_PREFIX + "rpc.enable";
   public static final boolean DFS_ROUTER_RPC_ENABLE_DEFAULT = true;
+  public static final String DFS_ROUTER_RPC_ENABLE_ASYNC =
+      FEDERATION_ROUTER_PREFIX + "rpc.async.enable";
+  public static final boolean DFS_ROUTER_RPC_ENABLE_ASYNC_DEFAULT = false;
+  public static final String DFS_ROUTER_RPC_ASYNC_HANDLER_COUNT =
+      FEDERATION_ROUTER_PREFIX + "rpc.async.handler.count";
+  public static final int DFS_ROUTER_RPC_ASYNC_HANDLER_COUNT_DEFAULT = 2;
+  public static final String DFS_ROUTER_RPC_ASYNC_RESPONDER_COUNT =
+      FEDERATION_ROUTER_PREFIX + "rpc.async.responder.count";
+  public static final int DFS_ROUTER_RPC_ASYNC_RESPONDER_COUNT_DEFAULT = 10;
 
   public static final String DFS_ROUTER_METRICS_ENABLE =
       FEDERATION_ROUTER_PREFIX + "metrics.enable";

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

@@ -0,0 +1,75 @@
+/**
+ * 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.hadoop.hdfs.server.federation.router;
+
+import org.apache.hadoop.fs.CacheFlag;
+import org.apache.hadoop.hdfs.protocol.CacheDirectiveEntry;
+import org.apache.hadoop.hdfs.protocol.CacheDirectiveInfo;
+import org.apache.hadoop.hdfs.protocol.CachePoolEntry;
+import org.apache.hadoop.fs.BatchedRemoteIterator.BatchedEntries;
+import org.apache.hadoop.hdfs.server.federation.resolver.FederationNamespaceInfo;
+import org.apache.hadoop.hdfs.server.federation.resolver.RemoteLocation;
+import org.apache.hadoop.hdfs.server.federation.router.async.ApplyFunction;
+
+import java.io.IOException;
+import java.util.EnumSet;
+import java.util.Map;
+
+import static org.apache.hadoop.hdfs.server.federation.router.async.AsyncUtil.asyncApply;
+import static org.apache.hadoop.hdfs.server.federation.router.async.AsyncUtil.asyncReturn;
+
+/**
+ * Module that implements all the asynchronous RPC calls in
+ * {@link org.apache.hadoop.hdfs.protocol.ClientProtocol} related to Cache Admin
+ * in the {@link RouterRpcServer}.
+ */
+public class RouterAsyncCacheAdmin extends RouterCacheAdmin{
+
+  public RouterAsyncCacheAdmin(RouterRpcServer server) {
+    super(server);
+  }
+
+  @Override
+  public long addCacheDirective(
+      CacheDirectiveInfo path, EnumSet<CacheFlag> flags) throws IOException {
+    invokeAddCacheDirective(path, flags);
+    asyncApply((ApplyFunction<Map<RemoteLocation, Long>, Long>) 
+        response -> response.values().iterator().next());
+    return asyncReturn(Long.class);
+  }
+
+  @Override
+  public BatchedEntries<CacheDirectiveEntry> listCacheDirectives(
+      long prevId, CacheDirectiveInfo filter) throws IOException {
+    invokeListCacheDirectives(prevId, filter);
+    asyncApply((ApplyFunction<Map,
+        BatchedEntries<CacheDirectiveEntry>>)
+        response -> (BatchedEntries<CacheDirectiveEntry>) response.values().iterator().next());
+    return asyncReturn(BatchedEntries.class);
+  }
+
+  @Override
+  public BatchedEntries<CachePoolEntry> listCachePools(String prevKey) throws IOException {
+    invokeListCachePools(prevKey);
+    asyncApply((ApplyFunction<Map<FederationNamespaceInfo, BatchedEntries>,
+        BatchedEntries<CachePoolEntry>>)
+        results -> results.values().iterator().next());
+    return asyncReturn(BatchedEntries.class);
+  }
+}