Migrate JMS instrumentation to Decorator

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/decorator/ClientDecorator.java

@@ -8,13 +8,17 @@ public abstract class ClientDecorator extends BaseDecorator {
 
   protected abstract String service();
 
+  protected String spanKind() {
+    return Tags.SPAN_KIND_CLIENT;
+  }
+
   @Override
   public Span afterStart(final Span span) {
     assert span != null;
     if (service() != null) {
       span.setTag(DDTags.SERVICE_NAME, service());
     }
-    Tags.SPAN_KIND.set(span, Tags.SPAN_KIND_CLIENT);
+    Tags.SPAN_KIND.set(span, spanKind());
     return super.afterStart(span);
   }
 }

dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/JMSDecorator.java

@@ -0,0 +1,112 @@
+package datadog.trace.instrumentation.jms;
+
+import datadog.trace.agent.decorator.ClientDecorator;
+import datadog.trace.api.DDSpanTypes;
+import datadog.trace.api.DDTags;
+import io.opentracing.Scope;
+import io.opentracing.Span;
+import io.opentracing.tag.Tags;
+import java.lang.reflect.Method;
+import javax.jms.Destination;
+import javax.jms.Message;
+import javax.jms.Queue;
+import javax.jms.TemporaryQueue;
+import javax.jms.TemporaryTopic;
+import javax.jms.Topic;
+
+public abstract class JMSDecorator extends ClientDecorator {
+  public static final JMSDecorator PRODUCER_DECORATE =
+      new JMSDecorator() {
+        @Override
+        protected String spanKind() {
+          return Tags.SPAN_KIND_PRODUCER;
+        }
+
+        @Override
+        protected String spanType() {
+          return DDSpanTypes.MESSAGE_PRODUCER;
+        }
+      };
+
+  public static final JMSDecorator CONSUMER_DECORATE =
+      new JMSDecorator() {
+        @Override
+        protected String spanKind() {
+          return Tags.SPAN_KIND_CONSUMER;
+        }
+
+        @Override
+        protected String spanType() {
+          return DDSpanTypes.MESSAGE_CONSUMER;
+        }
+      };
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"jms", "jms-1", "jms-2"};
+  }
+
+  @Override
+  protected String service() {
+    return "jms";
+  }
+
+  @Override
+  protected String component() {
+    return "jms";
+  }
+
+  @Override
+  protected abstract String spanKind();
+
+  public void onConsume(final Span span, final Message message) {
+    span.setTag(DDTags.RESOURCE_NAME, "Consumed from " + toResourceName(message, null));
+  }
+
+  public void onReceive(final Span span, final Method method) {
+    span.setTag(DDTags.RESOURCE_NAME, "JMS " + method.getName());
+  }
+
+  public void onReceive(final Scope scope, final Message message) {
+    scope.span().setTag(DDTags.RESOURCE_NAME, "Received from " + toResourceName(message, null));
+  }
+
+  public void onProduce(final Scope scope, final Message message, final Destination destination) {
+    scope
+        .span()
+        .setTag(DDTags.RESOURCE_NAME, "Produced for " + toResourceName(message, destination));
+  }
+
+  private static final String TIBCO_TMP_PREFIX = "$TMP$";
+
+  public static String toResourceName(final Message message, final Destination destination) {
+    Destination jmsDestination = null;
+    try {
+      jmsDestination = message.getJMSDestination();
+    } catch (final Exception e) {
+    }
+    if (jmsDestination == null) {
+      jmsDestination = destination;
+    }
+    try {
+      if (jmsDestination instanceof Queue) {
+        final String queueName = ((Queue) jmsDestination).getQueueName();
+        if (jmsDestination instanceof TemporaryQueue || queueName.startsWith(TIBCO_TMP_PREFIX)) {
+          return "Temporary Queue";
+        } else {
+          return "Queue " + queueName;
+        }
+      }
+      if (jmsDestination instanceof Topic) {
+        final String topicName = ((Topic) jmsDestination).getTopicName();
+        if (jmsDestination instanceof TemporaryTopic || topicName.startsWith(TIBCO_TMP_PREFIX)) {
+          return "Temporary Topic";
+        } else {
+          return "Topic " + topicName;
+        }
+      }
+    } catch (final Exception e) {
+    }
+    return "Destination";
+  }
+}

dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java

@@ -1,8 +1,7 @@
 package datadog.trace.instrumentation.jms;
 
 import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
-import static datadog.trace.instrumentation.jms.JmsUtil.toResourceName;
+import static datadog.trace.instrumentation.jms.JMSDecorator.CONSUMER_DECORATE;
-import static io.opentracing.log.Fields.ERROR_OBJECT;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -11,17 +10,12 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.api.DDSpanTypes;
-import datadog.trace.api.DDTags;
-import io.opentracing.Scope;
 import io.opentracing.Span;
 import io.opentracing.SpanContext;
 import io.opentracing.Tracer;
 import io.opentracing.propagation.Format;
-import io.opentracing.tag.Tags;
 import io.opentracing.util.GlobalTracer;
 import java.lang.reflect.Method;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
@@ -46,7 +40,14 @@ public final class JMSMessageConsumerInstrumentation extends Instrumenter.Defaul
 
   @Override
   public String[] helperClassNames() {
-    return new String[] {packageName + ".JmsUtil", packageName + ".MessagePropertyTextMap"};
+    return new String[] {
+      "datadog.trace.agent.decorator.BaseDecorator",
+      "datadog.trace.agent.decorator.ClientDecorator",
+      packageName + ".JMSDecorator",
+      packageName + ".JMSDecorator$1",
+      packageName + ".JMSDecorator$2",
+      packageName + ".MessagePropertyTextMap",
+    };
   }
 
   @Override
@@ -78,20 +79,10 @@ public final class JMSMessageConsumerInstrumentation extends Instrumenter.Defaul
       Tracer.SpanBuilder spanBuilder =
           GlobalTracer.get()
               .buildSpan("jms.consume")
-              .withTag(DDTags.SERVICE_NAME, "jms")
-              .withTag(DDTags.SPAN_TYPE, DDSpanTypes.MESSAGE_CONSUMER)
-              .withTag(Tags.COMPONENT.getKey(), "jms")
-              .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CONSUMER)
               .withTag("span.origin.type", consumer.getClass().getName())
               .withStartTimestamp(TimeUnit.MILLISECONDS.toMicros(startTime));
 
-      if (message == null) {
+      if (message != null) {
-        spanBuilder = spanBuilder.withTag(DDTags.RESOURCE_NAME, "JMS " + method.getName());
-      } else {
-        spanBuilder =
-            spanBuilder.withTag(
-                DDTags.RESOURCE_NAME, "Consumed from " + toResourceName(message, null));
-
         final SpanContext extractedContext =
             GlobalTracer.get()
                 .extract(Format.Builtin.TEXT_MAP, new MessagePropertyTextMap(message));
@@ -100,15 +91,16 @@ public final class JMSMessageConsumerInstrumentation extends Instrumenter.Defaul
         }
       }
 
-      final Scope scope = spanBuilder.startActive(true);
+      final Span span = spanBuilder.start();
-      final Span span = scope.span();
+      CONSUMER_DECORATE.afterStart(span);
-
+      if (message == null) {
-      if (throwable != null) {
+        CONSUMER_DECORATE.onReceive(span, method);
-        Tags.ERROR.set(span, Boolean.TRUE);
+      } else {
-        span.log(Collections.singletonMap(ERROR_OBJECT, throwable));
+        CONSUMER_DECORATE.onConsume(span, message);
       }
-
+      CONSUMER_DECORATE.onError(span, throwable);
-      scope.close();
+      CONSUMER_DECORATE.beforeFinish(span);
+      span.finish();
     }
   }
 }

dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/JMSMessageListenerInstrumentation.java

@@ -1,8 +1,7 @@
 package datadog.trace.instrumentation.jms;
 
 import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
-import static datadog.trace.instrumentation.jms.JmsUtil.toResourceName;
+import static datadog.trace.instrumentation.jms.JMSDecorator.CONSUMER_DECORATE;
-import static io.opentracing.log.Fields.ERROR_OBJECT;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -12,16 +11,11 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.api.DDSpanTypes;
-import datadog.trace.api.DDTags;
 import datadog.trace.context.TraceScope;
 import io.opentracing.Scope;
-import io.opentracing.Span;
 import io.opentracing.SpanContext;
 import io.opentracing.propagation.Format;
