Added package attribute extractor (#60)

Co-authored-by: oleksandrh <alex.hrebeniuk@gmail.com>
digma-ai · Apr 19, 2024 · 33ca94e · 33ca94e
1 parent ae47739
commit 33ca94e
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 4 deletions.
diff --git a/...a/otel/javaagent/extension/instrumentation/methods/DigmaMethodsInstrumentationModule.java b/...a/otel/javaagent/extension/instrumentation/methods/DigmaMethodsInstrumentationModule.java
@@ -71,6 +71,7 @@ public List<String> getAdditionalHelperClassNames() {
         list.add("io.opentelemetry.instrumentation.api.instrumenter.util.ClassAndMethodAttributesGetter");
         list.add("io.opentelemetry.instrumentation.api.instrumenter.code.CodeSpanNameExtractor");
         list.add("io.opentelemetry.semconv.trace.attributes.SemanticAttributes");
+        list.add("com.digma.otel.javaagent.extension.instrumentation.methods.PackageExtractor");
         return list;
     }
 

diff --git a/...ain/java/com/digma/otel/javaagent/extension/instrumentation/methods/MethodSingletons.java b/...ain/java/com/digma/otel/javaagent/extension/instrumentation/methods/MethodSingletons.java
@@ -29,6 +29,7 @@ public final class MethodSingletons {
                         .addAttributesExtractor(CodeAttributesExtractor.create(codeAttributesGetter))
                         .setInstrumentationVersion(DigmaExtensionVersion.VERSION)
                         .setErrorCauseExtractor(ErrorCauseExtractor.getDefault())
+                        .addAttributesExtractor(PackageExtractor.create(codeAttributesGetter))
                         .buildInstrumenter(SpanKindExtractor.alwaysInternal());
     }
 

diff --git a/...ain/java/com/digma/otel/javaagent/extension/instrumentation/methods/PackageExtractor.java b/...ain/java/com/digma/otel/javaagent/extension/instrumentation/methods/PackageExtractor.java
@@ -0,0 +1,38 @@
+package com.digma.otel.javaagent.extension.instrumentation.methods;
+
+import com.digma.otel.instrumentation.common.DigmaSemanticAttributes;
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.common.AttributesBuilder;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
+import io.opentelemetry.instrumentation.api.instrumenter.code.CodeAttributesGetter;
+import io.opentelemetry.instrumentation.api.internal.AttributesExtractorUtil;
+
+
+import javax.annotation.Nullable;
+
+public final class PackageExtractor<REQUEST, RESPONSE> implements AttributesExtractor<REQUEST, RESPONSE>
+{
+    private final CodeAttributesGetter<REQUEST> getter;
+
+    private PackageExtractor(CodeAttributesGetter<REQUEST> getter) {
+        this.getter = getter;
+    }
+
+    public static <REQUEST, RESPONSE> PackageExtractor<REQUEST, RESPONSE> create(CodeAttributesGetter<REQUEST> getter) {
+        return new PackageExtractor(getter);
+    }
+
+    @Override
+    public void onStart(AttributesBuilder attributes, Context context, REQUEST request) {
+        Class<?> cls = this.getter.codeClass(request);
+        if (cls != null) {
+            AttributesExtractorUtil.internalSet(attributes, AttributeKey.stringKey("digma.instrumentation.extended.package"), cls.getPackage().getName());
+            AttributesExtractorUtil.internalSet(attributes, AttributeKey.stringKey("digma.instrumentation.extended.enabled"), "true");
+        }
+    }
+
+    @Override
+    public void onEnd(AttributesBuilder attributesBuilder, Context context, REQUEST request, @Nullable RESPONSE response, @Nullable Throwable throwable) {
+    }
+}
diff --git a/...m/digma/otel/javaagent/extension/instrumentation/methods/MethodsInstrumentationTests.java b/...m/digma/otel/javaagent/extension/instrumentation/methods/MethodsInstrumentationTests.java
@@ -4,6 +4,7 @@
 import com.digma.otel.javaagent.extension.instrumentation.methods.test.*;
 import com.digma.otel.javaagent.extension.instrumentation.methods.test2.MyClassInOtherPackage;
 import com.digma.otel.javaagent.extension.version.DigmaExtensionVersion;
+import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
@@ -24,6 +25,8 @@ class MethodsInstrumentationTests {
 
     @RegisterExtension
     static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
+    static final AttributeKey<String> EXTENDED_PACKAGE = AttributeKey.stringKey("digma.instrumentation.extended.package");
+    static final AttributeKey<String> EXTENDED_ENABLED = AttributeKey.stringKey("digma.instrumentation.extended.enabled");
 
     @Test
     void methodTraced() {
@@ -45,7 +48,9 @@ void methodTraced() {
                                                 .hasKind(SpanKind.INTERNAL)
                                                 .hasAttributesSatisfyingExactly(
                                                         equalTo(CODE_NAMESPACE, ConfigTracedCallable.class.getName()),
-                                                        equalTo(CODE_FUNCTION, "call"))));
+                                                        equalTo(CODE_FUNCTION, "call"),
+                                                        equalTo(EXTENDED_PACKAGE, "com.digma.otel.javaagent.extension.instrumentation.methods.test"),
+                                                        equalTo(EXTENDED_ENABLED, "true"))));
 
         Assertions.assertThat(testing.spans()).size().isEqualTo(1);
         testing.spans().forEach(spanData -> Assertions.assertThat(spanData.getInstrumentationScopeInfo().getName()).isEqualTo("digma.io.opentelemetry.methods"));
