From ba7065e2cccaa9573f5bef29d6f1324f947dd1f6 Mon Sep 17 00:00:00 2001 From: Cedric Champeau Date: Fri, 29 Mar 2024 18:11:33 +0100 Subject: [PATCH] Fix AOT for GraalVM 22 This commit updates how Micronaut AOT generates code for GraalVM 22. In 22, the new default is to forbid build time initialization of classes. Each class which is initialized at build time has to be explicitly declared as such. This commit makes it so that classes we generate and know are initialized at build time are added to the `--initialize-at-build-time` option. While this fixes a lot of cases, we cannot, unfortunately, kwow upfront all the dependencies of the types we initialize at build time, which may _also_ need build time initialization. This means that in some cases, users will have to explicitly add some types to the list (e.g via a `buildArg` in the native binaries Gradle extension). Ideally, we should move AOT off build time initialization. However, this is not currently doable, since precisely Micronaut Core and optimization registration is designed to use static fields. An option would be to use service loading instead, but it's a breaking change, and that would remove the ability to have a service loading optimization. While removing the service loading optimization may not be an issue for native, it is, however, a problem for JIT optimizations. There's therefore no good solution that I'm aware of to this problem. --- .../io/micronaut/aot/core/AOTContext.java | 19 ++++++++ .../AbstractSingleClassFileGenerator.java | 2 + ...ApplicationContextConfigurerGenerator.java | 1 + .../DelegatingSourceGenerationContext.java | 16 +++++++ .../DefaultSourceGenerationContext.java | 36 +++++++++++---- .../AbstractSourceGeneratorSpec.groovy | 3 +- ...actStaticServiceLoaderSourceGenerator.java | 25 ++++++++--- ...lVMOptimizationFeatureSourceGenerator.java | 45 ++++++++++++------- ...imizationFeatureSourceGeneratorTest.groovy | 6 ++- 9 files changed, 121 insertions(+), 32 deletions(-) diff --git a/aot-core/src/main/java/io/micronaut/aot/core/AOTContext.java b/aot-core/src/main/java/io/micronaut/aot/core/AOTContext.java index b7a368e6..cc287f29 100644 --- a/aot-core/src/main/java/io/micronaut/aot/core/AOTContext.java +++ b/aot-core/src/main/java/io/micronaut/aot/core/AOTContext.java @@ -26,6 +26,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.function.Consumer; /** @@ -115,6 +116,12 @@ public interface AOTContext { */ void registerClassNeededAtCompileTime(@NonNull Class clazz); + /** + * Registers a type as a requiring initialization at build time. + * @param className the type + */ + void registerBuildTimeInit(@NonNull String className); + /** * Generates a java file spec. * @param typeSpec the type spec of the main class @@ -165,4 +172,16 @@ public interface AOTContext { */ @NonNull Runtime getRuntime(); + + /** + * Returns the set of classes which require build time initialization + * @return the set of classes needing build time init + */ + Set getBuildTimeInitClasses(); + + /** + * Performs actions which have to be done as late as possible during + * source generation. + */ + void finish(); } diff --git a/aot-core/src/main/java/io/micronaut/aot/core/codegen/AbstractSingleClassFileGenerator.java b/aot-core/src/main/java/io/micronaut/aot/core/codegen/AbstractSingleClassFileGenerator.java index 09467d6b..539bb206 100644 --- a/aot-core/src/main/java/io/micronaut/aot/core/codegen/AbstractSingleClassFileGenerator.java +++ b/aot-core/src/main/java/io/micronaut/aot/core/codegen/AbstractSingleClassFileGenerator.java @@ -34,6 +34,8 @@ public void generate(@NonNull AOTContext context) { this.context = context; JavaFile javaFile = generate(); context.registerGeneratedSourceFile(javaFile); + context.registerBuildTimeInit(javaFile.packageName + "." + javaFile.typeSpec.name); + context.registerBuildTimeInit(javaFile.packageName + "." + javaFile.typeSpec.name + "$1"); } protected final AOTContext getContext() { diff --git a/aot-core/src/main/java/io/micronaut/aot/core/codegen/ApplicationContextConfigurerGenerator.java b/aot-core/src/main/java/io/micronaut/aot/core/codegen/ApplicationContextConfigurerGenerator.java index af1a1605..ccc3acb9 100644 --- a/aot-core/src/main/java/io/micronaut/aot/core/codegen/ApplicationContextConfigurerGenerator.java +++ b/aot-core/src/main/java/io/micronaut/aot/core/codegen/ApplicationContextConfigurerGenerator.java @@ -69,6 +69,7 @@ public void generate(@NonNull AOTContext context) { optimizedEntryPoint.addStaticBlock(staticInitializer.build()); context.registerGeneratedSourceFile(context.javaFile(optimizedEntryPoint.build())); context.registerServiceImplementation(ApplicationContextConfigurer.class, CUSTOMIZER_CLASS_NAME); + context.finish(); } private void addDiagnostics(AOTContext context, TypeSpec.Builder optimizedEntryPoint) { diff --git a/aot-core/src/main/java/io/micronaut/aot/core/codegen/DelegatingSourceGenerationContext.java b/aot-core/src/main/java/io/micronaut/aot/core/codegen/DelegatingSourceGenerationContext.java index a0c974d5..c2c213c6 100644 --- a/aot-core/src/main/java/io/micronaut/aot/core/codegen/DelegatingSourceGenerationContext.java +++ b/aot-core/src/main/java/io/micronaut/aot/core/codegen/DelegatingSourceGenerationContext.java @@ -29,6 +29,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.function.Consumer; /** @@ -110,6 +111,11 @@ public void registerClassNeededAtCompileTime(@NonNull Class clazz) { delegate.registerClassNeededAtCompileTime(clazz); } + @Override + public void registerBuildTimeInit(String className) { + delegate.registerBuildTimeInit(className); + } + @Override @NonNull public JavaFile javaFile(TypeSpec typeSpec) { @@ -144,4 +150,14 @@ public Map> getDiagnostics() { public Runtime getRuntime() { return delegate.getRuntime(); } + + @Override + public Set getBuildTimeInitClasses() { + return delegate.getBuildTimeInitClasses(); + } + + @Override + public void finish() { + delegate.finish(); + } } diff --git a/aot-core/src/main/java/io/micronaut/aot/core/context/DefaultSourceGenerationContext.java b/aot-core/src/main/java/io/micronaut/aot/core/context/DefaultSourceGenerationContext.java index 85a80653..3801e74d 100644 --- a/aot-core/src/main/java/io/micronaut/aot/core/context/DefaultSourceGenerationContext.java +++ b/aot-core/src/main/java/io/micronaut/aot/core/context/DefaultSourceGenerationContext.java @@ -77,6 +77,8 @@ public final class DefaultSourceGenerationContext implements AOTContext { private final List generatedJavaFiles = new ArrayList<>(); private final List initializers = new ArrayList<>(); private final Path generatedResourcesDirectory; + private final Set buildTimeInitClasses = new HashSet<>(); + private final List deferredOperations = new ArrayList<>(); public DefaultSourceGenerationContext(String packageName, ApplicationContextAnalyzer analyzer, @@ -163,6 +165,7 @@ public void registerStaticOptimization(String className, Class optimizati .addSuperinterface(ParameterizedTypeName.get(StaticOptimizations.Loader.class, optimizationKind)) .addMethod(method) .build(); + registerBuildTimeInit(optimizationKind.getName()); registerGeneratedSourceFile(javaFile(generatedType)); registerServiceImplementation(StaticOptimizations.Loader.class, className); } @@ -190,14 +193,16 @@ public List getGeneratedStaticInitializers() { @Override public void registerGeneratedResource(@NonNull String path, Consumer consumer) { LOGGER.debug("Registering generated resource file: {}", path); - Path relative = generatedResourcesDirectory.resolve(path); - File resourceFile = relative.toFile(); - File parent = resourceFile.getParentFile(); - if (parent.exists() || parent.mkdirs()) { - consumer.accept(resourceFile); - } else { - throw new RuntimeException("Unable to create parent file " + parent + " for resource " + path); - } + deferredOperations.add(() -> { + Path relative = generatedResourcesDirectory.resolve(path); + File resourceFile = relative.toFile(); + File parent = resourceFile.getParentFile(); + if (parent.exists() || parent.mkdirs()) { + consumer.accept(resourceFile); + } else { + throw new RuntimeException("Unable to create parent file " + parent + " for resource " + path); + } + }); } @NonNull @@ -217,6 +222,11 @@ public List getExtraClasspath() { .collect(Collectors.toList()); } + @Override + public void registerBuildTimeInit(String className) { + buildTimeInitClasses.add(className); + } + /** * Returns the list of resources to be excluded from * the binary. @@ -258,4 +268,14 @@ public Optional get(@NonNull Class type) { public Map> getDiagnostics() { return diagnostics; } + + @Override + public Set getBuildTimeInitClasses() { + return Collections.unmodifiableSet(buildTimeInitClasses); + } + + @Override + public void finish() { + deferredOperations.forEach(Runnable::run); + } } diff --git a/aot-core/src/testFixtures/groovy/io/micronaut/aot/core/codegen/AbstractSourceGeneratorSpec.groovy b/aot-core/src/testFixtures/groovy/io/micronaut/aot/core/codegen/AbstractSourceGeneratorSpec.groovy index a69a9bdf..581771f1 100644 --- a/aot-core/src/testFixtures/groovy/io/micronaut/aot/core/codegen/AbstractSourceGeneratorSpec.groovy +++ b/aot-core/src/testFixtures/groovy/io/micronaut/aot/core/codegen/AbstractSourceGeneratorSpec.groovy @@ -63,9 +63,10 @@ abstract class AbstractSourceGeneratorSpec extends Specification { abstract AOTCodeGenerator newGenerator() - void generate() { + final void generate() { def sourceGenerator = newGenerator() sourceGenerator.generate(context) + context.finish() def sources = context.getGeneratedJavaFiles().collectEntries([:]) { def writer = new StringWriter() it.writeTo(writer) diff --git a/aot-std-optimizers/src/main/java/io/micronaut/aot/std/sourcegen/AbstractStaticServiceLoaderSourceGenerator.java b/aot-std-optimizers/src/main/java/io/micronaut/aot/std/sourcegen/AbstractStaticServiceLoaderSourceGenerator.java index 78bbfa8f..1d2cd750 100644 --- a/aot-std-optimizers/src/main/java/io/micronaut/aot/std/sourcegen/AbstractStaticServiceLoaderSourceGenerator.java +++ b/aot-std-optimizers/src/main/java/io/micronaut/aot/std/sourcegen/AbstractStaticServiceLoaderSourceGenerator.java @@ -79,7 +79,7 @@ public abstract class AbstractStaticServiceLoaderSourceGenerator extends Abstrac private Map substitutions; private Set forceInclude; private final Substitutes substitutes = new Substitutes(); - private final Map staticServiceClasses = new HashMap<>(); + private final Map staticServiceClasses = new HashMap<>(); private final Set disabledConfigurations = Collections.synchronizedSet(new HashSet<>()); private final Map>> serviceClasses = new HashMap<>(); private final Set> disabledServices = new HashSet<>(); @@ -144,7 +144,10 @@ public void generate(@NonNull AOTContext context) { LOGGER.debug("Generated static {} service loader substitutions", substitutes.values().size()); staticServiceClasses.values() .stream() - .map(context::javaFile) + .map(generatedType -> { + context.registerBuildTimeInit(generatedType.className()); + return context.javaFile(generatedType.typeSpec()); + }) .forEach(context::registerGeneratedSourceFile); context.registerStaticOptimization("StaticServicesLoader", SoftServiceLoader.Optimizations.class, this::buildOptimization); } @@ -166,7 +169,7 @@ private void generateServiceLoader() { serviceName, serviceType, factory); - staticServiceClasses.put(serviceName, factory.build()); + staticServiceClasses.put(serviceName, new GeneratedType(context.getPackageName() + "." + factoryNameFor(serviceName), factory.build())); } } @@ -234,7 +237,7 @@ protected abstract void generateFindAllMethod(Stream> serviceClasses, TypeSpec.Builder factory); private TypeSpec.Builder prepareServiceLoaderType(String serviceName, Class serviceType) { - String name = simpleNameOf(serviceName) + "Factory"; + String name = factoryNameFor(serviceName); TypeSpec.Builder factory = TypeSpec.classBuilder(name) .addModifiers(PUBLIC) .addAnnotation(Generated.class) @@ -242,6 +245,10 @@ private TypeSpec.Builder prepareServiceLoaderType(String serviceName, Class s return factory; } + private static String factoryNameFor(String serviceName) { + return simpleNameOf(serviceName) + "Factory"; + } + private void buildOptimization(CodeBlock.Builder body) { ParameterizedTypeName serviceLoaderType = ParameterizedTypeName.get( ClassName.get(SoftServiceLoader.StaticServiceLoader.class), WildcardTypeName.subtypeOf(Object.class)); @@ -249,8 +256,8 @@ private void buildOptimization(CodeBlock.Builder body) { ParameterizedTypeName.get(ClassName.get(Map.class), ClassName.get(String.class), serviceLoaderType), ParameterizedTypeName.get(ClassName.get(HashMap.class), ClassName.get(String.class), serviceLoaderType)); - for (Map.Entry entry : staticServiceClasses.entrySet()) { - body.addStatement("staticServices.put($S, new $T())", entry.getKey(), ClassName.bestGuess(entry.getValue().name)); + for (Map.Entry entry : staticServiceClasses.entrySet()) { + body.addStatement("staticServices.put($S, new $T())", entry.getKey(), ClassName.bestGuess(entry.getValue().typeSpec().name)); } body.addStatement("return new $T(staticServices)", SoftServiceLoader.Optimizations.class); } @@ -324,5 +331,11 @@ private boolean skipService(Class clazz, Throwable e) { } } + private record GeneratedType( + String className, + TypeSpec typeSpec + ) { + + } } diff --git a/aot-std-optimizers/src/main/java/io/micronaut/aot/std/sourcegen/GraalVMOptimizationFeatureSourceGenerator.java b/aot-std-optimizers/src/main/java/io/micronaut/aot/std/sourcegen/GraalVMOptimizationFeatureSourceGenerator.java index b89805f2..27f499d9 100644 --- a/aot-std-optimizers/src/main/java/io/micronaut/aot/std/sourcegen/GraalVMOptimizationFeatureSourceGenerator.java +++ b/aot-std-optimizers/src/main/java/io/micronaut/aot/std/sourcegen/GraalVMOptimizationFeatureSourceGenerator.java @@ -28,6 +28,7 @@ import java.io.IOException; import java.io.PrintWriter; import java.util.List; +import java.util.stream.Collectors; /** * Generates the GraalVM configuration file which is going to configure @@ -35,33 +36,47 @@ * the optimized entry point at build time. */ @AOTModule( - id = GraalVMOptimizationFeatureSourceGenerator.ID, - description = GraalVMOptimizationFeatureSourceGenerator.DESCRIPTION, - options = { - @Option( - key = "service.types", - description = "The list of service types to be scanned (comma separated)", - sampleValue = "io.micronaut.Service1,io.micronaut.Service2" - ) - }, - enabledOn = Runtime.NATIVE + id = GraalVMOptimizationFeatureSourceGenerator.ID, + description = GraalVMOptimizationFeatureSourceGenerator.DESCRIPTION, + options = { + @Option( + key = "service.types", + description = "The list of service types to be scanned (comma separated)", + sampleValue = "io.micronaut.Service1,io.micronaut.Service2" + ) + }, + enabledOn = Runtime.NATIVE ) public class GraalVMOptimizationFeatureSourceGenerator extends AbstractCodeGenerator { public static final String ID = "graalvm.config"; - public static final String DESCRIPTION = "Generates GraalVM configuration files required to load the AOT optimizations"; + public static final String DESCRIPTION = + "Generates GraalVM configuration files required to load the AOT optimizations"; private static final String NEXT_LINE = " \\"; - private static final Option OPTION = MetadataUtils.findOption(GraalVMOptimizationFeatureSourceGenerator.class, "service.types"); + private static final Option OPTION = + MetadataUtils.findOption(GraalVMOptimizationFeatureSourceGenerator.class, "service.types"); @Override public void generate(@NonNull AOTContext context) { List serviceTypes = context.getConfiguration().stringList(OPTION.key()); - String path = "META-INF/native-image/" + context.getPackageName() + "/native-image.properties"; + String path = + "META-INF/native-image/" + context.getPackageName() + "/native-image.properties"; context.registerGeneratedResource(path, propertiesFile -> { try (PrintWriter wrt = new PrintWriter(new FileWriter(propertiesFile))) { wrt.print("Args="); - wrt.println("--initialize-at-build-time=" + context.getPackageName() + "." + ApplicationContextConfigurerGenerator.CUSTOMIZER_CLASS_NAME + NEXT_LINE); - if (context.getConfiguration().isFeatureEnabled(NativeStaticServiceLoaderSourceGenerator.ID)) { + wrt.println("--initialize-at-build-time=io.micronaut.context.ApplicationContextConfigurer$1" + NEXT_LINE); + wrt.println(" --initialize-at-build-time=" + context.getPackageName() + "." + + ApplicationContextConfigurerGenerator.CUSTOMIZER_CLASS_NAME + + NEXT_LINE); + var buildTimeInit = context.getBuildTimeInitClasses() + .stream() + .map(clazz -> " --initialize-at-build-time=" + clazz) + .collect(Collectors.joining(NEXT_LINE + "\n")); + if (!buildTimeInit.isEmpty()) { + wrt.println(buildTimeInit); + } + if (context.getConfiguration() + .isFeatureEnabled(NativeStaticServiceLoaderSourceGenerator.ID)) { for (int i = 0; i < serviceTypes.size(); i++) { String serviceType = serviceTypes.get(i); wrt.print(" -H:ServiceLoaderFeatureExcludeServices=" + serviceType); diff --git a/aot-std-optimizers/src/test/groovy/io/micronaut/aot/std/sourcegen/GraalVMOptimizationFeatureSourceGeneratorTest.groovy b/aot-std-optimizers/src/test/groovy/io/micronaut/aot/std/sourcegen/GraalVMOptimizationFeatureSourceGeneratorTest.groovy index 8f12d014..01382325 100644 --- a/aot-std-optimizers/src/test/groovy/io/micronaut/aot/std/sourcegen/GraalVMOptimizationFeatureSourceGeneratorTest.groovy +++ b/aot-std-optimizers/src/test/groovy/io/micronaut/aot/std/sourcegen/GraalVMOptimizationFeatureSourceGeneratorTest.groovy @@ -18,7 +18,8 @@ class GraalVMOptimizationFeatureSourceGeneratorTest extends AbstractSourceGenera assertThatGeneratedSources { doesNotCreateInitializer() generatesMetaInfResource("native-image/$packageName/native-image.properties", """ -Args=--initialize-at-build-time=io.micronaut.test.AOTApplicationContextConfigurer \\ +Args=--initialize-at-build-time=io.micronaut.context.ApplicationContextConfigurer\$1 \\ + --initialize-at-build-time=io.micronaut.test.AOTApplicationContextConfigurer \\ """) } } @@ -32,7 +33,8 @@ Args=--initialize-at-build-time=io.micronaut.test.AOTApplicationContextConfigure assertThatGeneratedSources { doesNotCreateInitializer() generatesMetaInfResource("native-image/$packageName/native-image.properties", """ -Args=--initialize-at-build-time=io.micronaut.test.AOTApplicationContextConfigurer \\ +Args=--initialize-at-build-time=io.micronaut.context.ApplicationContextConfigurer\$1 \\ + --initialize-at-build-time=io.micronaut.test.AOTApplicationContextConfigurer \\ -H:ServiceLoaderFeatureExcludeServices=A \\ -H:ServiceLoaderFeatureExcludeServices=B \\ -H:ServiceLoaderFeatureExcludeServices=C