-import io.opentracing.tag.Tags;
 import io.opentracing.util.GlobalTracer;
-import java.util.Collections;
 import java.util.Map;
 import javax.jms.Message;
 import javax.jms.MessageListener;
@@ -44,7 +38,14 @@ public final class JMSMessageListenerInstrumentation extends Instrumenter.Defaul
 
   @Override
   public String[] helperClassNames() {
-    return new String[] {packageName + ".JmsUtil", packageName + ".MessagePropertyTextMap"};
+    return new String[] {
+      "datadog.trace.agent.decorator.BaseDecorator",
+      "datadog.trace.agent.decorator.ClientDecorator",
+      packageName + ".JMSDecorator",
+      packageName + ".JMSDecorator$1",
+      packageName + ".JMSDecorator$2",
+      packageName + ".MessagePropertyTextMap",
+    };
   }
 
   @Override
@@ -67,13 +68,10 @@ public final class JMSMessageListenerInstrumentation extends Instrumenter.Defaul
           GlobalTracer.get()
               .buildSpan("jms.onMessage")
               .asChildOf(extractedContext)
-              .withTag(DDTags.SERVICE_NAME, "jms")
-              .withTag(DDTags.SPAN_TYPE, DDSpanTypes.MESSAGE_CONSUMER)
-              .withTag(DDTags.RESOURCE_NAME, "Received from " + toResourceName(message, null))
-              .withTag(Tags.COMPONENT.getKey(), "jms")
-              .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CONSUMER)
               .withTag("span.origin.type", listener.getClass().getName())
               .startActive(true);
+      CONSUMER_DECORATE.afterStart(scope);
+      CONSUMER_DECORATE.onReceive(scope, message);
 
       if (scope instanceof TraceScope) {
         ((TraceScope) scope).setAsyncPropagation(true);
@@ -85,13 +83,9 @@ public final class JMSMessageListenerInstrumentation extends Instrumenter.Defaul
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(
         @Advice.Enter final Scope scope, @Advice.Thrown final Throwable throwable) {
-
       if (scope != null) {
-        if (throwable != null) {
+        CONSUMER_DECORATE.onError(scope, throwable);
-          final Span span = scope.span();
+        CONSUMER_DECORATE.beforeFinish(scope);
-          Tags.ERROR.set(span, Boolean.TRUE);
-          span.log(Collections.singletonMap(ERROR_OBJECT, throwable));
-        }
         scope.close();
       }
     }

dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/JMSMessageProducerInstrumentation.java

@@ -1,8 +1,7 @@
 package datadog.trace.instrumentation.jms;
 
 import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
-import static datadog.trace.instrumentation.jms.JmsUtil.toResourceName;
+import static datadog.trace.instrumentation.jms.JMSDecorator.PRODUCER_DECORATE;
-import static io.opentracing.log.Fields.ERROR_OBJECT;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -11,15 +10,10 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.api.DDSpanTypes;
-import datadog.trace.api.DDTags;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import io.opentracing.Scope;
-import io.opentracing.Span;
 import io.opentracing.propagation.Format;
-import io.opentracing.tag.Tags;
 import io.opentracing.util.GlobalTracer;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import javax.jms.Destination;
@@ -45,7 +39,14 @@ public final class JMSMessageProducerInstrumentation extends Instrumenter.Defaul
 
   @Override
   public String[] helperClassNames() {
-    return new String[] {packageName + ".JmsUtil", packageName + ".MessagePropertyTextMap"};
+    return new String[] {
+      "datadog.trace.agent.decorator.BaseDecorator",
+      "datadog.trace.agent.decorator.ClientDecorator",
+      packageName + ".JMSDecorator",
+      packageName + ".JMSDecorator$1",
+      packageName + ".JMSDecorator$2",
+      packageName + ".MessagePropertyTextMap",
+    };
   }
 
   @Override
@@ -79,18 +80,14 @@ public final class JMSMessageProducerInstrumentation extends Instrumenter.Defaul
       } catch (final JMSException e) {
         defaultDestination = null;
       }
+
       final Scope scope =
           GlobalTracer.get()
               .buildSpan("jms.produce")
