mdcLogLabel functionality

logback-appender/src/main/java/pl/tkowalcz/tjahzi/logback/LokiAppender.java

@@ -11,6 +11,8 @@ import pl.tkowalcz.tjahzi.stats.MonitoringModule;
 import pl.tkowalcz.tjahzi.stats.MutableMonitoringModuleWrapper;
 
 import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.Map;
 import java.util.function.Function;
 
 public class LokiAppender extends LokiAppenderConfigurator {
@@ -23,6 +25,7 @@ public class LokiAppender extends LokiAppenderConfigurator {
 
     private TjahziLogger logger;
     private String logLevelLabel;
+    private List<String> mdcLogLabels;
     private MutableMonitoringModuleWrapper monitoringModuleWrapper;
 
     public EfficientPatternLayout getEfficientLayout() {
@@ -58,11 +61,11 @@ public class LokiAppender extends LokiAppenderConfigurator {
     protected void append(ILoggingEvent event) {
         String logLevel = event.getLevel().toString();
         ByteBuffer logLine = actualEncoder.apply(event);
+        Map<String, String> mdcPropertyMap = event.getMDCPropertyMap();
 
         LabelSerializer labelSerializer = LabelSerializers.threadLocal();
-        if (logLevelLabel != null) {
-            labelSerializer.appendLabel(logLevelLabel, logLevel);
-        }
+        appendLogLabel(labelSerializer, logLevel);
+        appendMdcLogLabels(labelSerializer, mdcPropertyMap);
 
         logger.log(
                 event.getTimeStamp(),
@@ -71,6 +74,21 @@ public class LokiAppender extends LokiAppenderConfigurator {
         );
     }
 
+    private void appendLogLabel(LabelSerializer labelSerializer, String logLevel) {
+        if (logLevelLabel != null) {
+            labelSerializer.appendLabel(logLevelLabel, logLevel);
+        }
+    }
+
+    private void appendMdcLogLabels(LabelSerializer serializer,
+                                    Map<String, String> mdcPropertyMap) {
+        mdcLogLabels.forEach(mdcLogLabel -> {
+            if (mdcPropertyMap.containsKey(mdcLogLabel)) {
+                serializer.appendLabel(mdcLogLabel, mdcPropertyMap.get(mdcLogLabel));
+            }
+        });
+    }
+
     @Override
     public void start() {
         if (encoder == null) {
@@ -88,6 +106,7 @@ public class LokiAppender extends LokiAppenderConfigurator {
         LokiAppenderFactory lokiAppenderFactory = new LokiAppenderFactory(this);
         loggingSystem = lokiAppenderFactory.createAppender();
         logLevelLabel = lokiAppenderFactory.getLogLevelLabel();
+        mdcLogLabels = lokiAppenderFactory.getMdcLogLabels();
         monitoringModuleWrapper = lokiAppenderFactory.getMonitoringModuleWrapper();
 
         logger = loggingSystem.createLogger();

logback-appender/src/main/java/pl/tkowalcz/tjahzi/logback/LokiAppenderConfigurator.java

@@ -21,6 +21,7 @@ public abstract class LokiAppenderConfigurator extends UnsynchronizedAppenderBas
     private boolean useOffHeapBuffer = true;
 
     private String logLevelLabel;
+    private List<String> mdcLogLabels = new ArrayList<>();
 
     private long batchSize = 10_2400;
     private long batchWait = 5;
@@ -133,4 +134,12 @@ public abstract class LokiAppenderConfigurator extends UnsynchronizedAppenderBas
     public List<Label> getLabels() {
         return labels;
     }
+
+    public List<String> getMdcLogLabels() {
+        return mdcLogLabels;
+    }
+
+    public void addMdcLogLabel(String mdcLogLabel) {
+        this.mdcLogLabels.add(mdcLogLabel);
+    }
 }

logback-appender/src/main/java/pl/tkowalcz/tjahzi/logback/LokiAppenderFactory.java

@@ -10,6 +10,7 @@ import pl.tkowalcz.tjahzi.stats.MutableMonitoringModuleWrapper;
 import pl.tkowalcz.tjahzi.stats.StandardMonitoringModule;
 
 import java.util.HashMap;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Stream;
 
@@ -19,6 +20,7 @@ public class LokiAppenderFactory {
 
     private final HashMap<String, String> lokiLabels;
     private final String logLevelLabel;
+    private final List<String> mdcLogLabels;
     private MutableMonitoringModuleWrapper monitoringModuleWrapper;
 
     public LokiAppenderFactory(LokiAppenderConfigurator configurator) {
@@ -32,6 +34,7 @@ public class LokiAppenderFactory {
 
         lokiLabels = labelFactory.convertLabelsDroppingInvalid();
         logLevelLabel = labelFactory.validateLogLevelLabel(lokiLabels);
+        mdcLogLabels = configurator.getMdcLogLabels();
     }
 
     public LoggingSystem createAppender() {
@@ -88,4 +91,8 @@ public class LokiAppenderFactory {
     public MutableMonitoringModuleWrapper getMonitoringModuleWrapper() {
         return monitoringModuleWrapper;
     }
+
+    public List<String> getMdcLogLabels() {
+        return mdcLogLabels;
+    }
 }

logback-appender/src/test/java/pl/tkowalcz/tjahzi/logback/MdcLogLabelConfigurationTest.java

@@ -0,0 +1,99 @@
+package pl.tkowalcz.tjahzi.logback;
+
+import ch.qos.logback.classic.Logger;
+import ch.qos.logback.classic.LoggerContext;
+import ch.qos.logback.classic.joran.JoranConfigurator;
+import io.restassured.RestAssured;
+import io.restassured.http.ContentType;
+import io.restassured.parsing.Parser;
+import org.awaitility.Awaitility;
+import org.awaitility.Durations;
+import org.junit.jupiter.api.Test;
+import org.slf4j.LoggerFactory;
+import org.slf4j.MDC;
+import org.testcontainers.containers.BindMode;
+import org.testcontainers.containers.GenericContainer;
+import org.testcontainers.containers.wait.strategy.Wait;
+import org.testcontainers.junit.jupiter.Container;
+import org.testcontainers.junit.jupiter.Testcontainers;
+
+import java.net.URI;
+import java.util.UUID;
+
+import static io.restassured.RestAssured.given;
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.CoreMatchers.nullValue;
+
+@Testcontainers
+class MdcLogLabelConfigurationTest {
+
+    @Container
+    public GenericContainer loki = new GenericContainer("grafana/loki:latest")
+            .withCommand("-config.file=/etc/loki-config.yaml")
+            .withClasspathResourceMapping("loki-config.yaml",
+                    "/etc/loki-config.yaml",
+                    BindMode.READ_ONLY)
+            .waitingFor(
+                    Wait.forHttp("/ready")
+                            .forPort(3100)
+            )
+            .withExposedPorts(3100);
+
+    @Test
+    void shouldSendMdcLogLabelAsConfigured() throws Exception {
+        // Given
+        System.setProperty("loki.host", loki.getHost());
+        System.setProperty("loki.port", loki.getFirstMappedPort().toString());
+
+        URI uri = getClass()
+                .getClassLoader()
+                .getResource("appender-test-with-md-log-label-set.xml")
+                .toURI();
+
+        LoggerContext context = (LoggerContext) LoggerFactory.getILoggerFactory();
+
+        JoranConfigurator configurator = new JoranConfigurator();
+        configurator.setContext(context);
+
+        context.reset();
+        configurator.doConfigure(uri.toURL());
+
+        // When
+        String traceIdValue = UUID.randomUUID().toString();
+        String spanIdValue = UUID.randomUUID().toString();
+
+        MDC.put("trace_id", traceIdValue);
+        MDC.put("span_id", spanIdValue);
+
+        Logger logger = context.getLogger(LokiAppenderLargeBatchesTest.class);
+        logger.info("Test");
+
+        // Then
+        RestAssured.port = loki.getFirstMappedPort();
+        RestAssured.baseURI = "http://" + loki.getHost();
+        RestAssured.registerParser("text/plain", Parser.JSON);
+
+        Awaitility
+                .await()
+                .atMost(Durations.TEN_SECONDS)
+                .pollInterval(Durations.ONE_SECOND)
+                .ignoreExceptions()
+                .untilAsserted(() -> {
+                    given()
+                            .contentType(ContentType.URLENC)
+                            .urlEncodingEnabled(false)
+                            .formParam("query=%7Bserver%3D%22127.0.0.1%22%7D")
+                            .when()
+                            .get("/loki/api/v1/query_range")
+                            .then()
+                            .log()
+                            .all()
+                            .statusCode(200)
+                            .body("status", equalTo("success"))
+                            .body("data.result.size()", equalTo(1))
+                            .body("data.result[0].stream.server", equalTo("127.0.0.1"))
+                            .body("data.result[0].stream.trace_id", equalTo(traceIdValue))
+                            .body("data.result[0].stream.span_id", equalTo(spanIdValue));
+                });
+    }
+}

logback-appender/src/test/resources/appender-test-with-md-log-label-set.xml

@@ -0,0 +1,32 @@
+<configuration debug="true">
+    <appender name="Loki" class="pl.tkowalcz.tjahzi.logback.LokiAppender">
+        <host>${loki.host}</host>
+        <port>${loki.port}</port>
+
+        <efficientLayout>
+            <pattern>%-4relative [%thread] %-5level %logger{35} - %msg%n</pattern>
+        </efficientLayout>
+
+        <header>
+            <name>X-Org-Id</name>
+            <value>Circus</value>
+        </header>
+
+        <label>
+            <name>server</name>
+            <value>127.0.0.1</value>
+        </label>
+
+        <mdcLogLabel>
+            trace_id
+        </mdcLogLabel>
+
+        <mdcLogLabel>
+            span_id
+        </mdcLogLabel>
+    </appender>
+
+    <root level="debug">
+        <appender-ref ref="Loki"/>
+    </root>
+</configuration>