Capture log4j2 context data (mdc) attributes (#4959)

* Capture log4j2 context data (mdc) attributes

* Spotless

* Remove system.out.println

* Switch to allow-list

* Spotless

instrumentation/log4j/log4j-appender-2.16/javaagent/build.gradle.kts

@@ -33,3 +33,8 @@ tasks {
     dependsOn(testAsync)
   }
 }
+
+tasks.withType<Test>().configureEach {
+  // TODO run tests both with and without experimental log attributes
+  jvmArgs("-Dotel.instrumentation.log4j-appender.experimental.capture-context-data-attributes=*")
+}

instrumentation/log4j/log4j-appender-2.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/log4j/appender/v2_16/Log4jHelper.java

@@ -7,24 +7,43 @@
 
 import io.opentelemetry.instrumentation.api.appender.GlobalLogEmitterProvider;
 import io.opentelemetry.instrumentation.api.appender.LogBuilder;
+import io.opentelemetry.instrumentation.log4j.appender.v2_16.internal.ContextDataAccessor;
 import io.opentelemetry.instrumentation.log4j.appender.v2_16.internal.LogEventMapper;
+import java.util.Map;
+import java.util.function.BiConsumer;
+import javax.annotation.Nullable;
 import org.apache.logging.log4j.Level;
 import org.apache.logging.log4j.Logger;
+import org.apache.logging.log4j.ThreadContext;
 import org.apache.logging.log4j.message.Message;
 
 public final class Log4jHelper {
 
+  private static final LogEventMapper<Map<String, String>> mapper =
+      new LogEventMapper<>(ContextDataAccessorImpl.INSTANCE);
+
   public static void capture(Logger logger, Level level, Message message, Throwable throwable) {
 
     LogBuilder builder =
         GlobalLogEmitterProvider.get().logEmitterBuilder(logger.getName()).build().logBuilder();
+    Map<String, String> contextData = ThreadContext.getImmutableContext();
+    mapper.mapLogEvent(builder, message, level, throwable, null, contextData);
+    builder.emit();
+  }
 
-    LogEventMapper.setBody(builder, message);
-    LogEventMapper.setSeverity(builder, level);
-    LogEventMapper.setThrowable(builder, throwable);
-    LogEventMapper.setContext(builder);
+  private enum ContextDataAccessorImpl implements ContextDataAccessor<Map<String, String>> {
+    INSTANCE;
 
-    builder.emit();
+    @Override
+    @Nullable
+    public Object getValue(Map<String, String> contextData, String key) {
+      return contextData.get(key);
+    }
+
+    @Override
+    public void forEach(Map<String, String> contextData, BiConsumer<String, Object> action) {
+      contextData.forEach(action);
+    }
   }
 
   private Log4jHelper() {}

instrumentation/log4j/log4j-appender-2.16/javaagent/src/test/groovy/Log4j2Test.groovy

@@ -3,11 +3,13 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import io.opentelemetry.api.common.AttributeKey
 import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
 import io.opentelemetry.sdk.logs.data.Severity
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes
 import org.apache.logging.log4j.LogManager
 import org.apache.logging.log4j.Logger
+import org.apache.logging.log4j.ThreadContext
 import spock.lang.Unroll
 
 import static io.opentelemetry.instrumentation.test.utils.TraceUtils.runUnderTrace
@@ -88,4 +90,31 @@ class Log4j2Test extends AgentInstrumentationSpecification {
     severity = args[1]
     severityText = args[2]
   }
+
+  def "test context data"() {
+    when:
+    ThreadContext.put("key1", "val1");
+    ThreadContext.put("key2", "val2");
+    try {
+      logger.info("xyz")
+    } finally {
+      ThreadContext.clearMap();
+    }
+
+    then:
+
+    await()
+      .untilAsserted(
+        () -> {
+          assertThat(logs).hasSize(1)
+        })
+    def log = logs.get(0)
+    assertThat(log.getBody().asString()).isEqualTo("xyz")
+    assertThat(log.getInstrumentationLibraryInfo().getName()).isEqualTo("abc")
+    assertThat(log.getSeverity()).isEqualTo(Severity.INFO)
+    assertThat(log.getSeverityText()).isEqualTo("INFO")
+    assertThat(log.getAttributes().size()).isEqualTo(2)
+    assertThat(log.getAttributes().get(AttributeKey.stringKey("log4j.context_data.key1"))).isEqualTo("val1")
+    assertThat(log.getAttributes().get(AttributeKey.stringKey("log4j.context_data.key2"))).isEqualTo("val2")
+  }
 }

instrumentation/log4j/log4j-appender-2.16/library/build.gradle.kts

@@ -9,6 +9,12 @@ dependencies {
 
   testImplementation(project(":instrumentation-sdk-appender"))
   testImplementation("io.opentelemetry:opentelemetry-sdk-logs")
+  testImplementation("io.opentelemetry:opentelemetry-sdk-testing")
 
   testImplementation("org.mockito:mockito-core")
 }
+
+tasks.withType<Test>().configureEach {
+  // TODO run tests both with and without experimental log attributes
+  jvmArgs("-Dotel.instrumentation.log4j-appender.experimental.capture-context-data-attributes=*")
+}

