Fix gRPC plugin not working for server side in some case

devkanro · devkanro · commit e4a7db10b02d · 2023-02-14T15:19:38.000+08:00
diff --git a/CHANGES.md b/CHANGES.md
@@ -12,8 +12,8 @@ Release Notes.
 * Remove Powermock entirely from the test cases.
 * Fix H2 instrumentation point
 * Refactor pipeline in jedis-plugin.
-* Enhance kotlin coroutine plugin for stack tracing.
 * Add plugin to support ClickHouse JDBC driver (0.3.2.*).
+* Change gRPC instrumentation point to fix plugin not working for server side.
 
 #### Documentation
 * Update docs of Tracing APIs, reorganize the API docs into six parts.
diff --git a/apm-sniffer/apm-sdk-plugin/grpc-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/grpc/v1/define/AbstractServerImplBuilderInstrumentation.java b/apm-sniffer/apm-sdk-plugin/grpc-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/grpc/v1/define/AbstractServerImplBuilderInstrumentation.java
@@ -26,15 +26,14 @@
 import org.apache.skywalking.apm.agent.core.plugin.match.ClassMatch;
 
 import static net.bytebuddy.matcher.ElementMatchers.named;
-import static org.apache.skywalking.apm.agent.core.plugin.bytebuddy.ArgumentTypeNameMatch.takesArgumentWithType;
+import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments;
 import static org.apache.skywalking.apm.agent.core.plugin.match.NameMatch.byName;
 
 public class AbstractServerImplBuilderInstrumentation extends ClassInstanceMethodsEnhancePluginDefine {
 
     public static final String ENHANCE_CLASS = "io.grpc.internal.AbstractServerImplBuilder";
-    public static final String ENHANCE_METHOD = "addService";
+    public static final String ENHANCE_METHOD = "build";
     public static final String INTERCEPT_CLASS = "org.apache.skywalking.apm.plugin.grpc.v1.server.AbstractServerImplBuilderInterceptor";
-    public static final String ARGUMENT_TYPE = "io.grpc.ServerServiceDefinition";
 
     @Override
     public ConstructorInterceptPoint[] getConstructorsInterceptPoints() {
@@ -47,7 +46,7 @@ public InstanceMethodsInterceptPoint[] getInstanceMethodsInterceptPoints() {
             new InstanceMethodsInterceptPoint() {
                 @Override
                 public ElementMatcher<MethodDescription> getMethodsMatcher() {
-                    return named(ENHANCE_METHOD).and(takesArgumentWithType(0, ARGUMENT_TYPE));
+                    return named(ENHANCE_METHOD).and(takesNoArguments());
                 }
 
                 @Override
diff --git a/apm-sniffer/apm-sdk-plugin/grpc-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/grpc/v1/server/AbstractServerImplBuilderInterceptor.java b/apm-sniffer/apm-sdk-plugin/grpc-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/grpc/v1/server/AbstractServerImplBuilderInterceptor.java
@@ -18,8 +18,7 @@
 
 package org.apache.skywalking.apm.plugin.grpc.v1.server;
 
-import io.grpc.ServerInterceptors;
-import io.grpc.ServerServiceDefinition;
+import io.grpc.ServerBuilder;
 
 import java.lang.reflect.Method;
 
@@ -34,7 +33,12 @@ public class AbstractServerImplBuilderInterceptor implements InstanceMethodsArou
     @Override
     public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
         MethodInterceptResult result) {
-        allArguments[0] = ServerInterceptors.intercept((ServerServiceDefinition) allArguments[0], new ServerInterceptor());
+        if (objInst.getSkyWalkingDynamicField() == null) {
+            ServerBuilder<?> builder = (ServerBuilder) objInst;
+            ServerInterceptor interceptor = new ServerInterceptor();
+            builder.intercept(interceptor);
+            objInst.setSkyWalkingDynamicField(interceptor);
+        }
     }
 
     @Override
diff --git a/apm-sniffer/apm-sdk-plugin/grpc-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/grpc/v1/server/ServerInterceptor.java b/apm-sniffer/apm-sdk-plugin/grpc-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/grpc/v1/server/ServerInterceptor.java
@@ -18,14 +18,25 @@
 
 package org.apache.skywalking.apm.plugin.grpc.v1.server;
 
+import io.grpc.Context;
+import io.grpc.Contexts;
 import io.grpc.Metadata;
 import io.grpc.ServerCall;
 import io.grpc.ServerCallHandler;
 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;
+import org.apache.skywalking.apm.agent.core.context.ContextSnapshot;
+import org.apache.skywalking.apm.agent.core.context.trace.AbstractSpan;
+import org.apache.skywalking.apm.agent.core.context.trace.SpanLayer;
+import org.apache.skywalking.apm.network.trace.component.ComponentsDefine;
 import org.apache.skywalking.apm.util.StringUtil;
 
 public class ServerInterceptor implements io.grpc.ServerInterceptor {
+
+    static final Context.Key<ContextSnapshot> CONTEXT_SNAPSHOT_KEY = Context.key("skywalking-grpc-context-snapshot");
+    static final Context.Key<AbstractSpan> ACTIVE_SPAN_KEY = Context.key("skywalking-grpc-active-span");
+
     @Override
     public <REQUEST, RESPONSE> ServerCall.Listener<REQUEST> interceptCall(ServerCall<REQUEST, RESPONSE> call,
         Metadata headers, ServerCallHandler<REQUEST, RESPONSE> handler) {
@@ -38,7 +49,26 @@ public <REQUEST, RESPONSE> ServerCall.Listener<REQUEST> interceptCall(ServerCall
                 next.setHeadValue(contextValue);
             }
         }
-        return new TracingServerCallListener<>(handler.startCall(new TracingServerCall<>(call), headers), call
-                .getMethodDescriptor(), contextCarrier);
+
+        final AbstractSpan span = ContextManager
+                .createEntrySpan(call.getMethodDescriptor().getFullMethodName(), contextCarrier);
+        span.setComponent(ComponentsDefine.GRPC);
+        span.setLayer(SpanLayer.RPC_FRAMEWORK);
+        ContextSnapshot contextSnapshot = ContextManager.capture();
+        AbstractSpan asyncSpan = span.prepareForAsync();
+
+        Context context = Context.current().withValues(CONTEXT_SNAPSHOT_KEY, contextSnapshot, ACTIVE_SPAN_KEY, asyncSpan);
+
+        ServerCall.Listener<REQUEST> listener = Contexts.interceptCall(
+                context,
+                new TracingServerCall<>(call),
+                headers,
+                (serverCall, metadata) -> new TracingServerCallListener<>(
+                        handler.startCall(serverCall, metadata),
+                        serverCall.getMethodDescriptor()
+                )
+        );
+        ContextManager.stopSpan(asyncSpan);
+        return listener;
     }
 }
diff --git a/apm-sniffer/apm-sdk-plugin/grpc-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/grpc/v1/server/TracingServerCall.java b/apm-sniffer/apm-sdk-plugin/grpc-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/grpc/v1/server/TracingServerCall.java
@@ -49,7 +49,7 @@ public void sendMessage(RESPONSE message) {
             final AbstractSpan span = ContextManager.createLocalSpan(operationPrefix + RESPONSE_ON_MESSAGE_OPERATION_NAME);
             span.setComponent(ComponentsDefine.GRPC);
             span.setLayer(SpanLayer.RPC_FRAMEWORK);
-
+            ContextManager.continued(ServerInterceptor.CONTEXT_SNAPSHOT_KEY.get());
             try {
                 super.sendMessage(message);
             } catch (Throwable t) {
@@ -68,6 +68,7 @@ public void close(Status status, Metadata trailers) {
         final AbstractSpan span = ContextManager.createLocalSpan(operationPrefix + RESPONSE_ON_CLOSE_OPERATION_NAME);
         span.setComponent(ComponentsDefine.GRPC);
         span.setLayer(SpanLayer.RPC_FRAMEWORK);
+        ContextManager.continued(ServerInterceptor.CONTEXT_SNAPSHOT_KEY.get());
         switch (status.getCode()) {
             case OK:
                 break;
@@ -93,6 +94,7 @@ public void close(Status status, Metadata trailers) {
                 break;
         }
         Tags.RPC_RESPONSE_STATUS_CODE.set(span, status.getCode().name());
+        ServerInterceptor.ACTIVE_SPAN_KEY.get().asyncFinish();
 
         try {
             super.close(status, trailers);
diff --git a/apm-sniffer/apm-sdk-plugin/grpc-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/grpc/v1/server/TracingServerCallListener.java b/apm-sniffer/apm-sdk-plugin/grpc-1.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/grpc/v1/server/TracingServerCallListener.java
@@ -21,9 +21,7 @@
 import io.grpc.ForwardingServerCallListener;
 import io.grpc.MethodDescriptor;
 import io.grpc.ServerCall;
-import org.apache.skywalking.apm.agent.core.context.ContextCarrier;
 import org.apache.skywalking.apm.agent.core.context.ContextManager;
-import org.apache.skywalking.apm.agent.core.context.ContextSnapshot;
 import org.apache.skywalking.apm.agent.core.context.trace.AbstractSpan;
 import org.apache.skywalking.apm.agent.core.context.trace.SpanLayer;
 import org.apache.skywalking.apm.network.trace.component.ComponentsDefine;
@@ -36,18 +34,11 @@
 public class TracingServerCallListener<REQUEST> extends ForwardingServerCallListener.SimpleForwardingServerCallListener<REQUEST> {
     private final MethodDescriptor.MethodType methodType;
     private final String operationPrefix;
-    private final String operation;
-    private final ContextCarrier contextCarrier;
 
-    private AbstractSpan asyncSpan;
-    private ContextSnapshot contextSnapshot;
-
-    protected TracingServerCallListener(ServerCall.Listener<REQUEST> delegate, MethodDescriptor<REQUEST, ?> descriptor, ContextCarrier contextCarrier) {
+    protected TracingServerCallListener(ServerCall.Listener<REQUEST> delegate, MethodDescriptor<REQUEST, ?> descriptor) {
         super(delegate);
         this.methodType = descriptor.getType();
         this.operationPrefix = OperationNameFormatUtil.formatOperationName(descriptor) + SERVER;
-        this.operation = OperationNameFormatUtil.formatOperationName(descriptor);
-        this.contextCarrier = contextCarrier;
     }
 
     @Override
@@ -57,7 +48,7 @@ public void onMessage(REQUEST message) {
             final AbstractSpan span = ContextManager.createLocalSpan(operationPrefix + REQUEST_ON_MESSAGE_OPERATION_NAME);
             span.setComponent(ComponentsDefine.GRPC);
             span.setLayer(SpanLayer.RPC_FRAMEWORK);
-            ContextManager.continued(contextSnapshot);
+            ContextManager.continued(ServerInterceptor.CONTEXT_SNAPSHOT_KEY.get());
             try {
                 super.onMessage(message);
             } catch (Throwable t) {
@@ -73,21 +64,18 @@ public void onMessage(REQUEST message) {
 
     @Override
     public void onCancel() {
-        if (contextSnapshot == null) {
-            return;
-        }
         final AbstractSpan span = ContextManager.createLocalSpan(operationPrefix + REQUEST_ON_CANCEL_OPERATION_NAME);
         span.setComponent(ComponentsDefine.GRPC);
         span.setLayer(SpanLayer.RPC_FRAMEWORK);
-        ContextManager.continued(contextSnapshot);
+        ContextManager.continued(ServerInterceptor.CONTEXT_SNAPSHOT_KEY.get());
         try {
             super.onCancel();
         } catch (Throwable t) {
             ContextManager.activeSpan().log(t);
             throw t;
         } finally {
             ContextManager.stopSpan();
-            asyncSpan.asyncFinish();
+            ServerInterceptor.ACTIVE_SPAN_KEY.get().asyncFinish();
         }
     }
 
@@ -96,7 +84,7 @@ public void onHalfClose() {
         final AbstractSpan span = ContextManager.createLocalSpan(operationPrefix + REQUEST_ON_HALF_CLOSE_OPERATION_NAME);
         span.setComponent(ComponentsDefine.GRPC);
         span.setLayer(SpanLayer.RPC_FRAMEWORK);
-        ContextManager.continued(contextSnapshot);
+        ContextManager.continued(ServerInterceptor.CONTEXT_SNAPSHOT_KEY.get());
         try {
             super.onHalfClose();
         } catch (Throwable t) {
@@ -110,18 +98,10 @@ public void onHalfClose() {
     @Override
     public void onComplete() {
         super.onComplete();
-        asyncSpan.asyncFinish();
     }
 
     @Override
     public void onReady() {
-        final AbstractSpan span = ContextManager.createEntrySpan(operation, contextCarrier);
-        span.setComponent(ComponentsDefine.GRPC);
-        span.setLayer(SpanLayer.RPC_FRAMEWORK);
-        contextSnapshot = ContextManager.capture();
-        asyncSpan = span.prepareForAsync();
-        ContextManager.stopSpan(asyncSpan);
-
         super.onReady();
     }
 }
