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.