Reimplementing log correlation

Author: eyalkoren

apm-agent-plugins/apm-log-shader-plugin/apm-log-shader-plugin-common/src/main/java/co/elastic/apm/agent/log/shader/AbstractLogCorrelationHelper.java

@@ -0,0 +1,61 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License 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 co.elastic.apm.agent.log.shader;
+
+import co.elastic.apm.agent.impl.transaction.Transaction;
+import co.elastic.apm.agent.sdk.state.CallDepth;
+
+import javax.annotation.Nullable;
+
+public abstract class AbstractLogCorrelationHelper {
+
+    private static final CallDepth callDepth = CallDepth.get(AbstractLogCorrelationHelper.class);
+
+    public static final String TRACE_ID_MDC_KEY = "trace.id";
+    public static final String TRANSACTION_ID_MDC_KEY = "transaction.id";
+
+    /**
+     * Adds the active transaction's ID and trace ID to the MDC in the outmost logging API call
+     * @param activeTransaction the currently active transaction, or {@code null} if there is no such
+     * @return returns {@code true} if the transaction IDs were added to the MDC
+     */
+    public boolean beforeLoggingApiCall(@Nullable Transaction activeTransaction) {
+        if (callDepth.isNestedCallAndIncrement() || activeTransaction == null) {
+            return false;
+        }
+        addToMdc(TRACE_ID_MDC_KEY, activeTransaction.getTraceContext().getTraceId().toString());
+        addToMdc(TRANSACTION_ID_MDC_KEY, activeTransaction.getTraceContext().getTransactionId().toString());
+        return true;
+    }
+
+    /**
+     * Clears transaction and trace ID from the MDC if required
+     * @param added should reflect the value returned from {@link #beforeLoggingApiCall(Transaction)} for the corresponding API call
+     */
+    public void afterLoggingApi(boolean added) {
+        if (callDepth.isNestedCallAndDecrement() && added) {
+            removeFromMdc(TRACE_ID_MDC_KEY);
+            removeFromMdc(TRANSACTION_ID_MDC_KEY);
+        }
+    }
+
+    protected abstract void addToMdc(String key, String value);
+
+    protected abstract void removeFromMdc(String key);
+}

apm-agent-plugins/apm-log-shader-plugin/apm-log-shader-plugin-common/src/main/java/co/elastic/apm/agent/log/shader/AbstractLogCorrelationInstrumentation.java

@@ -0,0 +1,52 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License 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 co.elastic.apm.agent.log.shader;
+
+
+import co.elastic.apm.agent.bci.TracerAwareInstrumentation;
+import co.elastic.apm.agent.bci.bytebuddy.CustomElementMatchers;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+import static net.bytebuddy.matcher.ElementMatchers.isBootstrapClassLoader;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+
+public abstract class AbstractLogCorrelationInstrumentation extends TracerAwareInstrumentation {
+
+    @Override
+    public ElementMatcher.Junction<ClassLoader> getClassLoaderMatcher() {
+        return not(isBootstrapClassLoader()).and(not(CustomElementMatchers.isAgentClassLoader()));
+    }
+
+    @Override
+    public ElementMatcher<? super MethodDescription> getMethodMatcher() {
+        return named("trace").or(
+            named("debug").or(
+                named("info").or(
+                    named("warn").or(
+                        named("error").or(
+                            named("log")
+                        )
+                    )
+                )
+            )
+        );
+    }
+}

apm-agent-plugins/apm-log-shader-plugin/apm-log-shader-plugin-common/src/main/java/co/elastic/apm/agent/log/shader/AbstractLogShadingInstrumentation.java renamed to apm-agent-plugins/apm-log-shader-plugin/apm-log-shader-plugin-common/src/main/java/co/elastic/apm/agent/log/shader/AbstractLogIntegrationInstrumentation.java

