Skip to content

Commit 81b2d71

Browse files
wilkinsonasnicoll
andcommitted
Introduce Hook-based AOT processing
Closes gh- Co-authored-by: Stephane Nicoll <[email protected]>
1 parent 9778fe0 commit 81b2d71

File tree

3 files changed

+330
-0
lines changed

3 files changed

+330
-0
lines changed
Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
/*
2+
* Copyright 2012-2022 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.boot;
18+
19+
import org.springframework.boot.SpringApplicationHooks.Hook;
20+
import org.springframework.context.ConfigurableApplicationContext;
21+
import org.springframework.context.support.GenericApplicationContext;
22+
import org.springframework.util.Assert;
23+
24+
/**
25+
* A {@link Hook} used to prevent standard refresh of the application's context, ready for
26+
* subsequent {@link GenericApplicationContext#refreshForAotProcessing() AOT processing}.
27+
*
28+
* @author Andy Wilkinson
29+
*/
30+
class AotProcessingHook implements Hook {
31+
32+
private GenericApplicationContext context;
33+
34+
@Override
35+
public boolean preRefresh(SpringApplication application, ConfigurableApplicationContext context) {
36+
Assert.isInstanceOf(GenericApplicationContext.class, context,
37+
() -> "AOT processing requires a GenericApplicationContext but got at " + context.getClass().getName());
38+
this.context = (GenericApplicationContext) context;
39+
return false;
40+
}
41+
42+
GenericApplicationContext getApplicationContext() {
43+
return this.context;
44+
}
45+
46+
}
Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
/*
2+
* Copyright 2012-2022 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.boot;
18+
19+
import java.io.IOException;
20+
import java.nio.file.Files;
21+
import java.nio.file.Path;
22+
import java.nio.file.Paths;
23+
import java.util.ArrayList;
24+
import java.util.Arrays;
25+
import java.util.Collections;
26+
import java.util.List;
27+
28+
import org.springframework.aot.generator.DefaultGeneratedTypeContext;
29+
import org.springframework.aot.generator.GeneratedType;
30+
import org.springframework.aot.generator.GeneratedTypeReference;
31+
import org.springframework.aot.hint.ExecutableMode;
32+
import org.springframework.aot.hint.RuntimeHints;
33+
import org.springframework.aot.hint.TypeReference;
34+
import org.springframework.aot.nativex.FileNativeConfigurationWriter;
35+
import org.springframework.context.generator.ApplicationContextAotGenerator;
36+
import org.springframework.context.support.GenericApplicationContext;
37+
import org.springframework.javapoet.ClassName;
38+
import org.springframework.javapoet.JavaFile;
39+
import org.springframework.util.Assert;
40+
41+
/**
42+
* Entry point for AOT processing of a {@link SpringApplication}.
43+
* <p>
44+
* <strong>For internal use only.</strong>
45+
*
46+
* @author Stephane Nicoll
47+
* @author Andy Wilkinson
48+
* @since 3.0
49+
*/
50+
public class AotProcessor {
51+
52+
private final Class<?> application;
53+
54+
private final String[] applicationArgs;
55+
56+
private final Path sourceOutput;
57+
58+
private final Path resourceOutput;
59+
60+
/**
61+
* Create a new processor for the specified application and settings.
62+
* @param application the application main class
63+
* @param applicationArgs the arguments to provide to the main method
64+
* @param sourceOutput the location of generated sources
65+
* @param resourceOutput the location of generated resources
66+
*/
67+
public AotProcessor(Class<?> application, String[] applicationArgs, Path sourceOutput, Path resourceOutput) {
68+
this.application = application;
69+
this.applicationArgs = applicationArgs;
70+
this.sourceOutput = sourceOutput;
71+
this.resourceOutput = resourceOutput;
72+
}
73+
74+
/**
75+
* Trigger the processing of the application managed by this instance.
76+
*/
77+
public void process() {
78+
AotProcessingHook hook = new AotProcessingHook();
79+
SpringApplicationHooks.withHook(hook, this::callApplicationMainMethod);
80+
GenericApplicationContext applicationContext = hook.getApplicationContext();
81+
Assert.notNull(applicationContext, "No application context available after calling main method of '"
82+
+ this.application.getName() + "'. Does it run a SpringApplication?");
83+
performAotProcessing(applicationContext);
84+
}
85+
86+
private void callApplicationMainMethod() {
87+
try {
88+
this.application.getMethod("main", String[].class).invoke(null, new Object[] { this.applicationArgs });
89+
}
90+
catch (Exception ex) {
91+
throw new RuntimeException(ex);
92+
}
93+
}
94+
95+
private void performAotProcessing(GenericApplicationContext applicationContext) {
96+
DefaultGeneratedTypeContext generationContext = new DefaultGeneratedTypeContext(
97+
this.application.getPackageName(), (packageName) -> GeneratedType.of(ClassName.get(packageName,
98+
this.application.getSimpleName() + "__ApplicationContextInitializer")));
99+
ApplicationContextAotGenerator generator = new ApplicationContextAotGenerator();
100+
generator.generateApplicationContext(applicationContext, generationContext);
101+
102+
// Register reflection hint for entry point as we access it via reflection
103+
generationContext.runtimeHints().reflection()
104+
.registerType(GeneratedTypeReference.of(generationContext.getMainGeneratedType().getClassName()),
105+
(hint) -> hint.onReachableType(TypeReference.of(this.application)).withConstructor(
106+
Collections.emptyList(),
107+
(constructorHint) -> constructorHint.setModes(ExecutableMode.INVOKE)));
108+
109+
writeGeneratedSources(generationContext.toJavaFiles());
110+
writeGeneratedResources(generationContext.runtimeHints());
111+
writeNativeImageProperties();
112+
}
113+
114+
private void writeGeneratedSources(List<JavaFile> sources) {
115+
for (JavaFile source : sources) {
116+
try {
117+
source.writeTo(this.sourceOutput);
118+
}
119+
catch (IOException ex) {
120+
throw new IllegalStateException("Failed to write " + source.typeSpec.name, ex);
121+
}
122+
}
123+
}
124+
125+
private void writeGeneratedResources(RuntimeHints hints) {
126+
FileNativeConfigurationWriter writer = new FileNativeConfigurationWriter(this.resourceOutput);
127+
writer.write(hints);
128+
}
129+
130+
private void writeNativeImageProperties() {
131+
List<String> args = new ArrayList<>();
132+
args.add("-H:Class=" + this.application.getName());
133+
args.add("--allow-incomplete-classpath");
134+
args.add("--report-unsupported-elements-at-runtime");
135+
args.add("--no-fallback");
136+
args.add("--install-exit-handlers");
137+
StringBuilder sb = new StringBuilder();
138+
sb.append("Args = ");
139+
sb.append(String.join(String.format(" \\%n"), args));
140+
Path file = this.resourceOutput.resolve("META-INF/native-image/native-image.properties");
141+
try {
142+
if (!Files.exists(file)) {
143+
Files.createDirectories(file.getParent());
144+
Files.createFile(file);
145+
}
146+
Files.writeString(file, sb.toString());
147+
}
148+
catch (IOException ex) {
149+
throw new IllegalStateException("Failed to write native-image properties", ex);
150+
}
151+
}
152+
153+
public static void main(String[] args) throws Exception {
154+
if (args.length < 3) {
155+
throw new IllegalArgumentException("Usage: " + AotProcessor.class.getName()
156+
+ " <applicationName> <sourceOutput> <resourceOutput> <originalArgs...>");
157+
}
158+
String applicationName = args[0];
159+
Path sourceOutput = Paths.get(args[1]);
160+
Path resourceOutput = Paths.get(args[2]);
161+
String[] applicationArgs = (args.length > 3) ? Arrays.copyOfRange(args, 3, args.length) : new String[0];
162+
163+
Class<?> application = Class.forName(applicationName);
164+
AotProcessor aotProcess = new AotProcessor(application, applicationArgs, sourceOutput, resourceOutput);
165+
aotProcess.process();
166+
}
167+
168+
}
Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
/*
2+
* Copyright 2012-2022 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.boot;
18+
19+
import java.nio.file.Path;
20+
import java.util.function.Consumer;
21+
22+
import org.junit.jupiter.api.BeforeEach;
23+
import org.junit.jupiter.api.Test;
24+
import org.junit.jupiter.api.io.TempDir;
25+
26+
import org.springframework.context.annotation.Configuration;
27+
28+
import static org.assertj.core.api.Assertions.assertThat;
29+
import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
30+
31+
/**
32+
* Tests for {@link AotProcessor}.
33+
*
34+
* @author Stephane Nicoll
35+
*/
36+
class AotProcessorTests {
37+
38+
@BeforeEach
39+
void setup() {
40+
SampleApplication.argsHolder = null;
41+
}
42+
43+
@Test
44+
void processApplicationInvokesRunMethod(@TempDir Path directory) {
45+
String[] arguments = new String[] { "1", "2" };
46+
AotProcessor processor = new AotProcessor(SampleApplication.class, arguments, directory.resolve("source"),
47+
directory.resolve("resource"));
48+
processor.process();
49+
assertThat(SampleApplication.argsHolder).isEqualTo(arguments);
50+
assertThat(directory).satisfies(hasGeneratedAssetsForSampleApplication());
51+
}
52+
53+
@Test
54+
void processApplicationWithMainMethodThatDoesNotRun(@TempDir Path directory) {
55+
AotProcessor processor = new AotProcessor(BrokenApplication.class, new String[0], directory.resolve("source"),
56+
directory.resolve("resource"));
57+
assertThatIllegalArgumentException().isThrownBy(processor::process)
58+
.withMessageContaining("Does it run a SpringApplication?");
59+
assertThat(directory).isEmptyDirectory();
60+
}
61+
62+
@Test
63+
void invokeMainParseArgumentsAndInvokesRunMethod(@TempDir Path directory) throws Exception {
64+
String[] mainArguments = new String[] { SampleApplication.class.getName(),
65+
directory.resolve("source").toString(), directory.resolve("resource").toString(), "1", "2" };
66+
AotProcessor.main(mainArguments);
67+
assertThat(SampleApplication.argsHolder).containsExactly("1", "2");
68+
assertThat(directory).satisfies(hasGeneratedAssetsForSampleApplication());
69+
}
70+
71+
@Test
72+
void invokeMainWithMissingArguments() {
73+
assertThatIllegalArgumentException().isThrownBy(() -> AotProcessor.main(new String[] { "Test" }))
74+
.withMessageContaining("Usage:");
75+
}
76+
77+
private Consumer<Path> hasGeneratedAssetsForSampleApplication() {
78+
return (directory) -> {
79+
assertThat(directory
80+
.resolve("source/org/springframework/boot/SampleApplication__ApplicationContextInitializer.java"))
81+
.exists().isRegularFile();
82+
assertThat(directory.resolve("resource/META-INF/native-image/reflect-config.json")).exists()
83+
.isRegularFile();
84+
Path nativeImagePropertiesFile = directory
85+
.resolve("resource/META-INF/native-image/native-image.properties");
86+
assertThat(nativeImagePropertiesFile).exists().isRegularFile().hasContent("""
87+
Args = -H:Class=org.springframework.boot.AotProcessorTests$SampleApplication \\
88+
--allow-incomplete-classpath \\
89+
--report-unsupported-elements-at-runtime \\
90+
--no-fallback \\
91+
--install-exit-handlers
92+
""");
93+
};
94+
}
95+
96+
@Configuration(proxyBeanMethods = false)
97+
public static class SampleApplication {
98+
99+
public static String[] argsHolder;
100+
101+
public static void main(String[] args) {
102+
argsHolder = args;
103+
SpringApplication.run(SampleApplication.class, args);
104+
}
105+
106+
}
107+
108+
public static class BrokenApplication {
109+
110+
public static void main(String[] args) {
111+
// Does not run an application
112+
}
113+
114+
}
115+
116+
}

0 commit comments

Comments
 (0)