log4j-mdc-1.2

Author: SylvainJuge

instrumentation/log4j/log4j-mdc-1.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/log4j/mdc/v1_2/CategoryInstrumentation.java

@@ -11,8 +11,6 @@
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
-import io.opentelemetry.context.Context;
-import io.opentelemetry.instrumentation.api.util.VirtualField;
 import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
@@ -43,8 +41,7 @@ public static class CallAppendersAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void onEnter(@Advice.Argument(0) LoggingEvent event) {
-      VirtualField.find(LoggingEvent.class, Context.class)
-          .set(event, Java8BytecodeBridge.currentContext());
+      VirtualFieldHelper.CONTEXT.set(event, Java8BytecodeBridge.currentContext());
     }
   }
 }

instrumentation/log4j/log4j-mdc-1.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/log4j/mdc/v1_2/Log4j1InstrumentationModule.java

@@ -11,11 +11,13 @@
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
 import java.util.List;
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(InstrumentationModule.class)
-public class Log4j1InstrumentationModule extends InstrumentationModule {
+public class Log4j1InstrumentationModule extends InstrumentationModule
+    implements ExperimentalInstrumentationModule {
   public Log4j1InstrumentationModule() {
     super("log4j-mdc", "log4j-mdc-1.2");
   }
@@ -29,4 +31,9 @@ public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(new CategoryInstrumentation(), new LoggingEventInstrumentation());
   }
+
+  @Override
+  public boolean isIndyReady() {
+    return true;
+  }
 }

instrumentation/log4j/log4j-mdc-1.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/log4j/mdc/v1_2/LoggingEventInstrumentation.java

@@ -13,13 +13,14 @@
 
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.context.Context;
-import io.opentelemetry.instrumentation.api.util.VirtualField;
 import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.bootstrap.internal.AgentCommonConfig;
 import io.opentelemetry.javaagent.bootstrap.internal.ConfiguredResourceAttributesHolder;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import javax.annotation.Nullable;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.apache.log4j.spi.LoggingEvent;
@@ -44,41 +45,38 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class GetMdcAdvice {
 
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void onExit(
+    public static Object onExit(
         @Advice.This LoggingEvent event,
         @Advice.Argument(0) String key,
-        @Advice.Return(readOnly = false) Object value) {
-      if (AgentCommonConfig.get().getTraceIdKey().equals(key)
-          || AgentCommonConfig.get().getSpanIdKey().equals(key)
-          || AgentCommonConfig.get().getTraceFlagsKey().equals(key)) {
-        if (value != null) {
-          // Assume already instrumented event if traceId/spanId/sampled is present.
-          return;
-        }
+        @Advice.Return @Nullable Object returnValue) {
 
-        Context context = VirtualField.find(LoggingEvent.class, Context.class).get(event);
-        if (context == null) {
-          return;
-        }
-
-        SpanContext spanContext = Java8BytecodeBridge.spanFromContext(context).getSpanContext();
-        if (!spanContext.isValid()) {
-          return;
-        }
+      if (returnValue != null) {
+        return returnValue;
+      }
+      boolean traceId = AgentCommonConfig.get().getTraceIdKey().equals(key);
+      boolean spanId = AgentCommonConfig.get().getSpanIdKey().equals(key);
+      boolean traceFlags = AgentCommonConfig.get().getTraceFlagsKey().equals(key);
 
-        if (AgentCommonConfig.get().getTraceIdKey().equals(key)) {
-          value = spanContext.getTraceId();
-        }
-        if (AgentCommonConfig.get().getSpanIdKey().equals(key)) {
-          value = spanContext.getSpanId();
-        }
-        if (AgentCommonConfig.get().getTraceFlagsKey().equals(key)) {
-          value = spanContext.getTraceFlags().asHex();
-        }
-      } else if (value == null) {
-        value = ConfiguredResourceAttributesHolder.getAttributeValue(key);
+      if (!traceId && !spanId && !traceFlags) {
+        return ConfiguredResourceAttributesHolder.getAttributeValue(key);
+      }
+      Context context = VirtualFieldHelper.CONTEXT.get(event);
+      if (context == null) {
+        return null;
+      }
+      SpanContext spanContext = Java8BytecodeBridge.spanFromContext(context).getSpanContext();
+      if (!spanContext.isValid()) {
+        return null;
+      }
+      if (traceId) {
+        return spanContext.getTraceId();
+      }
+      if (spanId) {
+        return spanContext.getSpanId();
       }
+      return spanContext.getTraceFlags().asHex();
     }
   }
 }

instrumentation/log4j/log4j-mdc-1.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/log4j/mdc/v1_2/VirtualFieldHelper.java

@@ -0,0 +1,18 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.log4j.mdc.v1_2;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.util.VirtualField;
+import org.apache.log4j.spi.LoggingEvent;
+
+public class VirtualFieldHelper {
+
+  public static final VirtualField<LoggingEvent, Context> CONTEXT =
+      VirtualField.find(LoggingEvent.class, Context.class);
+
+  private VirtualFieldHelper() {}
+}