Skip to content

Commit

Permalink
Support Flux<ServerSentEvent<Fragment>> in WebFlux
Browse files Browse the repository at this point in the history 
Closes gh-33975
  • Loading branch information
@rstoyanchev
rstoyanchev committed Dec 11, 2024
1 parent 640e570 commit 3b95d2c
Show file tree
Hide file tree
Showing 3 changed files with 116 additions and 41 deletions.
85 changes: 59 additions & 26 deletions ...c/main/java/org/springframework/web/reactive/result/view/ViewResolutionResultHandler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatusCode;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
Expand Down Expand Up @@ -101,7 +102,7 @@ public class ViewResolutionResultHandler extends HandlerResultHandlerSupport imp

private final List<View> defaultViews = new ArrayList<>(4);

private final List<StreamHandler> streamHandlers = List.of(new SseStreamHandler());
private final SseStreamHandler sseHandler = new SseStreamHandler();


/**
Expand Down Expand Up @@ -175,7 +176,7 @@ public boolean supports(HandlerResult result) {
returnType = returnType.getNested(2);

if (adapter.isMultiValue()) {
return Fragment.class.isAssignableFrom(type);
return (Fragment.class.isAssignableFrom(type) || isSseFragmentStream(returnType));
}
}

Expand All @@ -194,8 +195,13 @@ private boolean hasModelAnnotation(MethodParameter parameter) {
}

private static boolean isFragmentCollection(ResolvableType returnType) {
Class<?> clazz = returnType.resolve(Object.class);
return (Collection.class.isAssignableFrom(clazz) && Fragment.class.equals(returnType.getNested(2).resolve()));
return (Collection.class.isAssignableFrom(returnType.resolve(Object.class)) &&
Fragment.class.equals(returnType.getNested(2).resolve()));
}

private static boolean isSseFragmentStream(ResolvableType returnType) {
return (ServerSentEvent.class.equals(returnType.resolve()) &&
Fragment.class.equals(returnType.getNested(2).resolve()));
}

@Override
Expand All @@ -204,9 +210,15 @@ public Mono<Void> handleResult(ServerWebExchange exchange, HandlerResult result)
Mono<Object> valueMono;
ResolvableType valueType;
ReactiveAdapter adapter = getAdapter(result);
BindingContext bindingContext = result.getBindingContext();
Locale locale = LocaleContextHolder.getLocale(exchange.getLocaleContext());

if (adapter != null) {
if (adapter.isMultiValue()) {
if (isSseFragmentStream(result.getReturnType().getNested(2))) {
return handleSseFragmentStream(exchange, result, adapter, locale, bindingContext);
}

valueMono = (result.getReturnValue() != null ?
Mono.just(FragmentsRendering.withPublisher(adapter.toPublisher(result.getReturnValue())).build()) :
Mono.empty());
Expand All @@ -233,8 +245,6 @@ public Mono<Void> handleResult(ServerWebExchange exchange, HandlerResult result)
Mono<List<View>> viewsMono;
Model model = result.getModel();
MethodParameter parameter = result.getReturnTypeSource();
BindingContext bindingContext = result.getBindingContext();
Locale locale = LocaleContextHolder.getLocale(exchange.getLocaleContext());

Class<?> clazz = valueType.toClass();
if (clazz == Object.class) {
Expand Down Expand Up @@ -277,13 +287,15 @@ else if (FragmentsRendering.class.isAssignableFrom(clazz)) {
response.getHeaders().putAll(render.headers());
bindingContext.updateModel(exchange);

StreamHandler streamHandler = getStreamHandler(exchange);
StreamHandler streamHandler =
(this.sseHandler.supports(exchange.getRequest()) ? this.sseHandler : null);

if (streamHandler != null) {
streamHandler.updateResponse(exchange);
}

Flux<Flux<DataBuffer>> renderFlux = render.fragments()
.concatMap(fragment -> renderFragment(fragment, streamHandler, locale, bindingContext, exchange))
.concatMap(fragment -> renderFragment(fragment, null, streamHandler, locale, bindingContext, exchange))
.doOnDiscard(DataBuffer.class, DataBufferUtils::release);

return response.writeAndFlushWith(renderFlux);
Expand Down Expand Up @@ -338,9 +350,29 @@ private Mono<List<View>> resolveViews(String viewName, Locale locale) {
});
}

private Mono<Void> handleSseFragmentStream(
ServerWebExchange exchange, HandlerResult result, ReactiveAdapter adapter, Locale locale,
BindingContext bindingContext) {

this.sseHandler.updateResponse(exchange);

Flux<ServerSentEvent<Fragment>> eventFlux =
Flux.from(adapter.toPublisher(result.getReturnValue()));

Flux<Flux<DataBuffer>> dataBufferFlux = eventFlux
.concatMap(event -> renderFragment(event.data(), event, this.sseHandler, locale, bindingContext, exchange))
.doOnDiscard(DataBuffer.class, DataBufferUtils::release);

return exchange.getResponse().writeAndFlushWith(dataBufferFlux);
}

private Mono<Flux<DataBuffer>> renderFragment(
Fragment fragment, @Nullable StreamHandler streamHandler, Locale locale,
BindingContext bindingContext, ServerWebExchange exchange) {
@Nullable Fragment fragment, @Nullable Object streamingHints, @Nullable StreamHandler streamHandler,
Locale locale, BindingContext bindingContext, ServerWebExchange exchange) {

if (fragment == null) {
return Mono.empty();
}

// Merge attributes from top-level model
fragment.mergeAttributes(bindingContext.getModel());
Expand All @@ -355,25 +387,18 @@ private Mono<Flux<DataBuffer>> renderFragment(
Map<String, Object> 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)));
return selectedViews
.flatMap(views ->
render(views, model, MediaType.TEXT_HTML, bindingContext, mutatedExchange))
.then(Mono.fromSupplier(() -> streamHandler.format(
response.getBodyFlux(), fragment, streamingHints, exchange)));
}
else {
return selectedViews.flatMap(views -> render(views, model, null, bindingContext, mutatedExchange))
.then(Mono.fromSupplier(response::getBodyFlux));
}
}

@Nullable
private StreamHandler getStreamHandler(ServerWebExchange exchange) {
for (StreamHandler handler : this.streamHandlers) {
if (handler.supports(exchange.getRequest())) {
return handler;
}
}
return null;
}

private String getNameForReturnValue(MethodParameter returnType) {
return Optional.ofNullable(returnType.getMethodAnnotation(ModelAttribute.class))
.filter(ann -> StringUtils.hasText(ann.value()))
Expand Down Expand Up @@ -499,10 +524,13 @@ private interface StreamHandler {
* Format the given fragment.
* @param fragmentContent the fragment serialized to data buffers
* @param fragment the fragment being rendered
* @param streamingHints extra hints for the stream format (e.g. ServerSentEvent wrapper)
* @param exchange the current exchange
* @return the formatted fragment
*/
Flux<DataBuffer> format(Flux<DataBuffer> fragmentContent, Fragment fragment, ServerWebExchange exchange);
Flux<DataBuffer> format(
Flux<DataBuffer> fragmentContent, Fragment fragment, @Nullable Object streamingHints,
ServerWebExchange exchange);
}


