Restore response after beforeCommit action errors

See gh-24186

Author: rstoyanchev

spring-web/src/main/java/org/springframework/http/server/reactive/AbstractServerHttpResponse.java

@@ -176,22 +176,25 @@ public boolean isCommitted() {
 	@Override
 	@SuppressWarnings("unchecked")
 	public final Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
-		// Write as Mono if possible as an optimization hint to Reactor Netty
-		// ChannelSendOperator not necessary for Mono
+		// For Mono we can avoid ChannelSendOperator and Reactor Netty is more optimized for Mono.
+		// We must resolve value first however, for a chance to handle potential error.
 		if (body instanceof Mono) {
-			return ((Mono<? extends DataBuffer>) body).flatMap(buffer ->
-					doCommit(() -> writeWithInternal(Mono.just(buffer)))
-							.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release))
-					.doOnError(t -> this.getHeaders().clearContentHeaders());
+			return ((Mono<? extends DataBuffer>) body)
+					.flatMap(buffer -> doCommit(() ->
+							writeWithInternal(Mono.fromCallable(() -> buffer)
+									.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release))))
+					.doOnError(t -> getHeaders().clearContentHeaders());
+		}
+		else {
+			return new ChannelSendOperator<>(body, inner -> doCommit(() -> writeWithInternal(inner)))
+					.doOnError(t -> getHeaders().clearContentHeaders());
 		}
-		return new ChannelSendOperator<>(body, inner -> doCommit(() -> writeWithInternal(inner)))
-				.doOnError(t -> this.getHeaders().clearContentHeaders());
 	}
 
 	@Override
 	public final Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
 		return new ChannelSendOperator<>(body, inner -> doCommit(() -> writeAndFlushWithInternal(inner)))
-				.doOnError(t -> this.getHeaders().clearContentHeaders());
+				.doOnError(t -> getHeaders().clearContentHeaders());
 	}
 
 	@Override
