Merge pull request #2970 from DataDog/rgs/servlet-checkpoints

emit thread migrations in TagSettingAsyncListener
DataDog · Jul 29, 2021 · 27100a0 · 27100a0
2 parents 87b9c98 + 92d96e6
commit 27100a0
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/...datadog/trace/instrumentation/java/concurrent/AsyncPropagatingDisableInstrumentation.java b/...datadog/trace/instrumentation/java/concurrent/AsyncPropagatingDisableInstrumentation.java
@@ -77,7 +77,8 @@ public ElementMatcher<? super TypeDescription> typeMatcher() {
             "org.elasticsearch.transport.netty4.Netty4TcpChannel",
             "org.springframework.cglib.core.internal.LoadingCache",
             "com.datastax.oss.driver.internal.core.channel.DefaultWriteCoalescer$Flusher",
-            "com.datastax.oss.driver.api.core.session.SessionBuilder")
+            "com.datastax.oss.driver.api.core.session.SessionBuilder",
+            "org.jvnet.hk2.internal.ServiceLocatorImpl")
         .or(RX_WORKERS)
         .or(GRPC_MANAGED_CHANNEL)
         .or(REACTOR_DISABLED_TYPE_INITIALIZERS);
@@ -135,6 +136,10 @@ public void adviceTransformations(AdviceTransformation transformation) {
         named("buildAsync")
             .and(isDeclaredBy(named("com.datastax.oss.driver.api.core.session.SessionBuilder"))),
         advice);
+    transformation.applyAdvice(
+        namedOneOf("getInjecteeDescriptor", "getService")
+            .and(isDeclaredBy(named("org.jvnet.hk2.internal.ServiceLocatorImpl"))),
+        advice);
     transformation.applyAdvice(
         isTypeInitializer().and(isDeclaredBy(REACTOR_DISABLED_TYPE_INITIALIZERS)), advice);
   }

diff --git a/...quest-3/src/main/java/datadog/trace/instrumentation/servlet3/TagSettingAsyncListener.java b/...quest-3/src/main/java/datadog/trace/instrumentation/servlet3/TagSettingAsyncListener.java
@@ -30,11 +30,13 @@ public TagSettingAsyncListener(
     this.activated = activated;
     this.span = span;
     this.isDispatch = isDispatch;
+    span.startThreadMigration();
   }
 
   @Override
   public void onComplete(final AsyncEvent event) throws IOException {
     if (activated.compareAndSet(false, true)) {
+      span.finishThreadMigration();
       if (!isDispatch) {
         DECORATE.onResponse(span, (HttpServletResponse) event.getSuppliedResponse());
       }
@@ -46,6 +48,7 @@ public void onComplete(final AsyncEvent event) throws IOException {
   @Override
   public void onTimeout(final AsyncEvent event) throws IOException {
     if (activated.compareAndSet(false, true)) {
+      span.finishThreadMigration();
       if (Config.get().isServletAsyncTimeoutError()) {
         span.setError(true);
       }
@@ -58,6 +61,7 @@ public void onTimeout(final AsyncEvent event) throws IOException {
   @Override
   public void onError(final AsyncEvent event) throws IOException {
     if (event.getThrowable() != null && activated.compareAndSet(false, true)) {
+      span.finishThreadMigration();
       if (!isDispatch) {
         DECORATE.onResponse(span, (HttpServletResponse) event.getSuppliedResponse());
         if (((HttpServletResponse) event.getSuppliedResponse()).getStatus()
@@ -75,6 +79,7 @@ public void onError(final AsyncEvent event) throws IOException {
   /** Transfer the listener over to the new context. */
   @Override
   public void onStartAsync(final AsyncEvent event) throws IOException {
+    span.startThreadMigration();
     event.getAsyncContext().addListener(this);
   }
 }