Skip to content

Commit

Permalink
Fix Spring AOT model relfection
Browse files Browse the repository at this point in the history
  • Loading branch information
rj93 committed Apr 3, 2024
1 parent d889d63 commit 00c9a40
Showing 1 changed file with 43 additions and 21 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import java.util.Set;
import org.jetbrains.annotations.NotNull;
import org.reflections.Reflections;
import org.reflections.scanners.Scanners;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.aot.hint.MemberCategory;
Expand Down Expand Up @@ -55,38 +56,52 @@ public BeanFactoryInitializationAotContribution processAheadOfTime(
@NotNull ConfigurableListableBeanFactory beanFactory) {
return (generationContext, beanFactoryInitializationCode) -> {
RuntimeHints hints = generationContext.getRuntimeHints();
String[] classNames =
new String[] {
"com.google.gson.JsonElement", //
"io.kubernetes.client.informer.cache.ProcessorListener", //
"io.kubernetes.client.extended.controller.Controller", //
"io.kubernetes.client.util.generic.GenericKubernetesApi$StatusPatch", //
"io.kubernetes.client.util.Watch$Response" //
};
for (String className : classNames) {
LOGGER.info("registering {} for reflection", className);
hints.reflection().registerType(TypeReference.of(className), allMemberCategories);
}
registerStaticClasses(hints);
registerModels(hints);
registerForPackage("io.kubernetes", hints);
Collection<String> packages = AutoConfigurationPackages.get(beanFactory);
for (String packageName : packages) {
registerForPackage(packageName, hints);
}
registerAutoconfigurationPackages(beanFactory, hints);
};
}

private void registerStaticClasses(RuntimeHints hints) {
Set<String> classNames = Set.of(
"com.google.gson.JsonElement",
"io.kubernetes.client.informer.cache.ProcessorListener",
"io.kubernetes.client.extended.controller.Controller",
"io.kubernetes.client.util.generic.GenericKubernetesApi$StatusPatch",
"io.kubernetes.client.util.Watch$Response"
);
for (String className : classNames) {
LOGGER.info("registering {} for reflection", className);
hints.reflection().registerType(TypeReference.of(className), allMemberCategories);
}
}

private void registerModels(RuntimeHints hints) {
Reflections reflections = new Reflections("io.kubernetes.client.openapi.models",
Scanners.SubTypes.filterResultsBy(s -> true));
Set<Class<?>> models = reflections.getSubTypesOf(Object.class);
LOGGER.info("Found {} models", models.size());
registerClassesForReflection(models, hints);
}

private void registerForPackage(String packageName, RuntimeHints hints) {
Reflections reflections = new Reflections(packageName);
Set<Class<?>> apiModels = reflections.getTypesAnnotatedWith(ApiModel.class);
Set<Class<? extends Controller>> controllers = reflections.getSubTypesOf(Controller.class);
LOGGER.info("Found {} controllers", controllers.size());
Set<Class<?>> jsonAdapters = findJsonAdapters(reflections);
LOGGER.info("Found {} jsonAdapters", jsonAdapters.size());
Set<Class<?>> all = new HashSet<>();
all.addAll(jsonAdapters);
all.addAll(controllers);
all.addAll(apiModels);
for (Class<?> clazz : all) {
LOGGER.info("registering {} for reflection", clazz.getName());
hints.reflection().registerType(clazz, allMemberCategories);
registerClassesForReflection(all, hints);
}

private void registerAutoconfigurationPackages(ConfigurableListableBeanFactory beanFactory,
RuntimeHints hints) {
Collection<String> packages = AutoConfigurationPackages.get(beanFactory);
for (String packageName : packages) {
registerForPackage(packageName, hints);
}
}

Expand All @@ -102,4 +117,11 @@ private <R extends Annotation> Set<Class<?>> findJsonAdapters(Reflections reflec
}
return classes;
}

private void registerClassesForReflection(Set<Class<?>> classes, RuntimeHints hints) {
for (Class<?> clazz : classes) {
LOGGER.info("registering {} for reflection", clazz);
hints.reflection().registerType(clazz, allMemberCategories);
}
}
}

0 comments on commit 00c9a40

Please sign in to comment.