Release retained response buffers on discard in NettyWriteResponseFilter

spring-cloud-gateway-server-webflux/src/main/java/org/springframework/cloud/gateway/filter/NettyWriteResponseFilter.java

@@ -19,6 +19,7 @@
 import java.util.List;
 
 import io.netty.buffer.ByteBuf;
+import io.netty.util.ReferenceCountUtil;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.jspecify.annotations.Nullable;
@@ -34,6 +35,7 @@
 import org.springframework.core.Ordered;
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
+import org.springframework.core.io.buffer.DataBufferUtils;
 import org.springframework.core.io.buffer.DefaultDataBufferFactory;
 import org.springframework.core.io.buffer.NettyDataBufferFactory;
 import org.springframework.http.MediaType;
@@ -114,7 +116,15 @@ public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
 							.inbound()
 							.receive()
 							.retain()
-							.map(byteBuf -> wrap(byteBuf, response));
+							.doOnDiscard(ByteBuf.class, ReferenceCountUtil::safeRelease)
+							.map(byteBuf -> wrap(byteBuf, response))
+							.doOnDiscard(DataBuffer.class, dataBuffer -> {
+								try {
+									DataBufferUtils.release(dataBuffer);
+								}
+								catch (Throwable ignored) {
+								}
+							});
 
 					MediaType contentType = null;
 					try {

spring-cloud-gateway-server-webflux/src/test/java/org/springframework/cloud/gateway/filter/NettyWriteResponseFilterTests.java

@@ -20,22 +20,33 @@
 import java.util.ArrayList;
 
 import io.netty.buffer.ByteBuf;
+import io.netty.channel.embedded.EmbeddedChannel;
 import org.junit.jupiter.api.Test;
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.netty.ByteBufFlux;
 import reactor.netty.Connection;
+import reactor.netty.NettyInbound;
 
 import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.DataBufferUtils;
 import org.springframework.core.io.buffer.NettyDataBufferFactory;
 import org.springframework.core.io.buffer.PooledDataBuffer;
+import org.springframework.http.server.reactive.ServerHttpResponse;
+import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
 import org.springframework.mock.http.server.reactive.MockServerHttpRequest;
 import org.springframework.mock.http.server.reactive.MockServerHttpResponse;
 import org.springframework.mock.web.server.MockServerWebExchange;
+import org.springframework.web.server.ServerWebExchange;
+import org.springframework.web.server.ServerWebExchangeDecorator;
 
 import static io.netty.buffer.PooledByteBufAllocator.DEFAULT;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.CLIENT_RESPONSE_CONN_ATTR;
 
 /**
@@ -68,6 +79,57 @@ public void committedResponseDisposesConnectionWithoutWriting() {
 		verify(connection, never()).inbound();
 	}
 
+	@Test
+	public void retainedBuffersAreReleasedOnDiscard() {
+		NettyWriteResponseFilter filter = new NettyWriteResponseFilter(new ArrayList<>(), null);
+
+		ByteBuf buffer1 = DEFAULT.buffer();
+		buffer1.writeCharSequence("test-1", Charset.defaultCharset());
+		ByteBuf buffer2 = DEFAULT.buffer();
+		buffer2.writeCharSequence("test-2", Charset.defaultCharset());
+		assertThat(buffer1.refCnt()).isEqualTo(1);
+		assertThat(buffer2.refCnt()).isEqualTo(1);
+
+		Connection connection = mock(Connection.class);
+		NettyInbound inbound = mock(NettyInbound.class);
+		EmbeddedChannel channel = new EmbeddedChannel();
+		when(connection.inbound()).thenReturn(inbound);
+		when(connection.channel()).thenReturn(channel);
+
+		Flux<ByteBuf> inboundFlux = Flux.create(sink -> {
+			sink.next(buffer1);
+			sink.next(buffer2);
+		});
+		when(inbound.receive()).thenReturn(ByteBufFlux.fromInbound(inboundFlux));
+
+		MockServerHttpRequest request = MockServerHttpRequest.get("/").build();
+		ServerHttpResponse response = new ServerHttpResponseDecorator(
+				new MockServerHttpResponse(new NettyDataBufferFactory(DEFAULT))) {
+			@Override
+			public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
+				return Flux.from(body).take(1).doOnNext(DataBufferUtils::release).then();
+			}
+		};
+
+		ServerWebExchange exchange = new ServerWebExchangeDecorator(MockServerWebExchange.from(request)) {
+			@Override
+			public ServerHttpResponse getResponse() {
+				return response;
+			}
+		};
+		exchange.getAttributes().put(CLIENT_RESPONSE_CONN_ATTR, connection);
+
+		filter.filter(exchange, ex -> Mono.empty()).block();
+
+		assertThat(buffer1.refCnt()).isEqualTo(1);
+		buffer1.release();
+		assertThat(buffer1.refCnt()).isEqualTo(0);
+
+		assertThat(buffer2.refCnt()).isEqualTo(0);
+
+		channel.finishAndReleaseAll();
+	}
+
 	private void doTestWrap(MockServerHttpResponse response) {
 		NettyWriteResponseFilter filter = new NettyWriteResponseFilter(new ArrayList<>(), null);