Additional fixes for discarding data buffers

Closes gh-26232
spring-projects · Dec 8, 2020 · 25101fb · 25101fb
1 parent cb44ae6
commit 25101fb
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 11 deletions.
diff --git a/spring-web/src/main/java/org/springframework/http/codec/EncoderHttpMessageWriter.java b/spring-web/src/main/java/org/springframework/http/codec/EncoderHttpMessageWriter.java
@@ -132,7 +132,8 @@ public Mono<Void> write(Publisher<? extends T> inputStream, ResolvableType eleme
 						message.getHeaders().setContentLength(buffer.readableByteCount());
 						return message.writeWith(Mono.just(buffer)
 								.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release));
-					});
+					})
+					.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release);
 		}
 
 		if (isStreamingMediaType(contentType)) {

diff --git a/...eb/src/main/java/org/springframework/http/server/reactive/AbstractServerHttpResponse.java b/...eb/src/main/java/org/springframework/http/server/reactive/AbstractServerHttpResponse.java
@@ -213,17 +213,27 @@ public final Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
 			return ((Mono<? extends DataBuffer>) body)
 					.flatMap(buffer -> {
 						touchDataBuffer(buffer);
-						return doCommit(() -> {
-							try {
-								return writeWithInternal(Mono.fromCallable(() -> buffer)
-										.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release));
-							}
-							catch (Throwable ex) {
-								return Mono.error(ex);
-							}
-						}).doOnError(ex -> DataBufferUtils.release(buffer));
+						AtomicReference<Boolean> subscribed = new AtomicReference<>(false);
+						return doCommit(
+								() -> {
+									try {
+										return writeWithInternal(Mono.fromCallable(() -> buffer)
+												.doOnSubscribe(s -> subscribed.set(true))
+												.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release));
+									}
+									catch (Throwable ex) {
+										return Mono.error(ex);
+									}
+								})
+								.doOnError(ex -> DataBufferUtils.release(buffer))
+								.doOnCancel(() -> {
+									if (!subscribed.get()) {
+										DataBufferUtils.release(buffer);
+									}
+								});
 					})
-					.doOnError(t -> getHeaders().clearContentHeaders());
+					.doOnError(t -> getHeaders().clearContentHeaders())
+					.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release);
 		}
 		else {
 			return new ChannelSendOperator<>(body, inner -> doCommit(() -> writeWithInternal(inner)))