Mdc aspect implementation

logging/declarative-logging/declarative-logging-annotation-processor/src/main/java/ru/tinkoff/kora/logging/aspect/mdc/MdcAspect.java

@@ -0,0 +1,147 @@
+package ru.tinkoff.kora.logging.aspect.mdc;
+
+import com.squareup.javapoet.ClassName;
+import com.squareup.javapoet.CodeBlock;
+import ru.tinkoff.kora.annotation.processor.common.CommonClassNames;
+import ru.tinkoff.kora.annotation.processor.common.MethodUtils;
+import ru.tinkoff.kora.annotation.processor.common.ProcessingErrorException;
+import ru.tinkoff.kora.aop.annotation.processor.KoraAspect;
+
+import javax.lang.model.element.AnnotationMirror;
+import javax.lang.model.element.ExecutableElement;
+import javax.lang.model.element.VariableElement;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.Future;
+
+import static ru.tinkoff.kora.annotation.processor.common.AnnotationUtils.findAnnotations;
+import static ru.tinkoff.kora.annotation.processor.common.AnnotationUtils.isAnnotationPresent;
+import static ru.tinkoff.kora.annotation.processor.common.AnnotationUtils.parseAnnotationValueWithoutDefault;
+import static ru.tinkoff.kora.logging.aspect.mdc.MdcAspectClassNames.mdc;
+import static ru.tinkoff.kora.logging.aspect.mdc.MdcAspectClassNames.mdcAnnotation;
+import static ru.tinkoff.kora.logging.aspect.mdc.MdcAspectClassNames.mdcContainerAnnotation;
+
+public class MdcAspect implements KoraAspect {
+
+    private static final String MDC_CONTEXT_VAR_NAME = "__mdcContext";
+
+    @Override
+    public Set<String> getSupportedAnnotationTypes() {
+        return Set.of(mdcAnnotation.canonicalName(), mdcContainerAnnotation.canonicalName());
+    }
+
+    @Override
+    public ApplyResult apply(ExecutableElement executableElement, String superCall, AspectContext aspectContext) {
+        final List<AnnotationMirror> methodAnnotations = findAnnotations(executableElement, mdcAnnotation, mdcContainerAnnotation);
+
+        final List<? extends VariableElement> parametersWithAnnotation = executableElement.getParameters()
+            .stream()
+            .filter(param -> isAnnotationPresent(param, mdcAnnotation))
+            .toList();
+
+        if (methodAnnotations.isEmpty() && parametersWithAnnotation.isEmpty()) {
+            final CodeBlock code = CodeBlock.builder()
+                .add(MethodUtils.isVoid(executableElement) ? "" : "return ")
+                .addStatement(KoraAspect.callSuper(executableElement, superCall))
+                .build();
+            return new ApplyResult.MethodBody(code);
+        }
+        if (MethodUtils.isFuture(executableElement)) {
+            throw new ProcessingErrorException("@Mdc can't be applied for types assignable from " + ClassName.get(Future.class), executableElement);
+        }
+        if (MethodUtils.isMono(executableElement) || MethodUtils.isFlux(executableElement)) {
+            throw new ProcessingErrorException("@Mdc can't be applied for types assignable from " + CommonClassNames.publisher, executableElement);
+        }
+
+        final CodeBlock.Builder currentContextBuilder = CodeBlock.builder();
+        currentContextBuilder.addStatement("var $N = $T.get().values()", MDC_CONTEXT_VAR_NAME, mdc);
+        final CodeBlock.Builder fillMdcBuilder = CodeBlock.builder();
+        final Set<String> methodKeys = fillMdcByMethodAnnotations(methodAnnotations, currentContextBuilder, fillMdcBuilder);
+        final Set<String> parametersKeys = fillMdcByParametersAnnotations(parametersWithAnnotation, currentContextBuilder, fillMdcBuilder);
+        final CodeBlock.Builder clearMdcBuilder = CodeBlock.builder();
+        clearMdc(methodKeys, clearMdcBuilder);
+        clearMdc(parametersKeys, clearMdcBuilder);
+
+        final CodeBlock code = CodeBlock.builder()
+            .add(currentContextBuilder.build())
+            .beginControlFlow("try")
+            .add(fillMdcBuilder.build())
+            .add(MethodUtils.isVoid(executableElement) ? "" : "return ")
+            .addStatement(KoraAspect.callSuper(executableElement, superCall))
+            .nextControlFlow("finally")
+            .add(clearMdcBuilder.build())
+            .endControlFlow()
+            .build();
+
+        return new ApplyResult.MethodBody(code);
+    }
+
+    private static Set<String> fillMdcByMethodAnnotations(List<AnnotationMirror> methodAnnotations, CodeBlock.Builder currentContextBuilder, CodeBlock.Builder fillMdcBuilder) {
+        final Set<String> keys = new HashSet<>();
+        for (AnnotationMirror annotation : methodAnnotations) {
+            final String key = extractStringParameter(annotation, "key")
+                .orElseThrow(() -> new ProcessingErrorException("@Mdc annotation must have 'key' attribute", annotation.getAnnotationType().asElement()));
+            final String value = extractStringParameter(annotation, "value")
+                .orElseThrow(() -> new ProcessingErrorException("@Mdc annotation must have 'value' attribute", annotation.getAnnotationType().asElement()));
+            final Boolean global = parseAnnotationValueWithoutDefault(annotation, "global");
+
+            if (global == null || !global) {
+                keys.add(key);
+                currentContextBuilder.addStatement("var __$N = $N.get($S)", key, MDC_CONTEXT_VAR_NAME, key);
+            }
+            if (value.startsWith("${") && value.endsWith("}")) {
+                fillMdcBuilder.addStatement("$T.put($S, $L)", mdc, key, value.substring(2, value.length() - 1));
+            } else {
+                fillMdcBuilder.addStatement("$T.put($S, $S)", mdc, key, value);
+            }
+        }
+        return keys;
+    }
+
+    private Set<String> fillMdcByParametersAnnotations(List<? extends VariableElement> parametersWithAnnotation, CodeBlock.Builder currentContextBuilder, CodeBlock.Builder fillMdcBuilder) {
+        final Set<String> keys = new HashSet<>();
+        for (VariableElement parameter : parametersWithAnnotation) {
+            final String parameterName = parameter.getSimpleName().toString();
+            final AnnotationMirror firstAnnotation = findAnnotations(parameter, mdcAnnotation, mdcContainerAnnotation)
+                .get(0);
+
+            final String key = extractStringParameter(firstAnnotation, "key")
+                .or(() -> extractStringParameter(firstAnnotation, "value"))
+                .orElse(parameterName);
+
+            final Boolean global = parseAnnotationValueWithoutDefault(firstAnnotation, "global");
+
+            fillMdcBuilder.addStatement(
+                "$T.put($S, $N)",
+                mdc,
+                key,
+                parameterName
+            );
+
+            if (global == null || !global) {
+                keys.add(key);
+                currentContextBuilder.addStatement("var __$N = $N.get($S)", key, MDC_CONTEXT_VAR_NAME, key);
+            }
+        }
+        return keys;
+    }
+
+    private static Optional<String> extractStringParameter(AnnotationMirror annotation, String name) {
+        final String value = parseAnnotationValueWithoutDefault(annotation, name);
+        return Optional.ofNullable(value)
+            .filter(s -> !s.isBlank());
+    }
+
+    private static void clearMdc(Set<String> keys, CodeBlock.Builder b) {
+        for (String key : keys) {
+            b.beginControlFlow("if (__$N != null)", key)
+                .addStatement("$T.put($S, __$N)", mdc, key, key)
+                .endControlFlow()
+                .beginControlFlow("else")
+                .addStatement("$T.remove($S)", mdc, key)
+                .endControlFlow();
+        }
+    }
+}

