Backport fixes for discarding data buffers

Closes gh-26232
This commit is contained in:
Rossen Stoyanchev
2020-12-08 22:00:28 +00:00
parent a8091b916b
commit 33476a2eae
3 changed files with 52 additions and 5 deletions

View File

@@ -128,7 +128,8 @@ public class EncoderHttpMessageWriter<T> implements HttpMessageWriter<T> {
message.getHeaders().setContentLength(buffer.readableByteCount());
return message.writeWith(Mono.just(buffer)
.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release));
});
})
.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release);
}
if (isStreamingMediaType(contentType)) {

View File

@@ -203,10 +203,28 @@ public abstract class AbstractServerHttpResponse implements ServerHttpResponse {
// 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.fromCallable(() -> buffer)
.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release))))
.doOnError(t -> getHeaders().clearContentHeaders());
.flatMap(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())
.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release);
}
else {
return new ChannelSendOperator<>(body, inner -> doCommit(() -> writeWithInternal(inner)))