Additional fixes for discarding data buffers

Closes gh-26232
This commit is contained in:
Rossen Stoyanchev 2020-12-08 21:33:50 +00:00
parent cb44ae62e9
commit 25101fb034
2 changed files with 22 additions and 11 deletions

View File

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

View File

@ -213,17 +213,27 @@ public abstract class AbstractServerHttpResponse implements ServerHttpResponse {
return ((Mono<? extends DataBuffer>) body) return ((Mono<? extends DataBuffer>) body)
.flatMap(buffer -> { .flatMap(buffer -> {
touchDataBuffer(buffer); touchDataBuffer(buffer);
return doCommit(() -> { AtomicReference<Boolean> subscribed = new AtomicReference<>(false);
try { return doCommit(
return writeWithInternal(Mono.fromCallable(() -> buffer) () -> {
.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release)); try {
} return writeWithInternal(Mono.fromCallable(() -> buffer)
catch (Throwable ex) { .doOnSubscribe(s -> subscribed.set(true))
return Mono.error(ex); .doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release));
} }
}).doOnError(ex -> DataBufferUtils.release(buffer)); 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 { else {
return new ChannelSendOperator<>(body, inner -> doCommit(() -> writeWithInternal(inner))) return new ChannelSendOperator<>(body, inner -> doCommit(() -> writeWithInternal(inner)))