logging/declarative-logging/declarative-logging-annotation-processor/src/main/java/ru/tinkoff/kora/logging/aspect/mdc/MdcAspectClassNames.java

@@ -0,0 +1,9 @@
+package ru.tinkoff.kora.logging.aspect.mdc;
+
+import com.squareup.javapoet.ClassName;
+
+public class MdcAspectClassNames {
+    public static final ClassName mdcAnnotation = ClassName.get("ru.tinkoff.kora.logging.common.annotation", "Mdc");
+    public static final ClassName mdcContainerAnnotation = ClassName.get("ru.tinkoff.kora.logging.common.annotation", "Mdc", "MdcContainer");
+    public static final ClassName mdc = ClassName.get("ru.tinkoff.kora.logging.common", "MDC");
+}

logging/declarative-logging/declarative-logging-annotation-processor/src/main/java/ru/tinkoff/kora/logging/aspect/mdc/MdcAspectFactory.java

@@ -0,0 +1,15 @@
+package ru.tinkoff.kora.logging.aspect.mdc;
+
+import ru.tinkoff.kora.aop.annotation.processor.KoraAspect;
+import ru.tinkoff.kora.aop.annotation.processor.KoraAspectFactory;
+
+import javax.annotation.processing.ProcessingEnvironment;
+import java.util.Optional;
+
+public class MdcAspectFactory implements KoraAspectFactory {
+
+    @Override
+    public Optional<KoraAspect> create(ProcessingEnvironment processingEnvironment) {
+        return Optional.of(new MdcAspect());
+    }
+}

