feat: support customization for MDC serializer (#293)

Author: aliaksandr-lavishak-tde

log4j2-ecs-layout/src/main/java/co/elastic/logging/log4j2/DefaultMdcSerializer.java

@@ -0,0 +1,72 @@
+/*-
+ * #%L
+ * Java ECS logging
+ * %%
+ * Copyright (C) 2019 - 2020 Elastic and contributors
+ * %%
+ * 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.
+ * #L%
+ */
+package co.elastic.logging.log4j2;
+
+import co.elastic.logging.EcsJsonSerializer;
+import co.elastic.logging.JsonUtils;
+import org.apache.logging.log4j.core.LogEvent;
+import org.apache.logging.log4j.util.TriConsumer;
+
+interface DefaultMdcSerializer extends MdcSerializer {
+
+    /**
+     * Garbage free MDC serialization for log4j2 2.7+
+     * Never reference directly in prod code so avoid linkage errors when TriConsumer or getContextData are not available
+     */
+    enum UsingContextData implements MdcSerializer {
+
+        @SuppressWarnings("unused")
+        INSTANCE;
+
+        private static final TriConsumer<String, Object, StringBuilder> WRITE_MDC = new TriConsumer<String, Object, StringBuilder>() {
+            @Override
+            public void accept(final String key, final Object value, final StringBuilder stringBuilder) {
+                stringBuilder.append('\"');
+                JsonUtils.quoteAsString(key, stringBuilder);
+                stringBuilder.append("\":\"");
+                JsonUtils.quoteAsString(EcsJsonSerializer.toNullSafeString(String.valueOf(value)), stringBuilder);
+                stringBuilder.append("\",");
+            }
+        };
+
+
+        @Override
+        public void serializeMdc(LogEvent event, StringBuilder builder) {
+            event.getContextData().forEach(WRITE_MDC, builder);
+        }
+    }
+
+    /**
+     * Fallback for log4j2 <= 2.6
+     */
+    enum UsingContextMap implements MdcSerializer {
+        INSTANCE;
+
+        @Override
+        public void serializeMdc(LogEvent event, StringBuilder builder) {
+            EcsJsonSerializer.serializeMDC(builder, event.getContextMap());
+        }
+    }
+}

log4j2-ecs-layout/src/main/java/co/elastic/logging/log4j2/EcsLayout.java

@@ -62,7 +62,6 @@ public class EcsLayout extends AbstractStringLayout {
 
     public static final Charset UTF_8 = Charset.forName("UTF-8");
     private static final ObjectMessageJacksonSerializer JACKSON_SERIALIZER = ObjectMessageJacksonSerializer.Resolver.resolve();
-    private static final MdcSerializer MDC_SERIALIZER = MdcSerializer.Resolver.resolve();
     private static final MultiFormatHandler MULTI_FORMAT_HANDLER = MultiFormatHandler.Resolver.resolve();
     private static final boolean FORMAT_MESSAGES_PATTERN_DISABLE_LOOKUPS = PropertiesUtil.getProperties().getBooleanProperty(
             "log4j2.formatMsgNoLookups", false);
@@ -79,9 +78,10 @@ public class EcsLayout extends AbstractStringLayout {
     private final boolean includeOrigin;
     private final PatternFormatter[] exceptionPatternFormatter;
     private final ConcurrentMap<Class<? extends MultiformatMessage>, Boolean> supportsJson = new ConcurrentHashMap<Class<? extends MultiformatMessage>, Boolean>();
+    private final MdcSerializer mdcSerializer;
 
     private EcsLayout(Configuration config, String serviceName, String serviceVersion, String serviceEnvironment, String serviceNodeName, String eventDataset, boolean includeMarkers,
-                      KeyValuePair[] additionalFields, boolean includeOrigin, String exceptionPattern, boolean stackTraceAsArray) {
+                      KeyValuePair[] additionalFields, boolean includeOrigin, String exceptionPattern, boolean stackTraceAsArray, String mdcSerializerFullClassName) {
         super(config, UTF_8, null, null);
         this.serviceName = serviceName;
         this.serviceVersion = serviceVersion;
@@ -109,6 +109,7 @@ private EcsLayout(Configuration config, String serviceName, String serviceVersio
         } else {
             exceptionPatternFormatter = null;
         }
+        mdcSerializer = MdcSerializerResolver.resolve(mdcSerializerFullClassName);
     }
 
     @PluginBuilderFactory
@@ -193,7 +194,7 @@ private void serializeAdditionalFieldsAndMDC(LogEvent event, StringBuilder build
                 }
             }
         }
-        MDC_SERIALIZER.serializeMdc(event, builder);
+        mdcSerializer.serializeMdc(event, builder);
     }
 
     private static void formatPattern(LogEvent event, PatternFormatter[] formatters, StringBuilder buffer) {
@@ -377,6 +378,8 @@ public static class Builder implements org.apache.logging.log4j.core.util.Builde
         private KeyValuePair[] additionalFields = new KeyValuePair[]{};
         @PluginBuilderAttribute("includeOrigin")
         private boolean includeOrigin = false;
+        @PluginBuilderAttribute("mdcSerializer")
+        private String mdcSerializerFullClassName = "";
 
         Builder() {
         }
@@ -428,6 +431,10 @@ public String getExceptionPattern() {
             return exceptionPattern;
         }
 
+        public String getMdcSerializerFullClassName() {
+            return mdcSerializerFullClassName;
+        }
+
         /**
          * Additional fields to set on each log event.
          *
@@ -483,11 +490,16 @@ public EcsLayout.Builder setExceptionPattern(String exceptionPattern) {
             return this;
         }
 
+        public EcsLayout.Builder setMdcSerializerFullClassName(String mdcSerializerFullClassName) {
+            this.mdcSerializerFullClassName = mdcSerializerFullClassName;
+            return this;
+        }
+
         @Override
         public EcsLayout build() {
             return new EcsLayout(getConfiguration(), serviceName, serviceVersion, serviceEnvironment, serviceNodeName,
                     EcsJsonSerializer.computeEventDataset(eventDataset, serviceName),
-                    includeMarkers, additionalFields, includeOrigin, exceptionPattern, stackTraceAsArray);
+                    includeMarkers, additionalFields, includeOrigin, exceptionPattern, stackTraceAsArray, mdcSerializerFullClassName);
         }
     }
 }

log4j2-ecs-layout/src/main/java/co/elastic/logging/log4j2/MdcSerializer.java

@@ -24,65 +24,23 @@
  */
 package co.elastic.logging.log4j2;
 
-import co.elastic.logging.EcsJsonSerializer;
-import co.elastic.logging.JsonUtils;
 import org.apache.logging.log4j.core.LogEvent;
-import org.apache.logging.log4j.util.TriConsumer;
-
-interface MdcSerializer {
-
-    void serializeMdc(LogEvent event, StringBuilder builder);
-
-    class Resolver {
-
-        public static MdcSerializer resolve() {
-            try {
-                LogEvent.class.getMethod("getContextData");
-                return (MdcSerializer) Class.forName("co.elastic.logging.log4j2.MdcSerializer$UsingContextData").getEnumConstants()[0];
-            } catch (Exception ignore) {
-            } catch (LinkageError ignore) {
-            }
-            return UsingContextMap.INSTANCE;
-        }
-
-    }
-
-    /**
-     * Garbage free MDC serialization for log4j2 2.7+
-     * Never reference directly in prod code so avoid linkage errors when TriConsumer or getContextData are not available
-     */
-    enum UsingContextData implements MdcSerializer {
-
-        @SuppressWarnings("unused")
-        INSTANCE;
-
-        private static final TriConsumer<String, Object, StringBuilder> WRITE_MDC = new TriConsumer<String, Object, StringBuilder>() {
-            @Override
-            public void accept(final String key, final Object value, final StringBuilder stringBuilder) {
-                stringBuilder.append('\"');
-                JsonUtils.quoteAsString(key, stringBuilder);
-                stringBuilder.append("\":\"");
-                JsonUtils.quoteAsString(EcsJsonSerializer.toNullSafeString(String.valueOf(value)), stringBuilder);
-                stringBuilder.append("\",");
-            }
-        };
-
-
-        @Override
-        public void serializeMdc(LogEvent event, StringBuilder builder) {
-            event.getContextData().forEach(WRITE_MDC, builder);
-        }
-    }
 
+/**
+ * Interface for serializing MDC (Mapped Diagnostic Context) data from a {@link LogEvent}.
+ * <p>
+ * Implementations must have a public no-argument constructor to allow dynamic instantiation.
+ * </p>
+ */
+public interface MdcSerializer {
     /**
-     * Fallback for log4j2 <= 2.6
+     * Add MDC data for the give log event to the provided output string builder. The output written to the string
+     * builder must be a valid JSON-object without the surrounding curly braces and with a trailing comma. If this MDC
+     * serializer does not append any content, no comma shall be added. For example, the serializer could output the
+     * following content: "foo":"bar","key":"value",
+     *
+     * @param event   the log event to write the MDC content for
+     * @param builder the output JSON string builder
      */
-    enum UsingContextMap implements MdcSerializer {
-        INSTANCE;
-
-        @Override
-        public void serializeMdc(LogEvent event, StringBuilder builder) {
-            EcsJsonSerializer.serializeMDC(builder, event.getContextMap());
-        }
-    }
+    void serializeMdc(LogEvent event, StringBuilder builder);
 }

log4j2-ecs-layout/src/main/java/co/elastic/logging/log4j2/MdcSerializerResolver.java

@@ -0,0 +1,58 @@
+/*-
+ * #%L
+ * Java ECS logging
+ * %%
+ * Copyright (C) 2019 - 2020 Elastic and contributors
+ * %%
+ * 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.
+ * #L%
+ */
+package co.elastic.logging.log4j2;
+
+import java.lang.reflect.InvocationTargetException;
+
+import org.apache.logging.log4j.core.LogEvent;
+
+import co.elastic.logging.log4j2.DefaultMdcSerializer.UsingContextMap;
+
+class MdcSerializerResolver {
+
+    static MdcSerializer resolve(String mdcSerializerFullClassName) {
+        if (mdcSerializerFullClassName == null || mdcSerializerFullClassName.isEmpty()) {
+            return resolveDefault();
+        }
+        try {
+            Class<?> clazz = Class.forName(mdcSerializerFullClassName);
+            return (MdcSerializer) clazz.getDeclaredConstructor().newInstance();
+        } catch (ClassNotFoundException | InstantiationException | IllegalAccessException | NoSuchMethodException |
+                 InvocationTargetException e) {
+            throw new IllegalArgumentException("Could not create MdcSerializer " + mdcSerializerFullClassName, e);
+        }
+    }
+
+    private static MdcSerializer resolveDefault() {
+        try {
+            LogEvent.class.getMethod("getContextData");
+            return (DefaultMdcSerializer) Class.forName(
+                    "co.elastic.logging.log4j2.DefaultMdcSerializer$UsingContextData").getEnumConstants()[0];
+        } catch (Exception | LinkageError ignore) {
+        }
+        return UsingContextMap.INSTANCE;
+    }
+
+}

log4j2-ecs-layout/src/test/java/co/elastic/logging/log4j2/AbstractLog4j2EcsLayoutTest.java

@@ -37,6 +37,7 @@
 
 import java.util.List;
 
+import static co.elastic.logging.log4j2.CustomMdcSerializer.CUSTOM_MDC_SERIALIZER_TEST_KEY;
 import static org.assertj.core.api.Assertions.assertThat;
 
 abstract class AbstractLog4j2EcsLayoutTest extends AbstractEcsLoggingTest {
@@ -52,6 +53,7 @@ void tearDown() throws Exception {
     void testAdditionalFieldsWithLookup() throws Exception {
         putMdc("trace.id", "foo");
         putMdc("foo", "bar");
+        putMdc(CUSTOM_MDC_SERIALIZER_TEST_KEY, "some_text_lower_case");
         debug("test");
         assertThat(getAndValidateLastLogLine().get("cluster.uuid").textValue()).isEqualTo("9fe9134b-20b0-465e-acf9-8cc09ac9053b");
         assertThat(getAndValidateLastLogLine().get("node.id").textValue()).isEqualTo("foo");
@@ -60,6 +62,7 @@ void testAdditionalFieldsWithLookup() throws Exception {
         assertThat(getAndValidateLastLogLine().get("clazz").textValue()).startsWith(getClass().getPackageName());
         assertThat(getAndValidateLastLogLine().get("404")).isNull();
         assertThat(getAndValidateLastLogLine().get("foo").textValue()).isEqualTo("bar");
+        assertThat(getAndValidateLastLogLine().get(CUSTOM_MDC_SERIALIZER_TEST_KEY).textValue()).isEqualTo("some_text_lower_case");
     }
 
     @Test

log4j2-ecs-layout/src/test/java/co/elastic/logging/log4j2/CustomMdcSerializer.java

@@ -0,0 +1,71 @@
+/*-
+ * #%L
+ * Java ECS logging
+ * %%
+ * Copyright (C) 2019 - 2020 Elastic and contributors
+ * %%
+ * 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.
+ * #L%
+ */
+package co.elastic.logging.log4j2;
+
+import org.apache.logging.log4j.core.LogEvent;
+import org.apache.logging.log4j.util.TriConsumer;
+
+import co.elastic.logging.EcsJsonSerializer;
+import co.elastic.logging.JsonUtils;
+
+public class CustomMdcSerializer implements MdcSerializer {
+
+    protected static final String  CUSTOM_MDC_SERIALIZER_TEST_KEY = "SPECIAL_TEST_CUSTOM_KEY";
+
+    @Override
+    public void serializeMdc(LogEvent event, StringBuilder builder) {
+        event.getContextData()
+                .forEach((key, value) -> getWriteFunctionForKey(key).accept(key, value, builder));
+    }
+
+    // Default function for serializing MDC entries
+    private static final TriConsumer<String, Object, StringBuilder> DEFAULT_WRITE_MDC_FUNCTION = (key, value, stringBuilder) -> {
+        stringBuilder.append('\"');
+        JsonUtils.quoteAsString(key, stringBuilder);
+        stringBuilder.append("\":\"");
+        JsonUtils.quoteAsString(EcsJsonSerializer.toNullSafeString(String.valueOf(value)), stringBuilder);
+        stringBuilder.append("\",");
+    };
+
+    // Custom function for handling a specific key
+    private static final TriConsumer<String, Object, StringBuilder> CUSTOM_KEY_WRITE_MDC_FUNCTION = (key, value, stringBuilder) -> DEFAULT_WRITE_MDC_FUNCTION.accept(
+            key,
+            value.toString().toUpperCase(),
+            stringBuilder
+    );
+
+    /**
+     * Returns the appropriate function to write an MDC entry based on the key.
+     *
+     * @param key MDC key.
+     * @return The function to serialize the MDC entry value.
+     */
+    private TriConsumer<String, Object, StringBuilder> getWriteFunctionForKey(String key) {
+        if (CUSTOM_MDC_SERIALIZER_TEST_KEY.equals(key)) {
+            return CUSTOM_KEY_WRITE_MDC_FUNCTION;
+        }
+        return DEFAULT_WRITE_MDC_FUNCTION;
+    }
+}