Update AbstractClientHttpRequest with server changes

This commit updates `AbstractClientHttpRequest` to make it more similar
to its server counterpart.

Author: bclozel

spring-web/src/main/java/org/springframework/http/client/reactive/AbstractClientHttpRequest.java

@@ -20,7 +20,9 @@
 import java.util.List;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import org.springframework.http.HttpCookie;
@@ -39,13 +41,21 @@
  */
 public abstract class AbstractClientHttpRequest implements ClientHttpRequest {
 
+	/**
+	 * COMMITTING -> COMMITTED is the period after doCommit is called but before
+	 * the response status and headers have been applied to the underlying
+	 * response during which time pre-commit actions can still make changes to
+	 * the response status and headers.
+	 */
+	private enum State {NEW, COMMITTING, COMMITTED}
+
 	private final HttpHeaders headers;
 
 	private final MultiValueMap<String, HttpCookie> cookies;
 
 	private AtomicReference<State> state = new AtomicReference<>(State.NEW);
 
-	private final List<Supplier<? extends Mono<Void>>> beforeCommitActions = new ArrayList<>(4);
+	private final List<Supplier<? extends Mono<Void>>> commitActions = new ArrayList<>(4);
 
 
 	public AbstractClientHttpRequest() {
@@ -61,50 +71,74 @@ public AbstractClientHttpRequest(HttpHeaders headers) {
 
 	@Override
 	public HttpHeaders getHeaders() {
-		if (State.COMITTED.equals(this.state.get())) {
+		if (State.COMMITTED.equals(this.state.get())) {
 			return HttpHeaders.readOnlyHttpHeaders(this.headers);
 		}
 		return this.headers;
 	}
 
 	@Override
 	public MultiValueMap<String, HttpCookie> getCookies() {
-		if (State.COMITTED.equals(this.state.get())) {
+		if (State.COMMITTED.equals(this.state.get())) {
 			return CollectionUtils.unmodifiableMultiValueMap(this.cookies);
 		}
 		return this.cookies;
 	}
 
-	protected Mono<Void> applyBeforeCommit() {
-		Mono<Void> mono = Mono.empty();
-		if (this.state.compareAndSet(State.NEW, State.COMMITTING)) {
-			for (Supplier<? extends Mono<Void>> action : this.beforeCommitActions) {
-				mono = mono.then(() -> action.get());
-			}
-			return mono
-					.otherwise(ex -> {
-						// Ignore errors from beforeCommit actions
-						return Mono.empty();
-					})
-					.then(() -> {
-						this.state.set(State.COMITTED);
-						writeHeaders();
-						writeCookies();
-						return Mono.empty();
-					});
+	/**
+	 * A variant of {@link #doCommit(Supplier)} for a request without body.
+	 * @return a completion publisher
+	 */
+	protected Mono<Void> doCommit() {
+		return (this.state.get() == State.NEW ? doCommit(null) : Mono.empty());
+	}
+
+	/**
+	 * Apply {@link #beforeCommit(Supplier) beforeCommit} actions, apply the
+	 * request headers/cookies, and write the request body.
+	 * @param writeAction the action to write the request body or {@code null}
+	 * @return a completion publisher
+	 */
+	protected Mono<Void> doCommit(Supplier<? extends Mono<Void>> writeAction) {
+
+		if (!this.state.compareAndSet(AbstractClientHttpRequest.State.NEW, AbstractClientHttpRequest.State.COMMITTING)) {
+			return Mono.empty();
+		}
+
+		this.commitActions.add(() -> {
+			applyHeaders();
+			applyCookies();
+			this.state.set(AbstractClientHttpRequest.State.COMMITTED);
+			return Mono.empty();
+		});
+
+		if (writeAction != null) {
+			this.commitActions.add(writeAction);
 		}
-		return mono;
+
+		List<? extends Mono<Void>> actions = this.commitActions.stream()
+				.map(Supplier::get).collect(Collectors.toList());
+
+		return Flux.concat(actions).next();
 	}
 
 	@Override
 	public void beforeCommit(Supplier<? extends Mono<Void>> action) {
 		Assert.notNull(action);
-		this.beforeCommitActions.add(action);
+		this.commitActions.add(action);
 	}
 
-	protected abstract void writeHeaders();
+	/**
+	 * Implement this method to apply header changes from {@link #getHeaders()}
+	 * to the underlying response. This method is called once only.
+	 */
+	protected abstract void applyHeaders();
+
+	/**
+	 * Implement this method to add cookies from {@link #getHeaders()} to the
+	 * underlying response. This method is called once only.
+	 */
+	protected abstract void applyCookies();
 
-	protected abstract void writeCookies();
 
-	private enum State {NEW, COMMITTING, COMITTED}
 }

spring-web/src/main/java/org/springframework/http/client/reactive/ReactorClientHttpRequest.java

@@ -75,16 +75,15 @@ public URI getURI() {
 
 	@Override
 	public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
-		return applyBeforeCommit().then(this.httpRequest
+		return doCommit(() -> this.httpRequest
 				.send(Flux.from(body).map(NettyDataBufferFactory::toByteBuf)).then());
 	}
 
 	@Override
 	public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
 		Publisher<Publisher<ByteBuf>> byteBufs = Flux.from(body).
 				map(ReactorClientHttpRequest::toByteBufs);
-		return applyBeforeCommit().then(this.httpRequest
-				.sendGroups(byteBufs).then());
+		return doCommit(() -> this.httpRequest.sendGroups(byteBufs).then());
 	}
 
 	private static Publisher<ByteBuf> toByteBufs(Publisher<? extends DataBuffer> dataBuffers) {
@@ -94,17 +93,17 @@ private static Publisher<ByteBuf> toByteBufs(Publisher<? extends DataBuffer> dat
 
 	@Override
 	public Mono<Void> setComplete() {
-		return applyBeforeCommit().then(httpRequest.sendHeaders().then());
+		return doCommit(() -> httpRequest.sendHeaders().then());
 	}
 
 	@Override
-	protected void writeHeaders() {
+	protected void applyHeaders() {
 		getHeaders().entrySet()
 				.forEach(e -> this.httpRequest.requestHeaders().set(e.getKey(), e.getValue()));
 	}
 
 	@Override
-	protected void writeCookies() {
+	protected void applyCookies() {
 		getCookies().values().stream().flatMap(Collection::stream)
 				.map(cookie -> new DefaultCookie(cookie.getName(), cookie.getValue()))
 				.forEach(this.httpRequest::addCookie);

spring-web/src/test/java/org/springframework/web/client/reactive/test/MockClientHttpRequest.java

@@ -92,13 +92,13 @@ public DataBufferFactory bufferFactory() {
 	@Override
 	public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
 		this.body = Flux.from(body);
-		return applyBeforeCommit().then(this.body.then());
+		return doCommit(() -> this.body.then());
 	}
 
 	@Override
 	public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
 		this.bodyWithFlushes = Flux.from(body).map(p -> Flux.from(p));
-		return applyBeforeCommit().then(this.bodyWithFlushes.then());
+		return doCommit(() -> this.bodyWithFlushes.then());
 	}
 
 	public Publisher<DataBuffer> getBody() {
@@ -111,12 +111,12 @@ public Publisher<Publisher<DataBuffer>> getBodyWithFlush() {
 
 	@Override
 	public Mono<Void> setComplete() {
-		return applyBeforeCommit().then();
+		return doCommit().then();
 	}
 
 	@Override
-	protected void writeHeaders() { }
+	protected void applyHeaders() { }
 
 	@Override
-	protected void writeCookies() { }
+	protected void applyCookies() { }
 }