logging/declarative-logging/declarative-logging-annotation-processor/src/main/resources/META-INF/services/ru.tinkoff.kora.aop.annotation.processor.KoraAspectFactory

@@ -1 +1,2 @@
 ru.tinkoff.kora.logging.aspect.LogAspectFactory
+ru.tinkoff.kora.logging.aspect.mdc.MdcAspectFactory

logging/declarative-logging/declarative-logging-annotation-processor/src/test/java/ru/tinkoff/kora/logging/aspect/mdc/AbstractMdcAspectTest.java

@@ -0,0 +1,26 @@
+package ru.tinkoff.kora.logging.aspect.mdc;
+
+import org.intellij.lang.annotations.Language;
+import ru.tinkoff.kora.annotation.processor.common.AbstractAnnotationProcessorTest;
+
+import java.util.List;
+
+public abstract class AbstractMdcAspectTest extends AbstractAnnotationProcessorTest {
+
+    @Override
+    protected String commonImports() {
+        return super.commonImports() + """
+            import ru.tinkoff.kora.logging.common.annotation.Mdc;
+            import ru.tinkoff.kora.logging.common.MDC;
+            import ru.tinkoff.kora.logging.aspect.mdc.MDCContextHolder;
+            import java.util.concurrent.CompletionStage;
+            import java.util.concurrent.CompletableFuture;
+            import reactor.core.publisher.Mono;
+            import reactor.core.publisher.Flux;
+            """;
+    }
+
+    protected static List<String> sources(@Language("java") String... sources) {
+        return List.of(sources);
+    }
+}

logging/declarative-logging/declarative-logging-annotation-processor/src/test/java/ru/tinkoff/kora/logging/aspect/mdc/MDCContextHolder.java

@@ -0,0 +1,21 @@
+package ru.tinkoff.kora.logging.aspect.mdc;
+
+import ru.tinkoff.kora.logging.common.arg.StructuredArgumentWriter;
+
+import java.util.Collections;
+import java.util.Map;
+
+public class MDCContextHolder {
+
+    private Map<String, StructuredArgumentWriter> mdcContext;
+
+    public Map<String, StructuredArgumentWriter> get() {
+        return mdcContext;
+    }
+
+    public void set(Map<String, StructuredArgumentWriter> mdcContext) {
+        this.mdcContext = mdcContext == null
+            ? Collections.emptyMap()
+            : Map.copyOf(mdcContext);
+    }
+}