JsonLayout.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF 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 net.morimekta.tiny.server.logback;

import ch.qos.logback.classic.spi.ILoggingEvent;
import ch.qos.logback.classic.spi.IThrowableProxy;
import ch.qos.logback.core.CoreConstants;
import ch.qos.logback.core.LayoutBase;
import com.fasterxml.jackson.databind.ObjectMapper;

import java.io.IOException;
import java.time.Instant;
import java.time.ZoneId;
import java.time.format.DateTimeFormatter;
import java.util.Locale;
import java.util.Map;
import java.util.Optional;
import java.util.TreeMap;
import java.util.TreeSet;

/**
 * {@link JsonLayout} for formatting logging messages to take use of stackdriver etc.
 * <p>
 * Usage is to configure an appender with an encoder using this layout.
 * See logback.xml example below:
 * <pre>{@code
 *   <configuration>
 *     <appender name="JSON-OUT" class="ch.qos.logback.core.ConsoleAppender">
 *         <encoder class="ch.qos.logback.core.encoder.LayoutWrappingEncoder">
 *             <layout class="net.morimekta.tiny.server.logback.JsonLayout">
 *                 <zoneId>UTC</zoneId>
 *                 <includeMdc>true</includeMdc>
 *                 <stackTraceFormat>full</stackTraceFormat>
 *                 <stackTraceIncludeShort>true</stackTraceIncludeShort>
 *                 <stackTraceFilter>
 *                     com.intellij,
 *                     com.sun.net.httpserver,
 *                     java.lang.reflect,
 *                     java.util.ArrayList.forEach,
 *                     java.util.concurrent,
 *                     java.util.stream,
 *                     jdk.httpserver,
 *                     jdk.internal.reflect,
 *                     org.apache.maven.surefire,
 *                     org.junit,
 *                     sun.net.httpserver,
 *                 </stackTraceFilter>
 *             </layout>
 *         </encoder>
 *     </appender>
 *     <root level="INFO">
 *         <appender-ref ref="JSON-OUT"/>
 *     </root>
 *     <logger name="ch.qos.logback" level="WARN"/>
 *     <logger name="net.morimekta.tiny.server" level="DEBUG"/>
 *   </configuration>
 * }</pre>
 */
public class JsonLayout extends LayoutBase<ILoggingEvent> {
    private enum Format {
        FULL,
        NORMAL,
        SHORT
    }

    private final ObjectMapper    mapper           = new ObjectMapper();
    private final TreeSet<String> stackTraceFilter = new TreeSet<>();

    private boolean includeMdc             = true;
    private Format  stackTraceFormat       = Format.NORMAL;
    private boolean stackTraceIncludeShort = false;
    private ZoneId  zoneId                 = ZoneId.systemDefault();

    public void setZoneId(String zoneId) {
        try {
            this.zoneId = ZoneId.of(zoneId);
        } catch (Exception e) {
            this.zoneId = ZoneId.of("UTC");
        }
    }

    public void setStackTraceIncludeShort(String bool) {
        stackTraceIncludeShort = Boolean.parseBoolean(bool);
    }

    public void setIncludeMdc(String bool) {
        includeMdc = Boolean.parseBoolean(bool);
    }

    public void setStackTraceFormat(String format) {
        switch (format.strip().toLowerCase(Locale.US)) {
            case "full":
                stackTraceFormat = Format.FULL;
                break;
            case "short":
                stackTraceFormat = Format.SHORT;
                break;
            default:
                stackTraceFormat = Format.NORMAL;
                break;
        }
    }

    public void setStackTraceFilter(String filter) {
        stackTraceFilter.clear();
        for (String s : filter.split("\\s*,\\s*")) {
            stackTraceFilter.add(s.strip());
        }
    }

    @Override
    public String doLayout(ILoggingEvent event) {
        if (!isStarted()) {
            return CoreConstants.EMPTY_STRING;
        }

        Instant timestamp = Optional
                .ofNullable(event.getInstant())
                .orElseGet(() -> Instant.ofEpochMilli(event.getTimeStamp()));

        Map<String, Object> entry = new TreeMap<>();
        if (includeMdc) {
            entry.putAll(event.getMDCPropertyMap());
        }
        entry.put("@timestamp", formatTimestamp(timestamp));
        entry.put("message", event.getFormattedMessage());
        entry.put("level", event.getLevel().levelStr);
        entry.put("level_value", event.getLevel().levelInt);
        entry.put("logger_name", event.getLoggerName());
        entry.put("thread_name", event.getThreadName());
        if (event.getThrowableProxy() != null) {
            entry.put("stack_trace", formatStackTrace(event.getThrowableProxy()));
            if (stackTraceIncludeShort && stackTraceFormat != Format.SHORT) {
                entry.put("stack_trace_short", formatStackTraceShort(event.getThrowableProxy()));
            }
        }
        try {
            return mapper.writeValueAsString(entry) + "\n";
        } catch (IOException e) {
            // should be impossible.
            return CoreConstants.EMPTY_STRING;
        }
    }

