From b734156f3237dfaeef34dfb86a3819f28f8d1c9d Mon Sep 17 00:00:00 2001 From: rstoyanchev Date: Mon, 5 Aug 2024 15:53:12 +0300 Subject: [PATCH] Update WebFlux fragment stream rendering Thymeleaf has its own special handling for SSE that gets in the way of fragment rendering. This is why we need to set the response content-type before streaming, and then pass text/html to the View for rendering each fragment. See gh-33194 --- .../view/ViewResolutionResultHandler.java | 109 ++++++++++++------ 1 file changed, 73 insertions(+), 36 deletions(-) diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/result/view/ViewResolutionResultHandler.java b/spring-webflux/src/main/java/org/springframework/web/reactive/result/view/ViewResolutionResultHandler.java index d2c3a5c5fb..f06e934fad 100644 --- a/spring-webflux/src/main/java/org/springframework/web/reactive/result/view/ViewResolutionResultHandler.java +++ b/spring-webflux/src/main/java/org/springframework/web/reactive/result/view/ViewResolutionResultHandler.java @@ -267,20 +267,24 @@ public class ViewResolutionResultHandler extends HandlerResultHandlerSupport imp Mono.just(Collections.singletonList((View) view))); } else if (FragmentsRendering.class.isAssignableFrom(clazz)) { + ServerHttpResponse response = exchange.getResponse(); FragmentsRendering render = (FragmentsRendering) returnValue; HttpStatusCode status = render.status(); if (status != null) { - exchange.getResponse().setStatusCode(status); + response.setStatusCode(status); } - exchange.getResponse().getHeaders().putAll(render.headers()); + response.getHeaders().putAll(render.headers()); bindingContext.updateModel(exchange); StreamHandler streamHandler = getStreamHandler(exchange); + if (streamHandler != null) { + streamHandler.updateResponse(exchange); + } Flux> renderFlux = render.fragments().concatMap(fragment -> renderFragment(fragment, streamHandler, locale, bindingContext, exchange)); - return exchange.getResponse().writeAndFlushWith(renderFlux); + return response.writeAndFlushWith(renderFlux); } else if (Model.class.isAssignableFrom(clazz)) { model.addAllAttributes(((Model) returnValue).asMap()); @@ -299,7 +303,7 @@ public class ViewResolutionResultHandler extends HandlerResultHandlerSupport imp viewsMono = resolveViews(getDefaultViewName(exchange), locale); } bindingContext.updateModel(exchange); - return viewsMono.flatMap(views -> render(views, model.asMap(), bindingContext, exchange)); + return viewsMono.flatMap(views -> render(views, model.asMap(), null, bindingContext, exchange)); }); } @@ -346,10 +350,16 @@ public class ViewResolutionResultHandler extends HandlerResultHandlerSupport imp Mono.just(List.of(fragment.view())) : resolveViews(fragment.viewName() != null ? fragment.viewName() : getDefaultViewName(exchange), locale)); - return selectedViews.flatMap(views -> render(views, fragment.model(), bindingContext, mutatedExchange)) - .then(Mono.fromSupplier(() -> (streamHandler != null ? - streamHandler.format(response.getBodyFlux(), fragment, exchange) : - response.getBodyFlux()))); + Map model = fragment.model(); + + if (streamHandler != null) { + return selectedViews.flatMap(views -> render(views, model, MediaType.TEXT_HTML, bindingContext, mutatedExchange)) + .then(Mono.fromSupplier(() -> streamHandler.format(response.getBodyFlux(), fragment, exchange))); + } + else { + return selectedViews.flatMap(views -> render(views, model, null, bindingContext, mutatedExchange)) + .then(Mono.fromSupplier(response::getBodyFlux)); + } } @Nullable @@ -369,7 +379,8 @@ public class ViewResolutionResultHandler extends HandlerResultHandlerSupport imp .orElseGet(() -> Conventions.getVariableNameForParameter(returnType)); } - private Mono render(List views, Map model, + private Mono render( + List views, Map model, @Nullable MediaType bestMediaType, BindingContext bindingContext, ServerWebExchange exchange) { for (View view : views) { @@ -378,19 +389,20 @@ public class ViewResolutionResultHandler extends HandlerResultHandlerSupport imp } } List mediaTypes = getMediaTypes(views); - MediaType bestMediaType; - try { - bestMediaType = selectMediaType(exchange, () -> mediaTypes); - } - catch (NotAcceptableStatusException ex) { - HttpStatusCode statusCode = exchange.getResponse().getStatusCode(); - if (statusCode != null && statusCode.isError()) { - if (logger.isDebugEnabled()) { - logger.debug("Ignoring error response content (if any). " + ex.getReason()); - } - return Mono.empty(); + if (bestMediaType == null) { + try { + bestMediaType = selectMediaType(exchange, () -> mediaTypes); + } + catch (NotAcceptableStatusException ex) { + HttpStatusCode statusCode = exchange.getResponse().getStatusCode(); + if (statusCode != null && statusCode.isError()) { + if (logger.isDebugEnabled()) { + logger.debug("Ignoring error response content (if any). " + ex.getReason()); + } + return Mono.empty(); + } + throw ex; } - throw ex; } if (bestMediaType != null) { for (View view : views) { @@ -427,15 +439,23 @@ public class ViewResolutionResultHandler extends HandlerResultHandlerSupport imp @Nullable private Flux bodyFlux; - private final HttpHeaders headers; + @Nullable + private HttpHeaders headers; BodySavingResponse(ServerHttpResponse delegate) { super(delegate); - this.headers = new HttpHeaders(delegate.getHeaders()); // Ignore header changes } @Override public HttpHeaders getHeaders() { + if (!super.getHeaders().containsKey(HttpHeaders.CONTENT_TYPE)) { + return super.getHeaders(); + } + // Content-type is set, ignore further updates + if (this.headers == null) { + this.headers = new HttpHeaders(); + this.headers.putAll(super.getHeaders()); + } return this.headers; } @@ -468,6 +488,11 @@ public class ViewResolutionResultHandler extends HandlerResultHandlerSupport imp */ boolean supports(ServerHttpRequest request); + /** + * Update the response before streaming, e.g. to set the content-type. + */ + void updateResponse(ServerWebExchange exchange); + /** * Format the given fragment. * @param fragmentContent the fragment serialized to data buffers @@ -476,7 +501,6 @@ public class ViewResolutionResultHandler extends HandlerResultHandlerSupport imp * @return the formatted fragment */ Flux format(Flux fragmentContent, Fragment fragment, ServerWebExchange exchange); - } @@ -492,20 +516,14 @@ public class ViewResolutionResultHandler extends HandlerResultHandlerSupport imp } @Override - public Flux format( - Flux fragmentContent, Fragment fragment, ServerWebExchange exchange) { - + public void updateResponse(ServerWebExchange exchange) { + MediaType mediaType = MediaType.TEXT_EVENT_STREAM; Charset charset = getCharset(exchange.getRequest()); - DataBufferFactory bufferFactory = exchange.getResponse().bufferFactory(); - - String eventLine = fragment.viewName() != null ? "event:" + fragment.viewName() + "\n" : ""; - - return Flux.concat( - Flux.just(encodeText(eventLine + "data:", charset, bufferFactory)), - fragmentContent, - Flux.just(encodeText("\n\n", charset, bufferFactory))); + mediaType = (charset != null ? new MediaType(mediaType, charset) : mediaType); + exchange.getResponse().getHeaders().setContentType(mediaType); } + @Nullable private Charset getCharset(ServerHttpRequest request) { for (MediaType mediaType : request.getHeaders().getAccept()) { if (mediaType.isCompatibleWith(MediaType.TEXT_EVENT_STREAM)) { @@ -515,7 +533,26 @@ public class ViewResolutionResultHandler extends HandlerResultHandlerSupport imp break; } } - return StandardCharsets.UTF_8; + return null; + } + + @Override + public Flux format( + Flux fragmentContent, Fragment fragment, ServerWebExchange exchange) { + + Charset charset = StandardCharsets.UTF_8; + MediaType contentType = exchange.getResponse().getHeaders().getContentType(); + if (contentType != null && contentType.getCharset() != null) { + charset = contentType.getCharset(); + } + + DataBufferFactory bufferFactory = exchange.getResponse().bufferFactory(); + + String eventLine = fragment.viewName() != null ? "event:" + fragment.viewName() + "\n" : ""; + DataBuffer prefix = encodeText(eventLine + "data:", charset, bufferFactory); + DataBuffer suffix = encodeText("\n\n", charset, bufferFactory); + + return Flux.concat(Flux.just(prefix), fragmentContent, Flux.just(suffix)); } private DataBuffer encodeText(String text, Charset charset, DataBufferFactory bufferFactory) {