Skip to content

Commit fde243d

Browse files
wilkinsonasnicoll
wilkinsona
and
snicoll
committed
Hacking on SpringApplicationHooks
Co-authored-by: Stephane Nicoll <[email protected]>
1 parent 5b05fd9 commit fde243d

File tree

4 files changed

+377
-10
lines changed

4 files changed

+377
-10
lines changed

‎spring-boot-project/spring-boot/src/main/java/org/springframework/boot/AotProcessingHook.java

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+
}

‎spring-boot-project/spring-boot/src/main/java/org/springframework/boot/AotProcessor.java

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
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.Path;
21+
import java.nio.file.Paths;
22+
import java.util.Arrays;
23+
import java.util.Collections;
24+
import java.util.List;
25+
26+
import org.springframework.aot.generator.DefaultGeneratedTypeContext;
27+
import org.springframework.aot.generator.GeneratedType;
28+
import org.springframework.aot.generator.GeneratedTypeReference;
29+
import org.springframework.aot.hint.ExecutableMode;
30+
import org.springframework.aot.hint.RuntimeHints;
31+
import org.springframework.aot.hint.TypeReference;
32+
import org.springframework.aot.nativex.FileNativeConfigurationGenerator;
33+
import org.springframework.context.generator.ApplicationContextAotGenerator;
34+
import org.springframework.context.support.GenericApplicationContext;
35+
import org.springframework.javapoet.ClassName;
36+
import org.springframework.javapoet.JavaFile;
37+
import org.springframework.util.Assert;
38+
39+
/**
40+
* Entry point for AOT processing of a {@link SpringApplication}. Used for internal
41+
* purposes only.
42+
*
43+
* @author Stephane Nicoll
44+
* @author Andy Wilkinson
45+
* @since 3.0
46+
*/
47+
public class AotProcessor {
48+
49+
private final Class<?> application;
50+
51+
private final String[] applicationArgs;
52+
53+
private final Path sourceOutput;
54+
55+
private final Path resourceOutput;
56+
57+
/**
58+
* Create a new processor for the specified application and settings.
59+
* @param application the application main class
60+
* @param applicationArgs the arguments to provide to the main method
61+
* @param sourceOutput the location of generated sources
62+
* @param resourceOutput the location of generated resources
63+
*/
64+
public AotProcessor(Class<?> application, String[] applicationArgs, Path sourceOutput, Path resourceOutput) {
65+
this.application = application;
66+
this.applicationArgs = applicationArgs;
67+
this.sourceOutput = sourceOutput;
68+
this.resourceOutput = resourceOutput;
69+
}
70+
71+
/**
72+
* Trigger the processing of the application managed by this instance.
73+
*/
74+
public void process() {
75+
AotProcessingHook hook = new AotProcessingHook();
76+
SpringApplicationHooks.withHook(hook, this::callApplicationMainMethod);
77+
GenericApplicationContext applicationContext = hook.getApplicationContext();
78+
Assert.notNull(applicationContext, "No application context available after calling main method of '"
79+
+ this.application.getName() + "'. Does it run a SpringApplication?");
80+
performAotProcessing(applicationContext);
81+
}
82+
83+
private void callApplicationMainMethod() {
84+
try {
85+
this.application.getMethod("main", String[].class).invoke(null, new Object[] { this.applicationArgs });
86+
}
87+
catch (Exception ex) {
88+
throw new RuntimeException(ex);
89+
}
90+
}
91+
92+
private void performAotProcessing(GenericApplicationContext applicationContext) {
93+
DefaultGeneratedTypeContext generationContext = new DefaultGeneratedTypeContext(
94+
this.application.getPackageName(), (packageName) -> GeneratedType.of(ClassName.get(packageName,
95+
this.application.getSimpleName() + "__ApplicationContextInitializer")));
96+
ApplicationContextAotGenerator generator = new ApplicationContextAotGenerator();
97+
generator.generateApplicationContext(applicationContext, generationContext);
98+
99+
// Register reflection hint for entry point as we access it via reflection
100+
generationContext.runtimeHints().reflection()
101+
.registerType(GeneratedTypeReference.of(generationContext.getMainGeneratedType().getClassName()),
102+
(hint) -> hint.onReachableType(TypeReference.of(this.application)).withConstructor(
103+
Collections.emptyList(),
104+
(constructorHint) -> constructorHint.setModes(ExecutableMode.INVOKE)));
105+
106+
writeGeneratedSources(generationContext.toJavaFiles());
107+
writeGeneratedResources(generationContext.runtimeHints());
108+
}
109+
110+
private void writeGeneratedSources(List<JavaFile> sources) {
111+
for (JavaFile source : sources) {
112+
try {
113+
source.writeTo(this.sourceOutput);
114+
}
115+
catch (IOException ex) {
116+
throw new IllegalStateException("Failed to write " + source.typeSpec.name, ex);
117+
}
118+
}
119+
}
120+
121+
private void writeGeneratedResources(RuntimeHints hints) {
122+
FileNativeConfigurationGenerator generator = new FileNativeConfigurationGenerator(this.resourceOutput);
123+
generator.generate(hints);
124+
}
125+
126+
public static void main(String[] args) throws Exception {
127+
if (args.length < 3) {
128+
throw new IllegalArgumentException("Usage: " + AotProcessor.class.getName()
129+
+ " <applicationName> <sourceOutput> <resourceOutput> <originalArgs...>");
130+
}
131+
String applicationName = args[0];
132+
Path sourceOutput = Paths.get(args[1]);
133+
Path resourceOutput = Paths.get(args[2]);
134+
String[] applicationArgs = (args.length > 3) ? Arrays.copyOfRange(args, 3, args.length) : new String[0];
135+
136+
Class<?> application = Class.forName(applicationName);
137+
AotProcessor aotProcess = new AotProcessor(application, applicationArgs, sourceOutput, resourceOutput);
138+
aotProcess.process();
139+
}
140+
141+
}

