Migrate gRPC instrumentation to Decorator

dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/client/GrpcClientBuilderInstrumentation.java

@@ -34,6 +34,9 @@ public class GrpcClientBuilderInstrumentation extends Instrumenter.Default {
       "datadog.trace.instrumentation.grpc.client.TracingClientInterceptor",
       "datadog.trace.instrumentation.grpc.client.TracingClientInterceptor$TracingClientCall",
       "datadog.trace.instrumentation.grpc.client.TracingClientInterceptor$TracingClientCallListener",
+      "datadog.trace.agent.decorator.BaseDecorator",
+      "datadog.trace.agent.decorator.ClientDecorator",
+      packageName + ".GrpcClientDecorator",
     };
   }
 

dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/client/GrpcClientDecorator.java

@@ -0,0 +1,44 @@
+package datadog.trace.instrumentation.grpc.client;
+
+import datadog.trace.agent.decorator.ClientDecorator;
+import datadog.trace.api.DDSpanTypes;
+import io.grpc.Status;
+import io.opentracing.Span;
+import io.opentracing.tag.Tags;
+
+public class GrpcClientDecorator extends ClientDecorator {
+  public static final GrpcClientDecorator DECORATE = new GrpcClientDecorator();
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"grpc", "grpc-client"};
+  }
+
+  @Override
+  protected String component() {
+    return "grpc-client";
+  }
+
+  @Override
+  protected String spanType() {
+    return DDSpanTypes.RPC;
+  }
+
+  @Override
+  protected String service() {
+    return null;
+  }
+
+  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/client/TracingClientInterceptor.java

@@ -1,8 +1,7 @@
 package datadog.trace.instrumentation.grpc.client;
 
-import static io.opentracing.log.Fields.ERROR_OBJECT;
+import static datadog.trace.instrumentation.grpc.client.GrpcClientDecorator.DECORATE;
 
-import datadog.trace.api.DDSpanTypes;
 import datadog.trace.api.DDTags;
 import io.grpc.CallOptions;
 import io.grpc.Channel;
@@ -17,8 +16,6 @@ import io.opentracing.Scope;
 import io.opentracing.Span;
 import io.opentracing.Tracer;
 import io.opentracing.propagation.Format;