instrumentation/log4j/log4j-appender-2.16/library/src/main/java/io/opentelemetry/instrumentation/log4j/appender/v2_16/OpenTelemetryAppender.java

@@ -7,8 +7,11 @@
 
 import io.opentelemetry.instrumentation.api.appender.GlobalLogEmitterProvider;
 import io.opentelemetry.instrumentation.api.appender.LogBuilder;
+import io.opentelemetry.instrumentation.log4j.appender.v2_16.internal.ContextDataAccessor;
 import io.opentelemetry.instrumentation.log4j.appender.v2_16.internal.LogEventMapper;
 import java.io.Serializable;
+import java.util.function.BiConsumer;
+import javax.annotation.Nullable;
 import org.apache.logging.log4j.core.Appender;
 import org.apache.logging.log4j.core.Core;
 import org.apache.logging.log4j.core.Filter;
@@ -18,6 +21,7 @@
 import org.apache.logging.log4j.core.config.Property;
 import org.apache.logging.log4j.core.config.plugins.Plugin;
 import org.apache.logging.log4j.core.config.plugins.PluginBuilderFactory;
+import org.apache.logging.log4j.util.ReadOnlyStringMap;
 
 @Plugin(
     name = OpenTelemetryAppender.PLUGIN_NAME,
@@ -27,6 +31,9 @@ public class OpenTelemetryAppender extends AbstractAppender {
 
   static final String PLUGIN_NAME = "OpenTelemetry";
 
+  private static final LogEventMapper<ReadOnlyStringMap> mapper =
+      new LogEventMapper<>(ContextDataAccessorImpl.INSTANCE);
+
   @PluginBuilderFactory
   public static <B extends Builder<B>> B builder() {
     return new Builder<B>().asBuilder();
@@ -58,7 +65,29 @@ public void append(LogEvent event) {
             .logEmitterBuilder(event.getLoggerName())
             .build()
             .logBuilder();
-    LogEventMapper.mapLogEvent(builder, event);
+    ReadOnlyStringMap contextData = event.getContextData();
+    mapper.mapLogEvent(
+        builder,
+        event.getMessage(),
+        event.getLevel(),
+        event.getThrown(),
+        event.getInstant(),
+        contextData);
     builder.emit();
   }
+
+  private enum ContextDataAccessorImpl implements ContextDataAccessor<ReadOnlyStringMap> {
+    INSTANCE;
+
+    @Override
+    @Nullable
+    public Object getValue(ReadOnlyStringMap contextData, String key) {
+      return contextData.getValue(key);
+    }
+
+    @Override
+    public void forEach(ReadOnlyStringMap contextData, BiConsumer<String, Object> action) {
+      contextData.forEach(action::accept);
+    }
+  }
 }

instrumentation/log4j/log4j-appender-2.16/library/src/main/java/io/opentelemetry/instrumentation/log4j/appender/v2_16/internal/ContextDataAccessor.java

@@ -0,0 +1,17 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.log4j.appender.v2_16.internal;
+
+import java.util.function.BiConsumer;
+import javax.annotation.Nullable;
+
+public interface ContextDataAccessor<T> {
+
+  @Nullable
+  Object getValue(T contextData, String key);
+
+  void forEach(T contextData, BiConsumer<String, Object> action);
+}

instrumentation/log4j/log4j-appender-2.16/library/src/main/java/io/opentelemetry/instrumentation/log4j/appender/v2_16/internal/LogEventMapper.java

@@ -5,21 +5,56 @@
 
 package io.opentelemetry.instrumentation.log4j.appender.v2_16.internal;
 
+import static java.util.Collections.emptyList;
+
+import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.common.AttributesBuilder;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.appender.LogBuilder;
 import io.opentelemetry.instrumentation.api.appender.Severity;
+import io.opentelemetry.instrumentation.api.cache.Cache;
+import io.opentelemetry.instrumentation.api.config.Config;
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
 import java.io.PrintWriter;
 import java.io.StringWriter;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
 import org.apache.logging.log4j.Level;
 import org.apache.logging.log4j.core.LogEvent;
 import org.apache.logging.log4j.core.time.Instant;
 import org.apache.logging.log4j.message.Message;
 
-public final class LogEventMapper {
+public final class LogEventMapper<T> {
+
+  private static final Cache<String, AttributeKey<String>> contextDataAttributeKeys =
+      Cache.bounded(100);
+
+  private final List<String> captureContextDataAttributes;
+
+  // cached as an optimization
+  private final boolean captureAllContextDataAttributes;
+
+  private final ContextDataAccessor<T> contextDataAccessor;
+
+  public LogEventMapper(ContextDataAccessor<T> contextDataAccessor) {
+    this(
+        contextDataAccessor,
+        Config.get()
+            .getList(
+                "otel.instrumentation.log4j-appender.experimental.capture-context-data-attributes",
+                emptyList()));
+  }
+
+  // visible for testing
+  LogEventMapper(
+      ContextDataAccessor<T> contextDataAccessor, List<String> captureContextDataAttributes) {
+    this.contextDataAccessor = contextDataAccessor;
+    this.captureContextDataAttributes = captureContextDataAttributes;
+    this.captureAllContextDataAttributes =
+        captureContextDataAttributes.size() == 1 && captureContextDataAttributes.get(0).equals("*");
+  }
 
   /**
    * Map the {@link LogEvent} data model onto the {@link LogBuilder}. Unmapped fields include:
@@ -31,55 +66,79 @@ public final class LogEventMapper {
    *   <li>Thread priority - {@link LogEvent#getThreadPriority()}
    *   <li>Marker - {@link LogEvent#getMarker()}
    *   <li>Nested diagnostic context - {@link LogEvent#getContextStack()}
-   *   <li>Mapped diagnostic context - {@link LogEvent#getContextData()}
    * </ul>
    */
-  public static void mapLogEvent(LogBuilder builder, LogEvent logEvent) {
-    setBody(builder, logEvent.getMessage());
-    setSeverity(builder, logEvent.getLevel());
-    setThrowable(builder, logEvent.getThrown());
-    setContext(builder);
-
-    // time
-    Instant instant = logEvent.getInstant();
-    if (instant != null) {
-      builder.setEpoch(
-          TimeUnit.MILLISECONDS.toNanos(instant.getEpochMillisecond())
-              + instant.getNanoOfMillisecond(),
-          TimeUnit.NANOSECONDS);
-    }
-  }
+  public void mapLogEvent(
+      LogBuilder builder,
+      Message message,
+      Level level,
+      @Nullable Throwable throwable,
+      @Nullable Instant timestamp,
+      T contextData) {
 
-  public static void setBody(LogBuilder builder, Message message) {
     if (message != null) {
       builder.setBody(message.getFormattedMessage());
     }
-  }
 
-  public static void setSeverity(LogBuilder builder, Level level) {
     if (level != null) {
       builder.setSeverity(levelToSeverity(level));
       builder.setSeverityText(level.name());
     }
-  }
 
-  public static void setThrowable(LogBuilder builder, Throwable throwable) {
+    AttributesBuilder attributes = Attributes.builder();
+
     if (throwable != null) {
-      AttributesBuilder attributes = Attributes.builder();
+      setThrowable(attributes, throwable);
+    }
+
+    captureContextDataAttributes(attributes, contextData);
 
-      // TODO (trask) extract method for recording exception into instrumentation-api-appender
-      attributes.put(SemanticAttributes.EXCEPTION_TYPE, throwable.getClass().getName());
-      attributes.put(SemanticAttributes.EXCEPTION_MESSAGE, throwable.getMessage());
-      StringWriter writer = new StringWriter();
-      throwable.printStackTrace(new PrintWriter(writer));
-      attributes.put(SemanticAttributes.EXCEPTION_STACKTRACE, writer.toString());
+    builder.setAttributes(attributes.build());
+
+    builder.setContext(Context.current());
 
-      builder.setAttributes(attributes.build());
+    if (timestamp != null) {
+      builder.setEpoch(
+          TimeUnit.MILLISECONDS.toNanos(timestamp.getEpochMillisecond())
+              + timestamp.getNanoOfMillisecond(),
+          TimeUnit.NANOSECONDS);
     }
   }
 
-  public static void setContext(LogBuilder builder) {
-    builder.setContext(Context.current());
+  // visible for testing
+  void captureContextDataAttributes(AttributesBuilder attributes, T contextData) {
+
+    if (captureAllContextDataAttributes) {
+      contextDataAccessor.forEach(
+          contextData,
+          (key, value) -> {
+            if (value != null) {
+              attributes.put(getContextDataAttributeKey(key), value.toString());
+            }
+          });
+      return;
+    }
+
+    for (String key : captureContextDataAttributes) {
+      Object value = contextDataAccessor.getValue(contextData, key);
+      if (value != null) {
+        attributes.put(getContextDataAttributeKey(key), value.toString());
+      }
+    }
+  }
+
+  public static AttributeKey<String> getContextDataAttributeKey(String key) {
+    return contextDataAttributeKeys.computeIfAbsent(
+        key, k -> AttributeKey.stringKey("log4j.context_data." + k));
+  }
+
+  private static void setThrowable(AttributesBuilder attributes, Throwable throwable) {
+    // TODO (trask) extract method for recording exception into instrumentation-api-appender
+    attributes.put(SemanticAttributes.EXCEPTION_TYPE, throwable.getClass().getName());
+    attributes.put(SemanticAttributes.EXCEPTION_MESSAGE, throwable.getMessage());
+    StringWriter writer = new StringWriter();
+    throwable.printStackTrace(new PrintWriter(writer));
+    attributes.put(SemanticAttributes.EXCEPTION_STACKTRACE, writer.toString());
   }
 
   private static Severity levelToSeverity(Level level) {
@@ -102,6 +161,4 @@ private static Severity levelToSeverity(Level level) {
     }
     return Severity.UNDEFINED_SEVERITY_NUMBER;
   }
-
-  private LogEventMapper() {}
 }