Migrate Jetty instrumentation to Decorator

dd-java-agent/instrumentation/jetty-8/src/main/java/datadog/trace/instrumentation/jetty8/HandlerInstrumentation.java

@@ -37,9 +37,13 @@ public final class HandlerInstrumentation extends Instrumenter.Default {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      "datadog.trace.instrumentation.jetty8.HttpServletRequestExtractAdapter",
+      "datadog.trace.agent.decorator.BaseDecorator",
-      "datadog.trace.instrumentation.jetty8.HttpServletRequestExtractAdapter$MultivaluedMapFlatIterator",
+      "datadog.trace.agent.decorator.ServerDecorator",
-      "datadog.trace.instrumentation.jetty8.TagSettingAsyncListener"
+      "datadog.trace.agent.decorator.HttpServerDecorator",
+      packageName + ".JettyDecorator",
+      packageName + ".HttpServletRequestExtractAdapter",
+      packageName + ".HttpServletRequestExtractAdapter$MultivaluedMapFlatIterator",
+      packageName + ".TagSettingAsyncListener"
     };
   }
 

dd-java-agent/instrumentation/jetty-8/src/main/java/datadog/trace/instrumentation/jetty8/JettyDecorator.java

@@ -0,0 +1,54 @@
+package datadog.trace.instrumentation.jetty8;
+
+import datadog.trace.agent.decorator.HttpServerDecorator;
+import io.opentracing.Span;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+
+public class JettyDecorator extends HttpServerDecorator<HttpServletRequest, HttpServletResponse> {
+  public static final JettyDecorator DECORATE = new JettyDecorator();
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"jetty", "jetty-8"};
+  }
+
+  @Override
+  protected String component() {
+    return "jetty-handler";
+  }
+
+  @Override
+  protected String method(final HttpServletRequest httpServletRequest) {
+    return httpServletRequest.getMethod();
+  }
+
+  @Override
+  protected String url(final HttpServletRequest httpServletRequest) {
+    return httpServletRequest.getRequestURL().toString();
+  }
+
+  @Override
+  protected String hostname(final HttpServletRequest httpServletRequest) {
+    return httpServletRequest.getServerName();
+  }
+
+  @Override
+  protected Integer port(final HttpServletRequest httpServletRequest) {
+    return httpServletRequest.getServerPort();
+  }
+
+  @Override
+  protected Integer status(final HttpServletResponse httpServletResponse) {
+    return httpServletResponse.getStatus();
+  }
+
+  @Override
+  public Span onRequest(final Span span, final HttpServletRequest request) {
+    assert span != null;
+    if (request != null) {
+      span.setTag("servlet.context", request.getContextPath());
+    }
+    return super.onRequest(span, request);
+  }
+}

dd-java-agent/instrumentation/jetty-8/src/main/java/datadog/trace/instrumentation/jetty8/JettyHandlerAdvice.java

@@ -1,8 +1,7 @@
 package datadog.trace.instrumentation.jetty8;
 
-import static io.opentracing.log.Fields.ERROR_OBJECT;
+import static datadog.trace.instrumentation.jetty8.JettyDecorator.DECORATE;
 
-import datadog.trace.api.DDSpanTypes;
 import datadog.trace.api.DDTags;
 import datadog.trace.context.TraceScope;
 import io.opentracing.Scope;
@@ -11,7 +10,6 @@ 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.concurrent.atomic.AtomicBoolean;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
@@ -31,29 +29,21 @@ public class JettyHandlerAdvice {
     final SpanContext extractedContext =
         GlobalTracer.get()
             .extract(Format.Builtin.HTTP_HEADERS, new HttpServletRequestExtractAdapter(req));
-    final String resourceName = req.getMethod() + " " + source.getClass().getName();
     final Scope scope =
         GlobalTracer.get()
             .buildSpan("jetty.request")
             .asChildOf(extractedContext)
-            .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_SERVER)
-            .withTag(DDTags.SPAN_TYPE, DDSpanTypes.HTTP_SERVER)
-            .withTag("servlet.context", req.getContextPath())
             .withTag("span.origin.type", source.getClass().getName())
             .startActive(false);
 
