Add instrumentation to inject x-ray formatted IDs into log4j and logb… (#49)

* Add instrumentation to inject x-ray formatted IDs into log4j and logback contexts.

* CLO

* Remove prefix

Author: anuraaga

awsagentprovider/build.gradle.kts

@@ -39,18 +39,5 @@ dependencies {
 tasks {
   shadowJar {
     archiveClassifier.set("")
-
-    exclude("**/module-info.class")
-
-    // rewrite dependencies calling Logger.getLogger
-    relocate("java.util.logging.Logger", "io.opentelemetry.javaagent.bootstrap.PatchLogger")
-
-    // prevents conflict with library instrumentation
-    relocate("io.opentelemetry.instrumentation.api", "io.opentelemetry.javaagent.shaded.instrumentation.api")
-
-    // relocate OpenTelemetry API usage
-    relocate("io.opentelemetry.api", "io.opentelemetry.javaagent.shaded.io.opentelemetry.api")
-    relocate("io.opentelemetry.context", "io.opentelemetry.javaagent.shaded.io.opentelemetry.context")
-    relocate("io.opentelemetry.spi", "io.opentelemetry.javaagent.shaded.io.opentelemetry.spi")
   }
 }

build.gradle.kts

@@ -13,6 +13,7 @@
  * permissions and limitations under the License.
  */
 
+import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar
 import com.github.jk1.license.render.InventoryMarkdownReportRenderer
 import nebula.plugin.release.git.opinion.Strategies
 import org.gradle.api.tasks.testing.logging.TestExceptionFormat
@@ -24,6 +25,8 @@ plugins {
   id("com.github.jk1.dependency-license-report")
   id("io.github.gradle-nexus.publish-plugin")
   id("nebula.release")
+
+  id("com.github.johnrengelman.shadow") apply false
 }
 
 release {
@@ -135,6 +138,32 @@ allprojects {
     }
   }
 
+  plugins.withId("com.github.johnrengelman.shadow") {
+    tasks {
+      named<ShadowJar>("shadowJar") {
+        exclude("**/module-info.class")
+
+        // rewrite library instrumentation dependencies
+        relocate("io.opentelemetry.instrumentation", "io.opentelemetry.javaagent.shaded.instrumentation")
+
+        // rewrite dependencies calling Logger.getLogger
+        relocate("java.util.logging.Logger", "io.opentelemetry.javaagent.bootstrap.PatchLogger")
+
+        // relocate OpenTelemetry API usage
+        relocate("io.opentelemetry.api", "io.opentelemetry.javaagent.shaded.io.opentelemetry.api")
+        relocate("io.opentelemetry.semconv", "io.opentelemetry.javaagent.shaded.io.opentelemetry.semconv")
+        relocate("io.opentelemetry.spi", "io.opentelemetry.javaagent.shaded.io.opentelemetry.spi")
+        relocate("io.opentelemetry.context", "io.opentelemetry.javaagent.shaded.io.opentelemetry.context")
+
+        // relocate the OpenTelemetry extensions that are used by instrumentation modules)
+        // these extensions live in the AgentClassLoader, and are injected into the user's class loader
+        // by the instrumentation modules that use them
+        relocate("io.opentelemetry.extension.aws", "io.opentelemetry.javaagent.shaded.io.opentelemetry.extension.aws")
+        relocate("io.opentelemetry.extension.kotlin", "io.opentelemetry.javaagent.shaded.io.opentelemetry.extension.kotlin")
+      }
+    }
+  }
+
   plugins.withId("maven-publish") {
     plugins.apply("signing")
 

dependencyManagement/build.gradle.kts

@@ -66,7 +66,8 @@ val DEPENDENCIES = listOf(
   "commons-logging:commons-logging:1.2",
   "com.sparkjava:spark-core:2.9.3",
   "com.squareup.okhttp3:okhttp:4.9.1",
-  "io.opentelemetry.javaagent:opentelemetry-javaagent:${if (!TEST_SNAPSHOTS) "1.0.1" else "1.1.0-SNAPSHOT"}"
+  "io.opentelemetry.javaagent:opentelemetry-javaagent:${if (!TEST_SNAPSHOTS) "1.0.1" else "1.1.0-SNAPSHOT"}",
+  "net.bytebuddy:byte-buddy:1.10.22"
 )
 
 javaPlatform {

instrumentation/log4j-2.13.2/build.gradle.kts

@@ -0,0 +1,29 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+plugins {
+  java
+  id("com.github.johnrengelman.shadow")
+}
+
+base.archivesBaseName = "aws-instrumentation-log4j-2.13.2"
+
+dependencies {
+  compileOnly("io.opentelemetry:opentelemetry-api")
+  compileOnly("io.opentelemetry.javaagent:opentelemetry-javaagent-tooling")
+  compileOnly("net.bytebuddy:byte-buddy")
+
+  compileOnly("org.apache.logging.log4j:log4j-core:2.13.2")
+}

instrumentation/log4j-2.13.2/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/log4j_2_13_2/AwsXrayContextDataProvider.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.opentelemetry.javaagent.instrumentation.log4j_2_13_2;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import java.util.Collections;
+import java.util.Map;
+import org.apache.logging.log4j.core.util.ContextDataProvider;
+
+/**
+ * A {@link ContextDataProvider} which injects the trace and span ID of the current {@link Span} in
+ * a format for consumption by AWS X-Ray and related services.
+ */
+public class AwsXrayContextDataProvider implements ContextDataProvider {
+  private static final String TRACE_ID_KEY = "AWS-XRAY-TRACE-ID";
+
+  @Override
+  public Map<String, String> supplyContextData() {
+    Span currentSpan = Span.current();
+    SpanContext spanContext = currentSpan.getSpanContext();
+    if (!spanContext.isValid()) {
+      return Collections.emptyMap();
+    }
+
+    String value =
+        "1-"
+            + spanContext.getTraceId().substring(0, 8)
+            + "-"
+            + spanContext.getTraceId().substring(8)
+            + "@"
+            + spanContext.getSpanId();
+    return Collections.singletonMap(TRACE_ID_KEY, value);
+  }
+}

instrumentation/log4j-2.13.2/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/log4j_2_13_2/AwsXrayLog4jInstrumentationModule.java

@@ -0,0 +1,73 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.opentelemetry.javaagent.instrumentation.log4j_2_13_2;
+
+import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.ClassLoaderMatcher.hasClassesNamed;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import io.opentelemetry.javaagent.tooling.InstrumentationModule;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class AwsXrayLog4jInstrumentationModule extends InstrumentationModule {
+
+  public AwsXrayLog4jInstrumentationModule() {
+    super("log4j", "log4j-2.13.2", "aws-log4j", "aws-log4j-2.13.2");
+  }
+
+  // The SPI will be merged with what's in the agent so we don't need to inject it, only our
+  // provider implementation.
+  @Override
+  protected String[] additionalHelperClassNames() {
+    return new String[] {
+      "software.amazon.opentelemetry.javaagent.instrumentation.log4j_2_13_2."
+          + "AwsXrayContextDataProvider"
+    };
+  }
+
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    return hasClassesNamed("org.apache.logging.log4j.core.util.ContextDataProvider");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return Collections.singletonList(new EmptyTypeInstrumentation());
+  }
+
+  public static class EmptyTypeInstrumentation implements TypeInstrumentation {
+    @Override
+    public ElementMatcher<? super TypeDescription> typeMatcher() {
+      // we cannot use ContextDataProvider here because one of the classes that we inject implements
+      // this interface, causing the interface to be loaded while it's being transformed, which
+      // leads
+      // to duplicate class definition error after the interface is transformed and the triggering
+      // class loader tries to load it.
+      return named("org.apache.logging.log4j.core.impl.ThreadContextDataInjector");
+    }
+
+    @Override
+    public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+      // Nothing to instrument, no methods to match
+      return Collections.emptyMap();
+    }
+  }
+}

instrumentation/log4j-2.13.2/src/main/resources/META-INF/services/io.opentelemetry.javaagent.tooling.InstrumentationModule

@@ -0,0 +1 @@
+software.amazon.opentelemetry.javaagent.instrumentation.log4j_2_13_2.AwsXrayLog4jInstrumentationModule

instrumentation/log4j-2.13.2/src/main/resources/META-INF/services/org.apache.logging.log4j.core.util.ContextDataProvider

@@ -0,0 +1 @@
+software.amazon.opentelemetry.javaagent.instrumentation.log4j_2_13_2.AwsXrayContextDataProvider

instrumentation/logback-1.0/build.gradle.kts

@@ -0,0 +1,31 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+plugins {
+  java
+  id("com.github.johnrengelman.shadow")
+}
+
+base.archivesBaseName = "aws-instrumentation-logback-1.0"
+
+dependencies {
+  compileOnly("io.opentelemetry:opentelemetry-api")
+  compileOnly("io.opentelemetry.instrumentation:opentelemetry-logback-1.0")
+  compileOnly("io.opentelemetry.javaagent:opentelemetry-javaagent-api")
+  compileOnly("io.opentelemetry.javaagent:opentelemetry-javaagent-tooling")
+  compileOnly("net.bytebuddy:byte-buddy")
+
+  compileOnly("ch.qos.logback:logback-classic:1.0.0")
+}

instrumentation/logback-1.0/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/logback_1_0/AwsXrayLogbackInstrumentationModule.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.opentelemetry.javaagent.instrumentation.logback_1_0;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.javaagent.tooling.InstrumentationModule;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+public class AwsXrayLogbackInstrumentationModule extends InstrumentationModule {
+  public AwsXrayLogbackInstrumentationModule() {
+    super("logback", "logback-1.0", "aws-logback", "aws-logback-1.0");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return Collections.singletonList(new AwsXrayLoggingEventInstrumentation());
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return Collections.singletonMap(
+        "ch.qos.logback.classic.spi.ILoggingEvent", Span.class.getName());
+  }
+}