Fix AOT for GraalVM 22 (#283)

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.

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<String> getBuildTimeInitClasses();
+
+    /**
+     * Performs actions which have to be done as late as possible during
+     * source generation.
+     */
+    void finish();
 }

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() {

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) {

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<String, List<String>> getDiagnostics() {
     public Runtime getRuntime() {
         return delegate.getRuntime();
     }
+
+    @Override
+    public Set<String> getBuildTimeInitClasses() {
+        return delegate.getBuildTimeInitClasses();
+    }
+
+    @Override
+    public void finish() {
+        delegate.finish();
+    }
 }

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<JavaFile> generatedJavaFiles = new ArrayList<>();
     private final List<MethodSpec> initializers = new ArrayList<>();
     private final Path generatedResourcesDirectory;
+    private final Set<String> buildTimeInitClasses = new HashSet<>();
+    private final List<Runnable> deferredOperations = new ArrayList<>();
 
     public DefaultSourceGenerationContext(String packageName,
                                           ApplicationContextAnalyzer analyzer,
@@ -163,6 +165,7 @@ public <T> void registerStaticOptimization(String className, Class<T> 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<MethodSpec> getGeneratedStaticInitializers() {
     @Override
     public void registerGeneratedResource(@NonNull String path, Consumer<? super File> 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<File> 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 <T> Optional<T> get(@NonNull Class<T> type) {
     public Map<String, List<String>> getDiagnostics() {
         return diagnostics;
     }
+
+    @Override
+    public Set<String> getBuildTimeInitClasses() {
+        return Collections.unmodifiableSet(buildTimeInitClasses);
+    }
+
+    @Override
+    public void finish() {
+        deferredOperations.forEach(Runnable::run);
+    }
 }

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)

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<String, AbstractCodeGenerator> substitutions;
     private Set<String> forceInclude;
     private final Substitutes substitutes = new Substitutes();
-    private final Map<String, TypeSpec> staticServiceClasses = new HashMap<>();
+    private final Map<String, GeneratedType> staticServiceClasses = new HashMap<>();
     private final Set<BeanConfiguration> disabledConfigurations = Collections.synchronizedSet(new HashSet<>());
     private final Map<String, List<Class<?>>> serviceClasses = new HashMap<>();
     private final Set<Class<?>> 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,23 +237,27 @@ protected abstract void generateFindAllMethod(Stream<Class<?>> 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)
                 .addSuperinterface(ParameterizedTypeName.get(SoftServiceLoader.StaticServiceLoader.class, serviceType));
         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));
             body.addStatement("$T staticServices = new $T()",
                     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<String, TypeSpec> entry : staticServiceClasses.entrySet()) {
-                body.addStatement("staticServices.put($S, new $T())", entry.getKey(), ClassName.bestGuess(entry.getValue().name));
+            for (Map.Entry<String, GeneratedType> 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
+    ) {
+
+    }
 
 }

aot-std-optimizers/src/main/java/io/micronaut/aot/std/sourcegen/GraalVMOptimizationFeatureSourceGenerator.java

@@ -28,40 +28,55 @@
 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
  * the native image code generation, typically asking to initialize
  * 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<String> 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);

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