diff --git a/logging/declarative-logging/declarative-logging-annotation-processor/src/main/java/ru/tinkoff/kora/logging/aspect/mdc/MdcAspect.java b/logging/declarative-logging/declarative-logging-annotation-processor/src/main/java/ru/tinkoff/kora/logging/aspect/mdc/MdcAspect.java new file mode 100644 index 000000000..15093d4ed --- /dev/null +++ b/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 getSupportedAnnotationTypes() { + return Set.of(mdcAnnotation.canonicalName(), mdcContainerAnnotation.canonicalName()); + } + + @Override + public ApplyResult apply(ExecutableElement executableElement, String superCall, AspectContext aspectContext) { + final List methodAnnotations = findAnnotations(executableElement, mdcAnnotation, mdcContainerAnnotation); + + final List 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 methodKeys = fillMdcByMethodAnnotations(methodAnnotations, currentContextBuilder, fillMdcBuilder); + final Set 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 fillMdcByMethodAnnotations(List methodAnnotations, CodeBlock.Builder currentContextBuilder, CodeBlock.Builder fillMdcBuilder) { + final Set 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 fillMdcByParametersAnnotations(List parametersWithAnnotation, CodeBlock.Builder currentContextBuilder, CodeBlock.Builder fillMdcBuilder) { + final Set 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 extractStringParameter(AnnotationMirror annotation, String name) { + final String value = parseAnnotationValueWithoutDefault(annotation, name); + return Optional.ofNullable(value) + .filter(s -> !s.isBlank()); + } + + private static void clearMdc(Set 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(); + } + } +} diff --git a/logging/declarative-logging/declarative-logging-annotation-processor/src/main/java/ru/tinkoff/kora/logging/aspect/mdc/MdcAspectClassNames.java b/logging/declarative-logging/declarative-logging-annotation-processor/src/main/java/ru/tinkoff/kora/logging/aspect/mdc/MdcAspectClassNames.java new file mode 100644 index 000000000..a449d8182 --- /dev/null +++ b/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"); +} diff --git a/logging/declarative-logging/declarative-logging-annotation-processor/src/main/java/ru/tinkoff/kora/logging/aspect/mdc/MdcAspectFactory.java b/logging/declarative-logging/declarative-logging-annotation-processor/src/main/java/ru/tinkoff/kora/logging/aspect/mdc/MdcAspectFactory.java new file mode 100644 index 000000000..eafbfbe15 --- /dev/null +++ b/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 create(ProcessingEnvironment processingEnvironment) { + return Optional.of(new MdcAspect()); + } +} diff --git a/logging/declarative-logging/declarative-logging-annotation-processor/src/main/resources/META-INF/services/ru.tinkoff.kora.aop.annotation.processor.KoraAspectFactory b/logging/declarative-logging/declarative-logging-annotation-processor/src/main/resources/META-INF/services/ru.tinkoff.kora.aop.annotation.processor.KoraAspectFactory index 26a2325ee..dc60c4374 100644 --- a/logging/declarative-logging/declarative-logging-annotation-processor/src/main/resources/META-INF/services/ru.tinkoff.kora.aop.annotation.processor.KoraAspectFactory +++ b/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 diff --git a/logging/declarative-logging/declarative-logging-annotation-processor/src/test/java/ru/tinkoff/kora/logging/aspect/mdc/AbstractMdcAspectTest.java b/logging/declarative-logging/declarative-logging-annotation-processor/src/test/java/ru/tinkoff/kora/logging/aspect/mdc/AbstractMdcAspectTest.java new file mode 100644 index 000000000..f6991d030 --- /dev/null +++ b/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 sources(@Language("java") String... sources) { + return List.of(sources); + } +} diff --git a/logging/declarative-logging/declarative-logging-annotation-processor/src/test/java/ru/tinkoff/kora/logging/aspect/mdc/MDCContextHolder.java b/logging/declarative-logging/declarative-logging-annotation-processor/src/test/java/ru/tinkoff/kora/logging/aspect/mdc/MDCContextHolder.java new file mode 100644 index 000000000..154e03c52 --- /dev/null +++ b/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 mdcContext; + + public Map get() { + return mdcContext; + } + + public void set(Map mdcContext) { + this.mdcContext = mdcContext == null + ? Collections.emptyMap() + : Map.copyOf(mdcContext); + } +} diff --git a/logging/declarative-logging/declarative-logging-annotation-processor/src/test/java/ru/tinkoff/kora/logging/aspect/mdc/MdcAspectTest.java b/logging/declarative-logging/declarative-logging-annotation-processor/src/test/java/ru/tinkoff/kora/logging/aspect/mdc/MdcAspectTest.java new file mode 100644 index 000000000..c822e3ed6 --- /dev/null +++ b/logging/declarative-logging/declarative-logging-annotation-processor/src/test/java/ru/tinkoff/kora/logging/aspect/mdc/MdcAspectTest.java @@ -0,0 +1,207 @@ +package ru.tinkoff.kora.logging.aspect.mdc; + +import org.intellij.lang.annotations.Language; +import org.jetbrains.annotations.NotNull; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import ru.tinkoff.kora.annotation.processor.common.CompileResult; +import ru.tinkoff.kora.aop.annotation.processor.AopAnnotationProcessor; +import ru.tinkoff.kora.logging.common.MDC; +import ru.tinkoff.kora.logging.common.arg.StructuredArgumentWriter; + +import java.lang.reflect.InvocationTargetException; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.stream.Stream; + +import static java.util.Collections.emptyMap; +import static java.util.stream.Collectors.toMap; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; + +class MdcAspectTest extends AbstractMdcAspectTest { + + private static final MDCContextHolder CONTEXT_HOLDER = new MDCContextHolder(); + + @ParameterizedTest + @MethodSource("provideTestCases") + void testMdc(@Language("java") String source) throws Exception { + var aopProxy = compile(List.of(new AopAnnotationProcessor()), source); + + invokeMethod(aopProxy); + + final Map context = extractMdcContextFromHolder(); + assertEquals(Map.of("key", "\"value\"", "key1", "\"value2\"", "123", "\"test\""), context); + assertEquals(emptyMap(), currentMdcContext()); + } + + @Test + void testMdcWithCode() throws Exception { + var aopProxy = compile( + List.of(new AopAnnotationProcessor()), + """ + public class TestMdc { + + private final MDCContextHolder mdcContextHolder; + + public TestMdc(MDCContextHolder mdcContextHolder) { + this.mdcContextHolder = mdcContextHolder; + } + + @Mdc(key = "key", value = "${java.util.UUID.randomUUID().toString()}") + public Integer test(String s) { + mdcContextHolder.set(MDC.get().values()); + return null; + } + } + """ + ); + + invokeMethod(aopProxy); + + final Map context = extractMdcContextFromHolder(); + final String value = context.get("key"); + assertNotNull(value); + assertDoesNotThrow(() -> UUID.fromString(value.substring(1, value.length() - 1))); + assertEquals(emptyMap(), currentMdcContext()); + } + + @Test + void testGlobalMdc() throws Exception { + var aopProxy = compile( + List.of(new AopAnnotationProcessor()), + """ + public class TestMdc { + + private final MDCContextHolder mdcContextHolder; + + public TestMdc(MDCContextHolder mdcContextHolder) { + this.mdcContextHolder = mdcContextHolder; + } + + @Mdc(key = "key", value = "value", global = true) + @Mdc(key = "key1", value = "value2") + public Integer test(@Mdc(key = "123", global = true) String s) { + mdcContextHolder.set(MDC.get().values()); + return null; + } + } + """ + ); + + invokeMethod(aopProxy); + + final Map context = extractMdcContextFromHolder(); + assertEquals(Map.of("key", "\"value\"", "key1", "\"value2\"", "123", "\"test\""), context); + assertEquals(Map.of("key", "\"value\"", "123", "\"test\""), currentMdcContext()); + clearCurrentContext(); + } + + @Test + void testMdcWithNotEmptyContext() throws Exception { + var aopProxy = compile( + List.of(new AopAnnotationProcessor()), + """ + public class TestMdc { + + private final MDCContextHolder mdcContextHolder; + + public TestMdc(MDCContextHolder mdcContextHolder) { + this.mdcContextHolder = mdcContextHolder; + } + + @Mdc(key = "key", value = "value") + @Mdc(key = "key1", value = "value2") + public Integer test(@Mdc(key = "123") String s) { + mdcContextHolder.set(MDC.get().values()); + return null; + } + } + """ + ); + + MDC.put("key", "special-value"); + MDC.put("123", "special-123"); + invokeMethod(aopProxy); + + final Map context = extractMdcContextFromHolder(); + assertEquals(Map.of("key", "\"value\"", "key1", "\"value2\"", "123", "\"test\""), context); + assertEquals(Map.of("key", "\"special-value\"", "123", "\"special-123\""), currentMdcContext()); + clearCurrentContext(); + } + + private static List provideTestCases() { + return sources( + """ + public class TestMdc { + + private final MDCContextHolder mdcContextHolder; + + public TestMdc(MDCContextHolder mdcContextHolder) { + this.mdcContextHolder = mdcContextHolder; + } + + @Mdc(key = "key", value = "value") + @Mdc(key = "key1", value = "value2") + public Integer test(@Mdc(key = "123") String s) { + mdcContextHolder.set(MDC.get().values()); + return null; + } + } + """, + """ + public class TestMdc { + + private final MDCContextHolder mdcContextHolder; + + public TestMdc(MDCContextHolder mdcContextHolder) { + this.mdcContextHolder = mdcContextHolder; + } + + @Mdc(key = "key", value = "value") + @Mdc(key = "key1", value = "value2") + public void test(@Mdc("123") String s) { + mdcContextHolder.set(MDC.get().values()); + } + } + """ + ); + } + + private static void invokeMethod(CompileResult aopProxy) throws Exception { + aopProxy.assertSuccess(); + + var generatedClass = aopProxy.loadClass("$TestMdc__AopProxy"); + var constructor = generatedClass.getConstructors()[0]; + var params = new Object[constructor.getParameterCount()]; + params[0] = CONTEXT_HOLDER; + final TestObject testObject = new TestObject(generatedClass, constructor.newInstance(params)); + + testObject.invoke("test", "test"); + } + + private static Map extractMdcContextFromHolder() { + return toMdcContext(CONTEXT_HOLDER.get()); + } + + private static Map currentMdcContext() { + return toMdcContext(MDC.get().values()); + } + + private static Map toMdcContext(Map values) { + return values.entrySet() + .stream() + .collect(toMap( + Map.Entry::getKey, + entry -> entry.getValue().writeToString() + )); + } + + private static void clearCurrentContext() { + MDC.get().values().keySet().forEach(MDC::remove); + } +} diff --git a/logging/declarative-logging/declarative-logging-annotation-processor/src/test/java/ru/tinkoff/kora/logging/aspect/mdc/MdcAspectUnsupportedTypesTest.java b/logging/declarative-logging/declarative-logging-annotation-processor/src/test/java/ru/tinkoff/kora/logging/aspect/mdc/MdcAspectUnsupportedTypesTest.java new file mode 100644 index 000000000..83a45ab01 --- /dev/null +++ b/logging/declarative-logging/declarative-logging-annotation-processor/src/test/java/ru/tinkoff/kora/logging/aspect/mdc/MdcAspectUnsupportedTypesTest.java @@ -0,0 +1,126 @@ +package ru.tinkoff.kora.logging.aspect.mdc; + +import org.intellij.lang.annotations.Language; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; +import ru.tinkoff.kora.annotation.processor.common.CompileResult; +import ru.tinkoff.kora.aop.annotation.processor.AopAnnotationProcessor; + +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class MdcAspectUnsupportedTypesTest extends AbstractMdcAspectTest { + + @ParameterizedTest + @MethodSource({"sourcesWithMdcAndFuture", "sourcesWithMdcAndMono", "sourcesWithMdcAndFlux"}) + void testMdc(@Language("java") String source) { + final CompileResult compileResult = compile( + List.of(new AopAnnotationProcessor()), + source + ); + + assertTrue(compileResult.isFailed()); + } + + private static List sourcesWithMdcAndFuture() { + return sources( + """ + public class TestMdc { + + @Mdc(key = "key", value = "value", global = true) + @Mdc(key = "key1", value = "value2") + public CompletionStage test(@Mdc(key = "123", value = "value3") String s) { + return CompletableFuture.completedFuture(1); + } + } + """, + """ + public class TestMdc { + + @Mdc(key = "key1", value = "value2") + public CompletionStage test(String s) { + return CompletableFuture.completedFuture(1); + } + } + """, + """ + public class TestMdc { + + public CompletionStage test(@Mdc(key = "123") String s) { + return CompletableFuture.completedFuture(1); + } + } + """ + ); + } + + private static List sourcesWithMdcAndMono() { + return sources( + """ + public class TestMdc { + + @Mdc(key = "key", value = "value", global = true) + @Mdc(key = "key1", value = "value2") + public Mono test(@Mdc(key = "123", value = "value3") String s) { + return Mono.just(1); + } + } + """, + """ + public class TestMdc { + + @Mdc(key = "key1", value = "value2") + public Mono test(String s) { + return Mono.just(1); + } + } + """, + """ + public class TestMdc { + + public Mono test(@Mdc(key = "123") String s) { + return Mono.just(1); + } + } + """ + ); + } + + private static List sourcesWithMdcAndFlux() { + return sources( + """ + public class TestMdc { + + @Mdc(key = "key", value = "value", global = true) + @Mdc(key = "key1", value = "value2") + public Flux test(@Mdc(key = "123", value = "value3") String s) { + return Mono.just(1) + .flux(); + } + } + """, + """ + public class TestMdc { + + @Mdc(key = "key1", value = "value2") + public Flux test(String s) { + return Mono.just(1) + .flux(); + } + } + """, + """ + public class TestMdc { + + public Flux test(@Mdc(key = "123") String s) { + return Mono.just(1) + .flux(); + } + } + """ + ); + } +} diff --git a/logging/declarative-logging/declarative-logging-symbol-processor/src/main/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/MdcKoraAspect.kt b/logging/declarative-logging/declarative-logging-symbol-processor/src/main/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/MdcKoraAspect.kt new file mode 100644 index 000000000..dbfc11801 --- /dev/null +++ b/logging/declarative-logging/declarative-logging-symbol-processor/src/main/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/MdcKoraAspect.kt @@ -0,0 +1,155 @@ +package ru.tinkoff.kora.logging.symbol.processor.aop.mdc + +import com.google.devtools.ksp.symbol.KSAnnotation +import com.google.devtools.ksp.symbol.KSFunctionDeclaration +import com.google.devtools.ksp.symbol.KSValueParameter +import com.squareup.kotlinpoet.ClassName +import com.squareup.kotlinpoet.CodeBlock +import ru.tinkoff.kora.aop.symbol.processor.KoraAspect +import ru.tinkoff.kora.ksp.common.AnnotationUtils.findValue +import ru.tinkoff.kora.ksp.common.AnnotationUtils.isAnnotationPresent +import ru.tinkoff.kora.ksp.common.CommonClassNames +import ru.tinkoff.kora.ksp.common.FunctionUtils.isCompletionStage +import ru.tinkoff.kora.ksp.common.FunctionUtils.isFlux +import ru.tinkoff.kora.ksp.common.FunctionUtils.isFuture +import ru.tinkoff.kora.ksp.common.FunctionUtils.isMono +import ru.tinkoff.kora.ksp.common.FunctionUtils.isSuspend +import ru.tinkoff.kora.ksp.common.FunctionUtils.isVoid +import ru.tinkoff.kora.ksp.common.KspCommonUtils.findRepeatableAnnotation +import ru.tinkoff.kora.ksp.common.exception.ProcessingErrorException +import java.util.concurrent.CompletionStage + +class MdcKoraAspect : KoraAspect { + + companion object { + const val MDC_CONTEXT_VAL_NAME = "__mdcContext" + val mdc = ClassName("ru.tinkoff.kora.logging.common", "MDC") + val mdcAnnotation = ClassName("ru.tinkoff.kora.logging.common.annotation", "Mdc") + val mdcContainerAnnotation = mdcAnnotation.nestedClass("MdcContainer") + } + + override fun getSupportedAnnotationTypes(): Set = setOf(mdcAnnotation.canonicalName, mdcContainerAnnotation.canonicalName) + + override fun apply( + ksFunction: KSFunctionDeclaration, + superCall: String, + aspectContext: KoraAspect.AspectContext + ): KoraAspect.ApplyResult { + val annotations = ksFunction.findRepeatableAnnotation(mdcAnnotation, mdcContainerAnnotation) + .toList() + + val parametersWithAnnotation = ksFunction.parameters + .filter { it.isAnnotationPresent(mdcAnnotation) } + + if (annotations.isEmpty() && parametersWithAnnotation.isEmpty()) { + return KoraAspect.ApplyResult.MethodBody(ksFunction.superCall(superCall)) + } + + if (ksFunction.isFuture()) { + throw ProcessingErrorException("@Mdc can't be applied for types assignable from ${CommonClassNames.future}", ksFunction) + } else if (ksFunction.isCompletionStage()) { + throw ProcessingErrorException("@Mdc can't be applied for types assignable from ${CompletionStage::class.java}", ksFunction) + } else if (ksFunction.isMono() || ksFunction.isFlux()) { + throw ProcessingErrorException("@Mdc can't be applied for types assignable from ${CommonClassNames.publisher}", ksFunction) + } + + val currentContextBuilder = CodeBlock.builder() + currentContextBuilder.addStatement("val %N = %T.get().values()", MDC_CONTEXT_VAL_NAME, mdc) + val fillMdcBuilder = CodeBlock.builder() + val methodKeys = fillMdcByMethodAnnotations(annotations, currentContextBuilder, fillMdcBuilder, !ksFunction.isSuspend()) + val parameterKeys = fillMdcByParametersAnnotations(parametersWithAnnotation, currentContextBuilder, fillMdcBuilder, !ksFunction.isSuspend()) + val clearMdcBuilder = CodeBlock.builder() + clearMdc(methodKeys, clearMdcBuilder) + clearMdc(parameterKeys, clearMdcBuilder) + + return CodeBlock.builder() + .add(currentContextBuilder.build()) + .add(if (ksFunction.isVoid()) "" else "return ") + .beginControlFlow("try") + .add(fillMdcBuilder.build()) + .addStatement("%L", ksFunction.superCall(superCall)) + .endControlFlow() + .beginControlFlow("finally") + .add(clearMdcBuilder.build()) + .endControlFlow() + .build() + .let { KoraAspect.ApplyResult.MethodBody(it) } + } + + private fun fillMdcByMethodAnnotations( + annotations: List, + currentContextBuilder: CodeBlock.Builder, + fillMdcBuilder: CodeBlock.Builder, + globalIsSupported: Boolean + ): Set { + val keys: MutableSet = HashSet() + for (annotation in annotations) { + val key: String = annotation.findValue("key") + .orEmpty() + .ifBlank { throw ProcessingErrorException("@Mdc annotation must have 'key' attribute", annotation.annotationType) } + val value: String = annotation.findValue("value") + .orEmpty() + .ifBlank { throw ProcessingErrorException("@Mdc annotation must have 'value' attribute", annotation.annotationType) } + val global = annotation.findValue("global") ?: false + + if (!global) { + keys.add(key) + currentContextBuilder.addStatement("val __%L = %N[%S]", key, MDC_CONTEXT_VAL_NAME, key) + } else if (!globalIsSupported) { + throw ProcessingErrorException("@Mdc annotation with 'global' attribute is not supported for this function", annotation.annotationType) + } + fillMdcBuilder.addStatement( + "%T.put(%S, %S)", + mdc, + key, + value + ) + } + return keys + } + + private fun fillMdcByParametersAnnotations( + parametersWithAnnotation: List, + currentContextBuilder: CodeBlock.Builder, + fillMdcBuilder: CodeBlock.Builder, + globalIsSupported: Boolean + ): Set { + val keys: MutableSet = HashSet() + for (parameter in parametersWithAnnotation) { + val parameterName = parameter.name?.asString() + val annotation = parameter.findRepeatableAnnotation(mdcAnnotation, mdcContainerAnnotation) + .first() + val key: String = annotation.findValue("key") + ?.ifBlank { annotation.findValue("value") } + ?.ifBlank { parameterName } + ?: throw ProcessingErrorException("@Mdc annotation must have key or value or parameter name", parameter) + + val global = annotation.findValue("global") ?: false + + fillMdcBuilder.addStatement( + "%T.put(%S, %N)", + mdc, + key, + parameterName + ) + + if (!global) { + keys.add(key) + currentContextBuilder.addStatement("val __%L = %N[%S]", key, MDC_CONTEXT_VAL_NAME, key) + } else if (!globalIsSupported) { + throw ProcessingErrorException("@Mdc annotation with 'global' attribute is not supported for this function", annotation.annotationType) + } + } + + return keys + } + + private fun clearMdc(keys: Set, b: CodeBlock.Builder) = keys.forEach { + b.beginControlFlow("if (__%L != null)", it) + .addStatement("%T.put(%S, __%L)", mdc, it, it) + .endControlFlow() + .beginControlFlow("else") + .addStatement("%T.remove(%S)", mdc, it) + .endControlFlow() + } +} diff --git a/logging/declarative-logging/declarative-logging-symbol-processor/src/main/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/MdcKoraAspectFactory.kt b/logging/declarative-logging/declarative-logging-symbol-processor/src/main/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/MdcKoraAspectFactory.kt new file mode 100644 index 000000000..3932cd1f4 --- /dev/null +++ b/logging/declarative-logging/declarative-logging-symbol-processor/src/main/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/MdcKoraAspectFactory.kt @@ -0,0 +1,9 @@ +package ru.tinkoff.kora.logging.symbol.processor.aop.mdc + +import com.google.devtools.ksp.processing.Resolver +import ru.tinkoff.kora.aop.symbol.processor.KoraAspect +import ru.tinkoff.kora.aop.symbol.processor.KoraAspectFactory + +class MdcKoraAspectFactory : KoraAspectFactory { + override fun create(resolver: Resolver): KoraAspect = MdcKoraAspect() +} diff --git a/logging/declarative-logging/declarative-logging-symbol-processor/src/main/resources/META-INF/services/ru.tinkoff.kora.aop.symbol.processor.KoraAspectFactory b/logging/declarative-logging/declarative-logging-symbol-processor/src/main/resources/META-INF/services/ru.tinkoff.kora.aop.symbol.processor.KoraAspectFactory index 494637110..6eb4cdb8d 100644 --- a/logging/declarative-logging/declarative-logging-symbol-processor/src/main/resources/META-INF/services/ru.tinkoff.kora.aop.symbol.processor.KoraAspectFactory +++ b/logging/declarative-logging/declarative-logging-symbol-processor/src/main/resources/META-INF/services/ru.tinkoff.kora.aop.symbol.processor.KoraAspectFactory @@ -1 +1,2 @@ ru.tinkoff.kora.logging.symbol.processor.aop.LogKoraAspectFactory +ru.tinkoff.kora.logging.symbol.processor.aop.mdc.MdcKoraAspectFactory diff --git a/logging/declarative-logging/declarative-logging-symbol-processor/src/test/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/AbstractMdcAspectTest.kt b/logging/declarative-logging/declarative-logging-symbol-processor/src/test/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/AbstractMdcAspectTest.kt new file mode 100644 index 000000000..f44c68ad2 --- /dev/null +++ b/logging/declarative-logging/declarative-logging-symbol-processor/src/test/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/AbstractMdcAspectTest.kt @@ -0,0 +1,18 @@ +package ru.tinkoff.kora.logging.symbol.processor.aop.mdc + +import ru.tinkoff.kora.ksp.common.AbstractSymbolProcessorTest + +abstract class AbstractMdcAspectTest : AbstractSymbolProcessorTest() { + + override fun commonImports(): String { + return super.commonImports() + """ + import ru.tinkoff.kora.logging.common.annotation.Mdc + import ru.tinkoff.kora.logging.common.MDC + import ru.tinkoff.kora.logging.symbol.processor.aop.mdc.MDCContextHolder + import java.util.concurrent.CompletionStage + import java.util.concurrent.CompletableFuture + import reactor.core.publisher.Mono + import reactor.core.publisher.Flux + """.trimIndent() + } +} diff --git a/logging/declarative-logging/declarative-logging-symbol-processor/src/test/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/MDCContextHolder.kt b/logging/declarative-logging/declarative-logging-symbol-processor/src/test/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/MDCContextHolder.kt new file mode 100644 index 000000000..b6fc14c9c --- /dev/null +++ b/logging/declarative-logging/declarative-logging-symbol-processor/src/test/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/MDCContextHolder.kt @@ -0,0 +1,16 @@ +package ru.tinkoff.kora.logging.symbol.processor.aop.mdc + +import ru.tinkoff.kora.logging.common.arg.StructuredArgumentWriter + +class MDCContextHolder { + + private var mdcContext: Map? = null + + fun get(): Map? { + return mdcContext + } + + fun set(mdcContext: Map?) { + this.mdcContext = mdcContext?.toMap() ?: emptyMap() + } +} diff --git a/logging/declarative-logging/declarative-logging-symbol-processor/src/test/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/MdcKoraAspectTest.kt b/logging/declarative-logging/declarative-logging-symbol-processor/src/test/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/MdcKoraAspectTest.kt new file mode 100644 index 000000000..82f2ec14b --- /dev/null +++ b/logging/declarative-logging/declarative-logging-symbol-processor/src/test/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/MdcKoraAspectTest.kt @@ -0,0 +1,241 @@ +package ru.tinkoff.kora.logging.symbol.processor.aop.mdc + +import org.intellij.lang.annotations.Language +import org.junit.jupiter.api.Assertions.assertDoesNotThrow +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.assertNotNull +import org.junit.jupiter.params.ParameterizedTest +import org.junit.jupiter.params.provider.MethodSource +import ru.tinkoff.kora.aop.symbol.processor.AopSymbolProcessorProvider +import ru.tinkoff.kora.logging.common.MDC +import java.util.UUID + +class MdcKoraAspectTest : AbstractMdcAspectTest() { + + private val contextHolder = MDCContextHolder() + + @ParameterizedTest + @MethodSource("provideTestCases") + fun testMdc(source: String) { + val aopProxy = compile0(listOf(AopSymbolProcessorProvider()), source.trimIndent()) + + invokeMethod(aopProxy) + + val context = contextHolder.get() + ?.mapValues { it.value.writeToString() } + + assertEquals(mapOf("key" to "\"value\"", "key1" to "\"value2\"", "123" to "\"test\""), context) + assertEquals(emptyMap(), currentContext()) + } + + @Test + fun testMdcWithCode() { + val aopProxy = compile0( + listOf(AopSymbolProcessorProvider()), + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + @Mdc(key = "key", value = "${UUID.randomUUID()}") + open fun test(s: String): Int? { + mdcContextHolder.set(MDC.get().values()) + return null + } + } + """.trimIndent() + ) + + invokeMethod(aopProxy) + + val context = contextHolder.get() + ?.mapValues { it.value.writeToString() } + + val value = context?.get("key") + assertNotNull(value) + assertDoesNotThrow { UUID.fromString(value.substring(1, value.length - 1)) } + assertEquals(emptyMap(), currentContext()) + } + + @Test + fun testGlobalMdc() { + val aopProxy = compile0( + listOf(AopSymbolProcessorProvider()), + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + @Mdc(key = "key", value = "value", global = true) + @Mdc(key = "key1", value = "value2") + open fun test(@Mdc(key = "123", global = true) s: String): Int? { + mdcContextHolder.set(MDC.get().values()) + return null + } + } + """.trimIndent() + ) + + invokeMethod(aopProxy) + + val context = contextHolder.get() + ?.mapValues { it.value.writeToString() } + + assertEquals(mapOf("key" to "\"value\"", "key1" to "\"value2\"", "123" to "\"test\""), context) + assertEquals(mapOf("key" to "\"value\"", "123" to "\"test\""), currentContext()) + clearCurrentContext() + } + + @ParameterizedTest + @MethodSource("provideSuspendTestCases") + fun testMdcWithCoroutines(source: String) { + val aopProxy = compile0(listOf(AopSymbolProcessorProvider()), source.trimIndent()) + + invokeMethod(aopProxy) + + val context = contextHolder.get() + ?.mapValues { it.value.writeToString() } + + assertEquals(mapOf("key" to "\"value\"", "key1" to "\"value2\"", "123" to "\"test\""), context) + assertEquals(emptyMap(), currentContext()) + } + + @ParameterizedTest + @MethodSource("provideGlobalSuspendTestCases") + fun testGlobalMdcWithCoroutines(source: String) { + val aopProxy = compile0(listOf(AopSymbolProcessorProvider()), source.trimIndent()) + assertTrue(aopProxy.isFailed()) + } + + @Test + fun testMdcWithNotEmptyContext() { + val aopProxy = compile0( + listOf(AopSymbolProcessorProvider()), + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + @Mdc(key = "key", value = "value") + @Mdc(key = "key1", value = "value2") + open fun test(@Mdc(key = "123") s: String): Int? { + mdcContextHolder.set(MDC.get().values()) + return null + } + } + """.trimIndent() + ) + + MDC.put("key", "special-value") + MDC.put("123", "special-123") + invokeMethod(aopProxy) + + val context = contextHolder.get() + ?.mapValues { it.value.writeToString() } + + assertEquals(mapOf("key" to "\"value\"", "key1" to "\"value2\"", "123" to "\"test\""), context) + assertEquals(mapOf("key" to "\"special-value\"", "123" to "\"special-123\""), currentContext()) + clearCurrentContext() + } + + private fun invokeMethod(aopProxy: CompileResult) { + aopProxy.assertSuccess() + + val generatedClass = aopProxy.loadClass("\$TestMdc__AopProxy") + val constructor = generatedClass.constructors.first() + + val testObject = TestObject(generatedClass.kotlin, constructor.newInstance(contextHolder)) + + testObject.invoke("test", "test") + } + + private fun currentContext() = MDC.get().values().mapValues { it.value.writeToString() } + + private fun clearCurrentContext() = MDC.get().values().keys.forEach { MDC.remove(it) } + + companion object { + + @JvmStatic + @Language("kotlin") + private fun provideTestCases() = listOf( + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + @Mdc(key = "key", value = "value") + @Mdc(key = "key1", value = "value2") + open fun test(@Mdc(key = "123") s: String): Int? { + mdcContextHolder.set(MDC.get().values()) + return null + } + } + """, + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + @Mdc(key = "key", value = "value") + @Mdc(key = "key1", value = "value2") + open fun test(@Mdc(key = "123") s: String) { + mdcContextHolder.set(MDC.get().values()) + } + } + """ + ) + + @JvmStatic + @Language("kotlin") + private fun provideSuspendTestCases() = listOf( + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + @Mdc(key = "key", value = "value") + @Mdc(key = "key1", value = "value2") + open suspend fun test(@Mdc(key = "123") s: String): Int? { + mdcContextHolder.set(MDC.get().values()) + return null + } + } + """, + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + @Mdc(key = "key", value = "value") + @Mdc(key = "key1", value = "value2") + open suspend fun test(@Mdc(key = "123") s: String) { + mdcContextHolder.set(MDC.get().values()) + } + } + """ + ) + + @JvmStatic + @Language("kotlin") + private fun provideGlobalSuspendTestCases() = listOf( + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + @Mdc(key = "key", value = "value", global = true) + @Mdc(key = "key1", value = "value2") + open suspend fun test(@Mdc(key = "123") s: String): Int? { + mdcContextHolder.set(MDC.get().values()) + return null + } + } + """, + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + @Mdc(key = "key", value = "value") + @Mdc(key = "key1", value = "value2") + open suspend fun test(@Mdc(key = "123", global = true) s: String) { + mdcContextHolder.set(MDC.get().values()) + } + } + """ + ) + } +} diff --git a/logging/declarative-logging/declarative-logging-symbol-processor/src/test/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/MdcKoraAspectUnsupportedTypesTest.kt b/logging/declarative-logging/declarative-logging-symbol-processor/src/test/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/MdcKoraAspectUnsupportedTypesTest.kt new file mode 100644 index 000000000..2a29fd7a3 --- /dev/null +++ b/logging/declarative-logging/declarative-logging-symbol-processor/src/test/kotlin/ru/tinkoff/kora/logging/symbol/processor/aop/mdc/MdcKoraAspectUnsupportedTypesTest.kt @@ -0,0 +1,135 @@ +package ru.tinkoff.kora.logging.symbol.processor.aop.mdc + +import org.intellij.lang.annotations.Language +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.params.ParameterizedTest +import org.junit.jupiter.params.provider.MethodSource +import ru.tinkoff.kora.aop.symbol.processor.AopSymbolProcessorProvider + +class MdcKoraAspectUnsupportedTypesTest : AbstractMdcAspectTest() { + + @ParameterizedTest + @MethodSource("sourcesWithMdcAndFuture", "sourcesWithMdcAndMono", "sourcesWithMdcAndFlux") + fun testMdc(source: String) { + val compileResult = compile0( + listOf(AopSymbolProcessorProvider()), + source + ) + + assertTrue(compileResult.isFailed()) + } + + companion object { + + @JvmStatic + @Language("kotlin") + private fun sourcesWithMdcAndFuture() = listOf( + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + @Mdc(key = "key", value = "value") + @Mdc(key = "key1", value = "value2") + open fun test(@Mdc(key = "123") s: String): CompletionStage<*> { + return CompletableFuture.completedFuture(1) + } + } + """, + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + @Mdc(key = "key1", value = "value2") + open fun test(s: String): CompletionStage<*> { + return CompletableFuture.completedFuture(1) + } + } + """, + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + + open fun test(@Mdc(key = "123") s: String): CompletionStage<*> { + return CompletableFuture.completedFuture(1) + } + } + """ + ) + + @JvmStatic + @Language("kotlin") + private fun sourcesWithMdcAndMono() = listOf( + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + @Mdc(key = "key", value = "value") + @Mdc(key = "key1", value = "value2") + open fun test(@Mdc(key = "123") s: String): Mono<*> { + return Mono.just(1) + } + } + """, + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + @Mdc(key = "key1", value = "value2") + open fun test(s: String): Mono<*> { + return Mono.just(1) + } + } + """, + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + + open fun test(@Mdc(key = "123") s: String): Mono<*> { + return Mono.just(1) + } + } + """ + ) + + @JvmStatic + @Language("kotlin") + private fun sourcesWithMdcAndFlux() = listOf( + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + @Mdc(key = "key", value = "value") + @Mdc(key = "key1", value = "value2") + open fun test(@Mdc(key = "123") s: String): Flux<*> { + return Mono.just(1) + .flux() + } + } + """, + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + @Mdc(key = "key1", value = "value2") + open fun test(s: String): Flux<*> { + return Mono.just(1) + .flux() + } + } + """, + """ + open class TestMdc( + private val mdcContextHolder: MDCContextHolder + ) { + + open fun test(@Mdc(key = "123") s: String): Flux<*> { + return Mono.just(1) + .flux() + } + } + """ + ) + } +} diff --git a/logging/logging-common/src/main/java/ru/tinkoff/kora/logging/common/annotation/Mdc.java b/logging/logging-common/src/main/java/ru/tinkoff/kora/logging/common/annotation/Mdc.java new file mode 100644 index 000000000..ad2ef815b --- /dev/null +++ b/logging/logging-common/src/main/java/ru/tinkoff/kora/logging/common/annotation/Mdc.java @@ -0,0 +1,31 @@ +package ru.tinkoff.kora.logging.common.annotation; + +import ru.tinkoff.kora.common.AopAnnotation; + +import java.lang.annotation.Repeatable; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; + +import static java.lang.annotation.ElementType.METHOD; +import static java.lang.annotation.ElementType.PARAMETER; + +@AopAnnotation +@Repeatable(Mdc.MdcContainer.class) +@Target({METHOD, PARAMETER}) +@Retention(RetentionPolicy.RUNTIME) +public @interface Mdc { + + String key() default ""; + + String value() default ""; + + boolean global() default false; + + @AopAnnotation + @Target(METHOD) + @Retention(RetentionPolicy.RUNTIME) + @interface MdcContainer { + Mdc[] value(); + } +}