emit thread migrations in TagSettingAsyncListener

DataDog · Jul 29, 2021 · 424e651 · 424e651
1 parent 09a790c
commit 424e651
Showing 1 changed file with 5 additions and 0 deletions.
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);
   }
 }