Add support for providing a POJO to an initializer registered by BeanFactoryInitializationCode

[Christopher-Chianelli]

I am adding support for Spring Boot native image for the Timefold Solver: TimefoldAI/timefold-solver#609. One of the things we do ahead of time is generate the SolverConfig, a POJO (a class with getter and setter for all fields (including inherited fields), who field types are all also POJOs (or a primitive/builtin/collection type)).

The Timefold Solver needs to register the generated SolverConfig as a bean, since it is used when constructing the other beans Timefold Solver provides (SolverFactory, SolverManager, etc). Outside of native mode, this is done by a BeanFactoryPostProcessor, which directly registers the SolverConfig to the ConfigurableListableBeanFactory. Inside native mode, this is done by a BeanFactoryInitializationAotProcessor, which generates a class with the generated SolverConfig and adds an initializer (that uses the generated class) that registers the SolverConfig to the ConfigurableListableBeanFactory.

The issue is, SolverConfig has many different fields, and have multiple configurations nested within it. This make it non-trivial to put inside a generated class. Normally, I would use XML serialization/deserialization to put it inside the generated class:

SolverConfigIO solverConfigIO = new SolverConfigIO();
for (Map.Entry<String, SolverConfig> solverConfigEntry : solverConfigMap.entrySet()) {
        StringWriter writer = new StringWriter();
        solverConfigIO.write(solverConfigEntry.getValue(), writer);
        solverConfigToXml.put(solverConfigEntry.getKey(), writer.toString());
}

GeneratedClass generatedClass = generationContext.getGeneratedClasses().addForFeature("timefold-aot",
    builder -> {
        final String SOLVER_CONFIG_MAP_FIELD = "solverConfigMap";
        builder.addField(Map.class, SOLVER_CONFIG_MAP_FIELD, Modifier.STATIC);
        CodeBlock.Builder initializer = CodeBlock.builder();
        initializer.add("$L = new $T();\n", SOLVER_CONFIG_MAP_FIELD, HashMap.class);
        initializer.add("$T solverConfigIO = new $T();\n", SolverConfigIO.class, SolverConfigIO.class);
        for (Map.Entry<String, String> solverConfigXmlEntry : solverConfigToXml.entrySet()) {
            initializer.add("$L.put($S, solverConfigIO.read(new $T($S)));\n", SOLVER_CONFIG_MAP_FIELD,
            solverConfigXmlEntry.getKey(),
            StringReader.class,
            solverConfigXmlEntry.getValue());
        }
        builder.addStaticBlock(initializer.build());
        // ...
  });

but this causes a ZipFile object to appear in the image heap, making native image compilation fail:

