TEZ-4708: TezTaskUmbilicalProtocol to use a protobuf-based protocol

tez-dag/src/main/java/org/apache/tez/dag/app/TezTaskCommunicatorImpl.java

@@ -1,15 +1,20 @@
 /*
- * Licensed 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
+ * 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
+ *   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.
+ * 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.tez.dag.app;
@@ -25,6 +30,7 @@
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.CommonConfigurationKeysPublic;
+import org.apache.hadoop.ipc.ProtobufRpcEngine2;
 import org.apache.hadoop.ipc.ProtocolSignature;
 import org.apache.hadoop.ipc.RPC;
 import org.apache.hadoop.ipc.Server;
@@ -53,6 +59,9 @@
 import org.apache.tez.runtime.api.impl.TaskSpec;
 import org.apache.tez.runtime.api.impl.TezHeartbeatRequest;
 import org.apache.tez.runtime.api.impl.TezHeartbeatResponse;
+import org.apache.tez.runtime.internals.protocolPB.TezTaskUmbilicalProtocolBlockingPB;
+import org.apache.tez.runtime.internals.protocolPB.TezTaskUmbilicalProtocolPBServerImpl;
+import org.apache.tez.runtime.internals.protocolPB.TezTaskUmbilicalProtocolProtos;
 import org.apache.tez.serviceplugins.api.ContainerEndReason;
 import org.apache.tez.serviceplugins.api.TaskAttemptEndReason;
 import org.apache.tez.serviceplugins.api.TaskCommunicator;
@@ -61,6 +70,7 @@
 import org.apache.tez.serviceplugins.api.TaskHeartbeatResponse;
 
 import com.google.common.collect.Maps;
+import com.google.protobuf.BlockingService;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -150,11 +160,19 @@ protected void startRpcServer() {
       JobTokenSecretManager jobTokenSecretManager = new JobTokenSecretManager(conf);
       jobTokenSecretManager.addTokenForJob(tokenIdentifier, sessionToken);
 
+      RPC.setProtocolEngine(
+          conf, TezTaskUmbilicalProtocolBlockingPB.class, ProtobufRpcEngine2.class);
+      TezTaskUmbilicalProtocolBlockingPB service =
+          new TezTaskUmbilicalProtocolPBServerImpl(taskUmbilical);
+      BlockingService umbilicalService =
+          TezTaskUmbilicalProtocolProtos.TezTaskUmbilicalProtocol.newReflectiveBlockingService(
+              service);
+
       server = new RPC.Builder(conf)
-          .setProtocol(TezTaskUmbilicalProtocol.class)
+          .setProtocol(TezTaskUmbilicalProtocolBlockingPB.class)
           .setBindAddress("0.0.0.0")
           .setPort(0)
-          .setInstance(taskUmbilical)
+          .setInstance(umbilicalService)
           .setNumHandlers(
               conf.getInt(TezConfiguration.TEZ_AM_TASK_LISTENER_THREAD_COUNT,
                   TezConfiguration.TEZ_AM_TASK_LISTENER_THREAD_COUNT_DEFAULT))
@@ -388,7 +406,6 @@ public TezHeartbeatResponse heartbeat(TezHeartbeatRequest request) throws IOExce
       return response;
     }
 
-
     // TODO Remove this method once we move to the Protobuf RPC engine
     @Override
     public long getProtocolVersion(String protocol, long clientVersion) throws IOException {

tez-dag/src/main/java/org/apache/tez/dag/app/security/authorize/TezAMPolicyProvider.java

@@ -1,4 +1,4 @@
-/**
+/*
  * 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
@@ -7,23 +7,24 @@
  * "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
+ *   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.
+ * 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.tez.dag.app.security.authorize;
 
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.classification.InterfaceStability;
 import org.apache.hadoop.security.authorize.PolicyProvider;
 import org.apache.hadoop.security.authorize.Service;
-import org.apache.tez.common.TezTaskUmbilicalProtocol;
 import org.apache.tez.dag.api.TezConstants;
 import org.apache.tez.dag.api.client.rpc.DAGClientAMProtocolBlockingPB;
+import org.apache.tez.runtime.internals.protocolPB.TezTaskUmbilicalProtocolBlockingPB;
 
 /**
  * {@link PolicyProvider} for YARN Tez client protocols.
@@ -36,7 +37,7 @@ public class TezAMPolicyProvider extends PolicyProvider {
       new Service[] {
     new Service(
         TezConstants.TEZ_AM_SECURITY_SERVICE_AUTHORIZATION_TASK_UMBILICAL,
-          TezTaskUmbilicalProtocol.class),
+          TezTaskUmbilicalProtocolBlockingPB.class),
     new Service(
     	TezConstants.TEZ_AM_SECURITY_SERVICE_AUTHORIZATION_CLIENT,
           DAGClientAMProtocolBlockingPB.class)

tez-ext-service-tests/src/test/java/org/apache/tez/service/impl/ContainerRunnerImpl.java

@@ -1,15 +1,20 @@
 /*
- * Licensed 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
+ * 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
+ *   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.
+ * 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.tez.service.impl;
@@ -34,6 +39,7 @@
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.io.DataInputBuffer;
+import org.apache.hadoop.ipc.ProtobufRpcEngine2;
 import org.apache.hadoop.ipc.RPC;
 import org.apache.hadoop.net.NetUtils;
 import org.apache.hadoop.security.Credentials;
@@ -57,6 +63,8 @@
 import org.apache.tez.runtime.api.impl.ExecutionContextImpl;
 import org.apache.tez.runtime.api.impl.TaskSpec;
 import org.apache.tez.runtime.common.objectregistry.ObjectRegistryImpl;
+import org.apache.tez.runtime.internals.protocolPB.TezTaskUmbilicalProtocolBlockingPB;
+import org.apache.tez.runtime.internals.protocolPB.TezTaskUmbilicalProtocolPBClientImpl;
 import org.apache.tez.runtime.task.TaskReporter;
 import org.apache.tez.runtime.task.TaskRunner2Result;
 import org.apache.tez.runtime.task.TezChild;
@@ -439,13 +447,20 @@ public ContainerExecutionResult call() throws Exception {
           NetUtils.createSocketAddrForHost(request.getAmHost(), request.getAmPort());
       SecurityUtil.setTokenService(jobToken, address);
       taskOwner.addToken(jobToken);
-      umbilical = taskOwner.doAs(new PrivilegedExceptionAction<TezTaskUmbilicalProtocol>() {
-        @Override
-        public TezTaskUmbilicalProtocol run() throws Exception {
-          return RPC.getProxy(TezTaskUmbilicalProtocol.class,
-              TezTaskUmbilicalProtocol.versionID, address, conf);
-        }
-      });
+      umbilical =
+          taskOwner.doAs(
+              (PrivilegedExceptionAction<TezTaskUmbilicalProtocol>)
+                  () -> {
+                    RPC.setProtocolEngine(
+                        conf, TezTaskUmbilicalProtocolBlockingPB.class, ProtobufRpcEngine2.class);
+                    TezTaskUmbilicalProtocolBlockingPB proxy =
+                        RPC.getProxy(
+                            TezTaskUmbilicalProtocolBlockingPB.class,
+                            RPC.getProtocolVersion(TezTaskUmbilicalProtocolBlockingPB.class),
+                            address,
+                            conf);
+                    return new TezTaskUmbilicalProtocolPBClientImpl(proxy);
+                  });
       // TODO Stop reading this on each request.
       taskReporter = new TaskReporter(
           umbilical,