Spring ConfigurationClassEnhancer does not use correct ClassLoader when called multiple times

Discovered whilst looking into https://github.com/spring-projects/spring-boot/issues/39733

The following test will pass with Spring Framework 5.3.36 but fail with Spring Framework 6.0.12:

package org.springframework.context.annotation;

import static org.assertj.core.api.Assertions.assertThat;

import java.io.IOException;
import java.io.InputStream;
import java.security.SecureClassLoader;

import org.junit.jupiter.api.Test;
import org.springframework.util.StreamUtils;

class ConfigurationClassEnhancerTests {

    @Test
    void test() throws Exception {
        ConfigurationClassEnhancer configurationClassEnhancer = new ConfigurationClassEnhancer();
        ClassLoader parentClassLoader = getClass().getClassLoader();
        MyClassLoader classLoader = new MyClassLoader(parentClassLoader);
        Class<?> myClass = parentClassLoader.loadClass(MyConfig.class.getName());
        configurationClassEnhancer.enhance(myClass, parentClassLoader);
        Class<?> myReloadedClass = classLoader.loadClass(MyConfig.class.getName());
        Class<?> enhancedReloadedClass = configurationClassEnhancer.enhance(myReloadedClass, classLoader);
        assertThat(enhancedReloadedClass.getClassLoader()).isEqualTo(classLoader);
    }

    @Configuration
    static class MyConfig {

        @Bean
        public String myBean() {
            return "bean";
        }

    }

    static class MyClassLoader extends SecureClassLoader {

        MyClassLoader(ClassLoader parent) {
            super(parent);
        }

        protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
            if (name.contains("MyConfig")) {
                String path = name.replace('.', '/').concat(".class");
                try (InputStream in = super.getResourceAsStream(path)) {
                    byte[] bytes = StreamUtils.copyToByteArray(in);
                    if (bytes.length > 0) {
                        return defineClass(name, bytes, 0, bytes.length);
                    }
                } catch (IOException ex) {
                    throw new IllegalStateException(ex);
                }
            }
            return super.loadClass(name, resolve);
        }

    }

}

I haven't managed to exactly track down why, but I suspect that b31a15851e7aaabf3629cc101d285e751e535927 might have introduced the problem. Perhaps this key?

Comment From: jhoeller

This is caused by the enhancer.setAttemptLoad(true) setting which we apply as of 6.0 for AOT purposes. Without that flag, the test passes on 6.x as well.

AttemptLoad seems to load a class with the same name irrespective of the specified ClassLoader. Maybe we need to conditionally set it to false depending on the ClassLoader arrangement.

Comment From: philwebb

I guess we're setting setAttemptLoad(true) to pick up any class files that we generated during the AOT phase. In this case, I think the first run is generating class bytecode in the parent loader and the second run is then picking that up. Ideally we want a way to pick up classes that were AOT generated but not those that were generated in-memory.

Comment From: jhoeller

Since we are primarily talking about reload scenarios, maybe we should simply set AttemptLoad to false if the current configuration class passes SmartClassLoader.isClassReloadable. That would cover Boot's scenario, wouldn't it?

In general, cleanly differentiating between pre-loaded classes and pre-defined classes is impossible for an arbitrary target ClassLoader. The ClassLoader API is trying hard to treat those the same, with only the protected findLoadedClass allowing to differentiate within the ClassLoader implementation itself. We got that problem in ContextTypeMatchClassLoader already, and it would be good to avoid any further variations of this... not least of it all: for compatibility with the module system.

Comment From: jhoeller

A modified test case where the custom ClassLoader implements SmartClassLoader.isClassReloadable passes fine with a corresponding check in ConfigurationClassEnhancer. I intend to go with this for 6.1.10 and 6.0.23.

Comment From: philwebb

Sounds good. Once that's in I can update Boot and test the read sample again.

Comment From: philwebb

Actually we already implement that method so we may not need any Boot changes.

Comment From: jhoeller

I would expect Boot's RestartClassLoader to be naturally compliant already, indeed. Alright, about to push this to 6.1.x here.

Comment From: philwebb

Tested and all looks good 👍