@@ -217,21 +220,30 @@ protected Mono<Void> doCommit(@Nullable Supplier<? extends Mono<Void>> writeActi
 		if (!this.state.compareAndSet(State.NEW, State.COMMITTING)) {
 			return Mono.empty();
 		}
-		this.commitActions.add(() ->
-				Mono.fromRunnable(() -> {
-					applyStatusCode();
-					applyHeaders();
-					applyCookies();
-					this.state.set(State.COMMITTED);
-				}));
-		if (writeAction != null) {
-			this.commitActions.add(writeAction);
+
+		Flux<Void> allActions = Flux.empty();
+
+		if (!this.commitActions.isEmpty()) {
+			allActions = Flux.concat(Flux.fromIterable(this.commitActions).map(Supplier::get))
+					.doOnError(ex -> {
+						if (this.state.compareAndSet(State.COMMITTING, State.NEW)) {
+							getHeaders().clearContentHeaders();
+						}
+					});
 		}
-		Flux<Void> commit = Flux.empty();
-		for (Supplier<? extends Mono<Void>> action : this.commitActions) {
-			commit = commit.concatWith(action.get());
+
+		allActions = allActions.concatWith(Mono.fromRunnable(() -> {
+			applyStatusCode();
+			applyHeaders();
+			applyCookies();
+			this.state.set(State.COMMITTED);
+		}));
+
+		if (writeAction != null) {
+			allActions = allActions.concatWith(writeAction.get());
 		}
-		return commit.then();
+
+		return allActions.then();
 	}
 
 

spring-web/src/test/java/org/springframework/http/server/reactive/ServerHttpResponseTests.java

@@ -20,11 +20,14 @@
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
 
 import org.junit.jupiter.api.Test;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
 
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DefaultDataBuffer;
@@ -36,14 +39,16 @@
 import static org.assertj.core.api.Assertions.assertThat;
 
 /**
+ * Unit tests for {@link AbstractServerHttpRequest}.
+ *
  * @author Rossen Stoyanchev
  * @author Sebastien Deleuze
  * @author Brian Clozel
  */
 public class ServerHttpResponseTests {
 
 	@Test
-	void writeWith() throws Exception {
+	void writeWith() {
 		TestServerHttpResponse response = new TestServerHttpResponse();
 		response.writeWith(Flux.just(wrap("a"), wrap("b"), wrap("c"))).block();
 
@@ -58,7 +63,7 @@ void writeWith() throws Exception {
 	}
 
 	@Test  // SPR-14952
-	void writeAndFlushWithFluxOfDefaultDataBuffer() throws Exception {
+	void writeAndFlushWithFluxOfDefaultDataBuffer() {
 		TestServerHttpResponse response = new TestServerHttpResponse();
 		Flux<Flux<DefaultDataBuffer>> flux = Flux.just(Flux.just(wrap("foo")));
 		response.writeAndFlushWith(flux).block();
@@ -72,18 +77,18 @@ void writeAndFlushWithFluxOfDefaultDataBuffer() throws Exception {
 	}
 
 	@Test
-	void writeWithFluxError() throws Exception {
+	void writeWithFluxError() {
 		IllegalStateException error = new IllegalStateException("boo");
 		writeWithError(Flux.error(error));
 	}
 
 	@Test
-	void writeWithMonoError() throws Exception {
+	void writeWithMonoError() {
 		IllegalStateException error = new IllegalStateException("boo");
 		writeWithError(Mono.error(error));
 	}
 
-	void writeWithError(Publisher<DataBuffer> body) throws Exception {
+	void writeWithError(Publisher<DataBuffer> body) {
 		TestServerHttpResponse response = new TestServerHttpResponse();
 		HttpHeaders headers = response.getHeaders();
 		headers.setContentType(MediaType.APPLICATION_JSON);
@@ -100,7 +105,7 @@ void writeWithError(Publisher<DataBuffer> body) throws Exception {
 	}
 
 	@Test
-	void setComplete() throws Exception {
+	void setComplete() {
 		TestServerHttpResponse response = new TestServerHttpResponse();
 		response.setComplete().block();
 
@@ -111,7 +116,7 @@ void setComplete() throws Exception {
 	}
 
 	@Test
-	void beforeCommitWithComplete() throws Exception {
+	void beforeCommitWithComplete() {
 		ResponseCookie cookie = ResponseCookie.from("ID", "123").build();
 		TestServerHttpResponse response = new TestServerHttpResponse();
 		response.beforeCommit(() -> Mono.fromRunnable(() -> response.getCookies().add(cookie.getName(), cookie)));
@@ -129,7 +134,7 @@ void beforeCommitWithComplete() throws Exception {
 	}
 
 	@Test
-	void beforeCommitActionWithSetComplete() throws Exception {
+	void beforeCommitActionWithSetComplete() {
 		ResponseCookie cookie = ResponseCookie.from("ID", "123").build();
 		TestServerHttpResponse response = new TestServerHttpResponse();
 		response.beforeCommit(() -> {
@@ -145,6 +150,32 @@ void beforeCommitActionWithSetComplete() throws Exception {
 		assertThat(response.getCookies().getFirst("ID")).isSameAs(cookie);
 	}
 
+	@Test // gh-24186
+	void beforeCommitErrorShouldLeaveResponseNotCommitted() {
+
+		Consumer<Supplier<Mono<Void>>> tester = preCommitAction -> {
+			TestServerHttpResponse response = new TestServerHttpResponse();
+			response.getHeaders().setContentType(MediaType.APPLICATION_JSON);
+			response.getHeaders().setContentLength(3);
+			response.beforeCommit(preCommitAction);
+
+			StepVerifier.create(response.writeWith(Flux.just(wrap("body"))))
+					.expectErrorMessage("Max sessions")
+					.verify();
+
+			assertThat(response.statusCodeWritten).isFalse();
+			assertThat(response.headersWritten).isFalse();
+			assertThat(response.cookiesWritten).isFalse();
+			assertThat(response.isCommitted()).isFalse();
+			assertThat(response.getHeaders()).isEmpty();
+		};
+
+		tester.accept(() -> Mono.error(new IllegalStateException("Max sessions")));
+		tester.accept(() -> {
+			throw new IllegalStateException("Max sessions");
+		});
+	}
+
 
 	private DefaultDataBuffer wrap(String a) {
 		return new DefaultDataBufferFactory().wrap(ByteBuffer.wrap(a.getBytes(StandardCharsets.UTF_8)));