From 51cc719c004da241c5a29f63d81c1203c2d00d26 Mon Sep 17 00:00:00 2001
From: Arjen Poutsma <apoutsma@pivotal.io>
Date: Thu, 28 May 2020 16:04:43 +0200
Subject: [PATCH] Remove UndertowDataBuffer

This commit removes the UndertowDataBuffer, in favor of using regular
DataBuffers from the DataBufferFactory. During the development of the
DefaultPartHttpMessageReader, it was determined that invoking various
slicing and releasing operators on the UndertowDataBuffer resulted in
memory leaks. This commit fixes that.
---
 .../reactive/UndertowServerHttpRequest.java   | 70 ++-----------------
 1 file changed, 4 insertions(+), 66 deletions(-)

diff --git a/spring-web/src/main/java/org/springframework/http/server/reactive/UndertowServerHttpRequest.java b/spring-web/src/main/java/org/springframework/http/server/reactive/UndertowServerHttpRequest.java
index ad05b7cd5b..030ac903b5 100644
--- a/spring-web/src/main/java/org/springframework/http/server/reactive/UndertowServerHttpRequest.java
+++ b/spring-web/src/main/java/org/springframework/http/server/reactive/UndertowServerHttpRequest.java
@@ -21,7 +21,6 @@ import java.net.InetSocketAddress;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.nio.ByteBuffer;
-import java.util.concurrent.atomic.AtomicInteger;
 
 import javax.net.ssl.SSLSession;
 
@@ -34,9 +33,6 @@ import reactor.core.publisher.Flux;
 
 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.DataBufferWrapper;
-import org.springframework.core.io.buffer.PooledDataBuffer;
 import org.springframework.http.HttpCookie;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
@@ -172,7 +168,6 @@ class UndertowServerHttpRequest extends AbstractServerHttpRequest {
 		@Nullable
 		protected DataBuffer read() throws IOException {
 			PooledByteBuffer pooledByteBuffer = this.byteBufferPool.allocate();
-			boolean release = true;
 			try {
 				ByteBuffer byteBuffer = pooledByteBuffer.getBuffer();
 				int read = this.channel.read(byteBuffer);
@@ -183,9 +178,9 @@ class UndertowServerHttpRequest extends AbstractServerHttpRequest {
 
 				if (read > 0) {
 					byteBuffer.flip();
-					DataBuffer dataBuffer = this.bufferFactory.wrap(byteBuffer);
-					release = false;
-					return new UndertowDataBuffer(dataBuffer, pooledByteBuffer);
+					DataBuffer dataBuffer = this.bufferFactory.allocateBuffer(read);
+					dataBuffer.write(byteBuffer);
+					return dataBuffer;
 				}
 				else if (read == -1) {
 					onAllDataRead();
@@ -193,9 +188,7 @@ class UndertowServerHttpRequest extends AbstractServerHttpRequest {
 				return null;
 			}
 			finally {
-				if (release && pooledByteBuffer.isOpen()) {
-					pooledByteBuffer.close();
-				}
+				pooledByteBuffer.close();
 			}
 		}
 
@@ -205,59 +198,4 @@ class UndertowServerHttpRequest extends AbstractServerHttpRequest {
 		}
 	}
 
-
-	private static class UndertowDataBuffer extends DataBufferWrapper implements PooledDataBuffer {
-
-		private final PooledByteBuffer pooledByteBuffer;
-
-		private final AtomicInteger refCount;
-
-
-		public UndertowDataBuffer(DataBuffer dataBuffer, PooledByteBuffer pooledByteBuffer) {
-			super(dataBuffer);
-			this.pooledByteBuffer = pooledByteBuffer;
-			this.refCount = new AtomicInteger(1);
-		}
-
-		private UndertowDataBuffer(DataBuffer dataBuffer, PooledByteBuffer pooledByteBuffer,
-				AtomicInteger refCount) {
-			super(dataBuffer);
-			this.refCount = refCount;
-			this.pooledByteBuffer = pooledByteBuffer;
-		}
-
-		@Override
-		public boolean isAllocated() {
-			return this.refCount.get() > 0;
-		}
-
-		@Override
-		public PooledDataBuffer retain() {
-			this.refCount.incrementAndGet();
-			DataBufferUtils.retain(dataBuffer());
-			return this;
-		}
-
-		@Override
-		public boolean release() {
-			int refCount = this.refCount.decrementAndGet();
-			if (refCount == 0) {
-				try {
-					return DataBufferUtils.release(dataBuffer());
-				}
-				finally {
-					this.pooledByteBuffer.close();
-				}
-			}
-			return false;
-		}
-
-		@Override
-		public DataBuffer slice(int index, int length) {
-			DataBuffer slice = dataBuffer().slice(index, length);
-			return new UndertowDataBuffer(slice, this.pooledByteBuffer, this.refCount);
-		}
-
-	}
-
 }