@@ -24,7 +24,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 
-public abstract class AbstractLogShadingInstrumentation extends TracerAwareInstrumentation {
+public abstract class AbstractLogIntegrationInstrumentation extends TracerAwareInstrumentation {
 
     @Override
     public Collection<String> getInstrumentationGroupNames() {

apm-agent-plugins/apm-log-shader-plugin/apm-log-shader-plugin-common/src/test/java/co/elastic/apm/agent/log/shader/LogShadingInstrumentationTest.java

@@ -20,6 +20,7 @@
 
 import co.elastic.apm.agent.AbstractInstrumentationTest;
 import co.elastic.apm.agent.configuration.CoreConfiguration;
+import co.elastic.apm.agent.impl.transaction.Transaction;
 import co.elastic.apm.agent.logging.LogEcsReformatting;
 import co.elastic.apm.agent.logging.LoggingConfiguration;
 import co.elastic.apm.agent.logging.TestUtils;
@@ -29,8 +30,6 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
-import javax.annotation.Nonnull;
-import javax.annotation.Nullable;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -42,11 +41,12 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.TimeZone;
-import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import static co.elastic.apm.agent.log.shader.AbstractLogCorrelationHelper.TRACE_ID_MDC_KEY;
+import static co.elastic.apm.agent.log.shader.AbstractLogCorrelationHelper.TRANSACTION_ID_MDC_KEY;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.when;
@@ -65,6 +65,7 @@ public abstract class LogShadingInstrumentationTest extends AbstractInstrumentat
     private final ObjectMapper objectMapper;
     private LoggingConfiguration loggingConfig;
     private String serviceName;
+    private Transaction transaction;
 
     public LogShadingInstrumentationTest() {
         logger = createLoggerFacade();
@@ -82,6 +83,8 @@ public void setup() throws Exception {
 
         // IMPORTANT: keep this last, so that it doesn't interfere with Mockito settings above
         serviceName = Objects.requireNonNull(tracer.getMetaDataFuture().get(2000, TimeUnit.MILLISECONDS).getService().getName());
+
+        transaction = Objects.requireNonNull(tracer.startRootTransaction(null)).activate();
     }
 
     private void setEcsReformattingConfig(LogEcsReformatting ecsReformattingConfig) {
@@ -96,6 +99,7 @@ private void initializeShadeDir(String dirName) throws IOException {
 
     @AfterEach
     public void closeLogger() {
+        transaction.deactivate().end();
         logger.close();
     }
 
@@ -109,16 +113,10 @@ public void testSimpleLogShading() throws Exception {
     }
 
     private void runSimpleScenario() throws Exception {
-        String traceId = UUID.randomUUID().toString();
-        logger.putTraceIdToMdc(traceId);
-        try {
-            logger.trace(TRACE_MESSAGE);
-            logger.debug(DEBUG_MESSAGE);
-            logger.warn(WARN_MESSAGE);
-            logger.error(ERROR_MESSAGE);
-        } finally {
-            logger.removeTraceIdFromMdc();
-        }
+        logger.trace(TRACE_MESSAGE);
+        logger.debug(DEBUG_MESSAGE);
+        logger.warn(WARN_MESSAGE);
+        logger.error(ERROR_MESSAGE);
 
         ArrayList<String[]> rawLogLines = readRawLogLines();
         assertThat(rawLogLines).hasSize(4);
@@ -127,7 +125,7 @@ private void runSimpleScenario() throws Exception {
         assertThat(ecsLogLines).hasSize(4);
 
         for (int i = 0; i < 4; i++) {
-            verifyEcsFormat(rawLogLines.get(i), ecsLogLines.get(i), traceId);
+            verifyEcsFormat(rawLogLines.get(i), ecsLogLines.get(i));
         }
     }
 
@@ -146,7 +144,7 @@ public void testMarkers() throws Exception {
             assertThat(ecsLogLines).hasSize(1);
             JsonNode ecsLogLine = ecsLogLines.get(0);
 
-            verifyEcsFormat(rawLogLine, ecsLogLine, null);
+            verifyEcsFormat(rawLogLine, ecsLogLine);
 
             JsonNode tagsJson = ecsLogLine.get("tags");
             assertThat(tagsJson.isArray()).isTrue();
@@ -181,8 +179,8 @@ public void testLazyShadeFileCreation() throws Exception {
 
         ArrayList<JsonNode> ecsLogLines = readShadeLogFile();
         assertThat(ecsLogLines).hasSize(2);
-        verifyEcsFormat(rawLogLines.get(2), ecsLogLines.get(0), null);
-        verifyEcsFormat(rawLogLines.get(3), ecsLogLines.get(1), null);
+        verifyEcsFormat(rawLogLines.get(2), ecsLogLines.get(0));
+        verifyEcsFormat(rawLogLines.get(3), ecsLogLines.get(1));
     }
 
     @Test
@@ -265,9 +263,9 @@ public void testDynamicConfiguration() throws Exception {
         assertThat(debugLogLine.get("log.level").textValue()).isEqualTo("DEBUG");
 
         // WARN messages should match content but not format
-        verifyEcsFormat(originalLogLines.get(2).split("\\s+"), ecsLogLines.get(0), null);
+        verifyEcsFormat(originalLogLines.get(2).split("\\s+"), ecsLogLines.get(0));
         assertThat(ecsLogLines.get(0).get("log.level").textValue()).isEqualTo("WARN");
-        verifyEcsFormat(originalLogLines.get(5).split("\\s+"), ecsLogLines.get(2), null);
+        verifyEcsFormat(originalLogLines.get(5).split("\\s+"), ecsLogLines.get(2));
         assertThat(ecsLogLines.get(2).get("log.level").textValue()).isEqualTo("WARN");
 
         // ERROR messages should be only in shade file in ECS format
@@ -288,14 +286,14 @@ private void verifyEcsLogLine(JsonNode ecsLogLineTree) {
         assertThat(ecsLogLineTree.get("event.dataset").textValue()).isEqualTo(serviceName + ".FILE");
         assertThat(ecsLogLineTree.get("service.version").textValue()).isEqualTo("v42");
         assertThat(ecsLogLineTree.get("some.field").textValue()).isEqualTo("some-value");
+        assertThat(ecsLogLineTree.get(TRACE_ID_MDC_KEY).textValue()).isEqualTo(transaction.getTraceContext().getTraceId().toString());
+        assertThat(ecsLogLineTree.get(TRANSACTION_ID_MDC_KEY).textValue()).isEqualTo(transaction.getTraceContext().getTransactionId().toString());
     }
 
-    @Nonnull
     private ArrayList<JsonNode> readShadeLogFile() throws IOException {
         return TestUtils.readJsonFile(getShadeLogFilePath());
     }
 
-    @Nonnull
     private ArrayList<String[]> readRawLogLines() throws IOException {
         ArrayList<String[]> rawLogLines;
         try (Stream<String> stream = Files.lines(getOriginalLogFilePath())) {
@@ -304,17 +302,15 @@ private ArrayList<String[]> readRawLogLines() throws IOException {
         return rawLogLines;
     }
 
-    @Nonnull
     private Path getOriginalLogFilePath() {
         return Paths.get(logger.getLogFilePath());
     }
 
-    @Nonnull
     protected String getShadeLogFilePath() {
         return Utils.computeShadeLogFilePath(logger.getLogFilePath(), loggingConfig.getLogEcsFormattingDestinationDir());
     }
 
-    private void verifyEcsFormat(String[] splitRawLogLine, JsonNode ecsLogLineTree, @Nullable String traceId) throws Exception {
+    private void verifyEcsFormat(String[] splitRawLogLine, JsonNode ecsLogLineTree) throws Exception {
         SimpleDateFormat timestampFormat = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSS'Z'");
         Date rawTimestamp = timestampFormat.parse(splitRawLogLine[0]);
         timestampFormat.setTimeZone(TimeZone.getTimeZone("UTC"));
@@ -329,13 +325,8 @@ private void verifyEcsFormat(String[] splitRawLogLine, JsonNode ecsLogLineTree,
         assertThat(ecsLogLineTree.get("event.dataset").textValue()).isEqualTo(serviceName + ".FILE");
         assertThat(ecsLogLineTree.get("service.version").textValue()).isEqualTo("v42");
         assertThat(ecsLogLineTree.get("some.field").textValue()).isEqualTo("some-value");
-        JsonNode jsonTraceId = ecsLogLineTree.get("trace.id");
-        if (traceId != null) {
-            assertThat(jsonTraceId).isNotNull();
-            assertThat(jsonTraceId.asText()).isEqualTo(traceId);
-        } else {
-            assertThat(jsonTraceId).isNull();
-        }
+        assertThat(ecsLogLineTree.get(TRACE_ID_MDC_KEY).textValue()).isEqualTo(transaction.getTraceContext().getTraceId().toString());
+        assertThat(ecsLogLineTree.get(TRANSACTION_ID_MDC_KEY).textValue()).isEqualTo(transaction.getTraceContext().getTransactionId().toString());
     }
 
     /**

apm-agent-plugins/apm-log-shader-plugin/apm-log4j1-plugin/src/main/java/co/elastic/apm/agent/log4j1/Log4j1LogCorrelationHelper.java

@@ -0,0 +1,34 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License 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 co.elastic.apm.agent.log4j1;
+
+import co.elastic.apm.agent.log.shader.AbstractLogCorrelationHelper;
+import org.apache.log4j.MDC;
+
+public class Log4j1LogCorrelationHelper extends AbstractLogCorrelationHelper {
+    @Override
+    protected void addToMdc(String key, String value) {
+        MDC.put(key, value);
+    }
+
+    @Override
+    protected void removeFromMdc(String key) {
+        MDC.remove(key);
+    }
+}

apm-agent-plugins/apm-log-shader-plugin/apm-log4j1-plugin/src/main/java/co/elastic/apm/agent/log4j1/Log4j1LogShadingInstrumentation.java

@@ -18,7 +18,7 @@
  */
 package co.elastic.apm.agent.log4j1;
 
-import co.elastic.apm.agent.log.shader.AbstractLogShadingInstrumentation;
+import co.elastic.apm.agent.log.shader.AbstractLogIntegrationInstrumentation;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -32,7 +32,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
-public abstract class Log4j1LogShadingInstrumentation extends AbstractLogShadingInstrumentation {
+public abstract class Log4j1LogShadingInstrumentation extends AbstractLogIntegrationInstrumentation {
 
     @Override
     public Collection<String> getInstrumentationGroupNames() {