[hotfix][network-tests] Simplify TestPooledBufferProvider

flink-runtime/src/test/java/org/apache/flink/runtime/io/network/util/TestBufferFactory.java

@@ -18,18 +18,17 @@
 
 package org.apache.flink.runtime.io.network.util;
 
-import org.apache.flink.core.memory.HeapMemorySegment;
-import org.apache.flink.core.memory.MemorySegment;
 import org.apache.flink.core.memory.MemorySegmentFactory;
 import org.apache.flink.runtime.io.network.buffer.Buffer;
 import org.apache.flink.runtime.io.network.buffer.BufferRecycler;
 import org.apache.flink.runtime.testutils.DiscardingRecycler;
 
-import java.util.concurrent.atomic.AtomicInteger;
+import javax.annotation.concurrent.ThreadSafe;
 
 import static org.apache.flink.util.Preconditions.checkArgument;
 import static org.apache.flink.util.Preconditions.checkNotNull;
 
+@ThreadSafe
 public class TestBufferFactory {
 
 	public static final int BUFFER_SIZE = 32 * 1024;
@@ -40,37 +39,31 @@ public class TestBufferFactory {
 
 	private final BufferRecycler bufferRecycler;
 
-	private AtomicInteger numberOfCreatedBuffers = new AtomicInteger();
+	private final int poolSize;
 
-	public TestBufferFactory() {
-		this(BUFFER_SIZE, RECYCLER);
-	}
-
-	public TestBufferFactory(int bufferSize) {
-		this(bufferSize, RECYCLER);
-	}
+	private int numberOfCreatedBuffers = 0;
 
-	public TestBufferFactory(int bufferSize, BufferRecycler bufferRecycler) {
+	public TestBufferFactory(int poolSize, int bufferSize, BufferRecycler bufferRecycler) {
 		checkArgument(bufferSize > 0);
+		this.poolSize = poolSize;
 		this.bufferSize = bufferSize;
 		this.bufferRecycler = checkNotNull(bufferRecycler);
 	}
 
-	public Buffer create() {
-		numberOfCreatedBuffers.incrementAndGet();
+	public synchronized Buffer create() {
+		if (numberOfCreatedBuffers >= poolSize) {
+			return null;
+		}
 
+		numberOfCreatedBuffers++;
 		return new Buffer(MemorySegmentFactory.allocateUnpooledSegment(bufferSize), bufferRecycler);
 	}
 
-	public Buffer createFrom(MemorySegment segment) {
-		return new Buffer(segment, bufferRecycler);
-	}
-
-	public int getNumberOfCreatedBuffers() {
-		return numberOfCreatedBuffers.get();
+	public synchronized int getNumberOfCreatedBuffers() {
+		return numberOfCreatedBuffers;
 	}
 
-	public int getBufferSize() {
+	public synchronized int getBufferSize() {
 		return bufferSize;
 	}
 

flink-runtime/src/test/java/org/apache/flink/runtime/io/network/util/TestPooledBufferProvider.java

@@ -36,58 +36,42 @@ import static org.apache.flink.util.Preconditions.checkArgument;
 
 public class TestPooledBufferProvider implements BufferProvider {
 
-	private final Object bufferCreationLock = new Object();
-
 	private final BlockingQueue<Buffer> buffers = new LinkedBlockingDeque<>();
 
 	private final TestBufferFactory bufferFactory;
 
 	private final PooledBufferProviderRecycler bufferRecycler;
 
-	private final int poolSize;
-
 	public TestPooledBufferProvider(int poolSize) {
 		checkArgument(poolSize > 0);
-		this.poolSize = poolSize;
 
 		this.bufferRecycler = new PooledBufferProviderRecycler(buffers);
-		this.bufferFactory = new TestBufferFactory(32 * 1024, bufferRecycler);
+		this.bufferFactory = new TestBufferFactory(poolSize, 32 * 1024, bufferRecycler);
 	}
 
 	@Override
 	public Buffer requestBuffer() throws IOException {
 		final Buffer buffer = buffers.poll();
-
 		if (buffer != null) {
 			return buffer;
 		}
-		else {
-			synchronized (bufferCreationLock) {
-				if (bufferFactory.getNumberOfCreatedBuffers() < poolSize) {
-					return bufferFactory.create();
-				}
-			}
 
-			return null;
-		}
+		return bufferFactory.create();
 	}
 
 	@Override
 	public Buffer requestBufferBlocking() throws IOException, InterruptedException {
-		final Buffer buffer = buffers.poll();
-
+		Buffer buffer = buffers.poll();
 		if (buffer != null) {
 			return buffer;
 		}
-		else {
-			synchronized (bufferCreationLock) {
-				if (bufferFactory.getNumberOfCreatedBuffers() < poolSize) {
-					return bufferFactory.create();
-				}
-			}
 
-			return buffers.take();
+		buffer = bufferFactory.create();
+		if (buffer != null) {
+			return buffer;
 		}
+
+		return buffers.take();
 	}
 
 	@Override