diff --git a/spring-web/src/main/java/org/springframework/http/server/reactive/HttpHeadResponseDecorator.java b/spring-web/src/main/java/org/springframework/http/server/reactive/HttpHeadResponseDecorator.java
index b52a38f722..04277cdead 100644
--- a/spring-web/src/main/java/org/springframework/http/server/reactive/HttpHeadResponseDecorator.java
+++ b/spring-web/src/main/java/org/springframework/http/server/reactive/HttpHeadResponseDecorator.java
@@ -41,18 +41,22 @@ public class HttpHeadResponseDecorator extends ServerHttpResponseDecorator {
/**
* Consume and release the body without writing.
*
If the headers contain neither Content-Length nor Transfer-Encoding,
- * count the bytes and set Content-Length.
+ * and the body is a {@link Mono}, count the bytes and set Content-Length.
*/
@Override
+ @SuppressWarnings("unchecked")
public final Mono writeWith(Publisher extends DataBuffer> body) {
- if (shouldSetContentLength()) {
- return Flux.from(body)
- .reduce(0, (current, buffer) -> {
- int next = current + buffer.readableByteCount();
- DataBufferUtils.release(buffer);
- return next;
+ if (shouldSetContentLength() && body instanceof Mono) {
+ return ((Mono extends DataBuffer>) body)
+ .doOnSuccess(buffer -> {
+ if (buffer != null) {
+ getHeaders().setContentLength(buffer.readableByteCount());
+ DataBufferUtils.release(buffer);
+ }
+ else {
+ getHeaders().setContentLength(0);
+ }
})
- .doOnNext(length -> getHeaders().setContentLength(length))
.then();
}
else {
diff --git a/spring-web/src/test/java/org/springframework/http/server/reactive/HttpHeadResponseDecoratorTests.java b/spring-web/src/test/java/org/springframework/http/server/reactive/HttpHeadResponseDecoratorTests.java
index d98b19356f..479bc3fc11 100644
--- a/spring-web/src/test/java/org/springframework/http/server/reactive/HttpHeadResponseDecoratorTests.java
+++ b/spring-web/src/test/java/org/springframework/http/server/reactive/HttpHeadResponseDecoratorTests.java
@@ -21,6 +21,7 @@ import io.netty.buffer.PooledByteBufAllocator;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
@@ -50,10 +51,17 @@ public class HttpHeadResponseDecoratorTests {
@Test
- public void write() {
+ public void writeWithFlux() {
Flux body = Flux.just(toDataBuffer("data1"), toDataBuffer("data2"));
this.response.writeWith(body).block();
- assertThat(this.response.getHeaders().getContentLength()).isEqualTo(10);
+ assertThat(this.response.getHeaders().getContentLength()).isEqualTo(-1);
+ }
+
+ @Test
+ public void writeWithMono() {
+ Mono body = Mono.just(toDataBuffer("data1,data2"));
+ this.response.writeWith(body).block();
+ assertThat(this.response.getHeaders().getContentLength()).isEqualTo(11);
}
@Test // gh-23484