-              .withTag(DDTags.SERVICE_NAME, "jms")
-              .withTag(
-                  DDTags.RESOURCE_NAME,
-                  "Produced for " + toResourceName(message, defaultDestination))
-              .withTag(Tags.COMPONENT.getKey(), "jms")
-              .withTag(DDTags.SPAN_TYPE, DDSpanTypes.MESSAGE_PRODUCER)
-              .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_PRODUCER)
               .withTag("span.origin.type", producer.getClass().getName())
               .startActive(true);
+      PRODUCER_DECORATE.afterStart(scope);
+      PRODUCER_DECORATE.onProduce(scope, message, defaultDestination);
 
       GlobalTracer.get()
           .inject(
@@ -104,11 +101,8 @@ public final class JMSMessageProducerInstrumentation extends Instrumenter.Defaul
         @Advice.Enter final Scope scope, @Advice.Thrown final Throwable throwable) {
 
       if (scope != null) {
-        if (throwable != null) {
+        PRODUCER_DECORATE.onError(scope, throwable);
-          final Span span = scope.span();
+        PRODUCER_DECORATE.beforeFinish(scope);
-          Tags.ERROR.set(span, Boolean.TRUE);
-          span.log(Collections.singletonMap(ERROR_OBJECT, throwable));
-        }
         scope.close();
         CallDepthThreadLocalMap.reset(MessageProducer.class);
       }
@@ -130,13 +124,10 @@ public final class JMSMessageProducerInstrumentation extends Instrumenter.Defaul
       final Scope scope =
           GlobalTracer.get()
               .buildSpan("jms.produce")
-              .withTag(DDTags.SERVICE_NAME, "jms")
-              .withTag(DDTags.SPAN_TYPE, DDSpanTypes.MESSAGE_PRODUCER)
-              .withTag(DDTags.RESOURCE_NAME, "Produced for " + toResourceName(message, destination))
-              .withTag(Tags.COMPONENT.getKey(), "jms")
-              .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_PRODUCER)
               .withTag("span.origin.type", producer.getClass().getName())
               .startActive(true);
+      PRODUCER_DECORATE.afterStart(scope);
+      PRODUCER_DECORATE.onProduce(scope, message, destination);
 
       GlobalTracer.get()
           .inject(
@@ -148,13 +139,9 @@ public final class JMSMessageProducerInstrumentation extends Instrumenter.Defaul
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(
         @Advice.Enter final Scope scope, @Advice.Thrown final Throwable throwable) {
-
       if (scope != null) {
-        if (throwable != null) {
+        PRODUCER_DECORATE.onError(scope, throwable);
-          final Span span = scope.span();
+        PRODUCER_DECORATE.beforeFinish(scope);
-          Tags.ERROR.set(span, Boolean.TRUE);
-          span.log(Collections.singletonMap(ERROR_OBJECT, throwable));
-        }
         scope.close();
         CallDepthThreadLocalMap.reset(MessageProducer.class);
       }

dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/JmsUtil.java

@@ -1,43 +0,0 @@
-package datadog.trace.instrumentation.jms;
-
-import javax.jms.Destination;
-import javax.jms.Message;
-import javax.jms.Queue;
-import javax.jms.TemporaryQueue;
-import javax.jms.TemporaryTopic;
-import javax.jms.Topic;
-
-public class JmsUtil {
-  private static final String TIBCO_TMP_PREFIX = "$TMP$";
-
-  public static String toResourceName(final Message message, final Destination destination) {
-    Destination jmsDestination = null;
-    try {
-      jmsDestination = message.getJMSDestination();
-    } catch (final Exception e) {
-    }
-    if (jmsDestination == null) {
-      jmsDestination = destination;
-    }
-    try {
-      if (jmsDestination instanceof Queue) {
-        final String queueName = ((Queue) jmsDestination).getQueueName();
-        if (jmsDestination instanceof TemporaryQueue || queueName.startsWith(TIBCO_TMP_PREFIX)) {
-          return "Temporary Queue";
-        } else {
-          return "Queue " + queueName;
-        }
-      }
-      if (jmsDestination instanceof Topic) {
-        final String topicName = ((Topic) jmsDestination).getTopicName();
-        if (jmsDestination instanceof TemporaryTopic || topicName.startsWith(TIBCO_TMP_PREFIX)) {
-          return "Temporary Topic";
-        } else {
-          return "Topic " + topicName;
-        }
-      }
-    } catch (final Exception e) {
-    }
-    return "Destination";
-  }
-}