@@ -78,7 +83,9 @@ void methodTracedWithAsyncStop() throws Exception {
                                                 .hasKind(SpanKind.INTERNAL)
                                                 .hasAttributesSatisfyingExactly(
                                                         equalTo(CODE_NAMESPACE, ConfigTracedCompletableFuture.class.getName()),
-                                                        equalTo(CODE_FUNCTION, "getResult"))));
+                                                        equalTo(CODE_FUNCTION, "getResult"),
+                                                        equalTo(EXTENDED_PACKAGE, "com.digma.otel.javaagent.extension.instrumentation.methods.test"),
+                                                        equalTo(EXTENDED_ENABLED, "true"))));
 
         Assertions.assertThat(testing.spans()).size().isEqualTo(1);
         testing.spans().forEach(spanData -> Assertions.assertThat(spanData.getInstrumentationScopeInfo().getName()).isEqualTo("digma.io.opentelemetry.methods"));
@@ -103,14 +110,18 @@ void methodTracedAnonymousAndLambda() {
                                         .hasKind(SpanKind.INTERNAL)
                                         .hasAttributesSatisfyingExactly(
                                                 equalTo(CODE_NAMESPACE, AnonymousTestClass.class.getName()),
-                                                equalTo(CODE_FUNCTION, "methodWithAnonymousClass"))),
+                                                equalTo(CODE_FUNCTION, "methodWithAnonymousClass"),
+                                                equalTo(EXTENDED_PACKAGE, "com.digma.otel.javaagent.extension.instrumentation.methods.test"),
+                                                equalTo(EXTENDED_ENABLED, "true"))),
                 trace -> trace.hasSpansSatisfyingExactly(
                         span ->
                                 span.hasName("AnonymousTestClass.methodWithLambda")
                                         .hasKind(SpanKind.INTERNAL)
                                         .hasAttributesSatisfyingExactly(
                                                 equalTo(CODE_NAMESPACE, AnonymousTestClass.class.getName()),
-                                                equalTo(CODE_FUNCTION, "methodWithLambda"))));
+                                                equalTo(CODE_FUNCTION, "methodWithLambda"),
+                                                equalTo(EXTENDED_PACKAGE, "com.digma.otel.javaagent.extension.instrumentation.methods.test"),
+                                                equalTo(EXTENDED_ENABLED, "true"))));
 
 
         Assertions.assertThat(testing.spans()).size().isEqualTo(2);