Expand Down Expand Up @@ -540,16 +568,21 @@ private Charset getCharset(ServerHttpRequest request) {

@Override
public Flux<DataBuffer> format(
Flux<DataBuffer> fragmentFlux, Fragment fragment, ServerWebExchange exchange) {
Flux<DataBuffer> fragmentFlux, Fragment fragment, @Nullable Object hints,
ServerWebExchange exchange) {

MediaType mediaType = exchange.getResponse().getHeaders().getContentType();
Charset charset = (mediaType != null && mediaType.getCharset() != null ?
mediaType.getCharset() : StandardCharsets.UTF_8);
Assert.state(hints == null || hints instanceof ServerSentEvent, "Expected ServerSentEvent");

DataBufferFactory bufferFactory = exchange.getResponse().bufferFactory();

String eventLine = (fragment.viewName() != null ? "event:" + fragment.viewName() + "\n" : "");
DataBuffer prefix = encodeText(eventLine + "data:", charset, bufferFactory);
ServerSentEvent<?> sse = (ServerSentEvent<?>) hints;
CharSequence eventText = (sse != null ? sse.format() :
(fragment.viewName() != null ? "event:" + fragment.viewName() + "\n" : "") + "data:");

DataBuffer prefix = encodeText(eventText.toString(), charset, bufferFactory);
DataBuffer suffix = encodeText("\n\n", charset, bufferFactory);

Mono<DataBuffer> content = DataBufferUtils.join(fragmentFlux)
Expand Down
67 changes: 52 additions & 15 deletions ...rg/springframework/web/reactive/result/view/FragmentViewResolutionResultHandlerTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,9 @@
import org.springframework.context.annotation.Configuration;
import org.springframework.context.support.ResourceBundleMessageSource;
import org.springframework.core.MethodParameter;
import org.springframework.core.ResolvableType;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.reactive.BindingContext;
import org.springframework.web.reactive.HandlerResult;
import org.springframework.web.reactive.accept.HeaderContentTypeResolver;
Expand Down Expand Up @@ -99,7 +101,51 @@ void render(Object returnValue, MethodParameter parameter) {
}

@Test
void renderSse() {
void renderFragmentStream() {

testSse(Flux.just(fragment1, fragment2),
on(Handler.class).resolveReturnType(Flux.class, Fragment.class),
"""
event:fragment1
data:<p>
data: Hello Foo
data:</p>
event:fragment2
data:<p>
data: Hello Bar
data:</p>
""");
}

@Test
void renderServerSentEventFragmentStream() {

ServerSentEvent<Fragment> event1 = ServerSentEvent.builder(fragment1).id("id1").event("event1").build();
ServerSentEvent<Fragment> event2 = ServerSentEvent.builder(fragment2).id("id2").event("event2").build();

MethodParameter returnType = on(Handler.class).resolveReturnType(
Flux.class, ResolvableType.forClassWithGenerics(ServerSentEvent.class, Fragment.class));

testSse(Flux.just(event1, event2), returnType,
"""
id:id1
event:event1
data:<p>
data: Hello Foo
data:</p>
id:id2
event:event2
data:<p>
data: Hello Bar
data:</p>
""");
}

private void testSse(Flux<?> dataFlux, MethodParameter returnType, String output) {
MockServerHttpRequest request = MockServerHttpRequest.get("/")
.accept(MediaType.TEXT_EVENT_STREAM)
.acceptLanguageAsLocales(Locale.ENGLISH)
Expand All @@ -110,27 +156,16 @@ void renderSse() {

HandlerResult result = new HandlerResult(
new Handler(),
Flux.just(fragment1, fragment2).subscribeOn(Schedulers.boundedElastic()),
on(Handler.class).resolveReturnType(Flux.class, Fragment.class),
dataFlux.subscribeOn(Schedulers.boundedElastic()),
returnType,
new BindingContext());

String body = initHandler().handleResult(exchange, result)
.then(Mono.defer(response::getBodyAsString))
.block(Duration.ofSeconds(60));

assertThat(response.getHeaders().getContentType()).isEqualTo(MediaType.TEXT_EVENT_STREAM);
assertThat(body).isEqualTo("""
event:fragment1
data:<p>
data: Hello Foo
data:</p>
event:fragment2
data:<p>
data: Hello Bar
data:</p>
""");
assertThat(body).isEqualTo(output);
}

private ViewResolutionResultHandler initHandler() {
Expand All @@ -155,6 +190,8 @@ private static class Handler {

Flux<Fragment> renderFlux() { return null; }

Flux<ServerSentEvent<Fragment>> renderSseFlux() { return null; }

List<Fragment> renderList() { return null; }

}
Expand Down
5 changes: 5 additions & 0 deletions ...t/java/org/springframework/web/reactive/result/view/ViewResolutionResultHandlerTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.lang.Nullable;
import org.springframework.ui.ConcurrentModel;
Expand Down Expand Up @@ -84,6 +85,9 @@ void supports() {

testSupports(on(Handler.class).resolveReturnType(FragmentsRendering.class));
testSupports(on(Handler.class).resolveReturnType(Flux.class, Fragment.class));
testSupports(on(Handler.class).resolveReturnType(
Flux.class, ResolvableType.forClassWithGenerics(ServerSentEvent.class, Fragment.class)));

testSupports(on(Handler.class).resolveReturnType(List.class, Fragment.class));
testSupports(on(Handler.class).resolveReturnType(
Mono.class, ResolvableType.forClassWithGenerics(List.class, Fragment.class)));
Expand Down Expand Up @@ -457,6 +461,7 @@ private static class Handler {

FragmentsRendering fragmentsRendering() { return null; }
Flux<Fragment> fragmentFlux() { return null; }
Flux<ServerSentEvent<Fragment>> fragmentServerSentEventFlux() { return null; }
Mono<List<Fragment>> monoFragmentList() { return null; }
List<Fragment> fragmentList() { return null; }

Expand Down

0 comments on commit 3b95d2c

Please sign in to comment.