Merge pull request #980 from marcoferrer/grpc-interceptor-error-tag

Fix span error tagging in grpc server interceptor

dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/server/GrpcServerBuilderInstrumentation.java

@@ -31,6 +31,7 @@ public class GrpcServerBuilderInstrumentation extends Instrumenter.Default {
   public String[] helperClassNames() {
     return new String[] {
       "datadog.trace.instrumentation.grpc.server.TracingServerInterceptor",
+      "datadog.trace.instrumentation.grpc.server.TracingServerInterceptor$TracingServerCall",
       "datadog.trace.instrumentation.grpc.server.TracingServerInterceptor$TracingServerCallListener",
       "datadog.trace.agent.decorator.BaseDecorator",
       "datadog.trace.agent.decorator.ServerDecorator",

dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/server/GrpcServerDecorator.java

@@ -2,6 +2,9 @@ package datadog.trace.instrumentation.grpc.server;
 
 import datadog.trace.agent.decorator.ServerDecorator;
 import datadog.trace.api.DDSpanTypes;
+import io.grpc.Status;
+import io.opentracing.Span;
+import io.opentracing.tag.Tags;
 
 public class GrpcServerDecorator extends ServerDecorator {
   public static final GrpcServerDecorator DECORATE = new GrpcServerDecorator();
@@ -20,4 +23,17 @@ public class GrpcServerDecorator extends ServerDecorator {
   protected String component() {
     return "grpc-server";
   }
+
+  public Span onClose(final Span span, final Status status) {
+
+    span.setTag("status.code", status.getCode().name());
+    span.setTag("status.description", status.getDescription());
+
+    onError(span, status.getCause());
+    if (!status.isOk()) {
+      Tags.ERROR.set(span, true);
+    }
+
+    return span;
+  }
 }

dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/server/TracingServerInterceptor.java

@@ -4,11 +4,13 @@ import static datadog.trace.instrumentation.grpc.server.GrpcServerDecorator.DECO
 
 import datadog.trace.api.DDTags;
 import datadog.trace.context.TraceScope;
+import io.grpc.ForwardingServerCall;
 import io.grpc.ForwardingServerCallListener;
 import io.grpc.Metadata;
 import io.grpc.ServerCall;
 import io.grpc.ServerCallHandler;
 import io.grpc.ServerInterceptor;
+import io.grpc.Status;
 import io.opentracing.Scope;
 import io.opentracing.Span;
 import io.opentracing.SpanContext;
@@ -60,8 +62,13 @@ public class TracingServerInterceptor implements ServerInterceptor {
 
     final ServerCall.Listener<ReqT> result;
     try {
+      // Wrap the server call so that we can decorate the span
+      // with the resulting status
+      TracingServerCall<ReqT, RespT> tracingServerCall =
+          new TracingServerCall<>(tracer, span, call);
+
       // call other interceptors
-      result = next.startCall(call, headers);
+      result = next.startCall(tracingServerCall, headers);
     } catch (final Throwable e) {
       DECORATE.onError(span, e);
       DECORATE.beforeFinish(span);
@@ -78,6 +85,36 @@ public class TracingServerInterceptor implements ServerInterceptor {
     return new TracingServerCallListener<>(tracer, span, result);
   }
 
+  static final class TracingServerCall<ReqT, RespT>
+      extends ForwardingServerCall.SimpleForwardingServerCall<ReqT, RespT> {
+    final Tracer tracer;
+    final Span span;
+
+    TracingServerCall(
+        final Tracer tracer, final Span span, final ServerCall<ReqT, RespT> delegate) {
+      super(delegate);
+      this.tracer = tracer;
+      this.span = span;
+    }
+
+    @Override
+    public void close(final Status status, final Metadata trailers) {
+      DECORATE.onClose(span, status);
+      try (final Scope scope = tracer.scopeManager().activate(span, false)) {
+        if (scope instanceof TraceScope) {
+          ((TraceScope) scope).setAsyncPropagation(true);
+        }
+        delegate().close(status, trailers);
+        if (scope instanceof TraceScope) {
+          ((TraceScope) scope).setAsyncPropagation(false);
+        }
+      } catch (final Throwable e) {
+        DECORATE.onError(span, e);
+        throw e;
+      }
+    }
+  }
+
   static final class TracingServerCallListener<ReqT>
       extends ForwardingServerCallListener.SimpleForwardingServerCallListener<ReqT> {
     final Tracer tracer;

dd-java-agent/instrumentation/grpc-1.5/src/test/groovy/GrpcStreamingTest.groovy

@@ -91,6 +91,7 @@ class GrpcStreamingTest extends AgentTestRunner {
           childOf trace(1).get(0)
           errored false
           tags {
+            "status.code" "OK"
             "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_SERVER
             "$Tags.COMPONENT.key" "grpc-server"
             defaultTags(true)

dd-java-agent/instrumentation/grpc-1.5/src/test/groovy/GrpcTest.groovy

@@ -47,6 +47,7 @@ class GrpcTest extends AgentTestRunner {
           childOf trace(1).get(0)
           errored false
           tags {
+            "status.code" "OK"
             "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_SERVER
             "$Tags.COMPONENT.key" "grpc-server"
             defaultTags(true)
@@ -136,10 +137,17 @@ class GrpcTest extends AgentTestRunner {
           resourceName "example.Greeter/SayHello"
           spanType DDSpanTypes.RPC
           childOf trace(1).get(0)
-          errored false
+          errored true
           tags {
+            "status.code" "${status.code.name()}"
+            "status.description" description
             "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_SERVER
             "$Tags.COMPONENT.key" "grpc-server"
+            if(status.cause != null){
+              errorTags status.cause.class, status.cause.message
+            }else{
+              tag "error", true
+            }
             defaultTags(true)
           }
         }