From 542691011cc69ce45730ca6bb094d59b82c5f76e Mon Sep 17 00:00:00 2001
From: ZS-Oliver <37582554+ZS-Oliver@users.noreply.github.com>
Date: Mon, 21 Dec 2020 21:50:52 +0800
Subject: [PATCH] Fix thrift trace broken and wrong arg collected. (#5989)

---
 CHANGES.md                                    |  3 ++-
 .../client/TServiceClientInterceptor.java     |  1 +
 .../wrapper/ServerInProtocolWrapper.java      | 22 +++++++++++++++++++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/CHANGES.md b/CHANGES.md
index 88317a10a..77c01f773 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -17,7 +17,8 @@ Release Notes.
 * Fix the unexpected RunningContext recreation in the Tomcat plugin.
 * Fix the potential NPE when trace_sql_parameters is enabled.
 * Update `byte-buddy` to 1.10.19.
-
+* Fix thrift plugin trace link broken when intermediate service does not mount agent
+* Fix thrift plugin collects wrong args when the method without parameter.
 
 #### OAP-Backend
 * Make meter receiver support MAL.
diff --git a/apm-sniffer/apm-sdk-plugin/thrift-plugin/src/main/java/org/apache/skywalking/apm/plugin/thrift/client/TServiceClientInterceptor.java b/apm-sniffer/apm-sdk-plugin/thrift-plugin/src/main/java/org/apache/skywalking/apm/plugin/thrift/client/TServiceClientInterceptor.java
index ae3a3c978..b9a32166f 100644
--- a/apm-sniffer/apm-sdk-plugin/thrift-plugin/src/main/java/org/apache/skywalking/apm/plugin/thrift/client/TServiceClientInterceptor.java
+++ b/apm-sniffer/apm-sdk-plugin/thrift-plugin/src/main/java/org/apache/skywalking/apm/plugin/thrift/client/TServiceClientInterceptor.java
@@ -100,6 +100,7 @@ public class TServiceClientInterceptor implements InstanceConstructorInterceptor
         while (true) {
             TFieldIdEnum field = base.fieldForId(++idx);
             if (field == null) {
+                idx--;
                 break;
             }
             buffer.append(field.getFieldName()).append(", ");
diff --git a/apm-sniffer/apm-sdk-plugin/thrift-plugin/src/main/java/org/apache/skywalking/apm/plugin/thrift/wrapper/ServerInProtocolWrapper.java b/apm-sniffer/apm-sdk-plugin/thrift-plugin/src/main/java/org/apache/skywalking/apm/plugin/thrift/wrapper/ServerInProtocolWrapper.java
index a05404641..3cde8b4c4 100644
--- a/apm-sniffer/apm-sdk-plugin/thrift-plugin/src/main/java/org/apache/skywalking/apm/plugin/thrift/wrapper/ServerInProtocolWrapper.java
+++ b/apm-sniffer/apm-sdk-plugin/thrift-plugin/src/main/java/org/apache/skywalking/apm/plugin/thrift/wrapper/ServerInProtocolWrapper.java
@@ -21,6 +21,7 @@ package org.apache.skywalking.apm.plugin.thrift.wrapper;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
+
 import org.apache.skywalking.apm.agent.core.context.CarrierItem;
 import org.apache.skywalking.apm.agent.core.context.ContextCarrier;
 import org.apache.skywalking.apm.agent.core.context.ContextManager;
@@ -45,6 +46,7 @@ public class ServerInProtocolWrapper extends AbstractProtocolWrapper {
     private static final ILog LOGGER = LogManager.getLogger(ServerInProtocolWrapper.class);
     private static final StringTag TAG_ARGS = new StringTag("args");
     private AbstractContext context;
+    private static final String HAVE_CREATED_SPAN = "HAVE_CREATED_SPAN";
 
     public ServerInProtocolWrapper(final TProtocol protocol) {
         super(protocol);
@@ -52,6 +54,7 @@ public class ServerInProtocolWrapper extends AbstractProtocolWrapper {
 
     public void initial(AbstractContext context) {
         this.context = context;
+        ContextManager.getRuntimeContext().put(HAVE_CREATED_SPAN, false);
     }
 
     @Override
@@ -72,6 +75,7 @@ public class ServerInProtocolWrapper extends AbstractProtocolWrapper {
                 span.tag(TAG_ARGS, context.getArguments());
                 span.setComponent(ComponentsDefine.THRIFT_SERVER);
                 SpanLayer.asRPCFramework(span);
+                ContextManager.getRuntimeContext().put(HAVE_CREATED_SPAN, true);
             } catch (Throwable throwable) {
                 LOGGER.error("Failed to resolve header or create EntrySpan.", throwable);
             } finally {
@@ -81,6 +85,24 @@ public class ServerInProtocolWrapper extends AbstractProtocolWrapper {
             }
             return readFieldBegin();
         }
+        if (field.type == TType.STOP) {
+            Boolean haveCreatedSpan =
+                    (Boolean) ContextManager.getRuntimeContext().get(HAVE_CREATED_SPAN);
+            if (haveCreatedSpan != null && !haveCreatedSpan) {
+                try {
+                    AbstractSpan span = ContextManager.createEntrySpan(
+                            context.getOperatorName(), createContextCarrier(null));
+                    span.start(context.startTime);
+                    span.tag(TAG_ARGS, context.getArguments());
+                    span.setComponent(ComponentsDefine.THRIFT_SERVER);
+                    SpanLayer.asRPCFramework(span);
+                } catch (Throwable throwable) {
+                    LOGGER.error("Failed to create EntrySpan.", throwable);
+                } finally {
+                    context = null;
+                }
+            }
+        }
         return field;
     }