From caf800d5290d8618003b764afb0b3ef8d9a5a0a8 Mon Sep 17 00:00:00 2001
From: Kihwal Lee <kihwal@apache.org>
Date: Tue, 9 Aug 2016 14:33:15 -0500
Subject: [PATCH] HADOOP-13473. Tracing in IPC Server is broken. Contributed by
 Daryn Sharp.

---
 .../apache/hadoop/ipc/ProtobufRpcEngine.java  | 36 ++++++++++++++++---
 1 file changed, 32 insertions(+), 4 deletions(-)

diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/ProtobufRpcEngine.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/ProtobufRpcEngine.java
index cce516618bf..eb30aa207e4 100644
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/ProtobufRpcEngine.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/ProtobufRpcEngine.java
@@ -69,7 +69,7 @@ public class ProtobufRpcEngine implements RpcEngine {
 
   static { // Register the rpcRequest deserializer for WritableRpcEngine 
     org.apache.hadoop.ipc.Server.registerProtocolEngine(
-        RPC.RpcKind.RPC_PROTOCOL_BUFFER, RpcWritable.Buffer.class,
+        RPC.RpcKind.RPC_PROTOCOL_BUFFER, RpcProtobufRequest.class,
         new Server.ProtoBufRpcInvoker());
   }
 
@@ -613,9 +613,8 @@ public class ProtobufRpcEngine implements RpcEngine {
        */
       public Writable call(RPC.Server server, String connectionProtocolName,
           Writable writableRequest, long receiveTime) throws Exception {
-        RpcWritable.Buffer request = (RpcWritable.Buffer) writableRequest;
-        RequestHeaderProto rpcRequest =
-            request.getValue(RequestHeaderProto.getDefaultInstance());
+        RpcProtobufRequest request = (RpcProtobufRequest) writableRequest;
+        RequestHeaderProto rpcRequest = request.getRequestHeader();
         String methodName = rpcRequest.getMethodName();
 
         /** 
@@ -687,4 +686,33 @@ public class ProtobufRpcEngine implements RpcEngine {
       }
     }
   }
+
+  // htrace in the ipc layer creates the span name based on toString()
+  // which uses the rpc header.  in the normal case we want to defer decoding
+  // the rpc header until needed by the rpc engine.
+  static class RpcProtobufRequest extends RpcWritable.Buffer {
+    private RequestHeaderProto lazyHeader;
+
+    public RpcProtobufRequest() {
+    }
+
+    synchronized RequestHeaderProto getRequestHeader() throws IOException {
+      if (lazyHeader == null) {
+        lazyHeader = getValue(RequestHeaderProto.getDefaultInstance());
+      }
+      return lazyHeader;
+    }
+
+    // this is used by htrace to name the span.
+    @Override
+    public String toString() {
+      try {
+        RequestHeaderProto header = getRequestHeader();
+        return header.getDeclaringClassProtocolName() + "." +
+               header.getMethodName();
+      } catch (IOException e) {
+        throw new IllegalArgumentException(e);
+      }
+    }
+  }
 }