+    DECORATE.afterStart(scope.span());
+    DECORATE.onRequest(scope.span(), req);
+    final String resourceName = req.getMethod() + " " + source.getClass().getName();
+    scope.span().setTag(DDTags.RESOURCE_NAME, resourceName);
+
     if (scope instanceof TraceScope) {
       ((TraceScope) scope).setAsyncPropagation(true);
     }
-
-    final Span span = scope.span();
-    Tags.COMPONENT.set(span, "jetty-handler");
-    Tags.HTTP_METHOD.set(span, req.getMethod());
-    Tags.HTTP_URL.set(span, req.getRequestURL().toString());
-    span.setTag(DDTags.RESOURCE_NAME, resourceName);
-    if (req.getUserPrincipal() != null) {
-      span.setTag(DDTags.USER_NAME, req.getUserPrincipal().getName());
-    }
     return scope;
   }
 
@@ -63,20 +53,19 @@ public class JettyHandlerAdvice {
       @Advice.Argument(3) final HttpServletResponse resp,
       @Advice.Enter final Scope scope,
       @Advice.Thrown final Throwable throwable) {
-
     if (scope != null) {
       final Span span = scope.span();
+      if (req.getUserPrincipal() != null) {
+        span.setTag(DDTags.USER_NAME, req.getUserPrincipal().getName());
+      }
       if (throwable != null) {
+        DECORATE.onResponse(span, resp);
         if (resp.getStatus() == HttpServletResponse.SC_OK) {
           // exception is thrown in filter chain, but status code is incorrect
           Tags.HTTP_STATUS.set(span, 500);
         }
-        Tags.ERROR.set(span, Boolean.TRUE);
+        DECORATE.onError(span, throwable);
-        span.log(Collections.singletonMap(ERROR_OBJECT, throwable));
+        DECORATE.beforeFinish(span);
-        if (scope instanceof TraceScope) {
-          ((TraceScope) scope).setAsyncPropagation(false);
-        }
-        scope.close();
         span.finish(); // Finish the span manually since finishSpanOnClose was false
       } else {
         final AtomicBoolean activated = new AtomicBoolean(false);
@@ -88,15 +77,14 @@ public class JettyHandlerAdvice {
             // finished after check above. We just ignore that exception and move on.
           }
         }
+        // Check again in case the request finished before adding the listener.
         if (!req.isAsyncStarted() && activated.compareAndSet(false, true)) {
-          Tags.HTTP_STATUS.set(span, resp.getStatus());
+          DECORATE.onResponse(span, resp);
-          if (scope instanceof TraceScope) {
+          DECORATE.beforeFinish(span);
-            ((TraceScope) scope).setAsyncPropagation(false);
-          }
           span.finish(); // Finish the span manually since finishSpanOnClose was false
         }
-        scope.close();
       }
+      scope.close();
     }
   }
 }

dd-java-agent/instrumentation/jetty-8/src/main/java/datadog/trace/instrumentation/jetty8/TagSettingAsyncListener.java

@@ -1,14 +1,10 @@
 package datadog.trace.instrumentation.jetty8;
 
-import static io.opentracing.log.Fields.ERROR_OBJECT;
+import static datadog.trace.instrumentation.jetty8.JettyDecorator.DECORATE;
 
-import datadog.trace.context.TraceScope;
-import io.opentracing.Scope;
 import io.opentracing.Span;
 import io.opentracing.tag.Tags;
-import io.opentracing.util.GlobalTracer;
 import java.io.IOException;
-import java.util.Collections;
 import java.util.concurrent.atomic.AtomicBoolean;
 import javax.servlet.AsyncEvent;
 import javax.servlet.AsyncListener;