    @Override
    public String getContentType() {
        return "application/json";
    }

    // Use this format to ensure always fully written out with millis.
    private static final DateTimeFormatter TIMESTAMP_FORMATTER = DateTimeFormatter.ofPattern(
            "yyyy-MM-dd'T'HH:mm:ss.SSSxxx");

    private String formatTimestamp(Instant instant) {
        return TIMESTAMP_FORMATTER.format(instant.atZone(zoneId));
    }

    private String formatStackTrace(IThrowableProxy throwableProxy) {
        var builder = new StringBuilder();
        formatStackTrace(stackTraceFormat, "", throwableProxy, builder);
        return builder.toString();
    }

    private String formatStackTraceShort(IThrowableProxy throwableProxy) {
        var builder = new StringBuilder();
        formatStackTrace(Format.SHORT, "", throwableProxy, builder);
        return builder.toString();
    }

    private boolean skipStackTraceElement(int index, StackTraceElement element) {
        if (index == 0) {
            // never skip the location the exception was thrown, even if filtered.
            return false;
        }
        var id = element.getClassName() + "." + element.getMethodName();
        if (stackTraceFilter.contains(id)) {
            return true;
        }
        var lower = stackTraceFilter.lower(id);
        return lower != null && id.startsWith(lower);
    }

    private void formatStackTrace(Format stackTraceFormat,
                                  String prefix,
                                  IThrowableProxy throwable,
                                  StringBuilder builder) {
        if (throwable.getCause() != null) {
            var cause = throwable.getCause();
            formatStackTrace(stackTraceFormat, prefix, cause, builder);
            builder.append(prefix).append(CoreConstants.WRAPPED_BY);
        }
        builder.append(throwable.getClassName())
               .append(": ")
               .append(throwable.getMessage())
               .append("\n");
        var trace = throwable.getStackTraceElementProxyArray();
        if (trace.length > 0) {
            int skippedElements = 0;
            // do not print common frames.
            var printNumElements =
                    stackTraceFormat != Format.SHORT
                    ? Math.max(1, trace.length - throwable.getCommonFrames())
                    : 1;
            for (int i = 0; i < printNumElements; ++i) {
                var element = trace[i].getStackTraceElement();
                if (skipStackTraceElement(i, element)) {
                    ++skippedElements;
                    continue;
                }
                if (skippedElements > 0) {
                    builder.append(prefix)
                           .append("\t[")
                           .append(skippedElements)
                           .append(" skipped]\n");
                    skippedElements = 0;
                }
                builder.append(prefix)
                       .append("\tat ");
                if (stackTraceFormat == Format.FULL && element.getModuleName() != null) {
                    builder.append(element.getModuleName());
                    if (element.getModuleVersion() != null) {
                        builder.append("@").append(element.getModuleVersion());
                    }
                    builder.append("/");
                }
                builder.append(element.getClassName())
                       .append(".")
                       .append(element.getMethodName())
                       .append("(");
                if (element.isNativeMethod()) {
                    builder.append("Native Method");
                } else {
                    builder.append(element.getFileName())
                           .append(":")
                           .append(element.getLineNumber());
                }
                builder.append(")\n");
            }
            if (skippedElements > 0) {
                builder.append(prefix)
                       .append("\t[")
                       .append(skippedElements)
                       .append(" skipped]\n");
            }
            if (stackTraceFormat != Format.SHORT) {
                if (throwable.getCommonFrames() > 0) {
                    builder.append(prefix)
                           .append("... ")
                           .append(throwable.getCommonFrames())
                           .append(" common frames omitted\n");
                }
            }
        }
        for (var suppressed : throwable.getSuppressed()) {
            builder.append(prefix)
                   .append("\t")
                   .append(CoreConstants.SUPPRESSED);
            formatStackTrace(stackTraceFormat, prefix + "\t", suppressed, builder);
        }
    }
}