-import io.opentracing.tag.Tags;
-import java.util.Collections;
 
 public class TracingClientInterceptor implements ClientInterceptor {
 
@@ -34,27 +31,22 @@ public class TracingClientInterceptor implements ClientInterceptor {
       final CallOptions callOptions,
       final Channel next) {
 
-    final Scope scope =
+    final Span span =
         tracer
             .buildSpan("grpc.client")
             .withTag(DDTags.RESOURCE_NAME, method.getFullMethodName())
-            .withTag(DDTags.SPAN_TYPE, DDSpanTypes.RPC)
+            .start();
-            .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT)
+    DECORATE.afterStart(span);
-            .withTag(Tags.COMPONENT.getKey(), "grpc-client")
-            .startActive(false);
-    final Span span = scope.span();
 
     final ClientCall<ReqT, RespT> result;
-    try {
+    try (final Scope ignore = tracer.scopeManager().activate(span, false)) {
       // call other interceptors
       result = next.newCall(method, callOptions);
     } catch (final Throwable e) {
-      Tags.ERROR.set(span, true);
+      DECORATE.onError(span, e);
-      span.log(Collections.singletonMap(ERROR_OBJECT, e));
+      DECORATE.beforeFinish(span);
       span.finish();
       throw e;
-    } finally {
-      scope.close();
     }
 
     return new TracingClientCall<>(tracer, span, result);
@@ -79,8 +71,8 @@ public class TracingClientInterceptor implements ClientInterceptor {
       try (final Scope ignored = tracer.scopeManager().activate(span, false)) {
         super.start(new TracingClientCallListener<>(tracer, span, responseListener), headers);
       } catch (final Throwable e) {
-        Tags.ERROR.set(span, true);
+        DECORATE.onError(span, e);
-        span.log(Collections.singletonMap(ERROR_OBJECT, e));
+        DECORATE.beforeFinish(span);
         span.finish();
         throw e;
       }
@@ -91,8 +83,8 @@ public class TracingClientInterceptor implements ClientInterceptor {
       try (final Scope ignored = tracer.scopeManager().activate(span, false)) {
         super.sendMessage(message);
       } catch (final Throwable e) {
-        Tags.ERROR.set(span, true);
+        DECORATE.onError(span, e);
-        span.log(Collections.singletonMap(ERROR_OBJECT, e));
+        DECORATE.beforeFinish(span);
         span.finish();
         throw e;
       }
@@ -118,41 +110,30 @@ public class TracingClientInterceptor implements ClientInterceptor {
               .buildSpan("grpc.message")
               .asChildOf(span)
               .withTag("message.type", message.getClass().getName())
-              .withTag(DDTags.SPAN_TYPE, DDSpanTypes.RPC)
-              .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT)
-              .withTag(Tags.COMPONENT.getKey(), "grpc-client")
               .startActive(true);
+      final Span messageSpan = scope.span();
+      DECORATE.afterStart(messageSpan);
       try {
         delegate().onMessage(message);
       } catch (final Throwable e) {
-        final Span span = scope.span();
+        DECORATE.onError(messageSpan, e);
-        Tags.ERROR.set(span, true);
-        this.span.log(Collections.singletonMap(ERROR_OBJECT, e));
-        this.span.finish();
         throw e;
       } finally {
+        DECORATE.beforeFinish(messageSpan);
         scope.close();
       }
     }
 
     @Override
     public void onClose(final Status status, final Metadata trailers) {
-      span.setTag("status.code", status.getCode().name());
+      DECORATE.onClose(span, status);
-      if (status.getDescription() != null) {
-        span.setTag("status.description", status.getDescription());
-      }
-      if (!status.isOk()) {
-        Tags.ERROR.set(span, true);
-      }
-      if (status.getCause() != null) {
-        span.log(Collections.singletonMap(ERROR_OBJECT, status.getCause()));
-      }
       // Finishes span.
       try (final Scope ignored = tracer.scopeManager().activate(span, true)) {
         delegate().onClose(status, trailers);
+        DECORATE.beforeFinish(span);
       } catch (final Throwable e) {
-        Tags.ERROR.set(span, true);
+        DECORATE.onError(span, e);
-        span.log(Collections.singletonMap(ERROR_OBJECT, e));
+        DECORATE.beforeFinish(span);
         span.finish();
         throw e;
       }
@@ -163,8 +144,8 @@ public class TracingClientInterceptor implements ClientInterceptor {
       try (final Scope ignored = tracer.scopeManager().activate(span, false)) {
         delegate().onReady();
       } catch (final Throwable e) {
-        Tags.ERROR.set(span, true);
+        DECORATE.onError(span, e);
-        span.log(Collections.singletonMap(ERROR_OBJECT, e));
+        DECORATE.beforeFinish(span);
         span.finish();
         throw e;
       }

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

@@ -32,6 +32,9 @@ public class GrpcServerBuilderInstrumentation extends Instrumenter.Default {
     return new String[] {
       "datadog.trace.instrumentation.grpc.server.TracingServerInterceptor",
       "datadog.trace.instrumentation.grpc.server.TracingServerInterceptor$TracingServerCallListener",
+      "datadog.trace.agent.decorator.BaseDecorator",
+      "datadog.trace.agent.decorator.ServerDecorator",
+      packageName + ".GrpcServerDecorator",
     };
   }
 

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

@@ -0,0 +1,23 @@
+package datadog.trace.instrumentation.grpc.server;
+
+import datadog.trace.agent.decorator.ServerDecorator;
+import datadog.trace.api.DDSpanTypes;
+
+public class GrpcServerDecorator extends ServerDecorator {
+  public static final GrpcServerDecorator DECORATE = new GrpcServerDecorator();
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"grpc", "grpc-server"};
+  }
+
+  @Override
+  protected String spanType() {
+    return DDSpanTypes.RPC;
+  }
+
+  @Override
+  protected String component() {
+    return "grpc-server";
+  }
+}

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

@@ -1,8 +1,7 @@
 package datadog.trace.instrumentation.grpc.server;
 
-import static io.opentracing.log.Fields.ERROR_OBJECT;
+import static datadog.trace.instrumentation.grpc.server.GrpcServerDecorator.DECORATE;
 
-import datadog.trace.api.DDSpanTypes;
 import datadog.trace.api.DDTags;
 import datadog.trace.context.TraceScope;
 import io.grpc.ForwardingServerCallListener;
@@ -16,8 +15,6 @@ import io.opentracing.SpanContext;
 import io.opentracing.Tracer;
 import io.opentracing.propagation.Format;
 import io.opentracing.propagation.TextMapExtractAdapter;
-import io.opentracing.tag.Tags;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -48,10 +45,7 @@ public class TracingServerInterceptor implements ServerInterceptor {
     final Tracer.SpanBuilder spanBuilder =
         tracer
             .buildSpan("grpc.server")
-            .withTag(DDTags.RESOURCE_NAME, call.getMethodDescriptor().getFullMethodName())
+            .withTag(DDTags.RESOURCE_NAME, call.getMethodDescriptor().getFullMethodName());
-            .withTag(DDTags.SPAN_TYPE, DDSpanTypes.RPC)
-            .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_SERVER)
-            .withTag(Tags.COMPONENT.getKey(), "grpc-server");
     if (spanContext != null) {
       spanBuilder.asChildOf(spanContext);
     }
@@ -62,14 +56,15 @@ public class TracingServerInterceptor implements ServerInterceptor {
     }
 
     final Span span = scope.span();
+    DECORATE.afterStart(span);
 
     final ServerCall.Listener<ReqT> result;
     try {
       // call other interceptors
       result = next.startCall(call, headers);
     } catch (final Throwable e) {
-      Tags.ERROR.set(span, true);
+      DECORATE.onError(span, e);
-      span.log(Collections.singletonMap(ERROR_OBJECT, e));
+      DECORATE.beforeFinish(span);
       span.finish();
       throw e;
     } finally {
@@ -102,10 +97,8 @@ public class TracingServerInterceptor implements ServerInterceptor {
               .buildSpan("grpc.message")
               .asChildOf(span)
               .withTag("message.type", message.getClass().getName())
-              .withTag(DDTags.SPAN_TYPE, DDSpanTypes.RPC)
-              .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_SERVER)
-              .withTag(Tags.COMPONENT.getKey(), "grpc-server")
               .startActive(true);
+      DECORATE.afterStart(scope.span());
       if (scope instanceof TraceScope) {
         ((TraceScope) scope).setAsyncPropagation(true);
       }
@@ -113,14 +106,15 @@ public class TracingServerInterceptor implements ServerInterceptor {
         delegate().onMessage(message);
       } catch (final Throwable e) {
         final Span span = scope.span();
-        Tags.ERROR.set(span, true);
+        DECORATE.onError(span, e);
-        this.span.log(Collections.singletonMap(ERROR_OBJECT, e));
+        DECORATE.beforeFinish(span);
         this.span.finish();
         throw e;
       } finally {
         if (scope instanceof TraceScope) {
           ((TraceScope) scope).setAsyncPropagation(false);
         }
+        DECORATE.afterStart(scope.span());
         scope.close();
       }
     }
@@ -136,8 +130,8 @@ public class TracingServerInterceptor implements ServerInterceptor {
           ((TraceScope) scope).setAsyncPropagation(false);
         }
       } catch (final Throwable e) {
-        Tags.ERROR.set(span, true);
+        DECORATE.onError(span, e);
-        span.log(Collections.singletonMap(ERROR_OBJECT, e));
+        DECORATE.beforeFinish(span);
         span.finish();
         throw e;
       }
@@ -155,9 +149,10 @@ public class TracingServerInterceptor implements ServerInterceptor {
         if (scope instanceof TraceScope) {
           ((TraceScope) scope).setAsyncPropagation(false);
         }
+        DECORATE.beforeFinish(span);
       } catch (final Throwable e) {
-        Tags.ERROR.set(span, true);
+        DECORATE.onError(span, e);
-        span.log(Collections.singletonMap(ERROR_OBJECT, e));
+        DECORATE.beforeFinish(span);
         span.finish();
         throw e;
       }
@@ -174,9 +169,10 @@ public class TracingServerInterceptor implements ServerInterceptor {
         if (scope instanceof TraceScope) {
           ((TraceScope) scope).setAsyncPropagation(false);
         }
+        DECORATE.beforeFinish(span);
       } catch (final Throwable e) {
-        Tags.ERROR.set(span, true);
+        DECORATE.onError(span, e);
-        span.log(Collections.singletonMap(ERROR_OBJECT, e));
+        DECORATE.beforeFinish(span);
         span.finish();
         throw e;
       }
@@ -193,8 +189,8 @@ public class TracingServerInterceptor implements ServerInterceptor {
           ((TraceScope) scope).setAsyncPropagation(false);
         }
       } catch (final Throwable e) {
-        Tags.ERROR.set(span, true);
+        DECORATE.onError(span, e);
-        span.log(Collections.singletonMap(ERROR_OBJECT, e));
+        DECORATE.beforeFinish(span);
         span.finish();
         throw e;
       }