Кейс использования Mapping Diagnostic Context и @Async

Речь в этой небольшой заметке пойдёт про то, как с пользой можно использовать MDC в Spring-проекте. Поводом написать статью послужила другая недавняя статья на Хабре.

Мы — небольшая команда backend-разработчиков, включая меня — работаем над проектом сервера для мобильных приложений для некой организации. Приложениями пользуются только её сотрудники и значительного highload-а у нас нет. Поэтому для сервера мы выбрали максимально привычный нам стек: Java и Spring Boot на servlet-контейнерах.

Здесь я хотел бы рассказать про свой подход к логированию бизнес-процессов при помощи MDC. Почему именно MDC? Фактически, имеется несколько реплик приложения, которые развёрнуты в Kubernetes, и все логи стекаются в единый агрегатор (Graylog). В конфигурации logback-а добавлен специальный appender, который отправляет все логи по нужному адресу, а также к каждому сообщению добавляет все имеющиеся в MDC поля:

<?xml version="1.0" encoding="UTF-8"?>
<configuration>
    <include resource="org/springframework/boot/logging/logback/defaults.xml"/>
    <include resource="org/springframework/boot/logging/logback/console-appender.xml"/>

    <springProperty name="graylog_environment"
                    scope="context"
                    source="logging.graylog.environment"
                    defaultValue="local"/>
    <springProperty name="graylog_host"
                    scope="context"
                    source="logging.graylog.host"
                    defaultValue="127.0.0.1"/>
    <springProperty name="graylog_port"
                    scope="context"
                    source="logging.graylog.port"
                    defaultValue="12201"/>
    <springProperty name="graylog_microservice"
                    scope="context"
                    source="logging.graylog.microservice"
                    defaultValue=""/>

    <appender name="UDP_GELF"
              class="biz.paluch.logging.gelf.logback.GelfLogbackAppender">
        <host>${graylog_host}</host>
        <port>${graylog_port}</port>

        <version>1.1</version>

        <extractStackTrace>true</extractStackTrace>
        <filterStackTrace>true</filterStackTrace>

        <includeFullMdc>true</includeFullMdc>
        <additionalFields>environment=${graylog_environment},microservice=${graylog_microservice}</additionalFields>
        <additionalFieldTypes>environment=String,microservice=String</additionalFieldTypes>

        <timestampPattern>yyyy-MM-dd HH:mm:ss,SSS</timestampPattern>
        <maximumMessageSize>8192</maximumMessageSize>
    </appender>

    <root level="DEBUG">
        <appender-ref ref="UDP_GELF"/>
        <appender-ref ref="CONSOLE"/>
    </root>

</configuration>

Это очень удобно — можно отправить название окружения (боевое или одно из тестовых), название конкретного микросервиса, Spring Cloud Sleuth добавляет в него данные трассировки (traceId и spanId), а мне захотелось отправлять туда такие поля, которые в последствии помогут искать в Graylog-е причины каких-то сбоев именно мне. Если на вашем проекте используется стек ELK, то скорее всего там всё аналогично, просто мне с ним сталкиваться не приходилось.

Подавляющая часть входящих запросов проходит через кастомный Security-фильтр, который по заголовкам запроса определяет пользователя. Именно в этот момент удобно подложить в MDC то, что хочется. Сперва описываем класс с теми полями, которые хочется подложить:

@UtilityClass
public class MdcKeys {

    /**
     * Значение HTTP-заголовка User-Agent из внешнего запроса.
     */
    public final String MDC_USER_AGENT = "user-agent";

    /**
     * Значение заголовка Authorization из внешнего запроса.
     */
    public final String MDC_USER_TOKEN = "authorization";

    /**
     * Логин пользователя, от имени которого поступил запрос.
     */
    public final String MDC_USER_LOGIN = "login";

    /**
     * URL, на который поступил внешний запрос.
     */
    public final String MDC_API_URL = "apiUrl";

    // ... и некоторые другие ...
}

Если просто устанавливать значения через MDC#put то они не будут удаляться из него после обработки запроса, что чревато проблемами: не все запросы должны быть авторизованы, поэтому часть проскакивает мимо вызова AuthenticationManager-а. Так как потоки, которые обрабатывают запросы, живут в servlet-контейнере и не пересоздаются, старые значения "замусоривают" нам картину. Выход простой — обернуть в try-catch и почистить за собой в блоке finally.

Далее возникает вопрос о том, что делать с методами, помеченными аннотацией @Async. Методы, помеченные ей, выполняются на отдельном пуле потоков, а так как значения в MDC локальны по отношению к текущему потоку, нужно их как-то туда скопировать и в конце метода не забыть удалить. Решение подсмотрено опять же у Spring Security. Находим место в конфигурации, где у нас создаётся пул потоков:

/**
 * TaskExecutor, используемый для запуска асинхронных методов.
 */
@Bean
@Qualifier("taskExecutor")
TaskExecutor taskExecutor() {
    ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();
    // ... код настройки taskExecutor-а опущен ...

    taskExecutor.setTaskDecorator(new AsyncTaskCustomDecorator());
    return taskExecutor;
}