@@ -26,49 +22,33 @@ public class TagSettingAsyncListener implements AsyncListener {
   @Override
   public void onComplete(final AsyncEvent event) throws IOException {
     if (activated.compareAndSet(false, true)) {
-      try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, true)) {
+      DECORATE.onResponse(span, (HttpServletResponse) event.getSuppliedResponse());
-        Tags.HTTP_STATUS.set(span, ((HttpServletResponse) event.getSuppliedResponse()).getStatus());
+      DECORATE.beforeFinish(span);
-
+      span.finish();
-        if (scope instanceof TraceScope) {
-          // This doesn't do anything because we're in a new scope, but just to be safe...
-          ((TraceScope) scope).setAsyncPropagation(false);
-        }
-      }
     }
   }
 
   @Override
   public void onTimeout(final AsyncEvent event) throws IOException {
     if (activated.compareAndSet(false, true)) {
-      try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, true)) {
+      Tags.ERROR.set(span, Boolean.TRUE);
-        Tags.ERROR.set(span, Boolean.TRUE);
+      span.setTag("timeout", event.getAsyncContext().getTimeout());
-        span.setTag("timeout", event.getAsyncContext().getTimeout());
+      DECORATE.beforeFinish(span);
-
+      span.finish();
-        if (scope instanceof TraceScope) {
-          // This doesn't do anything because we're in a new scope, but just to be safe...
-          ((TraceScope) scope).setAsyncPropagation(false);
-        }
-      }
     }
   }
 
   @Override
   public void onError(final AsyncEvent event) throws IOException {
     if (event.getThrowable() != null && activated.compareAndSet(false, true)) {
-      try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, true)) {
+      if (((HttpServletResponse) event.getSuppliedResponse()).getStatus()
-        if (((HttpServletResponse) event.getSuppliedResponse()).getStatus()
+          == HttpServletResponse.SC_OK) {
-            == HttpServletResponse.SC_OK) {
+        // exception is thrown in filter chain, but status code is incorrect
-          // exception is thrown in filter chain, but status code is incorrect
+        Tags.HTTP_STATUS.set(span, 500);
-          Tags.HTTP_STATUS.set(span, 500);
-        }
-        Tags.ERROR.set(span, Boolean.TRUE);
-        span.log(Collections.singletonMap(ERROR_OBJECT, event.getThrowable()));
-
-        if (scope instanceof TraceScope) {
-          // This doesn't do anything because we're in a new scope, but just to be safe...
-          ((TraceScope) scope).setAsyncPropagation(false);
-        }
       }
+      DECORATE.onError(span, event.getThrowable());
+      DECORATE.beforeFinish(span);
+      span.finish();
     }
   }
 

dd-java-agent/instrumentation/jetty-8/src/test/groovy/JettyHandlerTest.groovy

@@ -69,6 +69,8 @@ class JettyHandlerTest extends AgentTestRunner {
             "component" "jetty-handler"
             "span.origin.type" handler.class.name
             "http.status_code" 200
+            "peer.hostname" "localhost"
+            "peer.port" port
             defaultTags()
           }
         }
@@ -163,6 +165,8 @@ class JettyHandlerTest extends AgentTestRunner {
             "component" "jetty-handler"
             "span.origin.type" handler.class.name
             "http.status_code" 500
+            "peer.hostname" "localhost"
+            "peer.port" port
             errorTags RuntimeException
             defaultTags()
           }
@@ -184,6 +188,8 @@ class JettyHandlerTest extends AgentTestRunner {
               "component" "jetty-handler"
               "span.origin.type" handler.class.name
               "http.status_code" 500
+              "peer.hostname" "localhost"
+              "peer.port" port
               "error" true
               defaultTags()
             }

dd-java-agent/instrumentation/sparkjava-2.3/src/test/groovy/SparkJavaBasedTest.groovy

@@ -56,6 +56,8 @@ class SparkJavaBasedTest extends AgentTestRunner {
             "component" "jetty-handler"
             "span.origin.type" spark.embeddedserver.jetty.JettyHandler.name
             "http.status_code" 200
+            "peer.hostname" "localhost"
+            "peer.port" port
             defaultTags()
           }
         }