From a0f3ad7279483f0a33443962c084584fe5934217 Mon Sep 17 00:00:00 2001
From: Trask Stalnaker <trask.stalnaker@gmail.com>
Date: Mon, 7 Feb 2022 17:31:11 -0800
Subject: [PATCH] Add code attributes to spring-scheduling spans (#5306)

---
 ...ringSchedulingCodeAttributesExtractor.java | 46 +++++++++++++++++++
 .../SpringSchedulingSingletons.java           |  1 +
 .../test/groovy/SpringSchedulingTest.groovy   |  6 +++
 3 files changed, 53 insertions(+)
 create mode 100644 instrumentation/spring/spring-scheduling-3.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/scheduling/SpringSchedulingCodeAttributesExtractor.java

diff --git a/instrumentation/spring/spring-scheduling-3.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/scheduling/SpringSchedulingCodeAttributesExtractor.java b/instrumentation/spring/spring-scheduling-3.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/scheduling/SpringSchedulingCodeAttributesExtractor.java
new file mode 100644
index 000000000000..00d3a209db04
--- /dev/null
+++ b/instrumentation/spring/spring-scheduling-3.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/scheduling/SpringSchedulingCodeAttributesExtractor.java
@@ -0,0 +1,46 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.spring.scheduling;
+
+import io.opentelemetry.instrumentation.api.instrumenter.code.CodeAttributesExtractor;
+import javax.annotation.Nullable;
+import org.springframework.scheduling.support.ScheduledMethodRunnable;
+
+public class SpringSchedulingCodeAttributesExtractor
+    extends CodeAttributesExtractor<Runnable, Void> {
+
+  @Override
+  protected Class<?> codeClass(Runnable runnable) {
+    if (runnable instanceof ScheduledMethodRunnable) {
+      ScheduledMethodRunnable scheduledMethodRunnable = (ScheduledMethodRunnable) runnable;
+      return scheduledMethodRunnable.getTarget().getClass();
+    } else {
+      return runnable.getClass();
+    }
+  }
+
+  @Override
+  protected String methodName(Runnable runnable) {
+    if (runnable instanceof ScheduledMethodRunnable) {
+      ScheduledMethodRunnable scheduledMethodRunnable = (ScheduledMethodRunnable) runnable;
+      return scheduledMethodRunnable.getMethod().getName();
+    } else {
+      return "run";
+    }
+  }
+
+  @Override
+  @Nullable
+  protected String filePath(Runnable runnable) {
+    return null;
+  }
+
+  @Override
+  @Nullable
+  protected Long lineNumber(Runnable runnable) {
+    return null;
+  }
+}
diff --git a/instrumentation/spring/spring-scheduling-3.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/scheduling/SpringSchedulingSingletons.java b/instrumentation/spring/spring-scheduling-3.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/scheduling/SpringSchedulingSingletons.java
index e328dd3afaee..7a6177e82c31 100644
--- a/instrumentation/spring/spring-scheduling-3.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/scheduling/SpringSchedulingSingletons.java
+++ b/instrumentation/spring/spring-scheduling-3.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/scheduling/SpringSchedulingSingletons.java
@@ -17,6 +17,7 @@ public final class SpringSchedulingSingletons {
               GlobalOpenTelemetry.get(),
               "io.opentelemetry.spring-scheduling-3.1",
               SpringSchedulingSingletons::extractSpanName)
+          .addAttributesExtractor(new SpringSchedulingCodeAttributesExtractor())
           .newInstrumenter();
 
   private static String extractSpanName(Runnable runnable) {
diff --git a/instrumentation/spring/spring-scheduling-3.1/javaagent/src/test/groovy/SpringSchedulingTest.groovy b/instrumentation/spring/spring-scheduling-3.1/javaagent/src/test/groovy/SpringSchedulingTest.groovy
index e887440e9293..23dc08d87ac9 100644
--- a/instrumentation/spring/spring-scheduling-3.1/javaagent/src/test/groovy/SpringSchedulingTest.groovy
+++ b/instrumentation/spring/spring-scheduling-3.1/javaagent/src/test/groovy/SpringSchedulingTest.groovy
@@ -25,6 +25,8 @@ class SpringSchedulingTest extends AgentInstrumentationSpecification {
           name "TriggerTask.run"
           hasNoParent()
           attributes {
+            "code.namespace" "TriggerTask"
+            "code.function" "run"
           }
         }
       }
@@ -46,6 +48,8 @@ class SpringSchedulingTest extends AgentInstrumentationSpecification {
           name "IntervalTask.run"
           hasNoParent()
           attributes {
+            "code.namespace" "IntervalTask"
+            "code.function" "run"
           }
         }
       }
@@ -67,6 +71,8 @@ class SpringSchedulingTest extends AgentInstrumentationSpecification {
           nameContains "LambdaTaskConfigurer\$\$Lambda\$"
           hasNoParent()
           attributes {
+            "code.namespace" { it.contains("LambdaTaskConfigurer\$\$Lambda\$") }
+            "code.function" "run"
           }
         }
       }