From 54bbbb5422c812c1d7490eb29c6f1cad2dbd8ce1 Mon Sep 17 00:00:00 2001 From: Olga MaciaszekSharma Date: Wed, 24 Apr 2024 15:28:01 +0200 Subject: [PATCH] Manually register bean definitions in AOT. --- ...BeanFactoryInitializationAotProcessor.java | 156 ++++++++++++++++++ .../resources/META-INF/spring/aot.factories | 6 + 2 files changed, 162 insertions(+) create mode 100644 spring-cloud-config-server/src/main/java/org/springframework/cloud/config/server/aot/CompositeEnvironmentBeanFactoryInitializationAotProcessor.java diff --git a/spring-cloud-config-server/src/main/java/org/springframework/cloud/config/server/aot/CompositeEnvironmentBeanFactoryInitializationAotProcessor.java b/spring-cloud-config-server/src/main/java/org/springframework/cloud/config/server/aot/CompositeEnvironmentBeanFactoryInitializationAotProcessor.java new file mode 100644 index 00000000..f634865d --- /dev/null +++ b/spring-cloud-config-server/src/main/java/org/springframework/cloud/config/server/aot/CompositeEnvironmentBeanFactoryInitializationAotProcessor.java @@ -0,0 +1,156 @@ +/* + * Copyright 2014-2024 the original author or authors. + * + * Licensed 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 + * + * https://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 org.springframework.cloud.config.server.aot; + +import java.util.Arrays; +import java.util.Map; +import java.util.Objects; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +import javax.lang.model.element.Modifier; + +import org.springframework.aot.generate.GeneratedMethod; +import org.springframework.aot.generate.GenerationContext; +import org.springframework.aot.hint.MemberCategory; +import org.springframework.aot.hint.ReflectionHints; +import org.springframework.aot.hint.RuntimeHints; +import org.springframework.aot.hint.TypeReference; +import org.springframework.beans.factory.aot.BeanFactoryInitializationAotContribution; +import org.springframework.beans.factory.aot.BeanFactoryInitializationAotProcessor; +import org.springframework.beans.factory.aot.BeanFactoryInitializationCode; +import org.springframework.beans.factory.aot.BeanRegistrationExcludeFilter; +import org.springframework.beans.factory.config.BeanDefinition; +import org.springframework.beans.factory.config.ConfigurableListableBeanFactory; +import org.springframework.beans.factory.support.DefaultListableBeanFactory; +import org.springframework.beans.factory.support.RegisteredBean; +import org.springframework.cloud.config.server.environment.EnvironmentRepository; +import org.springframework.cloud.config.server.support.EnvironmentRepositoryProperties; +import org.springframework.javapoet.MethodSpec; +import org.springframework.util.ClassUtils; + +/** + * @author Olga Maciaszek-Sharma + */ +public class CompositeEnvironmentBeanFactoryInitializationAotProcessor + implements BeanFactoryInitializationAotProcessor, BeanRegistrationExcludeFilter { + + @Override + public BeanFactoryInitializationAotContribution processAheadOfTime(ConfigurableListableBeanFactory beanFactory) { + Map propertyBeanDefinitions = getCompositeEnvironmentBeanDefinitions(beanFactory, + "-env-repo-properties", EnvironmentRepositoryProperties.class); + Map repoBeanDefinitions = getCompositeEnvironmentBeanDefinitions(beanFactory, + "-env-repo", EnvironmentRepository.class); + + return new CompositeEnvironmentBeanFactoryInitializationAotContribution(propertyBeanDefinitions, repoBeanDefinitions); + } + + private static Map getCompositeEnvironmentBeanDefinitions(ConfigurableListableBeanFactory beanFactory, + String infix, Class beanClass) { + return Arrays.stream(beanFactory.getBeanDefinitionNames()) + .filter(beanName -> beanName.contains(infix)) + .map(beanName -> Map.entry(beanName, beanFactory.getBeanDefinition(beanName))) + .filter(entry -> { + try { + return beanClass.isAssignableFrom(Class.forName(entry.getValue() + .getBeanClassName())); + } + catch (ClassNotFoundException e) { + throw new RuntimeException("Class " + entry.getValue() + .getBeanClassName() + " could not be found", e); + } + }) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + } + + @Override + public boolean isExcludedFromAotProcessing(RegisteredBean registeredBean) { + return EnvironmentRepositoryProperties.class.isAssignableFrom(registeredBean.getBeanClass()) + && registeredBean.getBeanName().contains("-env-repo-properties") + || EnvironmentRepository.class.isAssignableFrom(registeredBean.getBeanClass()) + && registeredBean.getBeanName().contains("-env-repo"); + } + + private record CompositeEnvironmentBeanFactoryInitializationAotContribution( + Map propertyBeanDefinitions, + Map repoBeanDefinitions) implements BeanFactoryInitializationAotContribution { + + @Override + public void applyTo(GenerationContext generationContext, BeanFactoryInitializationCode beanFactoryInitializationCode) { + GeneratedMethod environmentRepositoryPropertiesGeneratedMethod = beanFactoryInitializationCode.getMethods() + .add("registerCompositeEnvironmentRepositoryPropertiesBeanDefinitions", + method -> generateRegisterBeanDefinitionsMethod(method, propertyBeanDefinitions, "EnvironmentRepositoryProperties")); + GeneratedMethod environmentRepositoriesGeneratedMethod = beanFactoryInitializationCode.getMethods() + .add("registerCompositeEnvironmentRepositoryBeanDefinitions", + method -> generateRegisterBeanDefinitionsMethod(method, propertyBeanDefinitions, "EnvironmentRepository")); + beanFactoryInitializationCode.addInitializer(environmentRepositoryPropertiesGeneratedMethod.toMethodReference()); + beanFactoryInitializationCode.addInitializer(environmentRepositoriesGeneratedMethod.toMethodReference()); + generateRuntimeHints(generationContext.getRuntimeHints()); + } + + private void generateRuntimeHints(RuntimeHints runtimeHints) { + ReflectionHints hints = runtimeHints.reflection(); + Stream.concat(propertyBeanDefinitions.values().stream() + , repoBeanDefinitions.values().stream()) + .map(BeanDefinition::getBeanClassName) + .filter(Objects::nonNull) + .map(beanClassName -> { + try { + return Class.forName(beanClassName); + } + catch (ClassNotFoundException e) { + throw new RuntimeException("Class " + beanClassName + " could not be found", e); + } + }) + .forEach(beanClassName -> { + hints.registerType(TypeReference.of(beanClassName), + MemberCategory.INTROSPECT_PUBLIC_METHODS, + MemberCategory.INTROSPECT_DECLARED_METHODS); + introspectPublicMethodsOnAllInterfaces(hints, beanClassName); + } + ); + + } + + private void generateRegisterBeanDefinitionsMethod(MethodSpec.Builder method, + Map beanDefinitions, String name) { + method.addJavadoc("Register the $S bean definitions for composite config data sources.", name); + method.addModifiers(Modifier.PUBLIC); + method.addParameter(DefaultListableBeanFactory.class, "beanFactory"); + beanDefinitions.keySet() + .forEach(beanName -> + method.addStatement("beanFactory.registerBeanDefinition($S, $L)", + beanName, beanDefinitions.get(beanName))); + + } + + // originally from Spring Framework BeanRegistrationsAotContribution + private void introspectPublicMethodsOnAllInterfaces(ReflectionHints hints, Class type) { + Class currentClass = type; + while (currentClass != null && currentClass != Object.class) { + for (Class interfaceType : currentClass.getInterfaces()) { + if (!ClassUtils.isJavaLanguageInterface(interfaceType)) { + hints.registerType(interfaceType, MemberCategory.INTROSPECT_PUBLIC_METHODS); + introspectPublicMethodsOnAllInterfaces(hints, interfaceType); + } + } + currentClass = currentClass.getSuperclass(); + } + } + } + +} diff --git a/spring-cloud-config-server/src/main/resources/META-INF/spring/aot.factories b/spring-cloud-config-server/src/main/resources/META-INF/spring/aot.factories index 35c4ba6b..682c7d78 100644 --- a/spring-cloud-config-server/src/main/resources/META-INF/spring/aot.factories +++ b/spring-cloud-config-server/src/main/resources/META-INF/spring/aot.factories @@ -1,2 +1,8 @@ org.springframework.aot.hint.RuntimeHintsRegistrar=\ org.springframework.cloud.config.server.config.ConfigServerRuntimeHints + +org.springframework.beans.factory.aot.BeanFactoryInitializationAotProcessor= \ +org.springframework.cloud.config.server.aot.CompositeEnvironmentBeanFactoryInitialisationAotProcessor + +org.springframework.beans.factory.aot.BeanRegistrationExcludeFilter=\ +org.springframework.cloud.config.server.aot.CompositeEnvironmentBeanFactoryInitializationAotProcessor