Tag @async @scheduled spring-scheduling spans as errored on exception (#12071)

Tag @async @scheduled spring-scheduling spans as errored on exception

Co-authored-by: devflow.devflow-routing-intake <devflow.devflow-routing-intake@kubernetes.us1.ddbuild.io>

Author: ValentinZakharov

dd-java-agent/instrumentation/spring/spring-scheduling-3.1/src/main/java/datadog/trace/instrumentation/springscheduling/SpannedMethodInvocation.java

@@ -50,8 +50,15 @@ private Object invokeWithContinuation(CharSequence spanName) throws Throwable {
 
   private Object invokeWithSpan(CharSequence spanName) throws Throwable {
     AgentSpan span = startSpan("spring-scheduling", spanName);
+    DECORATE.afterStart(span);
+    DECORATE.measureIfEnabled(span);
     try (ContextScope scope = activateSpan(span)) {
-      return delegate.proceed();
+      try {
+        return delegate.proceed();
+      } catch (Throwable throwable) {
+        DECORATE.onError(span, throwable);
+        throw throwable;
+      }
     } finally {
       span.finish();
     }

dd-java-agent/instrumentation/spring/spring-scheduling-3.1/src/main/java/datadog/trace/instrumentation/springscheduling/SpringSchedulingDecorator.java

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.springscheduling;
 
+import datadog.trace.api.Config;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
 import datadog.trace.bootstrap.instrumentation.decorator.BaseDecorator;
@@ -9,6 +10,8 @@ public class SpringSchedulingDecorator extends BaseDecorator {
   public static final CharSequence SCHEDULED_CALL = UTF8BytesString.create("scheduled.call");
   public static final SpringSchedulingDecorator DECORATE = new SpringSchedulingDecorator();
 
+  private static final boolean MEASURED = Config.get().isSpringSchedulingMeasuredEnabled();
+
   private SpringSchedulingDecorator() {}
 
   @Override
@@ -26,6 +29,12 @@ protected CharSequence component() {
     return "spring-scheduling";
   }
 
+  public void measureIfEnabled(final AgentSpan span) {
+    if (MEASURED) {
+      span.setMeasured(true);
+    }
+  }
+
   public void onRun(final AgentSpan span, final Runnable runnable) {
     if (runnable != null) {
       CharSequence resourceName = "";

dd-java-agent/instrumentation/spring/spring-scheduling-3.1/src/main/java/datadog/trace/instrumentation/springscheduling/SpringSchedulingRunnableWrapper.java

@@ -59,6 +59,7 @@ public void run() {
             ? startSpan("spring-scheduling", SCHEDULED_CALL)
             : startSpan("spring-scheduling", SCHEDULED_CALL, null);
     DECORATE.afterStart(span);
+    DECORATE.measureIfEnabled(span);
 
     try (final ContextScope scope = activateSpan(span)) {
       DECORATE.onRun(span, runnable);

dd-java-agent/instrumentation/spring/spring-scheduling-3.1/src/test/groovy/SpringAsyncTest.groovy

@@ -1,10 +1,57 @@
 import datadog.trace.agent.test.InstrumentationSpecification
+import datadog.trace.bootstrap.instrumentation.api.Tags
 import org.springframework.context.annotation.AnnotationConfigApplicationContext
 
 import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
 
 class SpringAsyncTest extends InstrumentationSpecification {
 
+  boolean asyncMeasured() {
+    false
+  }
+
+  @Override
+  protected void configurePreAgent() {
+    super.configurePreAgent()
+    if (asyncMeasured()) {
+      injectSysConfig("spring-scheduling.measured.enabled", "true")
+    }
+  }
+
+  def "exception in @async method tags span as errored"() {
+    setup:
+    def context = new AnnotationConfigApplicationContext(AsyncTaskConfig)
+    AsyncTask asyncTask = context.getBean(AsyncTask)
+
+    when:
+    Throwable thrown = null
+    try {
+      asyncTask.asyncThrow().join()
+    } catch (Throwable t) {
+      thrown = t
+    }
+
+    then:
+    thrown != null
+    assertTraces(1) {
+      trace(1) {
+        span {
+          resourceName "AsyncTask.asyncThrow"
+          errored true
+          measured asyncMeasured()
+          tags {
+            "$Tags.COMPONENT" "spring-scheduling"
+            errorTags(RuntimeException, "Datadog async repro")
+            defaultTags()
+          }
+        }
+      }
+    }
+
+    cleanup:
+    context.close()
+  }
+
   def "context propagated through @async annotation"() {
     setup:
     def context = new AnnotationConfigApplicationContext(AsyncTaskConfig)
@@ -54,3 +101,10 @@ class SpringAsyncTest extends InstrumentationSpecification {
     hasParent << [true, false]
   }
 }
+
+class SpringAsyncMeasuredForkedTest extends SpringAsyncTest {
+  @Override
+  boolean asyncMeasured() {
+    true
+  }
+}

dd-java-agent/instrumentation/spring/spring-scheduling-3.1/src/test/java/AsyncTask.java

@@ -10,6 +10,11 @@ public CompletableFuture<Integer> async() {
     return CompletableFuture.completedFuture(getInt());
   }
 
+  @Async
+  public CompletableFuture<Void> asyncThrow() {
+    throw new RuntimeException("Datadog async repro");
+  }
+
   @Trace
   public int getInt() {
     return ThreadLocalRandom.current().nextInt();

dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java

@@ -157,6 +157,9 @@ public final class TraceInstrumentationConfig {
   public static final String SPRING_DATA_REPOSITORY_INTERFACE_RESOURCE_NAME =
       "spring-data.repository.interface.resource-name";
 
+  public static final String SPRING_SCHEDULING_MEASURED_ENABLED =
+      "spring-scheduling.measured.enabled";
+
   public static final String INSTRUMENTATION_CONFIG_ID = "instrumentation_config_id";
 
   public static final String RESOLVER_CACHE_CONFIG = "resolver.cache.config";

internal-api/src/main/java/datadog/trace/api/Config.java

@@ -626,6 +626,7 @@
 import static datadog.trace.api.config.TraceInstrumentationConfig.SPARK_APP_NAME_AS_SERVICE;
 import static datadog.trace.api.config.TraceInstrumentationConfig.SPARK_TASK_HISTOGRAM_ENABLED;
 import static datadog.trace.api.config.TraceInstrumentationConfig.SPRING_DATA_REPOSITORY_INTERFACE_RESOURCE_NAME;
+import static datadog.trace.api.config.TraceInstrumentationConfig.SPRING_SCHEDULING_MEASURED_ENABLED;
 import static datadog.trace.api.config.TraceInstrumentationConfig.SQS_BODY_PROPAGATION_ENABLED;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_128_BIT_TRACEID_LOGGING_ENABLED;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_HTTP_CLIENT_TAG_QUERY_STRING;
@@ -1293,6 +1294,8 @@ public static String getHostName() {
   private final boolean hystrixTagsEnabled;
   private final boolean hystrixMeasuredEnabled;
 
+  private final boolean springSchedulingMeasuredEnabled;
+
   private final boolean resilience4jMeasuredEnabled;
   private final boolean resilience4jTagMetricsEnabled;
 
@@ -3052,6 +3055,9 @@ PROFILING_DATADOG_PROFILER_ENABLED, isDatadogProfilerSafeInCurrentEnvironment())
     hystrixTagsEnabled = configProvider.getBoolean(HYSTRIX_TAGS_ENABLED, false);
     hystrixMeasuredEnabled = configProvider.getBoolean(HYSTRIX_MEASURED_ENABLED, false);
 
+    springSchedulingMeasuredEnabled =
+        configProvider.getBoolean(SPRING_SCHEDULING_MEASURED_ENABLED, false);
+
     resilience4jMeasuredEnabled = configProvider.getBoolean(RESILIENCE4J_MEASURED_ENABLED, false);
     resilience4jTagMetricsEnabled =
         configProvider.getBoolean(RESILIENCE4J_TAG_METRICS_ENABLED, false);
@@ -4973,6 +4979,10 @@ public boolean isResilience4jMeasuredEnabled() {
     return resilience4jMeasuredEnabled;
   }
 
+  public boolean isSpringSchedulingMeasuredEnabled() {
+    return springSchedulingMeasuredEnabled;
+  }
+
   public boolean isResilience4jTagMetricsEnabled() {
     return resilience4jTagMetricsEnabled;
   }
@@ -6666,6 +6676,8 @@ public String toString() {
         + hystrixTagsEnabled
         + ", hystrixMeasuredEnabled="
         + hystrixMeasuredEnabled
+        + ", springSchedulingMeasuredEnabled="
+        + springSchedulingMeasuredEnabled
         + ", resilience4jMeasuredEnable="
         + resilience4jMeasuredEnabled
         + ", resilience4jTagMetricsEnabled="

metadata/supported-configurations.json

@@ -3897,6 +3897,14 @@
         "aliases": []
       }
     ],
+    "DD_SPRING_SCHEDULING_MEASURED_ENABLED": [
+      {
+        "version": "A",
+        "type": "boolean",
+        "default": "false",
+        "aliases": []
+      }
+    ],
     "DD_STACK_TRACE_LENGTH_LIMIT": [
       {
         "version": "A",