zhijiangW commented on a change in pull request #9993:
[FLINK-14498][runtime]Introduce NetworkBufferPool#isAvailable() for interacting
with LocalBufferPool.
URL: https://github.com/apache/flink/pull/9993#discussion_r340484430
##########
File path:
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/NetworkBufferPoolTest.java
##########
@@ -518,4 +527,204 @@ public void go() throws Exception {
globalPool.destroy();
}
}
+
+ /**
+ * Tests {@link NetworkBufferPool#isAvailable()}, verifying that the
buffer availability is correctly
+ * maintained after memory segments are requested by {@link
NetworkBufferPool#requestMemorySegment()}
+ * and recycled by {@link NetworkBufferPool#recycle(MemorySegment)}.
+ */
+ @Test
+ public void testIsAvailableAfterRequestMemorySegmentAndRecycle() throws
Exception {
+ final int numBuffers = 2;
+
+ final NetworkBufferPool globalPool = new
NetworkBufferPool(numBuffers, 128, 1);
+
+ try {
+ // the global pool should be in available state
initially
+ assertTrue(globalPool.isAvailable().isDone());
+
+ // request the first segment
+ final MemorySegment segment1 =
checkNotNull(globalPool.requestMemorySegment());
+ assertTrue(globalPool.isAvailable().isDone());
+
+ // request the second segment
+ final MemorySegment segment2 =
checkNotNull(globalPool.requestMemorySegment());
+ assertFalse(globalPool.isAvailable().isDone());
+
+ final CompletableFuture<?> availableFuture =
globalPool.isAvailable();
+
+ // recycle the first segment
+ globalPool.recycle(segment1);
+ assertTrue(availableFuture.isDone());
+ assertTrue(globalPool.isAvailable().isDone());
+
+ // recycle the second segment
+ globalPool.recycle(segment2);
+ assertTrue(globalPool.isAvailable().isDone());
+
+ } finally {
+ globalPool.destroy();
+ }
+ }
+
+ /**
+ * Tests {@link NetworkBufferPool#isAvailable()}, verifying that the
buffer availability is correctly
+ * maintained after memory segments are requested by {@link
NetworkBufferPool#requestMemorySegments()}
+ * and recycled by {@link
NetworkBufferPool#recycleMemorySegments(Collection)}.
+ */
+ @Test(timeout = 10000L)
+ public void testIsAvailableAfterRequestAndRecycleMultiSegments() throws
Exception {
+ final int numberOfSegmentsToRequest = 5;
+ final int numBuffers = 2 * numberOfSegmentsToRequest;
+
+ final NetworkBufferPool globalPool = new
NetworkBufferPool(numBuffers, 128, numberOfSegmentsToRequest);
+
+ try {
+ // the global pool should be in available state
initially
+ assertTrue(globalPool.isAvailable().isDone());
+
+ // request 5 segments
+ List<MemorySegment> segments1 =
globalPool.requestMemorySegments();
+ assertTrue(globalPool.isAvailable().isDone());
+ assertEquals(numberOfSegmentsToRequest,
segments1.size());
+
+ // request another 5 segments
+ List<MemorySegment> segments2 =
globalPool.requestMemorySegments();
+ assertFalse(globalPool.isAvailable().isDone());
+ assertEquals(numberOfSegmentsToRequest,
segments2.size());
+
+ CompletableFuture<?> availableFuture =
globalPool.isAvailable();
+
+ // request another 5 segments
+ final CountDownLatch latch = new CountDownLatch(1);
+ final List<MemorySegment> segments3 = new
ArrayList<>(numberOfSegmentsToRequest);
+ CheckedThread asyncRequest = new CheckedThread() {
+ @Override
+ public void go() throws Exception {
+ // this request should be blocked until
at least 5 segments are recycled
+
segments3.addAll(globalPool.requestMemorySegments());
+ latch.countDown();
+ }
+ };
+ asyncRequest.start();
+
+ // recycle 5 segments
+ globalPool.recycleMemorySegments(segments1);
+ assertTrue(availableFuture.isDone());
+
+ // wait util the third request is fulfilled
+ latch.await();
+ assertFalse(globalPool.isAvailable().isDone());
+ assertEquals(numberOfSegmentsToRequest,
segments3.size());
+
+ // recycle another 5 segments
+ globalPool.recycleMemorySegments(segments2);
+ assertTrue(globalPool.isAvailable().isDone());
+
+ // recycle the last 5 segments
+ globalPool.recycleMemorySegments(segments3);
+ assertTrue(globalPool.isAvailable().isDone());
+
+ } finally {
+ globalPool.destroy();
+ }
+ }
+
+ /**
+ * Tests that blocking request of multi local buffer pools can be
fulfilled by recycled segments
+ * to the global network buffer pool.
+ */
+ @Test(timeout = 10000L)
+ public void testBlockingRequestFromMultiLocalBufferPool() throws
Exception {
+ final int localPoolRequiredSize = 5;
+ final int localPoolMaxSize = 10;
+ final int numLocalBufferPool = 2;
+ final int numberOfSegmentsToRequest = 10;
+ final int numBuffers = numLocalBufferPool * localPoolMaxSize;
+
+ final ExecutorService executorService =
Executors.newFixedThreadPool(numLocalBufferPool);
+ final NetworkBufferPool globalPool = new
NetworkBufferPool(numBuffers, 128, numberOfSegmentsToRequest);
+ final List<BufferPool> localBufferPools = new
ArrayList<>(numLocalBufferPool);
+
+ try {
+ // create local buffer pools
+ for (int i = 0; i < numLocalBufferPool; ++i) {
+ final BufferPool localPool =
globalPool.createBufferPool(localPoolRequiredSize, localPoolMaxSize);
+ localBufferPools.add(localPool);
+ assertTrue(localPool.isAvailable().isDone());
+ }
+
+ // request some segments from the global pool in two
different ways
+ final List<MemorySegment> segments = new
ArrayList<>(numberOfSegmentsToRequest - 1);
+ for (int i = 0; i < numberOfSegmentsToRequest - 1; ++i)
{
+ segments.add(globalPool.requestMemorySegment());
+ }
+ final List<MemorySegment> exclusiveSegments =
globalPool.requestMemorySegments();
+ assertTrue(globalPool.isAvailable().isDone());
+ for (BufferPool localPool: localBufferPools) {
+ assertTrue(localPool.isAvailable().isDone());
+ }
+
+ // blocking request buffers form local buffer pools
+ final CountDownLatch latch = new
CountDownLatch(numLocalBufferPool);
+ final BlockingQueue<BufferBuilder> segmentsRequested =
new ArrayBlockingQueue<>(numBuffers);
+ for (final BufferPool localPool: localBufferPools) {
+ executorService.submit(() -> {
+ for (int num = localPoolMaxSize; num >
0; --num) {
+
ExceptionUtils.suppressExceptions(
+ () ->
segmentsRequested.add(localPool.requestBufferBuilderBlocking()));
+ }
+ latch.countDown();
+ });
+ }
+
+ // wait until all available buffers are requested
+ while (globalPool.getNumberOfAvailableMemorySegments()
> 0) {
+ Thread.sleep(100);
+ }
+
+ final CompletableFuture<?> globalPoolAvailableFuture =
globalPool.isAvailable();
+ assertFalse(globalPoolAvailableFuture.isDone());
+
+ final List<CompletableFuture<?>>
localPoolAvailableFutures = new ArrayList<>(numLocalBufferPool);
+ for (BufferPool localPool: localBufferPools) {
+ CompletableFuture<?> localPoolAvailableFuture =
localPool.isAvailable();
+
localPoolAvailableFutures.add(localPoolAvailableFuture);
+ assertFalse(localPool.isAvailable().isDone());
Review comment:
`localPool.isAvailable()` -> `localPoolAvailableFuture`
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
