Skip to content

Commit f76ac5b

Browse files
rstoyanchevrstoyanchev
committed
WebFlux support for @SessionAttributes
Issue: SPR-15887
1 parent bc470fc commit f76ac5b

File tree

10 files changed

+608
-69
lines changed

10 files changed

+608
-69
lines changed

‎spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/ControllerMethodResolver.java

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,8 @@ class ControllerMethodResolver {
9797
private final Map<ControllerAdviceBean, ExceptionHandlerMethodResolver> exceptionHandlerAdviceCache =
9898
new LinkedHashMap<>(64);
9999

100+
private final Map<Class<?>, SessionAttributesHandler> sessionAttributesHandlerCache = new ConcurrentHashMap<>(64);
101+
100102

101103
ControllerMethodResolver(ArgumentResolverConfigurer argumentResolvers,
102104
List<HttpMessageReader<?>> messageReaders, ReactiveAdapterRegistry reactiveRegistry,
@@ -154,6 +156,7 @@ private void addResolversTo(ArgumentResolverRegistrar registrar,
154156
registrar.addIfModelAttribute(() -> new ErrorsMethodArgumentResolver(reactiveRegistry));
155157
registrar.add(new ServerWebExchangeArgumentResolver(reactiveRegistry));
156158
registrar.add(new PrincipalArgumentResolver(reactiveRegistry));
159+
registrar.addIfRequestBody(readers -> new SessionStatusMethodArgumentResolver());
157160
registrar.add(new WebSessionArgumentResolver(reactiveRegistry));
158161

159162
// Custom...
@@ -315,6 +318,25 @@ public InvocableHandlerMethod getExceptionHandlerMethod(Throwable ex, HandlerMet
315318
return invocable;
316319
}
317320

321+
/**
322+
* Return the handler for the type-level {@code @SessionAttributes} annotation
323+
* based on the given controller method.
324+
*/
325+
public SessionAttributesHandler getSessionAttributesHandler(HandlerMethod handlerMethod) {
326+
Class<?> handlerType = handlerMethod.getBeanType();
327+
SessionAttributesHandler result = this.sessionAttributesHandlerCache.get(handlerType);
328+
if (result == null) {
329+
synchronized (this.sessionAttributesHandlerCache) {
330+
result = this.sessionAttributesHandlerCache.get(handlerType);
331+
if (result == null) {
332+
result = new SessionAttributesHandler(handlerType);
333+
this.sessionAttributesHandlerCache.put(handlerType, result);
334+
}
335+
}
336+
}
337+
return result;
338+
}
339+
318340

319341
/** Filter for {@link InitBinder @InitBinder} methods. */
320342
private static final ReflectionUtils.MethodFilter BINDER_METHODS = method ->
@@ -336,6 +358,7 @@ private static class ArgumentResolverRegistrar {
336358

337359
private final List<HandlerMethodArgumentResolver> result = new ArrayList<>();
338360

361+
339362
private ArgumentResolverRegistrar(ArgumentResolverConfigurer resolvers,
340363
List<HttpMessageReader<?>> messageReaders, boolean modelAttribute) {
341364

‎spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/InitBinderBindingContext.java

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,12 +23,15 @@
2323
import org.springframework.lang.Nullable;
2424
import org.springframework.util.Assert;
2525
import org.springframework.web.bind.annotation.InitBinder;
26+
import org.springframework.web.bind.support.SessionStatus;
27+
import org.springframework.web.bind.support.SimpleSessionStatus;
2628
import org.springframework.web.bind.support.WebBindingInitializer;
2729
import org.springframework.web.bind.support.WebExchangeDataBinder;
2830
import org.springframework.web.reactive.BindingContext;
2931
import org.springframework.web.reactive.HandlerResult;
3032
import org.springframework.web.reactive.result.method.SyncInvocableHandlerMethod;
3133
import org.springframework.web.server.ServerWebExchange;
34+
import org.springframework.web.server.WebSession;
3235

3336
/**
3437
* Extends {@link BindingContext} with {@code @InitBinder} method initialization.
@@ -43,6 +46,11 @@ class InitBinderBindingContext extends BindingContext {
4346
/* Simple BindingContext to help with the invoking @InitBinder methods */
4447
private final BindingContext binderMethodContext;
4548

49+
private final SessionStatus sessionStatus = new SimpleSessionStatus();
50+
51+
@Nullable
52+
private Runnable saveModelOperation;
53+
4654

4755
InitBinderBindingContext(@Nullable WebBindingInitializer initializer,
4856
List<SyncInvocableHandlerMethod> binderMethods) {
@@ -53,6 +61,15 @@ class InitBinderBindingContext extends BindingContext {
5361
}
5462

5563

64+
/**
65+
* Return the {@link SessionStatus} instance to use that can be used to
66+
* signal that session processing is complete.
67+
*/
68+
public SessionStatus getSessionStatus() {
69+
return this.sessionStatus;
70+
}
71+
72+
5673
@Override
5774
protected WebExchangeDataBinder initDataBinder(WebExchangeDataBinder dataBinder,
5875
ServerWebExchange exchange) {
@@ -87,4 +104,29 @@ private void invokeBinderMethod(WebExchangeDataBinder dataBinder,
87104
}
88105
}
89106

107+
/**
108+
* Provide the context required to apply {@link #saveModel()} after the
109+
* controller method has been invoked.
110+
*/
111+
public void setSessionContext(SessionAttributesHandler attributesHandler, WebSession session) {
112+
this.saveModelOperation = () -> {
113+
if (getSessionStatus().isComplete()) {
114+
attributesHandler.cleanupAttributes(session);
115+
}
116+
else {
117+
attributesHandler.storeAttributes(session, getModel().asMap());
118+
}
119+
};
120+
}
121+
122+
/**
123+
* Save model attributes in the session based on a type-level declarations
124+
* in an {@code @SessionAttributes} annotation.
125+
*/
126+
public void saveModel() {
127+
if (this.saveModelOperation != null) {
128+
this.saveModelOperation.run();
129+
}
130+
}
131+
90132
}

‎spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/ModelAttributeMethodArgumentResolver.java

Lines changed: 1 addition & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,11 @@
2121
import java.lang.reflect.Constructor;
2222
import java.util.List;
2323
import java.util.Map;
24-
import java.util.Optional;
2524

2625
import reactor.core.publisher.Mono;
2726
import reactor.core.publisher.MonoProcessor;
2827

2928
import org.springframework.beans.BeanUtils;
30-
import org.springframework.core.Conventions;
3129
import org.springframework.core.DefaultParameterNameDiscoverer;
3230
import org.springframework.core.MethodParameter;
3331
import org.springframework.core.ParameterNameDiscoverer;
@@ -39,7 +37,6 @@
3937
import org.springframework.ui.Model;
4038
import org.springframework.util.Assert;
4139
import org.springframework.util.ClassUtils;
42-
import org.springframework.util.StringUtils;
4340
import org.springframework.validation.BindingResult;
4441
import org.springframework.validation.Errors;
4542
import org.springframework.validation.annotation.Validated;
@@ -115,7 +112,7 @@ public Mono<Object> resolveArgument(
115112
() -> getClass().getSimpleName() + " does not support multi-value reactive type wrapper: " +
116113
parameter.getGenericParameterType());
117114

118-
String name = getAttributeName(parameter);
115+
String name = ModelInitializer.getNameForParameter(parameter);
119116
Mono<?> valueMono = prepareAttributeMono(name, valueType, context, exchange);
120117

121118
Map<String, Object> model = context.getModel().asMap();
@@ -150,13 +147,6 @@ public Mono<Object> resolveArgument(
150147
});
151148
}
152149

153-
private String getAttributeName(MethodParameter parameter) {
154-
return Optional.ofNullable(parameter.getParameterAnnotation(ModelAttribute.class))
155-
.filter(ann -> StringUtils.hasText(ann.value()))
156-
.map(ModelAttribute::value)
157-
.orElse(Conventions.getVariableNameForParameter(parameter));
158-
}
159-
160150
private Mono<?> prepareAttributeMono(String attributeName, ResolvableType attributeType,
161151
BindingContext context, ServerWebExchange exchange) {
162152

‎spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/ModelInitializer.java

Lines changed: 85 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,11 @@
1919
import java.util.ArrayList;
2020
import java.util.Arrays;
2121
import java.util.List;
22+
import java.util.Map;
2223
import java.util.Optional;
2324
import java.util.stream.Collectors;
2425

26+
import org.jetbrains.annotations.NotNull;
2527
import reactor.core.publisher.Mono;
2628

2729
import org.springframework.core.Conventions;
@@ -31,8 +33,10 @@
3133
import org.springframework.core.ResolvableType;
3234
import org.springframework.core.annotation.AnnotatedElementUtils;
3335
import org.springframework.lang.Nullable;
36+
import org.springframework.util.Assert;
3437
import org.springframework.util.StringUtils;
3538
import org.springframework.web.bind.annotation.ModelAttribute;
39+
import org.springframework.web.method.HandlerMethod;
3640
import org.springframework.web.reactive.BindingContext;
3741
import org.springframework.web.reactive.HandlerResult;
3842
import org.springframework.web.reactive.result.method.InvocableHandlerMethod;
@@ -47,40 +51,72 @@
4751
*/
4852
class ModelInitializer {
4953

54+
private final ControllerMethodResolver methodResolver;
55+
5056
private final ReactiveAdapterRegistry adapterRegistry;
5157

5258

53-
public ModelInitializer(ReactiveAdapterRegistry adapterRegistry) {
59+
public ModelInitializer(ControllerMethodResolver methodResolver, ReactiveAdapterRegistry adapterRegistry) {
60+
Assert.notNull(methodResolver, "ControllerMethodResolver is required");
61+
Assert.notNull(adapterRegistry, "ReactiveAdapterRegistry is required");
62+
this.methodResolver = methodResolver;
5463
this.adapterRegistry = adapterRegistry;
5564
}
5665

5766

5867
/**
59-
* Initialize the default model in the given {@code BindingContext} through
60-
* the {@code @ModelAttribute} methods and indicate when complete.
61-
* <p>This will wait for {@code @ModelAttribute} methods that return
62-
* {@code Mono<Void>} since those may be adding attributes asynchronously.
63-
* However if methods return async attributes, those will be added to the
64-
* model as-is and without waiting for them to be resolved.
65-
* @param bindingContext the BindingContext with the default model
66-
* @param attributeMethods the {@code @ModelAttribute} methods
68+
* Initialize the {@link org.springframework.ui.Model Model} based on a
69+
* (type-level) {@code @SessionAttributes} annotation and
70+
* {@code @ModelAttribute} methods.
71+
* @param handlerMethod the target controller method
72+
* @param bindingContext the context containing the model
6773
* @param exchange the current exchange
6874
* @return a {@code Mono} for when the model is populated.
6975
*/
7076
@SuppressWarnings("Convert2MethodRef")
71-
public Mono<Void> initModel(BindingContext bindingContext,
72-
List<InvocableHandlerMethod> attributeMethods, ServerWebExchange exchange) {
77+
public Mono<Void> initModel(HandlerMethod handlerMethod, InitBinderBindingContext bindingContext,
78+
ServerWebExchange exchange) {
79+
80+
List<InvocableHandlerMethod> modelMethods =
81+
this.methodResolver.getModelAttributeMethods(handlerMethod);
82+
83+
SessionAttributesHandler sessionAttributesHandler =
84+
this.methodResolver.getSessionAttributesHandler(handlerMethod);
85+
86+
if (!sessionAttributesHandler.hasSessionAttributes()) {
87+
return invokeModelAttributeMethods(bindingContext, modelMethods, exchange);
88+
}
89+
90+
return exchange.getSession()
91+
.flatMap(session -> {
92+
Map<String, Object> attributes = sessionAttributesHandler.retrieveAttributes(session);
93+
bindingContext.getModel().mergeAttributes(attributes);
94+
bindingContext.setSessionContext(sessionAttributesHandler, session);
95+
return invokeModelAttributeMethods(bindingContext, modelMethods, exchange)
96+
.doOnSuccess(aVoid -> {
97+
findModelAttributes(handlerMethod, sessionAttributesHandler).forEach(name -> {
98+
if (!bindingContext.getModel().containsAttribute(name)) {
99+
Object value = session.getRequiredAttribute(name);
100+
bindingContext.getModel().addAttribute(name, value);
101+
}
102+
});
103+
});
104+
});
105+
}
106+
107+
@NotNull
108+
private Mono<Void> invokeModelAttributeMethods(BindingContext bindingContext,
109+
List<InvocableHandlerMethod> modelMethods, ServerWebExchange exchange) {
73110

74111
List<Mono<HandlerResult>> resultList = new ArrayList<>();
75-
attributeMethods.forEach(invocable -> resultList.add(invocable.invoke(exchange, bindingContext)));
112+
modelMethods.forEach(invocable -> resultList.add(invocable.invoke(exchange, bindingContext)));
76113

77114
return Mono
78-
.zip(resultList, objectArray -> {
79-
return Arrays.stream(objectArray)
80-
.map(object -> handleResult(((HandlerResult) object), bindingContext))
81-
.collect(Collectors.toList());
82-
})
83-
.flatMap(completionList -> Mono.when(completionList));
115+
.zip(resultList, objectArray ->
116+
Arrays.stream(objectArray)
117+
.map(object -> handleResult(((HandlerResult) object), bindingContext))
118+
.collect(Collectors.toList()))
119+
.flatMap(Mono::when);
84120
}
85121

86122
private Mono<Void> handleResult(HandlerResult handlerResult, BindingContext bindingContext) {
@@ -109,4 +145,35 @@ private String getAttributeName(MethodParameter param) {
109145
.orElse(Conventions.getVariableNameForParameter(param));
110146
}
111147

148+
/** Find {@code @ModelAttribute} arguments also listed as {@code @SessionAttributes}. */
149+
private List<String> findModelAttributes(HandlerMethod handlerMethod,
150+
SessionAttributesHandler sessionAttributesHandler) {
151+
152+
List<String> result = new ArrayList<>();
153+
for (MethodParameter parameter : handlerMethod.getMethodParameters()) {
154+
if (parameter.hasParameterAnnotation(ModelAttribute.class)) {
155+
String name = getNameForParameter(parameter);
156+
Class<?> paramType = parameter.getParameterType();
157+
if (sessionAttributesHandler.isHandlerSessionAttribute(name, paramType)) {
158+
result.add(name);
159+
}
160+
}
161+
}
162+
return result;
163+
}
164+
165+
/**
166+
* Derive the model attribute name for the given method parameter based on
167+
* a {@code @ModelAttribute} parameter annotation (if present) or falling
168+
* back on parameter type based conventions.
169+
* @param parameter a descriptor for the method parameter
170+
* @return the derived name
171+
* @see Conventions#getVariableNameForParameter(MethodParameter)
172+
*/
173+
public static String getNameForParameter(MethodParameter parameter) {
174+
ModelAttribute ann = parameter.getParameterAnnotation(ModelAttribute.class);
175+
String name = (ann != null ? ann.value() : null);
176+
return (StringUtils.hasText(name) ? name : Conventions.getVariableNameForParameter(parameter));
177+
}
178+
112179
}

‎spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/RequestMappingHandlerAdapter.java

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -169,7 +169,7 @@ public void afterPropertiesSet() throws Exception {
169169
this.methodResolver = new ControllerMethodResolver(this.argumentResolverConfigurer,
170170
this.messageReaders, this.reactiveAdapterRegistry, this.applicationContext);
171171

172-
this.modelInitializer = new ModelInitializer(this.reactiveAdapterRegistry);
172+
this.modelInitializer = new ModelInitializer(this.methodResolver, this.reactiveAdapterRegistry);
173173
}
174174

175175

@@ -183,21 +183,20 @@ public Mono<HandlerResult> handle(ServerWebExchange exchange, Object handler) {
183183
HandlerMethod handlerMethod = (HandlerMethod) handler;
184184
Assert.state(this.methodResolver != null && this.modelInitializer != null, "Not initialized");
185185

186-
BindingContext bindingContext = new InitBinderBindingContext(
186+
InitBinderBindingContext bindingContext = new InitBinderBindingContext(
187187
getWebBindingInitializer(), this.methodResolver.getInitBinderMethods(handlerMethod));
188188

189-
List<InvocableHandlerMethod> modelAttributeMethods =
190-
this.methodResolver.getModelAttributeMethods(handlerMethod);
189+
InvocableHandlerMethod invocableMethod = this.methodResolver.getRequestMappingMethod(handlerMethod);
191190

192191
Function<Throwable, Mono<HandlerResult>> exceptionHandler =
193192
ex -> handleException(ex, handlerMethod, bindingContext, exchange);
194193

195194
return this.modelInitializer
196-
.initModel(bindingContext, modelAttributeMethods, exchange)
197-
.then(Mono.defer(() -> this.methodResolver.getRequestMappingMethod(handlerMethod)
198-
.invoke(exchange, bindingContext)
199-
.doOnNext(result -> result.setExceptionHandler(exceptionHandler))
200-
.onErrorResume(exceptionHandler)));
195+
.initModel(handlerMethod, bindingContext, exchange)
196+
.then(Mono.defer(() -> invocableMethod.invoke(exchange, bindingContext)))
197+
.doOnNext(result -> result.setExceptionHandler(exceptionHandler))
198+
.doOnNext(result -> bindingContext.saveModel())
199+
.onErrorResume(exceptionHandler);
201200
}
202201

203202
private Mono<HandlerResult> handleException(Throwable exception, HandlerMethod handlerMethod,

0 commit comments

Comments
 (0)