HADOOP-19864. (Followup) Migrate MiniRPCBenchmark to ProtobufRpcEngine2 (apache#8455)

pan3793 · web-flow · commit f6182e5680cf · 2026-04-24T19:37:27.000+08:00
Reviewed-by: Steve Loughran &lt;stevel@apache.org&gt;
Signed-off-by: Cheng Pan &lt;chengpan@apache.org&gt;
diff --git a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/ipc/MiniRPCBenchmark.java b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/ipc/MiniRPCBenchmark.java
@@ -30,6 +30,10 @@
 
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.io.Text;
+import org.apache.hadoop.ipc.protobuf.MiniRPCBenchmarkProtos.MiniDelegationTokenProto;
+import org.apache.hadoop.ipc.protobuf.MiniRPCBenchmarkProtos.MiniGetDelegationTokenRequestProto;
+import org.apache.hadoop.ipc.protobuf.MiniRPCBenchmarkProtos.MiniGetDelegationTokenResponseProto;
+import org.apache.hadoop.ipc.protobuf.MiniRPCBenchmarkProtos.MiniProtocolService;
 import org.apache.hadoop.net.NetUtils;
 import org.apache.hadoop.security.KerberosInfo;
 import org.apache.hadoop.security.SecurityUtil;
@@ -40,6 +44,10 @@
 import org.apache.hadoop.security.token.delegation.AbstractDelegationTokenSelector;
 import org.apache.hadoop.security.token.delegation.TestDelegationToken.TestDelegationTokenIdentifier;
 import org.apache.hadoop.security.token.delegation.TestDelegationToken.TestDelegationTokenSecretManager;
+import org.apache.hadoop.thirdparty.protobuf.BlockingService;
+import org.apache.hadoop.thirdparty.protobuf.ByteString;
+import org.apache.hadoop.thirdparty.protobuf.RpcController;
+import org.apache.hadoop.thirdparty.protobuf.ServiceException;
 import org.apache.hadoop.util.Time;
 import org.slf4j.event.Level;
 
@@ -101,18 +109,19 @@ protected TestDelegationTokenSelector() {
       super(new Text("MY KIND"));
     }    
   }