И устанавливаем декоратор, который будет оборачивать все запускаемые асинхронные методы своим кодом:

private static class AsyncTaskCustomDecorator implements TaskDecorator {

    @Override
    @NonNull
    public Runnable decorate(@NonNull Runnable runnable) {
        var runnableWithRestoredMDC = LoggingUtils.decorateMdcCopying(runnable);
        return new DelegatingSecurityContextRunnable(runnableWithRestoredMDC);
    }
}

Фактически, оборачивание происходит дважды: один раз нашим кодом (LoggingUtils#decorateMdcCopying) и один раз для самой Spring Security (иначе мы сломаем передачу контекста в SecurityContextHolder-е). Можно было бы "сделать всю работу" на месте, но мы вынесли код в утильный класс. Смотрим дальше:

@UtilityClass
public class LoggingUtils {

    private final Set<String> COPYABLE_MDC_FIELDS = Set.of(
            MdcKeys.MDC_USER_AGENT,
            MdcKeys.MDC_USER_TOKEN,
            MdcKeys.MDC_USER_LOGIN,
            MdcKeys.MDC_API_URL,
            MdcKeys.MDC_MOBILE_FEATURE);

    /**
     * Декорирует Runnable таким образом, что при его вызове в другом потоке
     * будут восстановлены значения нескольких определённых полей MDC из
     * текущего потока.
     */
    public Runnable decorateMdcCopying(Runnable runnable) {
        // Получаем значения, которые нужно будет восстановить в MDC.
        Map<String, String> mdcMap = getMdcMeaningfulMap();
        return () -> {
            // Восстанавливаем в значения полей MDC из потока-создателя.
            try (var ignored = mdcCloseable(mdcMap)) {
                // Вызываем оригинальный код.
                runnable.run();
            }
        };
    }

    private Map<String, String> getMdcMeaningfulMap() {
        return StreamEx.of(COPYABLE_MDC_FIELDS)
                .mapToEntry(MDC::get)
                .nonNullValues()
                .toMap();
    }
  
    public MdcCloseable mdcCloseable(Map<String, String> values) {
        // Если нечего устанавливать, вернём ничего не закрывающий singleton.
        if (MapUtils.isEmpty(values)) {
            return MdcCloseable.EMPTY;
        }

        // Вариант, когда нам не нужно будет ничего восстанавливать в MDC.
        var mdcMap = MapUtils.emptyIfNull(MDC.getCopyOfContextMap());
        if (MapUtils.isEmpty(mdcMap)) {
            return new MdcCloseable(values, Collections.emptyMap());
        }

        // Значения, которые нужно будет вернуть в MDC
        // (если они там были до нас).
        Map<String, String> original = EntryStream.of(mdcMap)
                .nonNullValues()
                .filterKeys(values::containsKey)
                .filterKeyValue((k, v) -> Objects.equals(v, mdcMap.get(k)))
                .toMap();
        return new MdcCloseable(values, original);
    }

    public MdcCloseable mdcCloseable(String key, String value) {
        return mdcCloseable(Map.of(key, value));
    }
}

И, да, мы написали собственную альтернативу MDC.MDCCloseable:

/**
 * Данный класс отличается от org.slf4j.MDC.MDCCloseable следующим:
 * <ol>
 *     <li>работает со множеством значений,</li>
 *     <li>умеет восстанавливать предыдущие значения.</li>
 * </ol>
 */
@RequiredArgsConstructor(access = AccessLevel.PRIVATE)
public class MdcCloseable implements Closeable {

    public static final MdcCloseable EMPTY = new MdcCloseable(
            Collections.emptySet(),
            Collections.emptyMap());

    private final Set<String> values;
    private final Map<String, String> original;

    MdcCloseable(Map<String, String> values, Map<String, String> original) {
        this(values.keySet(), original);
        values.forEach(MDC::put);
    }

    @Override
    public void close() {
        // Удалили все новые значения при закрытии.
        values.forEach(MDC::remove);
        // Восстановили все предыдущие значения.
        original.forEach(MDC::put);
    }
}

Данный класс можно использовать отдельно в коде приложения для установки каких-то дополнительных полей, которые помогут в дальнейшем искать логи в агрегаторе, например:

// ... какой-то код ...
try (var ignored = LoggingUtils.mdcCloseable(MdcKeys.SOME_EXT_SVC_URL, url) {
    /* Все логи в агрегаторе из этого блока кода будут содержать
       ключ MdcKeys.SOME_EXT_SVC_URL и значение url. */
}
// ... какой-то код ...

Всё вышеописанное может оказаться дичайшим over-engineering-ом.

Я не очень близко знаком с Reactor и WebFlux, поэтому мне кажется, что применить с ними аналогичный подход будет немного сложнее.

Lombok; var, Map.of, Set.of и другие фичи новых версий Java; StreamEx — это всё огонь.

Не бейте строго за первую статью на ресурсе.