‎spring-boot-project/spring-boot/src/main/java/org/springframework/boot/SpringApplication.java

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,7 @@ private Class<?> deduceMainApplicationClass() {
286286
* @return a running {@link ApplicationContext}
287287
*/
288288
public ConfigurableApplicationContext run(String... args) {
289+
SpringApplicationHooks.hooks().preRun(this);
289290
long startTime = System.nanoTime();
290291
DefaultBootstrapContext bootstrapContext = createBootstrapContext();
291292
ConfigurableApplicationContext context = null;
@@ -300,27 +301,32 @@ public ConfigurableApplicationContext run(String... args) {
300301
context = createApplicationContext();
301302
context.setApplicationStartup(this.applicationStartup);
302303
prepareContext(bootstrapContext, context, environment, listeners, applicationArguments, printedBanner);
303-
refreshContext(context);
304-
afterRefresh(context, applicationArguments);
305-
Duration timeTakenToStartup = Duration.ofNanos(System.nanoTime() - startTime);
306-
if (this.logStartupInfo) {
307-
new StartupInfoLogger(this.mainApplicationClass).logStarted(getApplicationLog(), timeTakenToStartup);
304+
if (refreshContext(context)) {
305+
afterRefresh(context, applicationArguments);
306+
Duration timeTakenToStartup = Duration.ofNanos(System.nanoTime() - startTime);
307+
if (this.logStartupInfo) {
308+
new StartupInfoLogger(this.mainApplicationClass).logStarted(getApplicationLog(),
309+
timeTakenToStartup);
310+
}
311+
listeners.started(context, timeTakenToStartup);
312+
callRunners(context, applicationArguments);
308313
}
309-
listeners.started(context, timeTakenToStartup);
310-
callRunners(context, applicationArguments);
311314
}
312315
catch (Throwable ex) {
313316
handleRunFailure(context, ex, listeners);
314317
throw new IllegalStateException(ex);
315318
}
316319
try {
317-
Duration timeTakenToReady = Duration.ofNanos(System.nanoTime() - startTime);
318-
listeners.ready(context, timeTakenToReady);
320+
if (context.isRunning()) {
321+
Duration timeTakenToReady = Duration.ofNanos(System.nanoTime() - startTime);
322+
listeners.ready(context, timeTakenToReady);
323+
}
319324
}
320325
catch (Throwable ex) {
321326
handleRunFailure(context, ex, null);
322327
throw new IllegalStateException(ex);
323328
}
329+
SpringApplicationHooks.hooks().postRun(this, context);
324330
return context;
325331
}
326332

@@ -395,11 +401,15 @@ private void prepareContext(DefaultBootstrapContext bootstrapContext, Configurab
395401
listeners.contextLoaded(context);
396402
}
397403

398-
private void refreshContext(ConfigurableApplicationContext context) {
404+
private boolean refreshContext(ConfigurableApplicationContext context) {
405+
if (!SpringApplicationHooks.hooks().preRefresh(this, context)) {
406+
return false;
407+
}
399408
if (this.registerShutdownHook) {
400409
shutdownHook.registerApplicationContext(context);
401410
}
402411
refresh(context);
412+
return true;
403413
}
404414

405415
private void configureHeadlessProperty() {

0 commit comments

Comments
 (0)