Trace: Object was reached by
  reading field jdk.internal.module.ModuleReferences$JarModuleReader.jf of constant 
    jdk.internal.module.ModuleReferences$JarModuleReader@498e3ac2: jdk.internal.module.ModuleReferences$JarModuleReader@498e3ac2
  reading field java.util.concurrent.ConcurrentHashMap$Node.val of constant 
    java.util.concurrent.ConcurrentHashMap$Node@17d8a5ed: [module org.graalvm.nativeimage.objectfile, location=...
  indexing into array java.util.concurrent.ConcurrentHashMap$Node[]@4776bc5a: [Ljava.util.concurrent.ConcurrentHashMap$Node;@4776bc5a
  reading field java.util.concurrent.ConcurrentHashMap.table of constant 
    java.util.concurrent.ConcurrentHashMap@e71ccfd: {[module org.graalvm.nativeimage.base, location=file:...
  reading field jdk.internal.loader.BuiltinClassLoader.moduleToReader of constant 
    jdk.internal.loader.ClassLoaders$AppClassLoader@c818063: jdk.internal.loader.ClassLoaders$AppClassLoader@c818063
  reading field com.oracle.svm.core.hub.DynamicHubCompanion.classLoader of constant 
    com.oracle.svm.core.hub.DynamicHubCompanion@1027bef: com.oracle.svm.core.hub.DynamicHubCompanion@1027bef
  reading field java.lang.Class.companion of constant 
    java.lang.Class@76e6eb32: class com.oracle.svm.core.code.CodeInfoDecoderCounters
  manually triggered rescan

To workaround this, I basically needed to write my own POJO serializer that serializes an arbitrary POJO into the static initialization block of the generated class: https://github.com/Christopher-Chianelli/timefold-solver/blob/spring-boot-native/spring-integration/spring-boot-autoconfigure/src/main/java/ai/timefold/solver/spring/boot/autoconfigure/util/PojoInliner.java. For instance, given

public class BasicPojo {
    BasicPojo parentPojo;
    int id;
    String name;

    public BasicPojo() {
    }

    public BasicPojo(BasicPojo parentPojo, int id, String name) {
        this.parentPojo = parentPojo;
        this.id = id;
        this.name = name;
    }

    public BasicPojo getParentPojo() {
        return parentPojo;
    }

    public void setParentPojo(BasicPojo parentPojo) {
        this.parentPojo = parentPojo;
    }

    public int getId() {
        return id;
    }

    public void setId(int id) {
        this.id = id;
    }

    public String getName() {
        return name;
    }

    public void setName(String name) {
        this.name = name;
    }
}

To inline BasicPojo myField = new BasicPojo(new BasicPojo(null, 0, "parent"), 1, "child"), it would generate the following code:

static BasicPojo myField;
static {
    Map $pojoMap = new HashMap();
    BasicPojo $obj0;
    $obj0 = new BasicPojo();
    $pojoMap.put("$obj0", $obj0);
    $obj0.setId(1);
    $obj0.setName("child");
    BasicPojo $obj1;
    $obj1 = new BasicPojo();
    $pojoMap.put("$obj1", $obj1);
    $obj1.setId(0);
    $obj1.setName("parent");
    $obj1.setParentPojo(null);
    $obj0.setParentPojo(((BasicPojo) $pojoMap.get("$obj1")));
    myField = $obj0;
}

What I would want to do instead is something like this:

static public void registerSolverConfigs(Environment environment, ConfigurableListableBeanFactory beanFactory, Map<String, SolverConfig> solverConfigMap) {
    // ...
}
// ...
beanFactoryInitializationCode.addInitializer(new DefaultMethodReference(
            MethodSpec.methodBuilder("registerSolverConfigs")
                    .addModifiers(Modifier.PUBLIC)
                    .addModifiers(Modifier.STATIC)
                    .addParameter(Environment.class, "environment")
                    .addParameter(ConfigurableListableBeanFactory.class, "beanFactory")
                    .addFixedParameter(Map.class, "solverConfigMap", solverConfigMap)
                    .build(), MyClass.getName()));

which would be translated to something like this:

GeneratedClass generatedClass = generationContext.getGeneratedClasses().addForFeature("timefold-aot",
        builder -> {
            PojoInliner.inlineFields(builder,
              PojoInliner.field(Map.class, "solverConfigMap", solverConfigMap)
            );
            CodeBlock.Builder registerSolverConfigsMethod = CodeBlock.builder();
            registerSolverConfigsMethod.add("$T.$L($L, $L, $L);", MyClass.class, "registerSolverConfigs",  "environment",  "beanFactory", "solverConfigMap");
            builder.addMethod(MethodSpec.methodBuilder("registerSolverConfigs")
                .addModifiers(Modifier.PUBLIC)
                .addModifiers(Modifier.STATIC)
                .addParameter(Environment.class, "environment")
                .addParameter(ConfigurableListableBeanFactory.class, "beanFactory")
                .addCode(registerSolverConfigsMethod.build())
                .build());

            builder.build();
});

beanFactoryInitializationCode.addInitializer(new DefaultMethodReference(
    MethodSpec.methodBuilder("registerSolverConfigs")
        .addModifiers(Modifier.PUBLIC)
        .addModifiers(Modifier.STATIC)
        .addParameter(Environment.class, "environment")
        .addParameter(ConfigurableListableBeanFactory.class, "beanFactory")
        .build(), generatedClass.getName()));

[snicoll]

@Christopher-Chianelli I am afraid I don't understand what you're trying to do and what is currently missing in the core framework. Tools to generate code is available, I need to insist that you should not generate your own code using Spring AOT unless absolutely necessary.

Can you try to rephrase and, in particular, what you think need to change here.

[Christopher-Chianelli]

The ability to pass an actual object instance to an initializer.
Per the BeanFactoryInitializationCode Javadoc (https://docs.spring.io/spring-framework/docs/6.0.x/javadoc-api/org/springframework/beans/factory/aot/BeanFactoryInitializationCode.html#addInitializer(org.springframework.aot.generate.MethodReference))

Add an initializer method call. An initializer can use a flexible signature, using any of the following:

DefaultListableBeanFactory, or ConfigurableListableBeanFactory to use the bean factory.
ConfigurableEnvironment or Environment to access the environment.
ResourceLoader to load resources.

This works for code that only needs access to a beanFactory, environment, or resourceLoader. But in my particular case, I have a pre-constructed object that I need to pass to the initializer so it can register it to the beanFactory (i.e. my AOT initializer needs to do something like this):

public static void registerSolverConfig(ConfigurableListableBeanFactory beanFactory, SolverConfig solverConfig) {
    beanFactory.registerSingleton(DEFAULT_SOLVER_CONFIG_NAME, solverConfig);
}

Since SolverConfig is a complicated object (many different fields and nested configurations), it is far from trivial to inline it in generated code. In my ideal solution, I would NOT need to generate any code at all, and instead just do:

beanFactoryInitializationCode.addInitializer(new DefaultMethodReference(
    MethodSpec.methodBuilder("registerSolverConfig")
        .addModifiers(Modifier.PUBLIC)
        .addModifiers(Modifier.STATIC)
        .addParameter(ConfigurableListableBeanFactory.class, "beanFactory")
        .addFixedParameter(SolverConfig.class, "solverConfig", solverConfig)
        .build(), MyClass.class));

In particular, there a new method to DefaultMethodReference called addFixedParameter, which allows objects generated AOT to be passed to an initializer.

This needs to be done as beans registered by postProcessBeanFactory (which is where SolverConfig is registered when not in native mode) are unavailable in native mode, so I need to create an initializer that duplicates the postProcessBeanFactory logic (but uses the precomputed SolverConfig) in native mode.

[ge0ffrey]

@snicoll We're going to create a simple, isolated reproducer in a separate repo to demonstrate the issue clearly, out of respect for your time.

Let us get back to you here when we have that.

[snicoll]

That would be very much appreciated @ge0ffrey as the second comment is not helping I am afraid. Before you do all of that, please keep in mind:

1. Code generation is strictly targeted to libraries that have a specific feature and either want to optimize it in native or is too difficult to implement with only hints. A good example is a library that scan the classpath for files can do it at build-time and store the result in generated code.
2. I saw on the linked issue that "If we don't know how to make it happen, we need to ask the Spring [team]. I can not accept responsibility for maintaining our own serialization framework". I am not sure that's right as Spring itself shouldn't be involved in "serialization". There are a number of 3rd party serialization frameworks out there that you can use with Spring and we're not generating code for any of them.

[triceo]

2. I saw on the linked issue that "If we don't know how to make it happen, we need to ask the Spring [team]. I can not accept responsibility for maintaining our own serialization framework". I am not sure that's right as Spring itself shouldn't be involved in "serialization". There are a number of 3rd party serialization frameworks out there that you can use with Spring and we're not generating code for any of them.

@snicoll Hello. It's me who you're quoting here. In fact, I was making the same point you are - the approach we came up with is not acceptable, and there has to be a standard approach that doesn't rely on any kind of serialization. If we can't find the correct approach ourselves, we should ask the Spring team. I was making the point that nobody should be doing serialization here. :-)

Anyway, we'll work on getting this down to a minimal reproducer.

[Christopher-Chianelli]

Reproducer of why we need to do AOT processing: https://github.com/Christopher-Chianelli/issue-reproducer/tree/spring-entity-scanner-native/project
The short of it is new EntityScanner(context).scan(MyAnnotation.class) fails to return anything in native mode, even when all the classes are registered for reflection.

[snicoll]

I am so confused. So we went from code generation needs a way to provide a POJO because of some sort of serialization problem to "claspath scanning does not work".

If you are figuring out things as you explore, I'd argue that it's not yet time to raise an issue here (@triceo). Ironically enough, I described what the reproducer shows in my last comment:

A good example is a library that scan the classpath for files can do it at build-time and store the result in generated code.

There's no classpath in a native image so there's no way for whatever code that does that to work. If Timefold Solver does classpath scanning then you need to do something about it. You can use PersistenceManagedTypesBeanRegistrationAotProcessor as an inspiration.

If you need more help understanding how GraalVM and native image with Java work, please refer to the GraalVM support or StackOverflow.

[ge0ffrey]

I am sorry to hear that the reproducer was rejected, because the EntityScanner is not designed to support native mode. We'll will evaluate the PersistenceManagedTypesBeanRegistrationAotProcessor proposal.

The issue consists out of multiple issues, the EntityScanner reproducer isolated only step 1. Sorry for the confusion that mixing multiple issues brought.

In essence (in the outer issue), for performance reasons we want to read the solverConfig.xml at build time into a javabean (SolverConfig) and pass that to the runtime (not just in native mode, but also in non-native mode), to avoid parsing XML at bootstrap time. For serverless deployments. In Quakus we do this through their Recorder class, which takes care of all the boilerplate code for us.

Currently we got Spring Native working 🎉 , although we serialize and deserialize (=parse) the config to pass it from build time to bootstrap time. See https://github.com/TimefoldAI/timefold-solver/pull/609/files We 'll run some benchmarks what's the startup cost for this and welcome ideas how to improve this going forward. Rome wasn't build in a day :)

[snicoll]

I am sorry to hear that the reproducer was rejected, because the EntityScanner is not designed to support native mode.

I think you misunderstood why the issue was closed. The issue was closed because it isn't actionable. We don't use our issue tracker as a support forum and this issue clearly looks like one to me. As for EntityScanner it can't be designed to support native mode, i.e. you can't just run classpath scanning inside one of your method and expect AOT to figure that out and replace things. That being said, I am happy to explore options where you could run classpath scanning in a place where we can detect it and provide the result back using generated code.

read the solverConfig.xml at build time into a javabean (SolverConfig) and pass that to the runtime (not just in native mode, but also in non-native mode

That's totally doable, and even more so if you want to do that in both JVM and native mode. That's said, Spring does not require any build-time transformation on the JVM at all. Ignoring that, and focusing on AOT, there's nothing I can see from that issue that prevents doing that. If we want to make progress, you could create an issue describing what you'd like to happen, potentially linking to what Quarkus does for reference.

[Christopher-Chianelli]

For what it is worth, I think the actual issue is upstream from Spring; I thought MethodSpec and CodeBlock were from Spring because of their package, and then I learned it was a vendored version of JavaPoet. This issue was essentially a duplicate of square/javapoet#968 (support for a formatSpec for constructing transparent objects). For what Quarkus does, see https://quarkus.io/guides/writing-extensions#bytecode-recording . It essentially allows you to do this in their equivalent of an AOT processor

MyObject1 a = new MyObject1(...);
MyObject2 b = new MyObject2(...);
recorder.myMethod(a, b);

and it generates this code either inside a static initializer or the main method

MyObject1 a = new MyObject1();
a.setFieldA(...);
a.setFieldB(...);
MyObject2 b = new MyObject2();
b.setFieldA(...);
b.setFieldB(...);
recorder.myMethod(a, b);

So what I wanted to see from Spring was a way to do this in a way that was as easy as Quarkus.
I did submit an implementation to JavaPoet (see square/javapoet#999).
With the JavaPoet PR, it does become almost as easy as Quarkus (although it still involve very minor code generator):

ObjectInliner.getDefault().trustEverything();
CodeBlock.of("$T.$N($V, $V)", MyRecord.class, "methodName", new MyObject1(...). new MyObject2(...));

which is transformed into something like this:

MyRecord.methodName(
((MyObject1)((java.util.function.Supplier)(() -> {
               MyComplexConfig $$javapoet$MyObject1 = new MyObject1();
               $$javapoet$MyObject1.setParameter1(1);
               $$javapoet$MyObject1.setParameter2("abc");
               // ...
               return $$javapoet$MyObject1;
            })).get()),
((MyObject2)((java.util.function.Supplier)(() -> {
               MyComplexConfig $$javapoet$MyObject2 = new MyObject2();
               $$javapoet$MyObject2.setParameter1(1);
               $$javapoet$MyObject2.setParameter2("abc");
               // ...
               return $$javapoet$MyObject2;
            })).get())
);

(the generated code might be ugly, but allows the value to be put inside a fragment/incomplete block).

[Christopher-Chianelli]

As for why I thought deserialization did not work in generated code, it because if you register Class.class for reflection, any usage of ObjectMapper or XML deserialization will cause a native image error (see https://github.com/Christopher-Chianelli/issue-reproducer/tree/spring-aot-cannot-pass-read-xml/project ). Class.class was register for reflection since I was registered the types needed for reflection recursively (i.e. like this):

    private static void registerForReflection(ReflectionHints reflectionHints, Class<?> type, Set<Class<?>> visited) {
        if (type == null || visited.contains(type)) {
            return;
        }
        visited.add(type);
        reflectionHints.registerType(type,
                MemberCategory.INTROSPECT_PUBLIC_METHODS,
                MemberCategory.INTROSPECT_DECLARED_METHODS,
                MemberCategory.INTROSPECT_DECLARED_CONSTRUCTORS,
                MemberCategory.INTROSPECT_PUBLIC_CONSTRUCTORS,
                MemberCategory.PUBLIC_FIELDS,
                MemberCategory.DECLARED_FIELDS,
                MemberCategory.INVOKE_DECLARED_CONSTRUCTORS,
                MemberCategory.INVOKE_PUBLIC_CONSTRUCTORS,
                MemberCategory.INVOKE_DECLARED_METHODS,
                MemberCategory.INVOKE_PUBLIC_METHODS);
        for (Field field : type.getDeclaredFields()) {
            registerForReflection(reflectionHints, field.getType(), visited);
        }
        registerForReflection(reflectionHints, type.getSuperclass(), visited);
    }

This is needed because for XML/JSON deserialization to work, the type of the fields of the root type must also be registered for reflection (i.e. it not enough to register A.class for reflection if it contains instances of B.class; both A.class and B.class need to be registered for reflection).

As long as you don't register Class.class or ClassLoader.class for reflection (i.e. by checking for them in the if), you can use ObjectMapper in generated code (to use JAXB currently requires additional work, see https://stackoverflow.com/questions/77969978/how-to-use-jaxb-in-native-image-without-using-agent)

[ge0ffrey]

The issue was closed because it isn't actionable.

Ok. We have a working way to do Spring Native support in Timefold, so all is well :)

If you think it's valuable us to extract latest Christopher's comments into clean, separate RFE issues, we're happy to do that to help out. But for us it's also fine to ignore them, as these are not blockers on our side and we are happy with what we have now: Spring Native works.

Thank you for your time, Stéphane. Much appreciated :)