-  
-  @KerberosInfo(
-      serverPrincipal=USER_NAME_KEY)
-  @TokenInfo(TestDelegationTokenSelector.class)
-  public static interface MiniProtocol extends VersionedProtocol {
-    public static final long versionID = 1L;
 
-    /**
-     * Get a Delegation Token.
-     */
-    public Token<TestDelegationTokenIdentifier> getDelegationToken(Text renewer) 
-        throws IOException;
+  /**
+   * Protobuf based MiniProtocol used by {@link MiniRPCBenchmark}.
+   * Replaces the legacy Writable based protocol now that
+   * {@code WritableRpcEngine} has been removed.
+   */
+  @KerberosInfo(serverPrincipal = USER_NAME_KEY)
+  @TokenInfo(TestDelegationTokenSelector.class)
+  @ProtocolInfo(
+      protocolName = "org.apache.hadoop.ipc.MiniRPCBenchmark$MiniProtocol",
+      protocolVersion = 1)
+  public interface MiniProtocol extends MiniProtocolService.BlockingInterface {
+    long versionID = 1L;
   }
 
   /**
@@ -125,34 +134,32 @@ static class MiniServer implements MiniProtocol {
     private TestDelegationTokenSecretManager secretManager;
     private Server rpcServer;
 
-    @Override // VersionedProtocol
-    public long getProtocolVersion(String protocol, 
-                                   long clientVersion) throws IOException {
-      if (protocol.equals(MiniProtocol.class.getName()))
-        return versionID;
-      throw new IOException("Unknown protocol: " + protocol);
-    }
-
-    @Override // VersionedProtocol
-    public ProtocolSignature getProtocolSignature(String protocol, 
-                                   long clientVersion,
-                                   int clientMethodsHashCode) throws IOException {
-      if (protocol.equals(MiniProtocol.class.getName()))
-        return new ProtocolSignature(versionID, null);
-      throw new IOException("Unknown protocol: " + protocol);
-    }
-
     @Override // MiniProtocol
-    public Token<TestDelegationTokenIdentifier> getDelegationToken(Text renewer) 
-    throws IOException {
-      String owner = UserGroupInformation.getCurrentUser().getUserName();
-      String realUser = 
-        UserGroupInformation.getCurrentUser().getRealUser() == null ? "":
-        UserGroupInformation.getCurrentUser().getRealUser().getUserName();
-      TestDelegationTokenIdentifier tokenId = 
-        new TestDelegationTokenIdentifier(
-            new Text(owner), renewer, new Text(realUser));
-      return new Token<TestDelegationTokenIdentifier>(tokenId, secretManager);
+    public MiniGetDelegationTokenResponseProto getDelegationToken(
+        RpcController controller,
+        MiniGetDelegationTokenRequestProto request) throws ServiceException {
+      try {
+        Text renewer = new Text(request.getRenewer());
+        String owner = UserGroupInformation.getCurrentUser().getUserName();
+        String realUser =
+            UserGroupInformation.getCurrentUser().getRealUser() == null ? "" :
+                UserGroupInformation.getCurrentUser().getRealUser().getUserName();
+        TestDelegationTokenIdentifier tokenId =
+            new TestDelegationTokenIdentifier(
+                new Text(owner), renewer, new Text(realUser));
+        Token<TestDelegationTokenIdentifier> token =
+            new Token<>(tokenId, secretManager);
+        return MiniGetDelegationTokenResponseProto.newBuilder()
+            .setToken(MiniDelegationTokenProto.newBuilder()
+                .setIdentifier(ByteString.copyFrom(token.getIdentifier()))
+                .setPassword(ByteString.copyFrom(token.getPassword()))
+                .setKind(token.getKind().toString())
+                .setService(token.getService().toString())
+                .build())
+            .build();
+      } catch (IOException ioe) {
+        throw new ServiceException(ioe);
+      }
     }
 
     /** Start RPC server */
@@ -164,8 +171,11 @@ public Token<TestDelegationTokenIdentifier> getDelegationToken(Text renewer)
         new TestDelegationTokenSecretManager(24*60*60*1000,
             7*24*60*60*1000,24*60*60*1000,3600000);
       secretManager.startThreads();
+      RPC.setProtocolEngine(conf, MiniProtocol.class, ProtobufRpcEngine2.class);
+      BlockingService service =
+          MiniProtocolService.newReflectiveBlockingService(this);
       rpcServer = new RPC.Builder(conf).setProtocol(MiniProtocol.class)
-          .setInstance(this).setBindAddress(DEFAULT_SERVER_ADDRESS).setPort(0)
+          .setInstance(service).setBindAddress(DEFAULT_SERVER_ADDRESS).setPort(0)
           .setNumHandlers(1).setVerbose(false).setSecretManager(secretManager)
           .build();
       rpcServer.start();
@@ -189,8 +199,8 @@ long connectToServer(Configuration conf, InetSocketAddress addr)
     MiniProtocol client = null;
     try {
       long start = Time.now();
-      client = RPC.getProxy(MiniProtocol.class,
-          MiniProtocol.versionID, addr, conf);
+      RPC.setProtocolEngine(conf, MiniProtocol.class, ProtobufRpcEngine2.class);
+      client = RPC.getProxy(MiniProtocol.class, MiniProtocol.versionID, addr, conf);
       long end = Time.now();
       return end - start;
     } finally {
@@ -211,14 +221,28 @@ void connectToServerAndGetDelegationToken(
         client =  proxyUserUgi.doAs(new PrivilegedExceptionAction<MiniProtocol>() {
           @Override
           public MiniProtocol run() throws IOException {
+            RPC.setProtocolEngine(conf, MiniProtocol.class,
+                ProtobufRpcEngine2.class);
             MiniProtocol p = RPC.getProxy(MiniProtocol.class,
                 MiniProtocol.versionID, addr, conf);
-            Token<TestDelegationTokenIdentifier> token;
-            token = p.getDelegationToken(new Text(RENEWER));
-            currentUgi = UserGroupInformation.createUserForTesting(MINI_USER, 
-                GROUP_NAMES);
-            SecurityUtil.setTokenService(token, addr);
-            currentUgi.addToken(token);
+            try {
+              MiniGetDelegationTokenResponseProto response =
+                  p.getDelegationToken(null,
+                      MiniGetDelegationTokenRequestProto.newBuilder()
+                          .setRenewer(RENEWER).build());
+              MiniDelegationTokenProto tokenProto = response.getToken();
+              Token<TestDelegationTokenIdentifier> token = new Token<>(
+                  tokenProto.getIdentifier().toByteArray(),
+                  tokenProto.getPassword().toByteArray(),
+                  new Text(tokenProto.getKind()),
+                  new Text(tokenProto.getService()));
+              currentUgi = UserGroupInformation.createUserForTesting(MINI_USER,
+                  GROUP_NAMES);
+              SecurityUtil.setTokenService(token, addr);
+              currentUgi.addToken(token);
+            } catch (ServiceException se) {
+              throw new IOException(se);
+            }
             return p;
           }
         });
@@ -239,6 +263,7 @@ long connectToServerUsingDelegationToken(
         client = currentUgi.doAs(new PrivilegedExceptionAction<MiniProtocol>() {
           @Override
           public MiniProtocol run() throws IOException {
+            RPC.setProtocolEngine(conf, MiniProtocol.class, ProtobufRpcEngine2.class);
             return RPC.getProxy(MiniProtocol.class,
                 MiniProtocol.versionID, addr, conf);
           }
diff --git a/hadoop-common-project/hadoop-common/src/test/proto/mini_rpc_benchmark.proto b/hadoop-common-project/hadoop-common/src/test/proto/mini_rpc_benchmark.proto
@@ -0,0 +1,54 @@
+/**
+ * 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.
+ */
+syntax = "proto2";
+option java_package = "org.apache.hadoop.ipc.protobuf";
+option java_outer_classname = "MiniRPCBenchmarkProtos";
+option java_generic_services = true;
+option java_generate_equals_and_hash = true;
+package hadoop.common;
+
+/**
+ * Protobuf service used by {@code MiniRPCBenchmark}.
+ *
+ * Mirror the legacy WritableRPCEngine-based MiniProtocol so the benchmark
+ * can exercise RPC connection establishment over ProtobufRpcEngine2.
+ *
+ * Messages are inlined (rather than importing Security.proto) to keep
+ * the test proto self contained and avoid additional proto path setup
+ * in the build configuration.
+ */
+
+message MiniGetDelegationTokenRequestProto {
+  required string renewer = 1;
+}
+
+message MiniDelegationTokenProto {
+  required bytes identifier = 1;
+  required bytes password = 2;
+  required string kind = 3;
+  required string service = 4;
+}
+
+message MiniGetDelegationTokenResponseProto {
+  optional MiniDelegationTokenProto token = 1;
+}
+
+service MiniProtocolService {
+  rpc getDelegationToken(MiniGetDelegationTokenRequestProto)
+      returns (MiniGetDelegationTokenResponseProto);
+}
